1# 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) 979