Allow reading and writing messages from sockets in async mode

capnp/helpers/capabilityHelper.cpp

@@ -203,7 +203,6 @@ void PyAsyncIoStream::shutdownWrite() {
   _asyncio_stream_shutdown_write(protocol->obj);
 }
 
-
 class TaskToPromiseAdapter {
 public:
   TaskToPromiseAdapter(kj::PromiseFulfiller<void>& fulfiller,
@@ -224,6 +223,19 @@ kj::Promise<void> taskToPromise(kj::Own<PyRefCounter> task, PyObject* callback)
   return kj::newAdaptedPromise<void, TaskToPromiseAdapter>(kj::mv(task), callback);
 }
 
+::kj::Promise<kj::Own<PyRefCounter>> tryReadMessage(kj::AsyncIoStream& stream, capnp::ReaderOptions opts) {
+  return capnp::tryReadMessage(stream, opts)
+    .then([](kj::Maybe<kj::Own<capnp::MessageReader>> maybeReader) -> kj::Promise<kj::Own<PyRefCounter>> {
+        KJ_IF_MAYBE(reader, maybeReader) {
+          PyObject* pyreader = make_async_message_reader(kj::mv(*reader));
+          check_py_error();
+          return kj::heap<PyRefCounter>(pyreader);
+        } else {
+          return kj::heap<PyRefCounter>(Py_None);
+        }
+      });
+}
+
 void init_capnp_api() {
     import_capnp__lib__capnp();
 }

capnp/helpers/capabilityHelper.h

@@ -2,6 +2,7 @@
 
 #include "capnp/dynamic.h"
 #include <kj/async-io.h>
+#include <capnp/serialize-async.h>
 #include <stdexcept>
 #include "Python.h"
 
@@ -147,4 +148,6 @@ inline kj::Exception makeException(kj::StringPtr message) {
 
 kj::Promise<void> taskToPromise(kj::Own<PyRefCounter> coroutine, PyObject* callback);
 
+::kj::Promise<kj::Own<PyRefCounter>> tryReadMessage(kj::AsyncIoStream& stream, capnp::ReaderOptions opts);
+
 void init_capnp_api();