diff options
Diffstat (limited to 'Lib/socket.py')
| -rw-r--r-- | Lib/socket.py | 148 |
1 files changed, 147 insertions, 1 deletions
diff --git a/Lib/socket.py b/Lib/socket.py index 6d67b3d..cbadff7 100644 --- a/Lib/socket.py +++ b/Lib/socket.py @@ -47,7 +47,7 @@ the setsockopt() and getsockopt() methods. import _socket from _socket import * -import os, sys, io +import os, sys, io, selectors from enum import IntEnum try: @@ -109,6 +109,9 @@ if sys.platform.lower().startswith("win"): __all__.append("errorTab") +class _GiveupOnSendfile(Exception): pass + + class socket(_socket.socket): """A subclass of _socket.socket adding the makefile() method.""" @@ -233,6 +236,149 @@ class socket(_socket.socket): text.mode = mode return text + if hasattr(os, 'sendfile'): + + def _sendfile_use_sendfile(self, file, offset=0, count=None): + 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: + raise _GiveupOnSendfile(err) # not a regular file + if not fsize: + return 0 # empty file + blocksize = fsize if not count else count + + 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 _socket.timeout('timed out') + if count: + blocksize = count - total_sent + 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) + else: + def _sendfile_use_sendfile(self, file, offset=0, count=None): + raise _GiveupOnSendfile( + "os.sendfile() not available on this platform") + + def _sendfile_use_send(self, file, offset=0, count=None): + self._check_sendfile_params(file, offset, count) + if self.gettimeout() == 0: + raise ValueError("non-blocking sockets are not supported") + if offset: + file.seek(offset) + blocksize = min(count, 8192) if count else 8192 + total_sent = 0 + # localize variable access to minimize overhead + file_read = file.read + sock_send = self.send + try: + while True: + if count: + blocksize = min(count - total_sent, blocksize) + if blocksize <= 0: + break + data = memoryview(file_read(blocksize)) + if not data: + break # EOF + while True: + try: + sent = sock_send(data) + except BlockingIOError: + continue + else: + total_sent += sent + if sent < len(data): + data = data[sent:] + else: + break + return total_sent + finally: + if total_sent > 0 and hasattr(file, 'seek'): + file.seek(offset + total_sent) + + def _check_sendfile_params(self, file, offset, count): + if 'b' not in getattr(file, 'mode', 'b'): + raise ValueError("file should be opened in binary mode") + if not self.type & SOCK_STREAM: + raise ValueError("only SOCK_STREAM type sockets are supported") + if count is not None: + if not isinstance(count, int): + raise TypeError( + "count must be a positive integer (got {!r})".format(count)) + if count <= 0: + raise ValueError( + "count must be a positive integer (got {!r})".format(count)) + + def sendfile(self, file, offset=0, count=None): + """sendfile(file[, offset[, count]]) -> sent + + Send a file until EOF is reached by using high-performance + os.sendfile() and return the total number of bytes which + were sent. + *file* must be a regular file object opened in binary mode. + If os.sendfile() is not available (e.g. Windows) or file is + not a regular file socket.send() will be used instead. + *offset* tells from where to start reading the file. + If specified, *count* is the total number of bytes to transmit + as opposed to sending the file until EOF is reached. + File position is updated on return or also in case of error in + which case file.tell() can be used to figure out the number of + bytes which were sent. + The socket must be of SOCK_STREAM type. + Non-blocking sockets are not supported. + """ + try: + return self._sendfile_use_sendfile(file, offset, count) + except _GiveupOnSendfile: + return self._sendfile_use_send(file, offset, count) + def _decref_socketios(self): if self._io_refs > 0: self._io_refs -= 1 |