问题背景

扫描器的基本功能包括对某个主机列表的端口列表做扫描，为了实现这种需求，曾经我写过类似下面的代码

def get_host():
    ret_host_list = []
    ...   # 从api中获取扫描主机ip
    return ret_host_list


def get_port():
    ret_port_list = []
    ...
    return ret_port_list


def generate_targets():
    ret_target_list = []
    for host in get_host():
        for port in get_port():
            ret_target_list.append(
                {
                    "host": host,
                    "port": port
                }
            )
    return ret_target_list


for target in generate_targets():
    ...

不知道你能不能看出来问题所在：上面的代码，当扫描的主机和端口都比较少时没什么问题，但是当主机和端口很多时，就会占用大量内存。

本文记录两个问题：

• 怎么改进上面的代码，避免内存占用过大的问题
• 研究为什么会占用大量内存

过程

• 复现

  我们先来写一个demo复现这个"内存占用"过大的问题

  import sys
  
  def generate_target():
    result = []
    for i in range(0, 256 * 256):
      for j in range(8000, 8050):
        result.append({"host":i, "port":j})
    return result
  
  result = generate_target()
  
  print("ok")
  sys.stdin.readline()
  

  运行上面的脚本，通过free -m命令可以观察到物理内存接近减少900M。

  如果range(8000,8050)修改成range(8000,9000)，也即扫描8000-9000端口时，内存至少减少12G(因为我的测试机器只有12G的物理内存，所以只能得到这个数字)。

• 怎么改进上面的代码，避免内存问题？

  看着像是因为生成大量的{"host":i, "port":j}的扫描对象，所以才占用很多内存。

  那么改进很简单，如果我们将列表改成"生成器"，就不用在generate_target函数生成所有{"host":i, "port":j}的扫描对象。

  比如在函数中用yield关键字

  import sys
  
  def generate_target():
    result = []
    for i in range(0, 256 * 256):
      for j in range(8000, 9000):
        yield {"host":i, "port":j}
  
  result = generate_target()
  
  print("ok")
  sys.stdin.readline()
  

  或者用生成器表达式^[1]

  import sys
  
  def generate_target():
    return ({"host":i, "port":j} for i in range(0, 256 * 256) for j in range(8000, 9000))
  
  result = generate_target()
  
  print("ok")
  sys.stdin.readline()
  

  关于"生成器"的概念，可以参考 廖雪峰的教程^[2]。

  如果对"生成器"的实现感兴趣，可以参考 重新认识生成器generator^[3]

  复现脚本消耗了接近900MB的物理内存，难道真的"区区几个"dict就能占用这么多内存吗？

• 为什么会占用大量内存？

  我们可以用pympler库来看看python程序中的对象都占用了多少内存，修改后的脚本如下。

  [root@instance-fj5pftdp tmp]# cat 20.py
  import sys
  from pympler import tracker, muppy, summary
  
  
  def print_mem():
    all_objects = muppy.get_objects()
    sum = summary.summarize(all_objects)
    summary.print_(sum)
  
  def generate_target():
    result = []
    for i in range(0, 256 * 256):
      for j in range(8000, 8050):
        result.append({"host":i, "port":j})
    return result
  
  result = generate_target()
  
  print_mem()
  sys.stdin.readline()
  

  执行后，可以看到有3278449个dict实例，总共占用750.76MB。

  3278449约等于256*256*50,和脚本中的循环次数吻合。

  那这里一个dict实例占用多少个字节呢？我们用sys.getsizeof函数可以看到，在python3.6中{}和{"host":"1","port":"1"}都占用了240字节

  这里一个dict实例占用240个字节，总共有3278449个实例，算一下确实会占用750MB，接近900MB。

  差不多我最开始的疑问都解开了，只有最后一个疑问。

  到这里我不知道你会不会和我一样奇怪：为啥{}啥也没存储，sys.getsizeof显示占用240字节，而{"host":"1","port":"1"}明显多了点字符串，sys.getsizeof为啥仍显示占用240字节。

