241 lines
8.2 KiB
Python
241 lines
8.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
hyper/http20/bufsocket.py
|
|
~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
|
|
This file implements a buffered socket wrapper.
|
|
|
|
The purpose of this is to avoid the overhead of unnecessary syscalls while
|
|
allowing small reads from the network. This represents a potentially massive
|
|
performance optimisation at the cost of burning some memory in the userspace
|
|
process.
|
|
"""
|
|
import select
|
|
from .exceptions import ConnectionResetError, LineTooLongError
|
|
|
|
|
|
class BufferedSocket(object):
|
|
"""
|
|
A buffered socket wrapper.
|
|
|
|
The purpose of this is to avoid the overhead of unnecessary syscalls while
|
|
allowing small reads from the network. This represents a potentially
|
|
massive performance optimisation at the cost of burning some memory in the
|
|
userspace process.
|
|
"""
|
|
def __init__(self, sck, buffer_size=1000):
|
|
"""
|
|
Create the buffered socket.
|
|
|
|
:param sck: The socket to wrap.
|
|
:param buffer_size: The size of the backing buffer in bytes. This
|
|
parameter should be set to an appropriate value for your use case.
|
|
Small values of ``buffer_size`` increase the overhead of buffer
|
|
management: large values cause more memory to be used.
|
|
"""
|
|
# The wrapped socket.
|
|
self._sck = sck
|
|
|
|
# The buffer we're using.
|
|
self._backing_buffer = bytearray(buffer_size)
|
|
self._buffer_view = memoryview(self._backing_buffer)
|
|
|
|
# The size of the buffer.
|
|
self._buffer_size = buffer_size
|
|
|
|
# The start index in the memory view.
|
|
self._index = 0
|
|
|
|
# The number of bytes in the buffer.
|
|
self._bytes_in_buffer = 0
|
|
|
|
@property
|
|
def _remaining_capacity(self):
|
|
"""
|
|
The maximum number of bytes the buffer could still contain.
|
|
"""
|
|
return self._buffer_size - self._index
|
|
|
|
@property
|
|
def _buffer_end(self):
|
|
"""
|
|
The index of the first free byte in the buffer.
|
|
"""
|
|
return self._index + self._bytes_in_buffer
|
|
|
|
@property
|
|
def can_read(self):
|
|
"""
|
|
Whether or not there is more data to read from the socket.
|
|
"""
|
|
read = select.select([self._sck], [], [], 0)[0]
|
|
if read:
|
|
return True
|
|
|
|
return False
|
|
|
|
@property
|
|
def buffer(self):
|
|
"""
|
|
Get access to the buffer itself.
|
|
"""
|
|
return self._buffer_view[self._index:self._buffer_end]
|
|
|
|
def advance_buffer(self, count):
|
|
"""
|
|
Advances the buffer by the amount of data consumed outside the socket.
|
|
"""
|
|
self._index += count
|
|
self._bytes_in_buffer -= count
|
|
|
|
def new_buffer(self):
|
|
"""
|
|
This method moves all the data in the backing buffer to the start of
|
|
a new, fresh buffer. This gives the ability to read much more data.
|
|
"""
|
|
def read_all_from_buffer():
|
|
end = self._index + self._bytes_in_buffer
|
|
return self._buffer_view[self._index:end]
|
|
|
|
new_buffer = bytearray(self._buffer_size)
|
|
new_buffer_view = memoryview(new_buffer)
|
|
new_buffer_view[0:self._bytes_in_buffer] = read_all_from_buffer()
|
|
|
|
self._index = 0
|
|
self._backing_buffer = new_buffer
|
|
self._buffer_view = new_buffer_view
|
|
|
|
return
|
|
|
|
def recv(self, amt):
|
|
"""
|
|
Read some data from the socket.
|
|
|
|
:param amt: The amount of data to read.
|
|
:returns: A ``memoryview`` object containing the appropriate number of
|
|
bytes. The data *must* be copied out by the caller before the next
|
|
call to this function.
|
|
"""
|
|
# In this implementation you can never read more than the number of
|
|
# bytes in the buffer.
|
|
if amt > self._buffer_size:
|
|
amt = self._buffer_size
|
|
|
|
# If the amount of data we've been asked to read is less than the
|
|
# remaining space in the buffer, we need to clear out the buffer and
|
|
# start over.
|
|
if amt > self._remaining_capacity:
|
|
self.new_buffer()
|
|
|
|
# If there's still some room in the buffer, opportunistically attempt
|
|
# to read into it.
|
|
# If we don't actually _need_ the data (i.e. there's enough in the
|
|
# buffer to satisfy the request), use select to work out if the read
|
|
# attempt will block. If it will, don't bother reading. If we need the
|
|
# data, always do the read.
|
|
if self._bytes_in_buffer >= amt:
|
|
should_read = select.select([self._sck], [], [], 0)[0]
|
|
else:
|
|
should_read = True
|
|
|
|
if (self._remaining_capacity > self._bytes_in_buffer and should_read):
|
|
count = self._sck.recv_into(self._buffer_view[self._buffer_end:])
|
|
|
|
# The socket just got closed. We should throw an exception if we
|
|
# were asked for more data than we can return.
|
|
if not count and amt > self._bytes_in_buffer:
|
|
raise ConnectionResetError()
|
|
self._bytes_in_buffer += count
|
|
|
|
# Read out the bytes and update the index.
|
|
amt = min(amt, self._bytes_in_buffer)
|
|
data = self._buffer_view[self._index:self._index+amt]
|
|
|
|
self._index += amt
|
|
self._bytes_in_buffer -= amt
|
|
|
|
return data
|
|
|
|
def fill(self):
|
|
"""
|
|
Attempts to fill the buffer as much as possible. It will block for at
|
|
most the time required to have *one* ``recv_into`` call return.
|
|
"""
|
|
if not self._remaining_capacity:
|
|
self.new_buffer()
|
|
|
|
count = self._sck.recv_into(self._buffer_view[self._buffer_end:])
|
|
if not count:
|
|
raise ConnectionResetError()
|
|
|
|
self._bytes_in_buffer += count
|
|
|
|
return
|
|
|
|
def readline(self):
|
|
"""
|
|
Read up to a newline from the network and returns it. The implicit
|
|
maximum line length is the buffer size of the buffered socket.
|
|
|
|
Note that, unlike recv, this method absolutely *does* block until it
|
|
can read the line.
|
|
|
|
:returns: A ``memoryview`` object containing the appropriate number of
|
|
bytes. The data *must* be copied out by the caller before the next
|
|
call to this function.
|
|
"""
|
|
# First, check if there's anything in the buffer. This is one of those
|
|
# rare circumstances where this will work correctly on all platforms.
|
|
index = self._backing_buffer.find(
|
|
b'\n',
|
|
self._index,
|
|
self._index + self._bytes_in_buffer
|
|
)
|
|
|
|
if index != -1:
|
|
length = index + 1 - self._index
|
|
data = self._buffer_view[self._index:self._index+length]
|
|
self._index += length
|
|
self._bytes_in_buffer -= length
|
|
return data
|
|
|
|
# In this case, we didn't find a newline in the buffer. To fix that,
|
|
# read some data into the buffer. To do our best to satisfy the read,
|
|
# we should shunt the data down in the buffer so that it's right at
|
|
# the start. We don't bother if we're already at the start of the
|
|
# buffer.
|
|
if self._index != 0:
|
|
self.new_buffer()
|
|
|
|
while self._bytes_in_buffer < self._buffer_size:
|
|
count = self._sck.recv_into(self._buffer_view[self._buffer_end:])
|
|
if not count:
|
|
raise ConnectionResetError()
|
|
|
|
# We have some more data. Again, look for a newline in that gap.
|
|
first_new_byte = self._buffer_end
|
|
self._bytes_in_buffer += count
|
|
index = self._backing_buffer.find(
|
|
b'\n',
|
|
first_new_byte,
|
|
first_new_byte + count,
|
|
)
|
|
|
|
if index != -1:
|
|
# The length of the buffer is the index into the
|
|
# buffer at which we found the newline plus 1, minus the start
|
|
# index of the buffer, which really should be zero.
|
|
assert not self._index
|
|
length = index + 1
|
|
data = self._buffer_view[:length]
|
|
self._index += length
|
|
self._bytes_in_buffer -= length
|
|
return data
|
|
|
|
# If we got here, it means we filled the buffer without ever getting
|
|
# a newline. Time to throw an exception.
|
|
raise LineTooLongError()
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._sck, name)
|