heap.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227
  1. #
  2. # Module which supports allocation of memory from an mmap
  3. #
  4. # processing/heap.py
  5. #
  6. # Copyright (c) 2007-2008, R Oudkerk --- see COPYING.txt
  7. #
  8. import bisect
  9. import mmap
  10. import tempfile
  11. import os
  12. import sys
  13. import thread
  14. import itertools
  15. from processing import _processing
  16. from processing.finalize import Finalize
  17. from processing.logger import info
  18. from processing.forking import assertSpawning
  19. __all__ = ['BufferWrapper']
  20. #
  21. # Inheirtable class which wraps an mmap, and from which blocks can be allocated
  22. #
  23. if sys.platform == 'win32':
  24. from processing._processing import win32
  25. class Arena(object):
  26. _nextid = itertools.count().next
  27. def __init__(self, size):
  28. self.size = size
  29. self.name = 'pym-%d-%d' % (os.getpid(), Arena._nextid())
  30. self.buffer = mmap.mmap(0, self.size, tagname=self.name)
  31. assert win32.GetLastError() == 0, 'tagname already in use'
  32. self._state = (self.size, self.name)
  33. def __getstate__(self):
  34. assertSpawning(self)
  35. return self._state
  36. def __setstate__(self, state):
  37. self.size, self.name = self._state = state
  38. self.buffer = mmap.mmap(0, self.size, tagname=self.name)
  39. assert win32.GetLastError() == win32.ERROR_ALREADY_EXISTS
  40. else:
  41. class Arena(object):
  42. def __init__(self, size):
  43. if sys.version_info >= (2, 5, 0):
  44. self.buffer = mmap.mmap(-1, size)
  45. self.size = size
  46. self.name = None
  47. else:
  48. fd, self.name = tempfile.mkstemp(prefix='pym-')
  49. self.size = remaining = size
  50. while remaining > 0:
  51. remaining -= os.write(fd, '\0' * remaining)
  52. self.buffer = mmap.mmap(fd, size)
  53. os.close(fd)
  54. if sys.platform == 'cygwin':
  55. # cannot unlink file until it is no longer in use
  56. def _finalize_heap(mmap, unlink, name):
  57. mmap.close()
  58. unlink(name)
  59. Finalize(
  60. self, _finalize_heap,
  61. args=(self.buffer, os.unlink, name),
  62. exitpriority=-10
  63. )
  64. else:
  65. os.unlink(self.name)
  66. def __getstate__(self):
  67. assertSpawning(self) # will fail
  68. #
  69. # Class allowing allocation of chunks of memory from arenas
  70. #
  71. class Heap(object):
  72. _alignment = 8
  73. def __init__(self, size=4096):
  74. self._lastpid = os.getpid()
  75. self._lock = thread.allocate_lock()
  76. self._size = size
  77. self._lengths = []
  78. self._len_to_seq = {}
  79. self._start_to_block = {}
  80. self._stop_to_block = {}
  81. self._allocated_blocks = set()
  82. self._arenas = []
  83. def _roundUp(self, n):
  84. mask = self._alignment - 1
  85. return (max(1, n) + mask) & ~mask
  86. def _malloc(self, size):
  87. # returns a large enough block -- it might be much larger
  88. i = bisect.bisect_left(self._lengths, size)
  89. if i == len(self._lengths):
  90. length = max(self._size, size)
  91. self._size *= 2
  92. info('allocating a new mmap of length %d', length)
  93. arena = Arena(length)
  94. self._arenas.append(arena)
  95. return (arena, 0, length)
  96. else:
  97. length = self._lengths[i]
  98. seq = self._len_to_seq[length]
  99. block = seq.pop()
  100. if not seq:
  101. del self._len_to_seq[length], self._lengths[i]
  102. (arena, start, stop) = block
  103. del self._start_to_block[(arena, start)]
  104. del self._stop_to_block[(arena, stop)]
  105. return block
  106. def _free(self, block):
  107. # free location and try to merge with neighbours
  108. (arena, start, stop) = block
  109. try:
  110. prev_block = self._stop_to_block[(arena, start)]
  111. except KeyError:
  112. pass
  113. else:
  114. start, _ = self._absorb(prev_block)
  115. try:
  116. next_block = self._start_to_block[(arena, stop)]
  117. except KeyError:
  118. pass
  119. else:
  120. _, stop = self._absorb(next_block)
  121. block = (arena, start, stop)
  122. length = stop - start
  123. try:
  124. self._len_to_seq[length].append(block)
  125. except KeyError:
  126. self._len_to_seq[length] = [block]
  127. bisect.insort(self._lengths, length)
  128. self._start_to_block[(arena, start)] = block
  129. self._stop_to_block[(arena, stop)] = block
  130. def _absorb(self, block):
  131. # deregister this block so it can be merged with a neighbour
  132. (arena, start, stop) = block
  133. del self._start_to_block[(arena, start)]
  134. del self._stop_to_block[(arena, stop)]
  135. length = stop - start
  136. seq = self._len_to_seq[length]
  137. seq.remove(block)
  138. if not seq:
  139. del self._len_to_seq[length]
  140. self._lengths.remove(length)
  141. return start, stop
  142. def free(self, block):
  143. # free a block returned by malloc()
  144. assert os.getpid() == self._lastpid
  145. self._lock.acquire()
  146. try:
  147. self._allocated_blocks.remove(block)
  148. self._free(block)
  149. finally:
  150. self._lock.release()
  151. def malloc(self, size):
  152. # return a block of right size (possibly rounded up)
  153. if os.getpid() != self._lastpid:
  154. self.__init__() # reinitialize after fork
  155. self._lock.acquire()
  156. try:
  157. size = self._roundUp(size)
  158. (arena, start, stop) = self._malloc(size)
  159. new_stop = start + size
  160. if new_stop < stop:
  161. self._free((arena, new_stop, stop))
  162. block = (arena, start, new_stop)
  163. self._allocated_blocks.add(block)
  164. return block
  165. finally:
  166. self._lock.release()
  167. #
  168. # Class representing a chunk of an mmap -- can be inherited
  169. #
  170. class BufferWrapper(object):
  171. _heap = Heap()
  172. def __init__(self, size):
  173. assert 0 <= size <= sys.maxint
  174. block = BufferWrapper._heap.malloc(size)
  175. self._state = (block, size)
  176. Finalize(self, BufferWrapper._heap.free, args=(block,))
  177. def getAddress(self):
  178. (arena, start, stop), size = self._state
  179. address, length = _processing.addressOfBuffer(arena.buffer)
  180. assert size <= length
  181. return address + start
  182. def getView(self):
  183. (arena, start, stop), size = self._state
  184. return _processing.readWriteBuffer(arena.buffer, start, size)
  185. def getSize(self):
  186. return self._state[1]