module which supports allocation of memory from an mmap multiprocessingheap py c 20062008 r oudkerk licensed to psf under a contributor agreement inheritable class which wraps an mmap and from which blocks can be allocated a shared memory area backed by anonymous memory windows we have reopened a preexisting mmap reopen existing mmap xxx temporarily preventing buildbot failures while determining xxx the correct longterm fix see issue 23060 assert winapi getlasterror winapi erroralreadyexists a shared memory area backed by a temporary file posix arena is created anew if fd 1 it means we re coming from rebuildarena below choose a nonstorage backed directory if possible to improve performance class allowing allocation of chunks of memory from arenas minimum malloc alignment current arena allocation size a sorted list of available block sizes in arenas free block management map each block size to a list of arena start stop blocks map arena start tuple to the arena start stop block starting at that offset map arena stop tuple to the arena start stop block ending at that offset map arenas to their arena start stop blocks in use list of pending blocks to free see comment in free below statistics alignment must be a power of 2 create a new arena with at least the given size we carve larger and larger arenas for efficiency until we reach a largeish size roughly l3 cachesized possibly delete the given unused arena reusing an existing arena is faster than creating a new one so we only reclaim space if it s large enough returns a large enough block it might be much larger make block available and try to merge with its neighbours in the arena deregister this block so it can be merged with a neighbour arena is entirely free discard it from this process free all the blocks in the pending list called with the lock held free a block returned by malloc since free can be called asynchronously by the gc it could happen that it s called while self lock is held in that case self lock acquire would deadlock issue 12352 to avoid that a trylock is used instead and if the lock can t be acquired immediately the block is added to a list of blocks to be freed synchronously sometimes later from malloc or free by calling freependingblocks appending and retrieving from a list is not strictly threadsafe but under cpython it s atomic thanks to the gil can t acquire the lock right now add the block to the list of pending blocks to free we hold the lock return a block of right size possibly rounded up allow pending blocks to be marked available if the returned block is larger than necessary mark the remainder available class wrapping a block allocated out of a heap can be inherited by child process module which supports allocation of memory from an mmap multiprocessing heap py c 2006 2008 r oudkerk licensed to psf under a contributor agreement inheritable class which wraps an mmap and from which blocks can be allocated a shared memory area backed by anonymous memory windows we have reopened a preexisting mmap reopen existing mmap xxx temporarily preventing buildbot failures while determining xxx the correct long term fix see issue 23060 assert _winapi getlasterror _winapi error_already_exists a shared memory area backed by a temporary file posix arena is created anew if fd 1 it means we re coming from rebuild_arena below choose a non storage backed directory if possible to improve performance enough free space class allowing allocation of chunks of memory from arenas minimum malloc alignment 4 mb current arena allocation size a sorted list of available block sizes in arenas free block management map each block size to a list of arena start stop blocks map arena start tuple to the arena start stop block starting at that offset map arena stop tuple to the arena start stop block ending at that offset map arenas to their arena start stop blocks in use list of pending blocks to free see comment in free below statistics alignment must be a power of 2 create a new arena with at least the given size we carve larger and larger arenas for efficiency until we reach a large ish size roughly l3 cache sized possibly delete the given unused arena reusing an existing arena is faster than creating a new one so we only reclaim space if it s large enough returns a large enough block it might be much larger make block available and try to merge with its neighbours in the arena deregister this block so it can be merged with a neighbour arena is entirely free discard it from this process free all the blocks in the pending list called with the lock held free a block returned by malloc since free can be called asynchronously by the gc it could happen that it s called while self _lock is held in that case self _lock acquire would deadlock issue 12352 to avoid that a trylock is used instead and if the lock can t be acquired immediately the block is added to a list of blocks to be freed synchronously sometimes later from malloc or free by calling _free_pending_blocks appending and retrieving from a list is not strictly thread safe but under cpython it s atomic thanks to the gil can t acquire the lock right now add the block to the list of pending blocks to free we hold the lock return a block of right size possibly rounded up reinitialize after fork allow pending blocks to be marked available if the returned block is larger than necessary mark the remainder available class wrapping a block allocated out of a heap can be inherited by child process
import bisect
from collections import defaultdict
import mmap
import os
import sys
import tempfile
import threading

