Major fixes and new features
All checks were successful
continuous-integration/drone/push Build is passing
All checks were successful
continuous-integration/drone/push Build is passing
This commit is contained in:
817
venv/lib/python3.12/site-packages/asyncpg/pgproto/buffer.pyx
Normal file
817
venv/lib/python3.12/site-packages/asyncpg/pgproto/buffer.pyx
Normal file
@@ -0,0 +1,817 @@
|
||||
# Copyright (C) 2016-present the asyncpg authors and contributors
|
||||
# <see AUTHORS file>
|
||||
#
|
||||
# This module is part of asyncpg and is released under
|
||||
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
|
||||
from libc.string cimport memcpy
|
||||
|
||||
import collections
|
||||
|
||||
class BufferError(Exception):
|
||||
pass
|
||||
|
||||
@cython.no_gc_clear
|
||||
@cython.final
|
||||
@cython.freelist(_BUFFER_FREELIST_SIZE)
|
||||
cdef class WriteBuffer:
|
||||
|
||||
def __cinit__(self):
|
||||
self._smallbuf_inuse = True
|
||||
self._buf = self._smallbuf
|
||||
self._size = _BUFFER_INITIAL_SIZE
|
||||
self._length = 0
|
||||
self._message_mode = 0
|
||||
|
||||
def __dealloc__(self):
|
||||
if self._buf is not NULL and not self._smallbuf_inuse:
|
||||
cpython.PyMem_Free(self._buf)
|
||||
self._buf = NULL
|
||||
self._size = 0
|
||||
|
||||
if self._view_count:
|
||||
raise BufferError(
|
||||
'Deallocating buffer with attached memoryviews')
|
||||
|
||||
def __getbuffer__(self, Py_buffer *buffer, int flags):
|
||||
self._view_count += 1
|
||||
|
||||
cpython.PyBuffer_FillInfo(
|
||||
buffer, self, self._buf, self._length,
|
||||
1, # read-only
|
||||
flags)
|
||||
|
||||
def __releasebuffer__(self, Py_buffer *buffer):
|
||||
self._view_count -= 1
|
||||
|
||||
cdef inline _check_readonly(self):
|
||||
if self._view_count:
|
||||
raise BufferError('the buffer is in read-only mode')
|
||||
|
||||
cdef inline _ensure_alloced(self, ssize_t extra_length):
|
||||
cdef ssize_t new_size = extra_length + self._length
|
||||
|
||||
if new_size > self._size:
|
||||
self._reallocate(new_size)
|
||||
|
||||
cdef _reallocate(self, ssize_t new_size):
|
||||
cdef char *new_buf
|
||||
|
||||
if new_size < _BUFFER_MAX_GROW:
|
||||
new_size = _BUFFER_MAX_GROW
|
||||
else:
|
||||
# Add a little extra
|
||||
new_size += _BUFFER_INITIAL_SIZE
|
||||
|
||||
if self._smallbuf_inuse:
|
||||
new_buf = <char*>cpython.PyMem_Malloc(
|
||||
sizeof(char) * <size_t>new_size)
|
||||
if new_buf is NULL:
|
||||
self._buf = NULL
|
||||
self._size = 0
|
||||
self._length = 0
|
||||
raise MemoryError
|
||||
memcpy(new_buf, self._buf, <size_t>self._size)
|
||||
self._size = new_size
|
||||
self._buf = new_buf
|
||||
self._smallbuf_inuse = False
|
||||
else:
|
||||
new_buf = <char*>cpython.PyMem_Realloc(
|
||||
<void*>self._buf, <size_t>new_size)
|
||||
if new_buf is NULL:
|
||||
cpython.PyMem_Free(self._buf)
|
||||
self._buf = NULL
|
||||
self._size = 0
|
||||
self._length = 0
|
||||
raise MemoryError
|
||||
self._buf = new_buf
|
||||
self._size = new_size
|
||||
|
||||
cdef inline start_message(self, char type):
|
||||
if self._length != 0:
|
||||
raise BufferError(
|
||||
'cannot start_message for a non-empty buffer')
|
||||
self._ensure_alloced(5)
|
||||
self._message_mode = 1
|
||||
self._buf[0] = type
|
||||
self._length = 5
|
||||
|
||||
cdef inline end_message(self):
|
||||
# "length-1" to exclude the message type byte
|
||||
cdef ssize_t mlen = self._length - 1
|
||||
|
||||
self._check_readonly()
|
||||
if not self._message_mode:
|
||||
raise BufferError(
|
||||
'end_message can only be called with start_message')
|
||||
if self._length < 5:
|
||||
raise BufferError('end_message: buffer is too small')
|
||||
if mlen > _MAXINT32:
|
||||
raise BufferError('end_message: message is too large')
|
||||
|
||||
hton.pack_int32(&self._buf[1], <int32_t>mlen)
|
||||
return self
|
||||
|
||||
cdef inline reset(self):
|
||||
self._length = 0
|
||||
self._message_mode = 0
|
||||
|
||||
cdef write_buffer(self, WriteBuffer buf):
|
||||
self._check_readonly()
|
||||
|
||||
if not buf._length:
|
||||
return
|
||||
|
||||
self._ensure_alloced(buf._length)
|
||||
memcpy(self._buf + self._length,
|
||||
<void*>buf._buf,
|
||||
<size_t>buf._length)
|
||||
self._length += buf._length
|
||||
|
||||
cdef write_byte(self, char b):
|
||||
self._check_readonly()
|
||||
|
||||
self._ensure_alloced(1)
|
||||
self._buf[self._length] = b
|
||||
self._length += 1
|
||||
|
||||
cdef write_bytes(self, bytes data):
|
||||
cdef char* buf
|
||||
cdef ssize_t len
|
||||
|
||||
cpython.PyBytes_AsStringAndSize(data, &buf, &len)
|
||||
self.write_cstr(buf, len)
|
||||
|
||||
cdef write_bytestring(self, bytes string):
|
||||
cdef char* buf
|
||||
cdef ssize_t len
|
||||
|
||||
cpython.PyBytes_AsStringAndSize(string, &buf, &len)
|
||||
# PyBytes_AsStringAndSize returns a null-terminated buffer,
|
||||
# but the null byte is not counted in len. hence the + 1
|
||||
self.write_cstr(buf, len + 1)
|
||||
|
||||
cdef write_str(self, str string, str encoding):
|
||||
self.write_bytestring(string.encode(encoding))
|
||||
|
||||
cdef write_len_prefixed_buffer(self, WriteBuffer buf):
|
||||
# Write a length-prefixed (not NULL-terminated) bytes sequence.
|
||||
self.write_int32(<int32_t>buf.len())
|
||||
self.write_buffer(buf)
|
||||
|
||||
cdef write_len_prefixed_bytes(self, bytes data):
|
||||
# Write a length-prefixed (not NULL-terminated) bytes sequence.
|
||||
cdef:
|
||||
char *buf
|
||||
ssize_t size
|
||||
|
||||
cpython.PyBytes_AsStringAndSize(data, &buf, &size)
|
||||
if size > _MAXINT32:
|
||||
raise BufferError('string is too large')
|
||||
# `size` does not account for the NULL at the end.
|
||||
self.write_int32(<int32_t>size)
|
||||
self.write_cstr(buf, size)
|
||||
|
||||
cdef write_frbuf(self, FRBuffer *buf):
|
||||
cdef:
|
||||
ssize_t buf_len = buf.len
|
||||
if buf_len > 0:
|
||||
self.write_cstr(frb_read_all(buf), buf_len)
|
||||
|
||||
cdef write_cstr(self, const char *data, ssize_t len):
|
||||
self._check_readonly()
|
||||
self._ensure_alloced(len)
|
||||
|
||||
memcpy(self._buf + self._length, <void*>data, <size_t>len)
|
||||
self._length += len
|
||||
|
||||
cdef write_int16(self, int16_t i):
|
||||
self._check_readonly()
|
||||
self._ensure_alloced(2)
|
||||
|
||||
hton.pack_int16(&self._buf[self._length], i)
|
||||
self._length += 2
|
||||
|
||||
cdef write_int32(self, int32_t i):
|
||||
self._check_readonly()
|
||||
self._ensure_alloced(4)
|
||||
|
||||
hton.pack_int32(&self._buf[self._length], i)
|
||||
self._length += 4
|
||||
|
||||
cdef write_int64(self, int64_t i):
|
||||
self._check_readonly()
|
||||
self._ensure_alloced(8)
|
||||
|
||||
hton.pack_int64(&self._buf[self._length], i)
|
||||
self._length += 8
|
||||
|
||||
cdef write_float(self, float f):
|
||||
self._check_readonly()
|
||||
self._ensure_alloced(4)
|
||||
|
||||
hton.pack_float(&self._buf[self._length], f)
|
||||
self._length += 4
|
||||
|
||||
cdef write_double(self, double d):
|
||||
self._check_readonly()
|
||||
self._ensure_alloced(8)
|
||||
|
||||
hton.pack_double(&self._buf[self._length], d)
|
||||
self._length += 8
|
||||
|
||||
@staticmethod
|
||||
cdef WriteBuffer new_message(char type):
|
||||
cdef WriteBuffer buf
|
||||
buf = WriteBuffer.__new__(WriteBuffer)
|
||||
buf.start_message(type)
|
||||
return buf
|
||||
|
||||
@staticmethod
|
||||
cdef WriteBuffer new():
|
||||
cdef WriteBuffer buf
|
||||
buf = WriteBuffer.__new__(WriteBuffer)
|
||||
return buf
|
||||
|
||||
|
||||
@cython.no_gc_clear
|
||||
@cython.final
|
||||
@cython.freelist(_BUFFER_FREELIST_SIZE)
|
||||
cdef class ReadBuffer:
|
||||
|
||||
def __cinit__(self):
|
||||
self._bufs = collections.deque()
|
||||
self._bufs_append = self._bufs.append
|
||||
self._bufs_popleft = self._bufs.popleft
|
||||
self._bufs_len = 0
|
||||
self._buf0 = None
|
||||
self._buf0_prev = None
|
||||
self._pos0 = 0
|
||||
self._len0 = 0
|
||||
self._length = 0
|
||||
|
||||
self._current_message_type = 0
|
||||
self._current_message_len = 0
|
||||
self._current_message_len_unread = 0
|
||||
self._current_message_ready = 0
|
||||
|
||||
cdef feed_data(self, data):
|
||||
cdef:
|
||||
ssize_t dlen
|
||||
bytes data_bytes
|
||||
|
||||
if not cpython.PyBytes_CheckExact(data):
|
||||
if cpythonx.PyByteArray_CheckExact(data):
|
||||
# ProactorEventLoop in Python 3.10+ seems to be sending
|
||||
# bytearray objects instead of bytes. Handle this here
|
||||
# to avoid duplicating this check in every data_received().
|
||||
data = bytes(data)
|
||||
else:
|
||||
raise BufferError(
|
||||
'feed_data: a bytes or bytearray object expected')
|
||||
|
||||
# Uncomment the below code to test code paths that
|
||||
# read single int/str/bytes sequences are split over
|
||||
# multiple received buffers.
|
||||
#
|
||||
# ll = 107
|
||||
# if len(data) > ll:
|
||||
# self.feed_data(data[:ll])
|
||||
# self.feed_data(data[ll:])
|
||||
# return
|
||||
|
||||
data_bytes = <bytes>data
|
||||
|
||||
dlen = cpython.Py_SIZE(data_bytes)
|
||||
if dlen == 0:
|
||||
# EOF?
|
||||
return
|
||||
|
||||
self._bufs_append(data_bytes)
|
||||
self._length += dlen
|
||||
|
||||
if self._bufs_len == 0:
|
||||
# First buffer
|
||||
self._len0 = dlen
|
||||
self._buf0 = data_bytes
|
||||
|
||||
self._bufs_len += 1
|
||||
|
||||
cdef inline _ensure_first_buf(self):
|
||||
if PG_DEBUG:
|
||||
if self._len0 == 0:
|
||||
raise BufferError('empty first buffer')
|
||||
if self._length == 0:
|
||||
raise BufferError('empty buffer')
|
||||
|
||||
if self._pos0 == self._len0:
|
||||
self._switch_to_next_buf()
|
||||
|
||||
cdef _switch_to_next_buf(self):
|
||||
# The first buffer is fully read, discard it
|
||||
self._bufs_popleft()
|
||||
self._bufs_len -= 1
|
||||
|
||||
# Shouldn't fail, since we've checked that `_length >= 1`
|
||||
# in _ensure_first_buf()
|
||||
self._buf0_prev = self._buf0
|
||||
self._buf0 = <bytes>self._bufs[0]
|
||||
|
||||
self._pos0 = 0
|
||||
self._len0 = len(self._buf0)
|
||||
|
||||
if PG_DEBUG:
|
||||
if self._len0 < 1:
|
||||
raise BufferError(
|
||||
'debug: second buffer of ReadBuffer is empty')
|
||||
|
||||
cdef inline const char* _try_read_bytes(self, ssize_t nbytes):
|
||||
# Try to read *nbytes* from the first buffer.
|
||||
#
|
||||
# Returns pointer to data if there is at least *nbytes*
|
||||
# in the buffer, NULL otherwise.
|
||||
#
|
||||
# Important: caller must call _ensure_first_buf() prior
|
||||
# to calling try_read_bytes, and must not overread
|
||||
|
||||
cdef:
|
||||
const char *result
|
||||
|
||||
if PG_DEBUG:
|
||||
if nbytes > self._length:
|
||||
return NULL
|
||||
|
||||
if self._current_message_ready:
|
||||
if self._current_message_len_unread < nbytes:
|
||||
return NULL
|
||||
|
||||
if self._pos0 + nbytes <= self._len0:
|
||||
result = cpython.PyBytes_AS_STRING(self._buf0)
|
||||
result += self._pos0
|
||||
self._pos0 += nbytes
|
||||
self._length -= nbytes
|
||||
if self._current_message_ready:
|
||||
self._current_message_len_unread -= nbytes
|
||||
return result
|
||||
else:
|
||||
return NULL
|
||||
|
||||
cdef inline _read_into(self, char *buf, ssize_t nbytes):
|
||||
cdef:
|
||||
ssize_t nread
|
||||
char *buf0
|
||||
|
||||
while True:
|
||||
buf0 = cpython.PyBytes_AS_STRING(self._buf0)
|
||||
|
||||
if self._pos0 + nbytes > self._len0:
|
||||
nread = self._len0 - self._pos0
|
||||
memcpy(buf, buf0 + self._pos0, <size_t>nread)
|
||||
self._pos0 = self._len0
|
||||
self._length -= nread
|
||||
nbytes -= nread
|
||||
buf += nread
|
||||
self._ensure_first_buf()
|
||||
|
||||
else:
|
||||
memcpy(buf, buf0 + self._pos0, <size_t>nbytes)
|
||||
self._pos0 += nbytes
|
||||
self._length -= nbytes
|
||||
break
|
||||
|
||||
cdef inline _read_and_discard(self, ssize_t nbytes):
|
||||
cdef:
|
||||
ssize_t nread
|
||||
|
||||
self._ensure_first_buf()
|
||||
while True:
|
||||
if self._pos0 + nbytes > self._len0:
|
||||
nread = self._len0 - self._pos0
|
||||
self._pos0 = self._len0
|
||||
self._length -= nread
|
||||
nbytes -= nread
|
||||
self._ensure_first_buf()
|
||||
|
||||
else:
|
||||
self._pos0 += nbytes
|
||||
self._length -= nbytes
|
||||
break
|
||||
|
||||
cdef bytes read_bytes(self, ssize_t nbytes):
|
||||
cdef:
|
||||
bytes result
|
||||
ssize_t nread
|
||||
const char *cbuf
|
||||
char *buf
|
||||
|
||||
self._ensure_first_buf()
|
||||
cbuf = self._try_read_bytes(nbytes)
|
||||
if cbuf != NULL:
|
||||
return cpython.PyBytes_FromStringAndSize(cbuf, nbytes)
|
||||
|
||||
if nbytes > self._length:
|
||||
raise BufferError(
|
||||
'not enough data to read {} bytes'.format(nbytes))
|
||||
|
||||
if self._current_message_ready:
|
||||
self._current_message_len_unread -= nbytes
|
||||
if self._current_message_len_unread < 0:
|
||||
raise BufferError('buffer overread')
|
||||
|
||||
result = cpython.PyBytes_FromStringAndSize(NULL, nbytes)
|
||||
buf = cpython.PyBytes_AS_STRING(result)
|
||||
self._read_into(buf, nbytes)
|
||||
return result
|
||||
|
||||
cdef bytes read_len_prefixed_bytes(self):
|
||||
cdef int32_t size = self.read_int32()
|
||||
if size < 0:
|
||||
raise BufferError(
|
||||
'negative length for a len-prefixed bytes value')
|
||||
if size == 0:
|
||||
return b''
|
||||
return self.read_bytes(size)
|
||||
|
||||
cdef str read_len_prefixed_utf8(self):
|
||||
cdef:
|
||||
int32_t size
|
||||
const char *cbuf
|
||||
|
||||
size = self.read_int32()
|
||||
if size < 0:
|
||||
raise BufferError(
|
||||
'negative length for a len-prefixed bytes value')
|
||||
|
||||
if size == 0:
|
||||
return ''
|
||||
|
||||
self._ensure_first_buf()
|
||||
cbuf = self._try_read_bytes(size)
|
||||
if cbuf != NULL:
|
||||
return cpython.PyUnicode_DecodeUTF8(cbuf, size, NULL)
|
||||
else:
|
||||
return self.read_bytes(size).decode('utf-8')
|
||||
|
||||
cdef read_uuid(self):
|
||||
cdef:
|
||||
bytes mem
|
||||
const char *cbuf
|
||||
|
||||
self._ensure_first_buf()
|
||||
cbuf = self._try_read_bytes(16)
|
||||
if cbuf != NULL:
|
||||
return pg_uuid_from_buf(cbuf)
|
||||
else:
|
||||
return pg_UUID(self.read_bytes(16))
|
||||
|
||||
cdef inline char read_byte(self) except? -1:
|
||||
cdef const char *first_byte
|
||||
|
||||
if PG_DEBUG:
|
||||
if not self._buf0:
|
||||
raise BufferError(
|
||||
'debug: first buffer of ReadBuffer is empty')
|
||||
|
||||
self._ensure_first_buf()
|
||||
first_byte = self._try_read_bytes(1)
|
||||
if first_byte is NULL:
|
||||
raise BufferError('not enough data to read one byte')
|
||||
|
||||
return first_byte[0]
|
||||
|
||||
cdef inline int64_t read_int64(self) except? -1:
|
||||
cdef:
|
||||
bytes mem
|
||||
const char *cbuf
|
||||
|
||||
self._ensure_first_buf()
|
||||
cbuf = self._try_read_bytes(8)
|
||||
if cbuf != NULL:
|
||||
return hton.unpack_int64(cbuf)
|
||||
else:
|
||||
mem = self.read_bytes(8)
|
||||
return hton.unpack_int64(cpython.PyBytes_AS_STRING(mem))
|
||||
|
||||
cdef inline int32_t read_int32(self) except? -1:
|
||||
cdef:
|
||||
bytes mem
|
||||
const char *cbuf
|
||||
|
||||
self._ensure_first_buf()
|
||||
cbuf = self._try_read_bytes(4)
|
||||
if cbuf != NULL:
|
||||
return hton.unpack_int32(cbuf)
|
||||
else:
|
||||
mem = self.read_bytes(4)
|
||||
return hton.unpack_int32(cpython.PyBytes_AS_STRING(mem))
|
||||
|
||||
cdef inline int16_t read_int16(self) except? -1:
|
||||
cdef:
|
||||
bytes mem
|
||||
const char *cbuf
|
||||
|
||||
self._ensure_first_buf()
|
||||
cbuf = self._try_read_bytes(2)
|
||||
if cbuf != NULL:
|
||||
return hton.unpack_int16(cbuf)
|
||||
else:
|
||||
mem = self.read_bytes(2)
|
||||
return hton.unpack_int16(cpython.PyBytes_AS_STRING(mem))
|
||||
|
||||
cdef inline read_null_str(self):
|
||||
if not self._current_message_ready:
|
||||
raise BufferError(
|
||||
'read_null_str only works when the message guaranteed '
|
||||
'to be in the buffer')
|
||||
|
||||
cdef:
|
||||
ssize_t pos
|
||||
ssize_t nread
|
||||
bytes result
|
||||
const char *buf
|
||||
const char *buf_start
|
||||
|
||||
self._ensure_first_buf()
|
||||
|
||||
buf_start = cpython.PyBytes_AS_STRING(self._buf0)
|
||||
buf = buf_start + self._pos0
|
||||
while buf - buf_start < self._len0:
|
||||
if buf[0] == 0:
|
||||
pos = buf - buf_start
|
||||
nread = pos - self._pos0
|
||||
buf = self._try_read_bytes(nread + 1)
|
||||
if buf != NULL:
|
||||
return cpython.PyBytes_FromStringAndSize(buf, nread)
|
||||
else:
|
||||
break
|
||||
else:
|
||||
buf += 1
|
||||
|
||||
result = b''
|
||||
while True:
|
||||
pos = self._buf0.find(b'\x00', self._pos0)
|
||||
if pos >= 0:
|
||||
result += self._buf0[self._pos0 : pos]
|
||||
nread = pos - self._pos0 + 1
|
||||
self._pos0 = pos + 1
|
||||
self._length -= nread
|
||||
|
||||
self._current_message_len_unread -= nread
|
||||
if self._current_message_len_unread < 0:
|
||||
raise BufferError(
|
||||
'read_null_str: buffer overread')
|
||||
|
||||
return result
|
||||
|
||||
else:
|
||||
result += self._buf0[self._pos0:]
|
||||
nread = self._len0 - self._pos0
|
||||
self._pos0 = self._len0
|
||||
self._length -= nread
|
||||
|
||||
self._current_message_len_unread -= nread
|
||||
if self._current_message_len_unread < 0:
|
||||
raise BufferError(
|
||||
'read_null_str: buffer overread')
|
||||
|
||||
self._ensure_first_buf()
|
||||
|
||||
cdef int32_t take_message(self) except -1:
|
||||
cdef:
|
||||
const char *cbuf
|
||||
|
||||
if self._current_message_ready:
|
||||
return 1
|
||||
|
||||
if self._current_message_type == 0:
|
||||
if self._length < 1:
|
||||
return 0
|
||||
self._ensure_first_buf()
|
||||
cbuf = self._try_read_bytes(1)
|
||||
if cbuf == NULL:
|
||||
raise BufferError(
|
||||
'failed to read one byte on a non-empty buffer')
|
||||
self._current_message_type = cbuf[0]
|
||||
|
||||
if self._current_message_len == 0:
|
||||
if self._length < 4:
|
||||
return 0
|
||||
|
||||
self._ensure_first_buf()
|
||||
cbuf = self._try_read_bytes(4)
|
||||
if cbuf != NULL:
|
||||
self._current_message_len = hton.unpack_int32(cbuf)
|
||||
else:
|
||||
self._current_message_len = self.read_int32()
|
||||
|
||||
self._current_message_len_unread = self._current_message_len - 4
|
||||
|
||||
if self._length < self._current_message_len_unread:
|
||||
return 0
|
||||
|
||||
self._current_message_ready = 1
|
||||
return 1
|
||||
|
||||
cdef inline int32_t take_message_type(self, char mtype) except -1:
|
||||
cdef const char *buf0
|
||||
|
||||
if self._current_message_ready:
|
||||
return self._current_message_type == mtype
|
||||
elif self._length >= 1:
|
||||
self._ensure_first_buf()
|
||||
buf0 = cpython.PyBytes_AS_STRING(self._buf0)
|
||||
|
||||
return buf0[self._pos0] == mtype and self.take_message()
|
||||
else:
|
||||
return 0
|
||||
|
||||
cdef int32_t put_message(self) except -1:
|
||||
if not self._current_message_ready:
|
||||
raise BufferError(
|
||||
'cannot put message: no message taken')
|
||||
self._current_message_ready = False
|
||||
return 0
|
||||
|
||||
cdef inline const char* try_consume_message(self, ssize_t* len):
|
||||
cdef:
|
||||
ssize_t buf_len
|
||||
const char *buf
|
||||
|
||||
if not self._current_message_ready:
|
||||
return NULL
|
||||
|
||||
self._ensure_first_buf()
|
||||
buf_len = self._current_message_len_unread
|
||||
buf = self._try_read_bytes(buf_len)
|
||||
if buf != NULL:
|
||||
len[0] = buf_len
|
||||
self._finish_message()
|
||||
return buf
|
||||
|
||||
cdef discard_message(self):
|
||||
if not self._current_message_ready:
|
||||
raise BufferError('no message to discard')
|
||||
if self._current_message_len_unread > 0:
|
||||
self._read_and_discard(self._current_message_len_unread)
|
||||
self._current_message_len_unread = 0
|
||||
self._finish_message()
|
||||
|
||||
cdef bytes consume_message(self):
|
||||
if not self._current_message_ready:
|
||||
raise BufferError('no message to consume')
|
||||
if self._current_message_len_unread > 0:
|
||||
mem = self.read_bytes(self._current_message_len_unread)
|
||||
else:
|
||||
mem = b''
|
||||
self._finish_message()
|
||||
return mem
|
||||
|
||||
cdef redirect_messages(self, WriteBuffer buf, char mtype,
|
||||
int stop_at=0):
|
||||
if not self._current_message_ready:
|
||||
raise BufferError(
|
||||
'consume_full_messages called on a buffer without a '
|
||||
'complete first message')
|
||||
if mtype != self._current_message_type:
|
||||
raise BufferError(
|
||||
'consume_full_messages called with a wrong mtype')
|
||||
if self._current_message_len_unread != self._current_message_len - 4:
|
||||
raise BufferError(
|
||||
'consume_full_messages called on a partially read message')
|
||||
|
||||
cdef:
|
||||
const char* cbuf
|
||||
ssize_t cbuf_len
|
||||
int32_t msg_len
|
||||
ssize_t new_pos0
|
||||
ssize_t pos_delta
|
||||
int32_t done
|
||||
|
||||
while True:
|
||||
buf.write_byte(mtype)
|
||||
buf.write_int32(self._current_message_len)
|
||||
|
||||
cbuf = self.try_consume_message(&cbuf_len)
|
||||
if cbuf != NULL:
|
||||
buf.write_cstr(cbuf, cbuf_len)
|
||||
else:
|
||||
buf.write_bytes(self.consume_message())
|
||||
|
||||
if self._length > 0:
|
||||
self._ensure_first_buf()
|
||||
else:
|
||||
return
|
||||
|
||||
if stop_at and buf._length >= stop_at:
|
||||
return
|
||||
|
||||
# Fast path: exhaust buf0 as efficiently as possible.
|
||||
if self._pos0 + 5 <= self._len0:
|
||||
cbuf = cpython.PyBytes_AS_STRING(self._buf0)
|
||||
new_pos0 = self._pos0
|
||||
cbuf_len = self._len0
|
||||
|
||||
done = 0
|
||||
# Scan the first buffer and find the position of the
|
||||
# end of the last "mtype" message.
|
||||
while new_pos0 + 5 <= cbuf_len:
|
||||
if (cbuf + new_pos0)[0] != mtype:
|
||||
done = 1
|
||||
break
|
||||
if (stop_at and
|
||||
(buf._length + new_pos0 - self._pos0) > stop_at):
|
||||
done = 1
|
||||
break
|
||||
msg_len = hton.unpack_int32(cbuf + new_pos0 + 1) + 1
|
||||
if new_pos0 + msg_len > cbuf_len:
|
||||
break
|
||||
new_pos0 += msg_len
|
||||
|
||||
if new_pos0 != self._pos0:
|
||||
assert self._pos0 < new_pos0 <= self._len0
|
||||
|
||||
pos_delta = new_pos0 - self._pos0
|
||||
buf.write_cstr(
|
||||
cbuf + self._pos0,
|
||||
pos_delta)
|
||||
|
||||
self._pos0 = new_pos0
|
||||
self._length -= pos_delta
|
||||
|
||||
assert self._length >= 0
|
||||
|
||||
if done:
|
||||
# The next message is of a different type.
|
||||
return
|
||||
|
||||
# Back to slow path.
|
||||
if not self.take_message_type(mtype):
|
||||
return
|
||||
|
||||
cdef bytearray consume_messages(self, char mtype):
|
||||
"""Consume consecutive messages of the same type."""
|
||||
cdef:
|
||||
char *buf
|
||||
ssize_t nbytes
|
||||
ssize_t total_bytes = 0
|
||||
bytearray result
|
||||
|
||||
if not self.take_message_type(mtype):
|
||||
return None
|
||||
|
||||
# consume_messages is a volume-oriented method, so
|
||||
# we assume that the remainder of the buffer will contain
|
||||
# messages of the requested type.
|
||||
result = cpythonx.PyByteArray_FromStringAndSize(NULL, self._length)
|
||||
buf = cpythonx.PyByteArray_AsString(result)
|
||||
|
||||
while self.take_message_type(mtype):
|
||||
self._ensure_first_buf()
|
||||
nbytes = self._current_message_len_unread
|
||||
self._read_into(buf, nbytes)
|
||||
buf += nbytes
|
||||
total_bytes += nbytes
|
||||
self._finish_message()
|
||||
|
||||
# Clamp the result to an actual size read.
|
||||
cpythonx.PyByteArray_Resize(result, total_bytes)
|
||||
|
||||
return result
|
||||
|
||||
cdef finish_message(self):
|
||||
if self._current_message_type == 0 or not self._current_message_ready:
|
||||
# The message has already been finished (e.g by consume_message()),
|
||||
# or has been put back by put_message().
|
||||
return
|
||||
|
||||
if self._current_message_len_unread:
|
||||
if PG_DEBUG:
|
||||
mtype = chr(self._current_message_type)
|
||||
|
||||
discarded = self.consume_message()
|
||||
|
||||
if PG_DEBUG:
|
||||
print('!!! discarding message {!r} unread data: {!r}'.format(
|
||||
mtype,
|
||||
discarded))
|
||||
|
||||
self._finish_message()
|
||||
|
||||
cdef inline _finish_message(self):
|
||||
self._current_message_type = 0
|
||||
self._current_message_len = 0
|
||||
self._current_message_ready = 0
|
||||
self._current_message_len_unread = 0
|
||||
|
||||
@staticmethod
|
||||
cdef ReadBuffer new_message_parser(object data):
|
||||
cdef ReadBuffer buf
|
||||
|
||||
buf = ReadBuffer.__new__(ReadBuffer)
|
||||
buf.feed_data(data)
|
||||
|
||||
buf._current_message_ready = 1
|
||||
buf._current_message_len_unread = buf._len0
|
||||
|
||||
return buf
|
||||
Reference in New Issue
Block a user