reduction.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218
  1. #
  2. # Module to support the pickling of different types of connection
  3. # objects and file objects so that they can be transferred between
  4. # different processes.
  5. #
  6. # processing/reduction.py
  7. #
  8. # Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
  9. #
  10. __all__ = []
  11. import os
  12. import sys
  13. import socket
  14. import threading
  15. import copy_reg
  16. import processing
  17. from processing import _processing
  18. from processing.logger import debug, subDebug, subWarning
  19. from processing.forking import thisThreadIsSpawning
  20. from processing.process import _registerAfterFork
  21. #
  22. #
  23. #
  24. connections_are_picklable = (
  25. sys.platform == 'win32' or hasattr(_processing, 'recvFd')
  26. )
  27. try:
  28. fromfd = socket.fromfd
  29. except AttributeError:
  30. def fromfd(fd, family, type, proto=0):
  31. s = socket._socket.socket()
  32. _processing.changeFd(s, fd, family, type, proto)
  33. return s
  34. #
  35. # Platform specific definitions
  36. #
  37. if sys.platform == 'win32':
  38. import _subprocess
  39. from processing._processing import win32
  40. closeHandle = win32.CloseHandle
  41. def duplicateHandle(handle):
  42. return _subprocess.DuplicateHandle(
  43. _subprocess.GetCurrentProcess(), handle,
  44. _subprocess.GetCurrentProcess(),
  45. 0, False, _subprocess.DUPLICATE_SAME_ACCESS
  46. ).Detach()
  47. def sendHandle(conn, handle, destination_pid):
  48. process_handle = win32.OpenProcess(
  49. win32.PROCESS_ALL_ACCESS, False, destination_pid
  50. )
  51. try:
  52. new_handle = _subprocess.DuplicateHandle(
  53. _subprocess.GetCurrentProcess(), handle,
  54. process_handle, 0, False, _subprocess.DUPLICATE_SAME_ACCESS
  55. )
  56. conn.send(new_handle.Detach())
  57. finally:
  58. win32.CloseHandle(process_handle)
  59. def recvHandle(conn):
  60. return conn.recv()
  61. def isInheritableHandle(handle):
  62. return (win32.GetHandleInformation(handle) & win32.HANDLE_FLAG_INHERIT)
  63. else:
  64. closeHandle = os.close
  65. duplicateHandle = os.dup
  66. def sendHandle(conn, handle, destination_pid):
  67. _processing.sendFd(conn.fileno(), handle)
  68. def recvHandle(conn):
  69. return _processing.recvFd(conn.fileno())
  70. def isInheritableHandle(handle):
  71. return True
  72. #
  73. # Support for a per-process server thread which caches pickled handles
  74. #
  75. _cache = set()
  76. def _reset(obj):
  77. global _lock, _listener, _cache
  78. for h in _cache:
  79. closeHandle(h)
  80. _cache.clear()
  81. _lock = threading.Lock()
  82. _listener = None
  83. _reset(None)
  84. _registerAfterFork(_reset, _reset)
  85. def _getListener():
  86. global _listener
  87. if _listener is None:
  88. _lock.acquire()
  89. try:
  90. if _listener is None:
  91. from processing.connection import Listener
  92. debug('starting listener and thread for sending handles')
  93. _listener = Listener(authenticate=True)
  94. t = threading.Thread(target=_serve)
  95. t.setDaemon(True)
  96. t.start()
  97. finally:
  98. _lock.release()
  99. return _listener
  100. def _serve():
  101. while 1:
  102. try:
  103. conn = _listener.accept()
  104. handle_wanted, destination_pid = conn.recv()
  105. _cache.remove(handle_wanted)
  106. sendHandle(conn, handle_wanted, destination_pid)
  107. closeHandle(handle_wanted)
  108. conn.close()
  109. except (SystemExit, KeyboardInterrupt):
  110. raise
  111. except:
  112. if not processing.currentProcess()._exiting:
  113. import traceback
  114. subWarning(
  115. 'thread for sharing handles raised exception :\n' +
  116. '-'*79 + '\n' + traceback.format_exc() + '-'*79
  117. )
  118. #
  119. # Functions to be used for pickling/unpickling objects with handles
  120. #
  121. def reduceHandle(handle):
  122. if thisThreadIsSpawning() and isInheritableHandle(handle):
  123. return (None, handle, True)
  124. dup_handle = duplicateHandle(handle)
  125. _cache.add(dup_handle)
  126. subDebug('reducing handle %d', handle)
  127. return (_getListener().address, dup_handle, False)
  128. def rebuildHandle(pickled_data):
  129. from processing.connection import Client
  130. address, handle, inherited = pickled_data
  131. if inherited:
  132. return handle
  133. subDebug('rebuilding handle %d', handle)
  134. conn = Client(address, authenticate=True)
  135. conn.send((handle, os.getpid()))
  136. new_handle = recvHandle(conn)
  137. conn.close()
  138. return new_handle
  139. #
  140. # Register `_processing.Connection` with `copy_reg`
  141. #
  142. def reduceConnection(conn):
  143. return rebuildConnection, (reduceHandle(conn.fileno()),)
  144. def rebuildConnection(reduced_handle):
  145. fd = rebuildHandle(reduced_handle)
  146. return _processing.Connection(fd, duplicate=False)
  147. copy_reg.pickle(_processing.Connection, reduceConnection)
  148. #
  149. # Register `socket.socket` with `copy_reg`
  150. #
  151. def reduceSocket(s):
  152. try:
  153. Family, Type, Proto = s.family, s.type, s.proto
  154. except AttributeError:
  155. # have to guess family, type, proto
  156. address = s.getsockname()
  157. Family = type(address) is str and socket.AF_UNIX or socket.AF_INET
  158. Type = s.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
  159. Proto = 0
  160. reduced_handle = reduceHandle(s.fileno())
  161. return rebuildSocket, (reduced_handle, Family, Type, Proto)
  162. def rebuildSocket(reduced_handle, family, type, proto):
  163. fd = rebuildHandle(reduced_handle)
  164. _sock = fromfd(fd, family, type, proto)
  165. closeHandle(fd)
  166. return socket.socket(_sock=_sock)
  167. copy_reg.pickle(socket.socket, reduceSocket)
  168. #
  169. # Register `_processing.PipeConnection` with `copy_reg`
  170. #
  171. if sys.platform == 'win32':
  172. def reducePipeConnection(conn):
  173. return rebuildPipeConnection, (reduceHandle(conn.fileno()),)
  174. def rebuildPipeConnection(reduced_handle):
  175. handle = rebuildHandle(reduced_handle)
  176. return _processing.PipeConnection(handle, duplicate=False)
  177. copy_reg.pickle(_processing.PipeConnection, reducePipeConnection)