Python的低级内存操作：locmem

最近，我想研究研究内存底层操作的强烈好奇心（绝不是因为闲得慌！），想看看能否在 Python 中实现更细粒度的内存管理，甚至动态执行代码，或者探索是否能通过这种方式优化某些特定场景下的内存使用。然后我整出了 locmem 这个库。

Github: https://github.com/moyanj/locmem
Pypi: https://pypi.org/project/locmem/

locmem 的基本使用：一个简单的加法示例

为了展示 locmem 的核心功能，我们从一个最简单的整数加法开始。在这里，我们不再依赖 Python 解释器默认的整型对象及其内存管理，而是手动在操作系统层面分配和管理内存来存储和操作整数。

from locmem import Int

# 创建两个 Int 对象，它们会在底层分配内存并存储整数值
a = Int(1)  # 相当于在操作系统内存中分配空间存储整数 1
b = Int(2)  # 相当于在操作系统内存中分配空间存储整数 2

# 执行加法操作。locmem 会处理读取底层内存中的值，进行计算，并将结果写入新的内存空间
c = a + b

# 访问计算结果。通过 .value 属性获取底层内存中存储的实际值
print(f"a 的值: {a.value}")  # 输出: a 的值: 1
print(f"b 的值: {b.value}")  # 输出: b 的值: 2
print(f"c (a + b) 的值: {c.value}") # 输出: c (a + b) 的值: 3

# 手动释放内存。这与 C 语言中的 free() 类似，将之前分配的内存归还给操作系统。
# 在实际应用中，管理好内存的生命周期至关重要，避免内存泄露。
a.free()
b.free()
c.free()

# 示例：使用 with 语句进行自动释放，更安全地管理内存
with Int(10) as x:
    with Int(20) as y:
        result = x + y
        print(f"使用 with 语句计算：{result.value}")
        result.free() # result 仍需手动释放，或者将其也封装在 with 语句中
# x 和 y 会在离开 with 块时自动调用 free() 释放内存

这是其等效的C代码

#include <stdlib.h>
#include <stdio.h>

int main() {
    int *a = malloc(sizeof(int));
    *a = 1;
    int *b = malloc(sizeof(int));
    *b = 2;
    int *c = malloc(sizeof(int));
    *c = *a + *b;
    printf("a的值%d\n", *a);
    printf("b的值%d\n", *b);
    printf("c的值%d\n", *c);
    free(a);
    free(b);
    free(c);
}

通过这个例子，我们可以看到 locmem 的基本工作流程：创建数据类型实例 -> 执行操作 -> 读取结果 -> 手动或自动释放内存。这与 Python 对象自动垃圾回收的机制截然不同，它赋予了开发者更多的控制权。

基本原理

你可能会觉得，进行低级内存操作非常复杂。但实际上，它的核心原理却异常简洁！在 Linux/Unix-like 系统上，locmem 主要是通过 mmap 系统调用来实现向操作系统申请和释放内存。

大概这样：

import ctypes
libc = ctypes.CDLL('libc.so.6') # Python里面加载libc

libc.mmap.argtypes = [
    ctypes.c_void_p,  # addr: 建议的映射起始地址，通常设为 0 表示由系统选择
    ctypes.c_size_t,  # len: 映射区域的长度（字节）
    ctypes.c_int,     # prot: 内存保护（读、写、执行）
    ctypes.c_int,     # flags: 映射类型（共享、私有、匿名等）
    ctypes.c_int,     # fd: 文件描述符（匿名映射时为 -1）
    ctypes.c_void_p   # offset: 文件偏移量（匿名映射时为 0）
]
libc.mmap.restype = ctypes.c_void_p # 返回映射到的内存地址

# 定义 munmap 函数的签名
libc.munmap.argtypes = [
    ctypes.c_void_p,  # addr: 要解除映射区域的起始地址
    ctypes.c_size_t   # len: 要解除映射区域的长度
]
libc.munmap.restype = ctypes.c_int # 返回状态码，0 表示成功

# 内存保护标志
PROT_READ = 0x1    # 可读
PROT_WRITE = 0x2   # 可写
PROT_EXEC = 0x4    # 可执行（注意：谨慎使用）

# 映射类型标志
MAP_PRIVATE = 0x0002 # 私有映射，对映射区域的修改不会影响文件内容
MAP_ANONYMOUS = 0x20 # 匿名映射，不与任何文件关联，通常结合 MAP_PRIVATE 使用

# 申请1K内存，并拥有读写权限
addr = libc.mmap(0, MEM_SIZE, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0)

