blob: 510e4b5aba44a6755c4a63fec0d8b297b2066563 [file] [log] [blame]
Olivier Deprezf4ef2d02021-04-20 13:36:24 +02001#
2# A higher level module for using sockets (or Windows named pipes)
3#
4# multiprocessing/connection.py
5#
6# Copyright (c) 2006-2008, R Oudkerk
7# Licensed to PSF under a Contributor Agreement.
8#
9
10__all__ = [ 'Client', 'Listener', 'Pipe', 'wait' ]
11
12import io
13import os
14import sys
15import socket
16import struct
17import time
18import tempfile
19import itertools
20
21import _multiprocessing
22
23from . import util
24
25from . import AuthenticationError, BufferTooShort
26from .context import reduction
27_ForkingPickler = reduction.ForkingPickler
28
29try:
30 import _winapi
31 from _winapi import WAIT_OBJECT_0, WAIT_ABANDONED_0, WAIT_TIMEOUT, INFINITE
32except ImportError:
33 if sys.platform == 'win32':
34 raise
35 _winapi = None
36
37#
38#
39#
40
41BUFSIZE = 8192
42# A very generous timeout when it comes to local connections...
43CONNECTION_TIMEOUT = 20.
44
45_mmap_counter = itertools.count()
46
47default_family = 'AF_INET'
48families = ['AF_INET']
49
50if hasattr(socket, 'AF_UNIX'):
51 default_family = 'AF_UNIX'
52 families += ['AF_UNIX']
53
54if sys.platform == 'win32':
55 default_family = 'AF_PIPE'
56 families += ['AF_PIPE']
57
58
59def _init_timeout(timeout=CONNECTION_TIMEOUT):
60 return time.monotonic() + timeout
61
62def _check_timeout(t):
63 return time.monotonic() > t
64
65#
66#
67#
68
69def arbitrary_address(family):
70 '''
71 Return an arbitrary free address for the given family
72 '''
73 if family == 'AF_INET':
74 return ('localhost', 0)
75 elif family == 'AF_UNIX':
76 # Prefer abstract sockets if possible to avoid problems with the address
77 # size. When coding portable applications, some implementations have
78 # sun_path as short as 92 bytes in the sockaddr_un struct.
79 if util.abstract_sockets_supported:
80 return f"\0listener-{os.getpid()}-{next(_mmap_counter)}"
81 return tempfile.mktemp(prefix='listener-', dir=util.get_temp_dir())
82 elif family == 'AF_PIPE':
83 return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
84 (os.getpid(), next(_mmap_counter)), dir="")
85 else:
86 raise ValueError('unrecognized family')
87
88def _validate_family(family):
89 '''
90 Checks if the family is valid for the current environment.
91 '''
92 if sys.platform != 'win32' and family == 'AF_PIPE':
93 raise ValueError('Family %s is not recognized.' % family)
94
95 if sys.platform == 'win32' and family == 'AF_UNIX':
96 # double check
97 if not hasattr(socket, family):
98 raise ValueError('Family %s is not recognized.' % family)
99
100def address_type(address):
101 '''
102 Return the types of the address
103
104 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
105 '''
106 if type(address) == tuple:
107 return 'AF_INET'
108 elif type(address) is str and address.startswith('\\\\'):
109 return 'AF_PIPE'
110 elif type(address) is str or util.is_abstract_socket_namespace(address):
111 return 'AF_UNIX'
112 else:
113 raise ValueError('address type of %r unrecognized' % address)
114
115#
116# Connection classes
117#
118
119class _ConnectionBase:
120 _handle = None
121
122 def __init__(self, handle, readable=True, writable=True):
123 handle = handle.__index__()
124 if handle < 0:
125 raise ValueError("invalid handle")
126 if not readable and not writable:
127 raise ValueError(
128 "at least one of `readable` and `writable` must be True")
129 self._handle = handle
130 self._readable = readable
131 self._writable = writable
132
133 # XXX should we use util.Finalize instead of a __del__?
134
135 def __del__(self):
136 if self._handle is not None:
137 self._close()
138
139 def _check_closed(self):
140 if self._handle is None:
141 raise OSError("handle is closed")
142
143 def _check_readable(self):
144 if not self._readable:
145 raise OSError("connection is write-only")
146
147 def _check_writable(self):
148 if not self._writable:
149 raise OSError("connection is read-only")
150
151 def _bad_message_length(self):
152 if self._writable:
153 self._readable = False
154 else:
155 self.close()
156 raise OSError("bad message length")
157
158 @property
159 def closed(self):
160 """True if the connection is closed"""
161 return self._handle is None
162
163 @property
164 def readable(self):
165 """True if the connection is readable"""
166 return self._readable
167
168 @property
169 def writable(self):
170 """True if the connection is writable"""
171 return self._writable
172
173 def fileno(self):
174 """File descriptor or handle of the connection"""
175 self._check_closed()
176 return self._handle
177
178 def close(self):
179 """Close the connection"""
180 if self._handle is not None:
181 try:
182 self._close()
183 finally:
184 self._handle = None
185
186 def send_bytes(self, buf, offset=0, size=None):
187 """Send the bytes data from a bytes-like object"""
188 self._check_closed()
189 self._check_writable()
190 m = memoryview(buf)
191 # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
192 if m.itemsize > 1:
193 m = memoryview(bytes(m))
194 n = len(m)
195 if offset < 0:
196 raise ValueError("offset is negative")
197 if n < offset:
198 raise ValueError("buffer length < offset")
199 if size is None:
200 size = n - offset
201 elif size < 0:
202 raise ValueError("size is negative")
203 elif offset + size > n:
204 raise ValueError("buffer length < offset + size")
205 self._send_bytes(m[offset:offset + size])
206
207 def send(self, obj):
208 """Send a (picklable) object"""
209 self._check_closed()
210 self._check_writable()
211 self._send_bytes(_ForkingPickler.dumps(obj))
212
213 def recv_bytes(self, maxlength=None):
214 """
215 Receive bytes data as a bytes object.
216 """
217 self._check_closed()
218 self._check_readable()
219 if maxlength is not None and maxlength < 0:
220 raise ValueError("negative maxlength")
221 buf = self._recv_bytes(maxlength)
222 if buf is None:
223 self._bad_message_length()
224 return buf.getvalue()
225
226 def recv_bytes_into(self, buf, offset=0):
227 """
228 Receive bytes data into a writeable bytes-like object.
229 Return the number of bytes read.
230 """
231 self._check_closed()
232 self._check_readable()
233 with memoryview(buf) as m:
234 # Get bytesize of arbitrary buffer
235 itemsize = m.itemsize
236 bytesize = itemsize * len(m)
237 if offset < 0:
238 raise ValueError("negative offset")
239 elif offset > bytesize:
240 raise ValueError("offset too large")
241 result = self._recv_bytes()
242 size = result.tell()
243 if bytesize < offset + size:
244 raise BufferTooShort(result.getvalue())
245 # Message can fit in dest
246 result.seek(0)
247 result.readinto(m[offset // itemsize :
248 (offset + size) // itemsize])
249 return size
250
251 def recv(self):
252 """Receive a (picklable) object"""
253 self._check_closed()
254 self._check_readable()
255 buf = self._recv_bytes()
256 return _ForkingPickler.loads(buf.getbuffer())
257
258 def poll(self, timeout=0.0):
259 """Whether there is any input available to be read"""
260 self._check_closed()
261 self._check_readable()
262 return self._poll(timeout)
263
264 def __enter__(self):
265 return self
266
267 def __exit__(self, exc_type, exc_value, exc_tb):
268 self.close()
269
270
271if _winapi:
272
273 class PipeConnection(_ConnectionBase):
274 """
275 Connection class based on a Windows named pipe.
276 Overlapped I/O is used, so the handles must have been created
277 with FILE_FLAG_OVERLAPPED.
278 """
279 _got_empty_message = False
280
281 def _close(self, _CloseHandle=_winapi.CloseHandle):
282 _CloseHandle(self._handle)
283
284 def _send_bytes(self, buf):
285 ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True)
286 try:
287 if err == _winapi.ERROR_IO_PENDING:
288 waitres = _winapi.WaitForMultipleObjects(
289 [ov.event], False, INFINITE)
290 assert waitres == WAIT_OBJECT_0
291 except:
292 ov.cancel()
293 raise
294 finally:
295 nwritten, err = ov.GetOverlappedResult(True)
296 assert err == 0
297 assert nwritten == len(buf)
298
299 def _recv_bytes(self, maxsize=None):
300 if self._got_empty_message:
301 self._got_empty_message = False
302 return io.BytesIO()
303 else:
304 bsize = 128 if maxsize is None else min(maxsize, 128)
305 try:
306 ov, err = _winapi.ReadFile(self._handle, bsize,
307 overlapped=True)
308 try:
309 if err == _winapi.ERROR_IO_PENDING:
310 waitres = _winapi.WaitForMultipleObjects(
311 [ov.event], False, INFINITE)
312 assert waitres == WAIT_OBJECT_0
313 except:
314 ov.cancel()
315 raise
316 finally:
317 nread, err = ov.GetOverlappedResult(True)
318 if err == 0:
319 f = io.BytesIO()
320 f.write(ov.getbuffer())
321 return f
322 elif err == _winapi.ERROR_MORE_DATA:
323 return self._get_more_data(ov, maxsize)
324 except OSError as e:
325 if e.winerror == _winapi.ERROR_BROKEN_PIPE:
326 raise EOFError
327 else:
328 raise
329 raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
330
331 def _poll(self, timeout):
332 if (self._got_empty_message or
333 _winapi.PeekNamedPipe(self._handle)[0] != 0):
334 return True
335 return bool(wait([self], timeout))
336
337 def _get_more_data(self, ov, maxsize):
338 buf = ov.getbuffer()
339 f = io.BytesIO()
340 f.write(buf)
341 left = _winapi.PeekNamedPipe(self._handle)[1]
342 assert left > 0
343 if maxsize is not None and len(buf) + left > maxsize:
344 self._bad_message_length()
345 ov, err = _winapi.ReadFile(self._handle, left, overlapped=True)
346 rbytes, err = ov.GetOverlappedResult(True)
347 assert err == 0
348 assert rbytes == left
349 f.write(ov.getbuffer())
350 return f
351
352
353class Connection(_ConnectionBase):
354 """
355 Connection class based on an arbitrary file descriptor (Unix only), or
356 a socket handle (Windows).
357 """
358
359 if _winapi:
360 def _close(self, _close=_multiprocessing.closesocket):
361 _close(self._handle)
362 _write = _multiprocessing.send
363 _read = _multiprocessing.recv
364 else:
365 def _close(self, _close=os.close):
366 _close(self._handle)
367 _write = os.write
368 _read = os.read
369
370 def _send(self, buf, write=_write):
371 remaining = len(buf)
372 while True:
373 n = write(self._handle, buf)
374 remaining -= n
375 if remaining == 0:
376 break
377 buf = buf[n:]
378
379 def _recv(self, size, read=_read):
380 buf = io.BytesIO()
381 handle = self._handle
382 remaining = size
383 while remaining > 0:
384 chunk = read(handle, remaining)
385 n = len(chunk)
386 if n == 0:
387 if remaining == size:
388 raise EOFError
389 else:
390 raise OSError("got end of file during message")
391 buf.write(chunk)
392 remaining -= n
393 return buf
394
395 def _send_bytes(self, buf):
396 n = len(buf)
397 if n > 0x7fffffff:
398 pre_header = struct.pack("!i", -1)
399 header = struct.pack("!Q", n)
400 self._send(pre_header)
401 self._send(header)
402 self._send(buf)
403 else:
404 # For wire compatibility with 3.7 and lower
405 header = struct.pack("!i", n)
406 if n > 16384:
407 # The payload is large so Nagle's algorithm won't be triggered
408 # and we'd better avoid the cost of concatenation.
409 self._send(header)
410 self._send(buf)
411 else:
412 # Issue #20540: concatenate before sending, to avoid delays due
413 # to Nagle's algorithm on a TCP socket.
414 # Also note we want to avoid sending a 0-length buffer separately,
415 # to avoid "broken pipe" errors if the other end closed the pipe.
416 self._send(header + buf)
417
418 def _recv_bytes(self, maxsize=None):
419 buf = self._recv(4)
420 size, = struct.unpack("!i", buf.getvalue())
421 if size == -1:
422 buf = self._recv(8)
423 size, = struct.unpack("!Q", buf.getvalue())
424 if maxsize is not None and size > maxsize:
425 return None
426 return self._recv(size)
427
428 def _poll(self, timeout):
429 r = wait([self], timeout)
430 return bool(r)
431
432
433#
434# Public functions
435#
436
437class Listener(object):
438 '''
439 Returns a listener object.
440
441 This is a wrapper for a bound socket which is 'listening' for
442 connections, or for a Windows named pipe.
443 '''
444 def __init__(self, address=None, family=None, backlog=1, authkey=None):
445 family = family or (address and address_type(address)) \
446 or default_family
447 address = address or arbitrary_address(family)
448
449 _validate_family(family)
450 if family == 'AF_PIPE':
451 self._listener = PipeListener(address, backlog)
452 else:
453 self._listener = SocketListener(address, family, backlog)
454
455 if authkey is not None and not isinstance(authkey, bytes):
456 raise TypeError('authkey should be a byte string')
457
458 self._authkey = authkey
459
460 def accept(self):
461 '''
462 Accept a connection on the bound socket or named pipe of `self`.
463
464 Returns a `Connection` object.
465 '''
466 if self._listener is None:
467 raise OSError('listener is closed')
468 c = self._listener.accept()
469 if self._authkey:
470 deliver_challenge(c, self._authkey)
471 answer_challenge(c, self._authkey)
472 return c
473
474 def close(self):
475 '''
476 Close the bound socket or named pipe of `self`.
477 '''
478 listener = self._listener
479 if listener is not None:
480 self._listener = None
481 listener.close()
482
483 @property
484 def address(self):
485 return self._listener._address
486
487 @property
488 def last_accepted(self):
489 return self._listener._last_accepted
490
491 def __enter__(self):
492 return self
493
494 def __exit__(self, exc_type, exc_value, exc_tb):
495 self.close()
496
497
498def Client(address, family=None, authkey=None):
499 '''
500 Returns a connection to the address of a `Listener`
501 '''
502 family = family or address_type(address)
503 _validate_family(family)
504 if family == 'AF_PIPE':
505 c = PipeClient(address)
506 else:
507 c = SocketClient(address)
508
509 if authkey is not None and not isinstance(authkey, bytes):
510 raise TypeError('authkey should be a byte string')
511
512 if authkey is not None:
513 answer_challenge(c, authkey)
514 deliver_challenge(c, authkey)
515
516 return c
517
518
519if sys.platform != 'win32':
520
521 def Pipe(duplex=True):
522 '''
523 Returns pair of connection objects at either end of a pipe
524 '''
525 if duplex:
526 s1, s2 = socket.socketpair()
527 s1.setblocking(True)
528 s2.setblocking(True)
529 c1 = Connection(s1.detach())
530 c2 = Connection(s2.detach())
531 else:
532 fd1, fd2 = os.pipe()
533 c1 = Connection(fd1, writable=False)
534 c2 = Connection(fd2, readable=False)
535
536 return c1, c2
537
538else:
539
540 def Pipe(duplex=True):
541 '''
542 Returns pair of connection objects at either end of a pipe
543 '''
544 address = arbitrary_address('AF_PIPE')
545 if duplex:
546 openmode = _winapi.PIPE_ACCESS_DUPLEX
547 access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
548 obsize, ibsize = BUFSIZE, BUFSIZE
549 else:
550 openmode = _winapi.PIPE_ACCESS_INBOUND
551 access = _winapi.GENERIC_WRITE
552 obsize, ibsize = 0, BUFSIZE
553
554 h1 = _winapi.CreateNamedPipe(
555 address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
556 _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
557 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
558 _winapi.PIPE_WAIT,
559 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER,
560 # default security descriptor: the handle cannot be inherited
561 _winapi.NULL
562 )
563 h2 = _winapi.CreateFile(
564 address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
565 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
566 )
567 _winapi.SetNamedPipeHandleState(
568 h2, _winapi.PIPE_READMODE_MESSAGE, None, None
569 )
570
571 overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
572 _, err = overlapped.GetOverlappedResult(True)
573 assert err == 0
574
575 c1 = PipeConnection(h1, writable=duplex)
576 c2 = PipeConnection(h2, readable=duplex)
577
578 return c1, c2
579
580#
581# Definitions for connections based on sockets
582#
583
584class SocketListener(object):
585 '''
586 Representation of a socket which is bound to an address and listening
587 '''
588 def __init__(self, address, family, backlog=1):
589 self._socket = socket.socket(getattr(socket, family))
590 try:
591 # SO_REUSEADDR has different semantics on Windows (issue #2550).
592 if os.name == 'posix':
593 self._socket.setsockopt(socket.SOL_SOCKET,
594 socket.SO_REUSEADDR, 1)
595 self._socket.setblocking(True)
596 self._socket.bind(address)
597 self._socket.listen(backlog)
598 self._address = self._socket.getsockname()
599 except OSError:
600 self._socket.close()
601 raise
602 self._family = family
603 self._last_accepted = None
604
605 if family == 'AF_UNIX' and not util.is_abstract_socket_namespace(address):
606 # Linux abstract socket namespaces do not need to be explicitly unlinked
607 self._unlink = util.Finalize(
608 self, os.unlink, args=(address,), exitpriority=0
609 )
610 else:
611 self._unlink = None
612
613 def accept(self):
614 s, self._last_accepted = self._socket.accept()
615 s.setblocking(True)
616 return Connection(s.detach())
617
618 def close(self):
619 try:
620 self._socket.close()
621 finally:
622 unlink = self._unlink
623 if unlink is not None:
624 self._unlink = None
625 unlink()
626
627
628def SocketClient(address):
629 '''
630 Return a connection object connected to the socket given by `address`
631 '''
632 family = address_type(address)
633 with socket.socket( getattr(socket, family) ) as s:
634 s.setblocking(True)
635 s.connect(address)
636 return Connection(s.detach())
637
638#
639# Definitions for connections based on named pipes
640#
641
642if sys.platform == 'win32':
643
644 class PipeListener(object):
645 '''
646 Representation of a named pipe
647 '''
648 def __init__(self, address, backlog=None):
649 self._address = address
650 self._handle_queue = [self._new_handle(first=True)]
651
652 self._last_accepted = None
653 util.sub_debug('listener created with address=%r', self._address)
654 self.close = util.Finalize(
655 self, PipeListener._finalize_pipe_listener,
656 args=(self._handle_queue, self._address), exitpriority=0
657 )
658
659 def _new_handle(self, first=False):
660 flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
661 if first:
662 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
663 return _winapi.CreateNamedPipe(
664 self._address, flags,
665 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
666 _winapi.PIPE_WAIT,
667 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
668 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
669 )
670
671 def accept(self):
672 self._handle_queue.append(self._new_handle())
673 handle = self._handle_queue.pop(0)
674 try:
675 ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
676 except OSError as e:
677 if e.winerror != _winapi.ERROR_NO_DATA:
678 raise
679 # ERROR_NO_DATA can occur if a client has already connected,
680 # written data and then disconnected -- see Issue 14725.
681 else:
682 try:
683 res = _winapi.WaitForMultipleObjects(
684 [ov.event], False, INFINITE)
685 except:
686 ov.cancel()
687 _winapi.CloseHandle(handle)
688 raise
689 finally:
690 _, err = ov.GetOverlappedResult(True)
691 assert err == 0
692 return PipeConnection(handle)
693
694 @staticmethod
695 def _finalize_pipe_listener(queue, address):
696 util.sub_debug('closing listener with address=%r', address)
697 for handle in queue:
698 _winapi.CloseHandle(handle)
699
700 def PipeClient(address):
701 '''
702 Return a connection object connected to the pipe given by `address`
703 '''
704 t = _init_timeout()
705 while 1:
706 try:
707 _winapi.WaitNamedPipe(address, 1000)
708 h = _winapi.CreateFile(
709 address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
710 0, _winapi.NULL, _winapi.OPEN_EXISTING,
711 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
712 )
713 except OSError as e:
714 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT,
715 _winapi.ERROR_PIPE_BUSY) or _check_timeout(t):
716 raise
717 else:
718 break
719 else:
720 raise
721
722 _winapi.SetNamedPipeHandleState(
723 h, _winapi.PIPE_READMODE_MESSAGE, None, None
724 )
725 return PipeConnection(h)
726
727#
728# Authentication stuff
729#
730
731MESSAGE_LENGTH = 20
732
733CHALLENGE = b'#CHALLENGE#'
734WELCOME = b'#WELCOME#'
735FAILURE = b'#FAILURE#'
736
737def deliver_challenge(connection, authkey):
738 import hmac
739 if not isinstance(authkey, bytes):
740 raise ValueError(
741 "Authkey must be bytes, not {0!s}".format(type(authkey)))
742 message = os.urandom(MESSAGE_LENGTH)
743 connection.send_bytes(CHALLENGE + message)
744 digest = hmac.new(authkey, message, 'md5').digest()
745 response = connection.recv_bytes(256) # reject large message
746 if response == digest:
747 connection.send_bytes(WELCOME)
748 else:
749 connection.send_bytes(FAILURE)
750 raise AuthenticationError('digest received was wrong')
751
752def answer_challenge(connection, authkey):
753 import hmac
754 if not isinstance(authkey, bytes):
755 raise ValueError(
756 "Authkey must be bytes, not {0!s}".format(type(authkey)))
757 message = connection.recv_bytes(256) # reject large message
758 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
759 message = message[len(CHALLENGE):]
760 digest = hmac.new(authkey, message, 'md5').digest()
761 connection.send_bytes(digest)
762 response = connection.recv_bytes(256) # reject large message
763 if response != WELCOME:
764 raise AuthenticationError('digest sent was rejected')
765
766#
767# Support for using xmlrpclib for serialization
768#
769
770class ConnectionWrapper(object):
771 def __init__(self, conn, dumps, loads):
772 self._conn = conn
773 self._dumps = dumps
774 self._loads = loads
775 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
776 obj = getattr(conn, attr)
777 setattr(self, attr, obj)
778 def send(self, obj):
779 s = self._dumps(obj)
780 self._conn.send_bytes(s)
781 def recv(self):
782 s = self._conn.recv_bytes()
783 return self._loads(s)
784
785def _xml_dumps(obj):
786 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
787
788def _xml_loads(s):
789 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
790 return obj
791
792class XmlListener(Listener):
793 def accept(self):
794 global xmlrpclib
795 import xmlrpc.client as xmlrpclib
796 obj = Listener.accept(self)
797 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
798
799def XmlClient(*args, **kwds):
800 global xmlrpclib
801 import xmlrpc.client as xmlrpclib
802 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
803
804#
805# Wait
806#
807
808if sys.platform == 'win32':
809
810 def _exhaustive_wait(handles, timeout):
811 # Return ALL handles which are currently signalled. (Only
812 # returning the first signalled might create starvation issues.)
813 L = list(handles)
814 ready = []
815 while L:
816 res = _winapi.WaitForMultipleObjects(L, False, timeout)
817 if res == WAIT_TIMEOUT:
818 break
819 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
820 res -= WAIT_OBJECT_0
821 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
822 res -= WAIT_ABANDONED_0
823 else:
824 raise RuntimeError('Should not get here')
825 ready.append(L[res])
826 L = L[res+1:]
827 timeout = 0
828 return ready
829
830 _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
831
832 def wait(object_list, timeout=None):
833 '''
834 Wait till an object in object_list is ready/readable.
835
836 Returns list of those objects in object_list which are ready/readable.
837 '''
838 if timeout is None:
839 timeout = INFINITE
840 elif timeout < 0:
841 timeout = 0
842 else:
843 timeout = int(timeout * 1000 + 0.5)
844
845 object_list = list(object_list)
846 waithandle_to_obj = {}
847 ov_list = []
848 ready_objects = set()
849 ready_handles = set()
850
851 try:
852 for o in object_list:
853 try:
854 fileno = getattr(o, 'fileno')
855 except AttributeError:
856 waithandle_to_obj[o.__index__()] = o
857 else:
858 # start an overlapped read of length zero
859 try:
860 ov, err = _winapi.ReadFile(fileno(), 0, True)
861 except OSError as e:
862 ov, err = None, e.winerror
863 if err not in _ready_errors:
864 raise
865 if err == _winapi.ERROR_IO_PENDING:
866 ov_list.append(ov)
867 waithandle_to_obj[ov.event] = o
868 else:
869 # If o.fileno() is an overlapped pipe handle and
870 # err == 0 then there is a zero length message
871 # in the pipe, but it HAS NOT been consumed...
872 if ov and sys.getwindowsversion()[:2] >= (6, 2):
873 # ... except on Windows 8 and later, where
874 # the message HAS been consumed.
875 try:
876 _, err = ov.GetOverlappedResult(False)
877 except OSError as e:
878 err = e.winerror
879 if not err and hasattr(o, '_got_empty_message'):
880 o._got_empty_message = True
881 ready_objects.add(o)
882 timeout = 0
883
884 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
885 finally:
886 # request that overlapped reads stop
887 for ov in ov_list:
888 ov.cancel()
889
890 # wait for all overlapped reads to stop
891 for ov in ov_list:
892 try:
893 _, err = ov.GetOverlappedResult(True)
894 except OSError as e:
895 err = e.winerror
896 if err not in _ready_errors:
897 raise
898 if err != _winapi.ERROR_OPERATION_ABORTED:
899 o = waithandle_to_obj[ov.event]
900 ready_objects.add(o)
901 if err == 0:
902 # If o.fileno() is an overlapped pipe handle then
903 # a zero length message HAS been consumed.
904 if hasattr(o, '_got_empty_message'):
905 o._got_empty_message = True
906
907 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
908 return [o for o in object_list if o in ready_objects]
909
910else:
911
912 import selectors
913
914 # poll/select have the advantage of not requiring any extra file
915 # descriptor, contrarily to epoll/kqueue (also, they require a single
916 # syscall).
917 if hasattr(selectors, 'PollSelector'):
918 _WaitSelector = selectors.PollSelector
919 else:
920 _WaitSelector = selectors.SelectSelector
921
922 def wait(object_list, timeout=None):
923 '''
924 Wait till an object in object_list is ready/readable.
925
926 Returns list of those objects in object_list which are ready/readable.
927 '''
928 with _WaitSelector() as selector:
929 for obj in object_list:
930 selector.register(obj, selectors.EVENT_READ)
931
932 if timeout is not None:
933 deadline = time.monotonic() + timeout
934
935 while True:
936 ready = selector.select(timeout)
937 if ready:
938 return [key.fileobj for (key, events) in ready]
939 else:
940 if timeout is not None:
941 timeout = deadline - time.monotonic()
942 if timeout < 0:
943 return ready
944
945#
946# Make connection and socket objects sharable if possible
947#
948
949if sys.platform == 'win32':
950 def reduce_connection(conn):
951 handle = conn.fileno()
952 with socket.fromfd(handle, socket.AF_INET, socket.SOCK_STREAM) as s:
953 from . import resource_sharer
954 ds = resource_sharer.DupSocket(s)
955 return rebuild_connection, (ds, conn.readable, conn.writable)
956 def rebuild_connection(ds, readable, writable):
957 sock = ds.detach()
958 return Connection(sock.detach(), readable, writable)
959 reduction.register(Connection, reduce_connection)
960
961 def reduce_pipe_connection(conn):
962 access = ((_winapi.FILE_GENERIC_READ if conn.readable else 0) |
963 (_winapi.FILE_GENERIC_WRITE if conn.writable else 0))
964 dh = reduction.DupHandle(conn.fileno(), access)
965 return rebuild_pipe_connection, (dh, conn.readable, conn.writable)
966 def rebuild_pipe_connection(dh, readable, writable):
967 handle = dh.detach()
968 return PipeConnection(handle, readable, writable)
969 reduction.register(PipeConnection, reduce_pipe_connection)
970
971else:
972 def reduce_connection(conn):
973 df = reduction.DupFd(conn.fileno())
974 return rebuild_connection, (df, conn.readable, conn.writable)
975 def rebuild_connection(df, readable, writable):
976 fd = df.detach()
977 return Connection(fd, readable, writable)
978 reduction.register(Connection, reduce_connection)