From 9f06f0f3e873206160201e0bc7b33e046d39d0bc Mon Sep 17 00:00:00 2001 From: Jason Paryani Date: Mon, 24 Feb 2014 14:06:28 -0800 Subject: [PATCH] Fix bug in to_bytes_packed when packed size > unpacked size. --- capnp/helpers/helpers.pxd | 2 +- capnp/helpers/serialize.h | 10 ++----- capnp/includes/capnp_cpp.pxd | 6 ++++ capnp/lib/capnp.pyx | 54 ++++++++++++++++++++++++++++++------ 4 files changed, 55 insertions(+), 17 deletions(-) diff --git a/capnp/helpers/helpers.pxd b/capnp/helpers/helpers.pxd index 5f1efa5..577506e 100644 --- a/capnp/helpers/helpers.pxd +++ b/capnp/helpers/helpers.pxd @@ -33,7 +33,7 @@ cdef extern from "../helpers/rpcHelper.h": PyPromise connectServer(TaskSet &, PyRestorer &, AsyncIoContext *, StringPtr) cdef extern from "../helpers/serialize.h": - ByteArray messageToPackedBytes(MessageBuilder &) + ByteArray messageToPackedBytes(MessageBuilder &, size_t wordCount) cdef extern from "../helpers/asyncHelper.h": void waitNeverDone(WaitScope&) diff --git a/capnp/helpers/serialize.h b/capnp/helpers/serialize.h index 8393c7d..7303856 100644 --- a/capnp/helpers/serialize.h +++ b/capnp/helpers/serialize.h @@ -4,16 +4,10 @@ #include "capnp/dynamic.h" #include "capnp/serialize-packed.h" -kj::Array< ::capnp::byte> messageToPackedBytes(capnp::MessageBuilder & message) +kj::Array< ::capnp::byte> messageToPackedBytes(capnp::MessageBuilder & message, size_t wordCount) { - auto segments = message.getSegmentsForOutput(); - size_t totalSize = segments.size() / 2 + 1; - for (auto& segment: segments) { - totalSize += segment.size(); - } - - kj::Array result = kj::heapArray(totalSize * 8); + kj::Array result = kj::heapArray(wordCount * 8); kj::ArrayOutputStream out(result.asPtr()); capnp::writePackedMessage(out, message); return heapArray(out.getArray()); // TODO: make this non-copying somehow diff --git a/capnp/includes/capnp_cpp.pxd b/capnp/includes/capnp_cpp.pxd index f486ca8..0a76a3d 100644 --- a/capnp/includes/capnp_cpp.pxd +++ b/capnp/includes/capnp_cpp.pxd @@ -11,6 +11,9 @@ from .capnp.includes.types cimport * cdef extern from "capnp/common.h" namespace " ::capnp": enum Void: VOID " ::capnp::VOID" + cdef cppclass MessageSize: + uint64_t wordCount + uint capCount cdef extern from "capnp/common.h": int CAPNP_VERSION_MAJOR @@ -214,11 +217,13 @@ cdef extern from "capnp/dynamic.h" namespace " ::capnp": bint has(char *) except +reraise_kj_exception StructSchema getSchema() Maybe[StructSchema.Field] which() + MessageSize totalSize() cppclass Pipeline: Pipeline() Pipeline(Pipeline &) DynamicValueForward.Pipeline get(char *) StructSchema getSchema() + cdef cppclass DynamicStruct_Builder" ::capnp::DynamicStruct::Builder": # Need to flatten this class out, since nested C++ classes cause havoc with cython fused types DynamicStruct_Builder() DynamicStruct_Builder(DynamicStruct_Builder &) @@ -232,6 +237,7 @@ cdef extern from "capnp/dynamic.h" namespace " ::capnp": void adopt(char *, DynamicOrphan) except +reraise_kj_exception DynamicOrphan disown(char *) DynamicStruct.Reader asReader() + MessageSize totalSize() cdef extern from "capnp/dynamic.h" namespace " ::capnp": cdef cppclass DynamicCapability: diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index 88bd7e9..564edbd 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -837,6 +837,14 @@ cdef class _DynamicEnumField: def __call__(self): return str(self) +cdef class _MessageSize: + cdef public uint64_t word_count + cdef public uint cap_count + + def __init__(self, uint64_t word_count, uint cap_count): + self.word_count = word_count + self.cap_count = cap_count + cdef class _DynamicStructReader: """Reads Cap'n Proto structs @@ -911,6 +919,11 @@ cdef class _DynamicStructReader: builder = _MallocMessageBuilder() return builder.set_root(self) + property total_size: + def __get__(self): + size = self.thisptr.totalSize() + return _MessageSize(size.wordCount, size.capCount) + cdef class _DynamicStructBuilder: """Builds Cap'n Proto structs @@ -992,12 +1005,26 @@ cdef class _DynamicStructBuilder: self._is_written = True return ret - cpdef to_bytes_packed(_DynamicStructBuilder self) except +reraise_kj_exception: - self._check_write() + cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) except +reraise_kj_exception: cdef _MessageBuilder builder = self._parent - array = helpers.messageToPackedBytes(deref(builder.thisptr)) + array = helpers.messageToPackedBytes(deref(builder.thisptr), word_count) cdef const char* ptr = array.begin() cdef bytes ret = ptr[:array.size()] + return ret + + cpdef to_bytes_packed(_DynamicStructBuilder self) except +reraise_kj_exception: + self._check_write() + word_count = self.total_size.word_count + 2 + + try: + ret = self._to_bytes_packed_helper(word_count) + except Exception as e: + if 'backing array was not large enough' in str(e): + word_count *= 2 + ret = self._to_bytes_packed_helper(word_count) + else: + raise + self._is_written = True return ret @@ -1147,6 +1174,11 @@ cdef class _DynamicStructBuilder: def to_dict(self, verbose=False): return _to_dict(self, verbose) + property total_size: + def __get__(self): + size = self.thisptr.totalSize() + return _MessageSize(size.wordCount, size.capCount) + cdef class _DynamicStructPipeline: """Reads Cap'n Proto structs @@ -1742,9 +1774,9 @@ cdef class _Restorer: cdef class _TwoPartyVatNetwork: cdef Own[C_TwoPartyVatNetwork] thisptr - cdef _FdAsyncIoStream stream + cdef _AsyncIoStream stream - cdef _init(self, _FdAsyncIoStream stream, Side side): + cdef _init(self, _AsyncIoStream stream, Side side): self.stream = stream self.thisptr = makeTwoPartyVatNetwork(deref(stream.thisptr), side) return self @@ -1769,7 +1801,7 @@ cdef class TwoPartyClient: cdef public _TwoPartyVatNetwork _network cdef public object _orig_stream cdef public _Restorer _restorer - cdef public _FdAsyncIoStream _stream + cdef public _AsyncIoStream _stream def __init__(self, socket, restorer=None): if isinstance(socket, basestring): @@ -1849,7 +1881,7 @@ cdef class TwoPartyServer: cdef public _TwoPartyVatNetwork _network cdef public object _orig_stream, _server_socket, _disconnect_promise cdef public _Restorer _restorer - cdef public _FdAsyncIoStream _stream + cdef public _AsyncIoStream _stream cdef object _port cdef public object port_promise cdef capnp.TaskSet * _task_set @@ -1908,8 +1940,10 @@ cdef class TwoPartyServer: # TODO: add restore functionality here? -cdef class _FdAsyncIoStream: +cdef class _AsyncIoStream: cdef Own[AsyncIoStream] thisptr + +cdef class _FdAsyncIoStream(_AsyncIoStream): cdef _EventLoop _event_loop def __init__(self, int fd): @@ -1919,6 +1953,10 @@ cdef class _FdAsyncIoStream: self._event_loop = C_DEFAULT_EVENT_LOOP_GETTER() self.thisptr = self._event_loop.wrapSocketFd(fd) +cdef class PyAsyncIoStream(_AsyncIoStream): + def __init__(self, int fd): + pass + cdef class PromiseFulfillerPair: cdef Own[C_PromiseFulfillerPair] thisptr cdef public bint is_consumed