from .context import reduction, assert_spawning
from . import util

__all__ = ['BufferWrapper']





if sys.platform == 'win32':

    import _winapi

    class Arena(object):
        

        _rand = tempfile._RandomNameSequence()

        def __init__(self, size):
            self.size = size
            for i in range(100):
                name = 'pym-%d-%s' % (os.getpid(), next(self._rand))
                buf = mmap.mmap(-1, size, tagname=name)
                if _winapi.GetLastError() == 0:
                    break
                
                buf.close()
            else:
                raise FileExistsError('Cannot find name for new mmap')
            self.name = name
            self.buffer = buf
            self._state = (self.size, self.name)

        def __getstate__(self):
            assert_spawning(self)
            return self._state

        def __setstate__(self, state):
            self.size, self.name = self._state = state
            
            self.buffer = mmap.mmap(-1, self.size, tagname=self.name)
            
            
            

else:

    class Arena(object):
        

        if sys.platform == 'linux':
            _dir_candidates = ['/dev/shm']
        else:
            _dir_candidates = []

        def __init__(self, size, fd=-1):
            self.size = size
            self.fd = fd
            if fd == -1:
                
                
                self.fd, name = tempfile.mkstemp(
                     prefix='pym-%d-'%os.getpid(),
                     dir=self._choose_dir(size))
                os.unlink(name)
                util.Finalize(self, os.close, (self.fd,))
                os.ftruncate(self.fd, size)
            self.buffer = mmap.mmap(self.fd, self.size)

        def _choose_dir(self, size):
            
            
            for d in self._dir_candidates:
                st = os.statvfs(d)
                if st.f_bavail * st.f_frsize >= size:  
                    return d
            return util.get_temp_dir()

    def reduce_arena(a):
        if a.fd == -1:
            raise ValueError('Arena is unpicklable because '
                             'forking was enabled when it was created')
        return rebuild_arena, (a.size, reduction.DupFd(a.fd))

    def rebuild_arena(size, dupfd):
        return Arena(size, dupfd.detach())

    reduction.register(Arena, reduce_arena)





