diff options
Diffstat (limited to 'Lib/socket.py')
| -rw-r--r-- | Lib/socket.py | 143 |
1 files changed, 76 insertions, 67 deletions
diff --git a/Lib/socket.py b/Lib/socket.py index 727b0e75f03..3073c012b19 100644 --- a/Lib/socket.py +++ b/Lib/socket.py @@ -56,6 +56,7 @@ import io import os import sys from enum import IntEnum, IntFlag +from functools import partial try: import errno @@ -348,75 +349,83 @@ class socket(_socket.socket): text.mode = mode return text - if hasattr(os, 'sendfile'): + def _sendfile_zerocopy(self, zerocopy_func, giveup_exc_type, file, + offset=0, count=None): + """ + Send a file using a zero-copy function. + """ + import selectors - def _sendfile_use_sendfile(self, file, offset=0, count=None): - # Lazy import to improve module import time - import selectors + self._check_sendfile_params(file, offset, count) + sockno = self.fileno() + try: + fileno = file.fileno() + except (AttributeError, io.UnsupportedOperation) as err: + raise giveup_exc_type(err) # not a regular file + try: + fsize = os.fstat(fileno).st_size + except OSError as err: + raise giveup_exc_type(err) # not a regular file + if not fsize: + return 0 # empty file + # Truncate to 1GiB to avoid OverflowError, see bpo-38319. + blocksize = min(count or fsize, 2 ** 30) + timeout = self.gettimeout() + if timeout == 0: + raise ValueError("non-blocking sockets are not supported") + # poll/select have the advantage of not requiring any + # extra file descriptor, contrarily to epoll/kqueue + # (also, they require a single syscall). + if hasattr(selectors, 'PollSelector'): + selector = selectors.PollSelector() + else: + selector = selectors.SelectSelector() + selector.register(sockno, selectors.EVENT_WRITE) - self._check_sendfile_params(file, offset, count) - sockno = self.fileno() - try: - fileno = file.fileno() - except (AttributeError, io.UnsupportedOperation) as err: - raise _GiveupOnSendfile(err) # not a regular file - try: - fsize = os.fstat(fileno).st_size - except OSError as err: - raise _GiveupOnSendfile(err) # not a regular file - if not fsize: - return 0 # empty file - # Truncate to 1GiB to avoid OverflowError, see bpo-38319. - blocksize = min(count or fsize, 2 ** 30) - timeout = self.gettimeout() - if timeout == 0: - raise ValueError("non-blocking sockets are not supported") - # poll/select have the advantage of not requiring any - # extra file descriptor, contrarily to epoll/kqueue - # (also, they require a single syscall). - if hasattr(selectors, 'PollSelector'): - selector = selectors.PollSelector() - else: - selector = selectors.SelectSelector() - selector.register(sockno, selectors.EVENT_WRITE) - - total_sent = 0 - # localize variable access to minimize overhead - selector_select = selector.select - os_sendfile = os.sendfile - try: - while True: - if timeout and not selector_select(timeout): - raise TimeoutError('timed out') - if count: - blocksize = min(count - total_sent, blocksize) - if blocksize <= 0: - break - try: - sent = os_sendfile(sockno, fileno, offset, blocksize) - except BlockingIOError: - if not timeout: - # Block until the socket is ready to send some - # data; avoids hogging CPU resources. - selector_select() - continue - except OSError as err: - if total_sent == 0: - # We can get here for different reasons, the main - # one being 'file' is not a regular mmap(2)-like - # file, in which case we'll fall back on using - # plain send(). - raise _GiveupOnSendfile(err) - raise err from None - else: - if sent == 0: - break # EOF - offset += sent - total_sent += sent - return total_sent - finally: - if total_sent > 0 and hasattr(file, 'seek'): - file.seek(offset) + total_sent = 0 + # localize variable access to minimize overhead + selector_select = selector.select + try: + while True: + if timeout and not selector_select(timeout): + raise TimeoutError('timed out') + if count: + blocksize = min(count - total_sent, blocksize) + if blocksize <= 0: + break + try: + sent = zerocopy_func(fileno, offset, blocksize) + except BlockingIOError: + if not timeout: + # Block until the socket is ready to send some + # data; avoids hogging CPU resources. + selector_select() + continue + except OSError as err: + if total_sent == 0: + # We can get here for different reasons, the main + # one being 'file' is not a regular mmap(2)-like + # file, in which case we'll fall back on using + # plain send(). + raise giveup_exc_type(err) + raise err from None + else: + if sent == 0: + break # EOF + offset += sent + total_sent += sent + return total_sent + finally: + if total_sent > 0 and hasattr(file, 'seek'): + file.seek(offset) + + if hasattr(os, 'sendfile'): + def _sendfile_use_sendfile(self, file, offset=0, count=None): + return self._sendfile_zerocopy( + partial(os.sendfile, self.fileno()), + _GiveupOnSendfile, + file, offset, count, + ) else: def _sendfile_use_sendfile(self, file, offset=0, count=None): raise _GiveupOnSendfile( |