# 释放内存
libc.munmap(addr, 1024)

通过 ctypes 调用 mmap，我们成功地在 Python 中直接向操作系统请求了一块内存空间。这就为 locmem 的所有功能奠定了基础——既然我们可以申请和释放内存，那么我们就可以在这块内存上进行各种自定义的数据存储和操作。

架构设计

1. 抽象层：跨平台内存接口

考虑到不同操作系统（如 POSIX 兼容系统和 Windows）有不同的内存管理 API，locmem 首先定义了一个抽象基类 BaseMemory，用于规范底层内存操作的接口。这样，我们就可以根据当前运行的平台选择不同的具体实现，从而实现跨平台兼容性。

from abc import ABC, abstractmethod


class BaseMemory(ABC):
    def __init__(self):
        pass

    @abstractmethod
    def get(self, size: int, executable: bool = False) -> int:
        raise NotImplementedError

    @abstractmethod
    def release(self, address: int, size: int):
        raise NotImplementedError

get 方法类似于 C 语言中的 malloc 函数，用于从操作系统请求一块内存；而 release 方法则对应 free，用于将内存归还。executable 参数的引入，使得我们可以控制内存的执行权限，这对于动态代码生成（JIT 编译）等高级场景非常有用。

2. 内存管理器

仅仅能够申请和释放内存是不够的，我们还需要一个高效的内存管理器来细粒度地管理这些大块的内存空间，满足应用程序频繁的小块内存请求。我在下面列出了一些经典的内存分配算法：

• 堆分配器 (Heap Allocator):

  • 原理: 最简单的实现方式之一，通常使用一个哈希表或其他数据结构来跟踪已分配内存块的起始地址和大小。当有新的分配请求时，他会直接向操作系统申请内存，并记录下分配的起始地址和大小。
  • 特点: 实现简单，但效率很低。

• 分箱分配器 (Segregated Fit Allocator / Bin Allocator):

  • 原理: 将内存请求按照大小划分到不同的“箱子”或“链表”中。每个箱子专门管理特定大小范围的内存块。
  • 特点: 非常适合处理频繁的小块内存分配和释放，能够显著减少内部碎片，并提高分配速度。

• Buddy 分配器 (Buddy System Allocator):

  • 原理: 一种基于二叉树结构的分配算法。它将内存块分成大小为 2 的幂次方的小块（如 1KB, 2KB, 4KB 等）。当请求一块内存时，它会寻找最小的能容纳该请求的 2 的幂次方的块，如果太大则将其递归地分裂成两个“伙伴（buddy）”块，直到找到合适的大小。释放时，如果相邻的伙伴块也空闲，它们会被合并回更大的块。
  • 特点: 可以实现快速的合并和分割操作，减少外部碎片，但由于只能分配 2 的幂次方大小的块，可能会导致内部碎片（例如请求 3KB 内存却分配了 4KB）。

最终，我在 locmem 中同时实现了堆分配器和分箱分配器，并将分箱分配器设为默认选项，以兼顾易用性和性能。

此外，locmem 还设计为可以轻松集成外部高性能内存分配器。例如，我们可以通过 ctypes 调用优秀的 C 库分配器，如 mimalloc 和 tcmalloc。这为 locmem 提供了更强大的性能扩展能力。

注意： jemalloc 是一个非常优秀的内存分配器，但它会接管整个进程的 malloc/free，这会与 Python 自身的内存管理机制冲突。

3. 上层封装接口：Pythonic 的数据类型

有了底层的内存操作能力和高效的内存管理器，剩下的工作就是向上层封装，提供 Python 开发者熟悉的接口。locmem 实现了一系列数据类型（如 Int, Float, Char, Bytes 等），它们在底层都使用 Pointer 类来管理实际的内存地址。

以 Int 类型的实现为例，它封装了 C 语言中 int 类型（通常是 4 字节）的内存操作：

import struct
from locmem.allocator import alloc, free
from locmem.core import Pointer, memread, memwrite

