Fix bug in to_bytes_packed when packed size > unpacked size.

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&)

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) {
+    kj::Array<capnp::byte> result = kj::heapArray<capnp::byte>(wordCount * 8);
-        totalSize += segment.size();
-    }
-
-    kj::Array<capnp::byte> result = kj::heapArray<capnp::byte>(totalSize * 8);
     kj::ArrayOutputStream out(result.asPtr());
     capnp::writePackedMessage(out, message);
     return heapArray(out.getArray()); // TODO: make this non-copying somehow

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:

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:
+    cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) except +reraise_kj_exception:
-        self._check_write()
         cdef _MessageBuilder builder = self._parent
-        array = helpers.messageToPackedBytes(deref(builder.thisptr))
+        array = helpers.messageToPackedBytes(deref(builder.thisptr), word_count)
         cdef const char* ptr = <const char *>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