Python code coverage for Lib/multiprocessing/resource_sharer.py

#	count	content
1	n/a	#
2	n/a	# We use a background thread for sharing fds on Unix, and for sharing sockets on
3	n/a	# Windows.
4	n/a	#
5	n/a	# A client which wants to pickle a resource registers it with the resource
6	n/a	# sharer and gets an identifier in return. The unpickling process will connect
7	n/a	# to the resource sharer, sends the identifier and its pid, and then receives
8	n/a	# the resource.
9	n/a	#
10	n/a
11	n/a	import os
12	n/a	import signal
13	n/a	import socket
14	n/a	import sys
15	n/a	import threading
16	n/a
17	n/a	from . import process
18	n/a	from .context import reduction
19	n/a	from . import util
20	n/a
21	n/a	__all__ = ['stop']
22	n/a
23	n/a
24	n/a	if sys.platform == 'win32':
25	n/a	__all__ += ['DupSocket']
26	n/a
27	n/a	class DupSocket(object):
28	n/a	'''Picklable wrapper for a socket.'''
29	n/a	def __init__(self, sock):
30	n/a	new_sock = sock.dup()
31	n/a	def send(conn, pid):
32	n/a	share = new_sock.share(pid)
33	n/a	conn.send_bytes(share)
34	n/a	self._id = _resource_sharer.register(send, new_sock.close)
35	n/a
36	n/a	def detach(self):
37	n/a	'''Get the socket. This should only be called once.'''
38	n/a	with _resource_sharer.get_connection(self._id) as conn:
39	n/a	share = conn.recv_bytes()
40	n/a	return socket.fromshare(share)
41	n/a
42	n/a	else:
43	n/a	__all__ += ['DupFd']
44	n/a
45	n/a	class DupFd(object):
46	n/a	'''Wrapper for fd which can be used at any time.'''
47	n/a	def __init__(self, fd):
48	n/a	new_fd = os.dup(fd)
49	n/a	def send(conn, pid):
50	n/a	reduction.send_handle(conn, new_fd, pid)
51	n/a	def close():
52	n/a	os.close(new_fd)
53	n/a	self._id = _resource_sharer.register(send, close)
54	n/a
55	n/a	def detach(self):
56	n/a	'''Get the fd. This should only be called once.'''
57	n/a	with _resource_sharer.get_connection(self._id) as conn:
58	n/a	return reduction.recv_handle(conn)
59	n/a
60	n/a
61	n/a	class _ResourceSharer(object):
62	n/a	'''Manager for resouces using background thread.'''
63	n/a	def __init__(self):
64	n/a	self._key = 0
65	n/a	self._cache = {}
66	n/a	self._old_locks = []
67	n/a	self._lock = threading.Lock()
68	n/a	self._listener = None
69	n/a	self._address = None
70	n/a	self._thread = None
71	n/a	util.register_after_fork(self, _ResourceSharer._afterfork)
72	n/a
73	n/a	def register(self, send, close):
74	n/a	'''Register resource, returning an identifier.'''
75	n/a	with self._lock:
76	n/a	if self._address is None:
77	n/a	self._start()
78	n/a	self._key += 1
79	n/a	self._cache[self._key] = (send, close)
80	n/a	return (self._address, self._key)
81	n/a
82	n/a	@staticmethod
83	n/a	def get_connection(ident):
84	n/a	'''Return connection from which to receive identified resource.'''
85	n/a	from .connection import Client
86	n/a	address, key = ident
87	n/a	c = Client(address, authkey=process.current_process().authkey)
88	n/a	c.send((key, os.getpid()))
89	n/a	return c
90	n/a
91	n/a	def stop(self, timeout=None):
92	n/a	'''Stop the background thread and clear registered resources.'''
93	n/a	from .connection import Client
94	n/a	with self._lock:
95	n/a	if self._address is not None:
96	n/a	c = Client(self._address,
97	n/a	authkey=process.current_process().authkey)
98	n/a	c.send(None)
99	n/a	c.close()
100	n/a	self._thread.join(timeout)
101	n/a	if self._thread.is_alive():
102	n/a	util.sub_warning('_ResourceSharer thread did '
103	n/a	'not stop when asked')
104	n/a	self._listener.close()
105	n/a	self._thread = None
106	n/a	self._address = None
107	n/a	self._listener = None
108	n/a	for key, (send, close) in self._cache.items():
109	n/a	close()
110	n/a	self._cache.clear()
111	n/a
112	n/a	def _afterfork(self):
113	n/a	for key, (send, close) in self._cache.items():
114	n/a	close()
115	n/a	self._cache.clear()
116	n/a	# If self._lock was locked at the time of the fork, it may be broken
117	n/a	# -- see issue 6721. Replace it without letting it be gc'ed.
118	n/a	self._old_locks.append(self._lock)
119	n/a	self._lock = threading.Lock()
120	n/a	if self._listener is not None:
121	n/a	self._listener.close()
122	n/a	self._listener = None
123	n/a	self._address = None
124	n/a	self._thread = None
125	n/a
126	n/a	def _start(self):
127	n/a	from .connection import Listener
128	n/a	assert self._listener is None
129	n/a	util.debug('starting listener and thread for sending handles')
130	n/a	self._listener = Listener(authkey=process.current_process().authkey)
131	n/a	self._address = self._listener.address
132	n/a	t = threading.Thread(target=self._serve)
133	n/a	t.daemon = True
134	n/a	t.start()
135	n/a	self._thread = t
136	n/a
137	n/a	def _serve(self):
138	n/a	if hasattr(signal, 'pthread_sigmask'):
139	n/a	signal.pthread_sigmask(signal.SIG_BLOCK, range(1, signal.NSIG))
140	n/a	while 1:
141	n/a	try:
142	n/a	with self._listener.accept() as conn:
143	n/a	msg = conn.recv()
144	n/a	if msg is None:
145	n/a	break
146	n/a	key, destination_pid = msg
147	n/a	send, close = self._cache.pop(key)
148	n/a	try:
149	n/a	send(conn, destination_pid)
150	n/a	finally:
151	n/a	close()
152	n/a	except:
153	n/a	if not util.is_exiting():
154	n/a	sys.excepthook(*sys.exc_info())
155	n/a
156	n/a
157	n/a	_resource_sharer = _ResourceSharer()
158	n/a	stop = _resource_sharer.stop