class Int:
    _fmt = "i"  # C-style int, 通常是 4 字节
    _size = struct.calcsize(_fmt)

    def __init__(self, initial_value: int):
        self.size = self._size
        # 通过全局的 alloc 函数分配内存，并封装为 Pointer 对象
        self.ptr: Pointer = alloc(self.size)
        self.set(initial_value)

    def get(self) -> int:
        """从内存中读取并解包数据。"""
        # 从 self.ptr 指向的内存位置读取 self.size 字节的数据
        packed_data = memread(self.ptr, self.size)
        # 使用 struct 模块按照指定格式解包数据
        return struct.unpack(self._fmt, packed_data)[0]

    def set(self, value: int) -> None:
        """将数据打包并写入内存。"""
        # 使用 struct 模块按照指定格式打包数据
        packed_data = struct.pack(self._fmt, value)
        # 将打包后的数据写入 self.ptr 指向的内存位置
        memwrite(self.ptr, packed_data)

    @property
    def value(self) -> int:
        """属性访问器，提供更 Pythonic 的方式获取存储的值。"""
        return self.get()

    @value.setter
    def value(self, new_value: int) -> None:
        """允许通过属性直接设置值。"""
        self.set(new_value)

    def __del__(self):
        """
        析构函数，当 Int 对象被垃圾回收时，自动释放底层内存。
        注意：Python 的垃圾回收机制不保证 __del__ 的精确调用时机，
        因此手动调用 free() 或使用 with 语句是更推荐的做法。
        """
        # 仅在内存尚未被释放的情况下调用 free()
        if self.ptr:
            self.free()

    def free(self) -> None:
        """手动释放底层内存。这是推荐的内存管理方式。"""
        if self.ptr:
            free(self.ptr)
            self.ptr = None  # 清除指针，避免重复释放

    def __enter__(self):
        """支持 with 语句的上下文管理器协议，进入时返回自身。"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持 with 语句的上下文管理器协议，离开时自动释放内存。"""
        self.free()

    @classmethod
    def from_bytes(cls, data: bytes):
        """
        类方法：从字节序列创建 Int 对象。
        这允许从外部二进制数据恢复 Int 实例。
        """
        if len(data) != cls._size:
            raise ValueError(f"Input bytes length {len(data)} does not match Int size {cls._size}")
        # 使用类的格式字符串解包数据
        value = struct.unpack(cls._fmt, data)[0]
        # 用解包后的值创建一个新的实例
        return cls(value)

在这个 Int 类的实现中，struct 模块扮演了关键角色，它负责将 Python 数据类型（如整数）与底层字节序列进行相互转换。memread 和 memwrite 是与底层内存管理器交互的接口函数，负责最终的内存读写。

通过实现__del__、free和上下文管理器协议（__enter__/__exit__），locmem 试图提供多种内存释放机制，让开发者能够灵活选择是手动管理还是通过 with 语句进行自动管理，从而避免内存泄露。

核心思想：

• 分配 (alloc): 调用底层的 mmap 或 VirtualAlloc 获取指定大小的内存块。
• 读写 (memread/memwrite): 使用 ctypes 或其他低级方法将 Python 对象的数据打包成字节，然后写入或读出到（和从）分配的内存地址。
• 释放 (free): 调用底层的 munmap 或 VirtualFree 将内存归还给操作系统。

---

locmem 的应用场景与展望

既然 locmem 提供了这种低级内存操作的能力，那么它能解决什么问题呢？

1. 高性能数据结构： 对于需要紧凑内存布局或频繁随机访问的超大型数据结构（例如，模拟一些硬件内存，或者需要像 C 语言结构体一样存储复杂数据），locmem 可以提供比纯 Python 列表/字典更可控的内存占用和潜在的访问效率。
2. 外部库互操作性： 当与 C/C++ 库进行深度集成时，locmem 可以方便地创建 Python 端的数据结构，其内存地址可以直接传递给 C 函数，无需额外的复制，从而提高数据传输效率。
3. 动态代码生成 (JIT): 如果内存块被标记为可执行 (PROT_EXEC)，可以将机器码写入这块内存，然后将其转换为函数指针并调用。这为在 Python 中实现简易的 JIT 编译器或运行时代码生成提供了基础（当然，这需要高级的汇编知识和谨慎的安全考虑）。
4. 内存池管理： 对于需要频繁分配和释放小块内存的应用，可以基于 locmem 构建更复杂的内存池，减少系统调用的开销，优化整体性能。
5. 教学与研究： 对于想深入理解操作系统内存管理、ctypes 用法或 Python 底层机制的开发者，locmem 提供了一个绝佳的实验平台。

locmem 并非旨在取代 Python 默认的内存管理和垃圾回收机制，它们在绝大多数场景下都非常高效和方便。然而，对于那些需要极致性能、底层控制或特定系统级交互的场景，locmem 提供了一个有趣的、有力的工具。我希望这个库能激发更多 Python 开发者对底层原理的探索热情！