如何在python中为线程或进程设置内存限制？

Question

我正在用Python编写程序，它能够在某种类型的沙箱中运行不受信任的Python代码。所以，我需要一种方法来限制不可信代码可以分配的内存量。现在，我可以通过覆盖我的沙盒环境中的默认python数据结构来限制范围（），列表，字典和其他的最大长度。

任何想法？

Answer 1

在Unix，可以使用resource.setrlimit(resource.RLIMIT_AS, ...)限制 “的可采取的进程的地址空间的最大面积（字节）”。

import sys 
import resource 
soft, hard = 10**7, 10**7 
# soft, hard = 10**8, 10**8 # uncommenting this allows program to finish 
resource.setrlimit(resource.RLIMIT_AS,(soft, hard)) 
memory_hog = {} 
try: 
    for x in range(10000): 
     print(x) 
     memory_hog[str(x)]='The sky is so blue' 
except MemoryError as err: 
    sys.exit('memory exceeded') 
    # memory exceeded 

Answer 2

我不知道它是如何在Windows下完成的。希望其他人可以提供解决方案的这一部分 。

下面是一些示例代码使用​​设置Windows上的限制...

import ctypes 

PROCESS_SET_QUOTA = 0x100 
PROCESS_TERMINATE = 0x1 
JobObjectExtendedLimitInformation = 9 
JOB_OBJECT_LIMIT_PROCESS_MEMORY = 0x100 


class IO_COUNTERS(ctypes.Structure): 
    _fields_ = [('ReadOperationCount', ctypes.c_uint64), 
       ('WriteOperationCount', ctypes.c_uint64), 
       ('OtherOperationCount', ctypes.c_uint64), 
       ('ReadTransferCount', ctypes.c_uint64), 
       ('WriteTransferCount', ctypes.c_uint64), 
       ('OtherTransferCount', ctypes.c_uint64)] 


class JOBOBJECT_BASIC_LIMIT_INFORMATION(ctypes.Structure): 
    _fields_ = [('PerProcessUserTimeLimit', ctypes.c_int64), 
       ('PerJobUserTimeLimit', ctypes.c_int64), 
       ('LimitFlags', ctypes.c_uint32), 
       ('MinimumWorkingSetSize', ctypes.c_void_p), 
       ('MaximumWorkingSetSize', ctypes.c_void_p), 
       ('ActiveProcessLimit', ctypes.c_uint32), 
       ('Affinity', ctypes.c_void_p), 
       ('PriorityClass', ctypes.c_uint32), 
       ('SchedulingClass', ctypes.c_uint32)] 


class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure): 
    _fields_ = [('BasicLimitInformation', JOBOBJECT_BASIC_LIMIT_INFORMATION), 
       ('IoInfo', IO_COUNTERS), 
       ('ProcessMemoryLimit', ctypes.c_void_p), 
       ('JobMemoryLimit', ctypes.c_void_p), 
       ('PeakProcessMemoryUsed', ctypes.c_void_p), 
       ('PeakJobMemoryUsed', ctypes.c_void_p)] 


# Set memory limit for process with specfied 'pid', to specified 'size' in bytes 
def set_limit(pid, size): 

    job_info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION() 
    out_size = ctypes.c_uint32() 

    job = ctypes.windll.kernel32.CreateJobObjectA(None, None) 
    assert job != 0 

    success = ctypes.windll.kernel32.QueryInformationJobObject(job, 
                   JobObjectExtendedLimitInformation, 
                   ctypes.POINTER(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)(job_info), 
                   ctypes.sizeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION), 
                   ctypes.POINTER(ctypes.c_uint32)(out_size)) 
    assert success 

    job_info.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_PROCESS_MEMORY 
    job_info.ProcessMemoryLimit = size 
    success = ctypes.windll.kernel32.SetInformationJobObject(job, 
                  JobObjectExtendedLimitInformation, 
                  ctypes.POINTER(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)(job_info), 
                  ctypes.sizeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)) 
    assert success 

    process = ctypes.windll.kernel32.OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, 
               False, pid) 
    assert process != 0 

    success = ctypes.windll.kernel32.AssignProcessToJobObject(job, process) 
    assert success 


    success = ctypes.windll.kernel32.CloseHandle(job) 
    assert success 

    success = ctypes.windll.kernel32.CloseHandle(process) 
    assert success 


if __name__ == '__main__': 

    import os 

    five_mb = 5 * 1024 * 1024 

    def can_we_allocate_five_mb(): 
     try: 
      s = 'x' * five_mb 
      return True 
     except MemoryError: 
      return False 

    print can_we_allocate_five_mb() 
    set_limit(os.getpid(), five_mb) 
    print can_we_allocate_five_mb() 

...虽然可能没有必要创建一个单独的作业为每个进程对象 - 你应该能够将所有受限制的流程与单个工作联系起来。