• 为啥sys.getsizeof告诉我们{}占用240字节？

  这个现象是分python版本的，比如python3.8版本如下

  我想如果知道sys.getsizeof是怎么计算内存占用的，我们就知道它的结果是什么意思。于是我就去翻文档和看源码。

  翻了下文档，没找到sys.getsizeof的计算过程，于是只好去看下CPython代码看下sys.getsizeof的实现。

  在Python/sysmodule.c中可以看出来：sys.getsizeof等于 __sizeof__() + GC头大小(16字节)

  size_t
  _PySys_GetSizeOf(PyObject *o)
  {
      PyObject *res = NULL;
      PyObject *method;
      Py_ssize_t size;
  
      ...
  
      method = _PyObject_LookupSpecial(o, &PyId___sizeof__);    # 对象的sizeof函数
      ...
      res = _PyObject_CallNoArg(method);
      ...
  
      size = PyLong_AsSsize_t(res);
      ...
      if (PyObject_IS_GC(o))    # 容器对象(list、dict)会有GC头，str、int等没有GC头。GC头用来做垃圾回收
          return ((size_t)size) + sizeof(PyGC_Head);
      return (size_t)size;
  }
  

  下面也可以验证上面的结论

  ➜  cpython-3.8 ./python.exe
  Python 3.8.12+ (default, Jan  1 2022, 12:15:13)
  ...
  >>> [].__sizeof__()
  40
  >>> sys.getsizeof([])
  56
  >>> {"host":"1"}.__sizeof__()
  216
  >>> sys.getsizeof({"host":"1"})
  232
  

  __sizeof__()是什么呢？每种类型的sizeof函数实现逻辑不同，dict类型的sizeof函数就是Objects/dictobject.c中的dict_sizeof函数。

  你可以动态调试，或者翻一翻文件，最终能看到计算过程，如下

  Py_ssize_t
  _PyDict_SizeOf(PyDictObject *mp)    // mp就是dict的实例
  {
      Py_ssize_t size, usable, res;
  
      size = DK_SIZE(mp->ma_keys);    // 哈希表的大小，也就是PyDictObject数据结构中dk_indices数组的大小。这个场景下是8字节
      usable = USABLE_FRACTION(size); // size的三分之二，也就是5
  
      res = _PyObject_SIZE(Py_TYPE(mp));  // PyDictObject数据结构的大小，48字节
      if (mp->ma_values)
          res += usable * sizeof(PyObject*);
      /* If the dictionary is split, the keys portion is accounted-for
         in the type object. */
      if (mp->ma_keys->dk_refcnt == 1)
          res += (sizeof(PyDictKeysObject)    // 除两个数组外有 5 个字段，共 40 字节
                  + DK_IXSIZE(mp->ma_keys) * size  // dk_indices索引数组占用的大小，这个场景下是8字节
                  + sizeof(PyDictKeyEntry) * usable);  // 键值对数组，长度为5 。每个 PyDictKeyEntry 结构体 24 字节，共 120 字节
      return res;
  }
  

  关于PyDictObject数据结构，你可以参考 dict 对象，高效的关联式容器^[4]。

  根据上面的内容，48+40+8+120+16刚好就是232字节，也就是python3.8下sys.getsizeof({"host":1})的结果。

总结

问题背景中的场景可能还有其他的编程方式来实现，这里我只是为了引出我学到的"生成器"和"Python内置对象的内存占用"两个知识点。这两个点都背后能扯到更多的点，比如dict容器的动态扩容、哈希表冲突的解决、迭代器设计模式，如果有兴趣，推荐你可以看文章中的参考资料。

想起我以前老听说python性能不好，只以为是解释运行得慢。通过这个案例和参考资料的学习，感觉到还可以从"内存"方面比较。和c相比，python对象的内存占用会多一点，比如空字符串c中就占用1个字节，python中占用49个字节。

生成器可以节约内存。
参考资料