class Heap(object):

    
    _alignment = 8

    _DISCARD_FREE_SPACE_LARGER_THAN = 4 * 1024 ** 2  
    _DOUBLE_ARENA_SIZE_UNTIL = 4 * 1024 ** 2

    def __init__(self, size=mmap.PAGESIZE):
        self._lastpid = os.getpid()
        self._lock = threading.Lock()
        
        self._size = size
        
        self._lengths = []

        
        
        self._len_to_seq = {}
        
        
        self._start_to_block = {}
        
        
        self._stop_to_block = {}

        
        self._allocated_blocks = defaultdict(set)
        self._arenas = []

        
        self._pending_free_blocks = []

        
        self._n_mallocs = 0
        self._n_frees = 0

    @staticmethod
    def _roundup(n, alignment):
        
        mask = alignment - 1
        return (n + mask) & ~mask

    def _new_arena(self, size):
        
        length = self._roundup(max(self._size, size), mmap.PAGESIZE)
        
        
        if self._size < self._DOUBLE_ARENA_SIZE_UNTIL:
            self._size *= 2
        util.info('allocating a new mmap of length %d', length)
        arena = Arena(length)
        self._arenas.append(arena)
        return (arena, 0, length)

    def _discard_arena(self, arena):
        
        length = arena.size
        
        
        if length < self._DISCARD_FREE_SPACE_LARGER_THAN:
            return
        blocks = self._allocated_blocks.pop(arena)
        assert not blocks
        del self._start_to_block[(arena, 0)]
        del self._stop_to_block[(arena, length)]
        self._arenas.remove(arena)
        seq = self._len_to_seq[length]
        seq.remove((arena, 0, length))
        if not seq:
            del self._len_to_seq[length]
            self._lengths.remove(length)

    def _malloc(self, size):
        
        i = bisect.bisect_left(self._lengths, size)
        if i == len(self._lengths):
            return self._new_arena(size)
        else:
            length = self._lengths[i]
            seq = self._len_to_seq[length]
            block = seq.pop()
            if not seq:
                del self._len_to_seq[length], self._lengths[i]

        (arena, start, stop) = block
        del self._start_to_block[(arena, start)]
        del self._stop_to_block[(arena, stop)]
        return block

    def _add_free_block(self, block):
        
        (arena, start, stop) = block

        try:
            prev_block = self._stop_to_block[(arena, start)]
        except KeyError:
            pass
        else:
            start, _ = self._absorb(prev_block)

        try:
            next_block = self._start_to_block[(arena, stop)]
        except KeyError:
            pass
        else:
            _, stop = self._absorb(next_block)

        block = (arena, start, stop)
        length = stop - start

        try:
            self._len_to_seq[length].append(block)
        except KeyError:
            self._len_to_seq[length] = [block]
            bisect.insort(self._lengths, length)

        self._start_to_block[(arena, start)] = block
        self._stop_to_block[(arena, stop)] = block

    def _absorb(self, block):
        
        (arena, start, stop) = block
        del self._start_to_block[(arena, start)]
        del self._stop_to_block[(arena, stop)]

        length = stop - start
        seq = self._len_to_seq[length]
        seq.remove(block)
        if not seq:
            del self._len_to_seq[length]
            self._lengths.remove(length)

        return start, stop

    def _remove_allocated_block(self, block):
        arena, start, stop = block
        blocks = self._allocated_blocks[arena]
        blocks.remove((start, stop))
        if not blocks:
            
            self._discard_arena(arena)

    def _free_pending_blocks(self):
        
        while True:
            try:
                block = self._pending_free_blocks.pop()
            except IndexError:
                break
            self._add_free_block(block)
            self._remove_allocated_block(block)

    def free(self, block):
        
        
        
        
        
        
        
        
        
        if os.getpid() != self._lastpid:
            raise ValueError(
                "My pid ({0:n}) is not last pid {1:n}".format(
                    os.getpid(),self._lastpid))
        if not self._lock.acquire(False):
            
            
            self._pending_free_blocks.append(block)
        else:
            
            try:
                self._n_frees += 1
                self._free_pending_blocks()
                self._add_free_block(block)
                self._remove_allocated_block(block)
            finally:
                self._lock.release()

    def malloc(self, size):
        
        if size < 0:
            raise ValueError("Size {0:n} out of range".format(size))
        if sys.maxsize <= size:
            raise OverflowError("Size {0:n} too large".format(size))
        if os.getpid() != self._lastpid:
            self.__init__()                     
        with self._lock:
            self._n_mallocs += 1
            
            self._free_pending_blocks()
            size = self._roundup(max(size, 1), self._alignment)
            (arena, start, stop) = self._malloc(size)
            real_stop = start + size
            if real_stop < stop:
                
                
                self._add_free_block((arena, real_stop, stop))
            self._allocated_blocks[arena].add((start, real_stop))
            return (arena, start, real_stop)





class BufferWrapper(object):

    _heap = Heap()

    def __init__(self, size):
        if size < 0:
            raise ValueError("Size {0:n} out of range".format(size))
        if sys.maxsize <= size:
            raise OverflowError("Size {0:n} too large".format(size))
        block = BufferWrapper._heap.malloc(size)
        self._state = (block, size)
        util.Finalize(self, BufferWrapper._heap.free, args=(block,))

    def create_memoryview(self):
        (arena, start, stop), size = self._state
        return memoryview(arena.buffer)[start:start+size]