fix loop blocking

Author: danlapid

src/rust/container/lib.rs

@@ -25,6 +25,7 @@ use tokio::runtime::Builder;
 pub mod io;
 use io::channel;
 use io::signo_as_string;
+use tokio::sync::Notify;
 
 #[derive(Debug, Error)]
 pub enum ContainerError {
@@ -99,6 +100,7 @@ impl ContainerService {
             return false;
         }
         self.r#impl.sender.try_send(data.to_vec()).unwrap();
+        self.r#impl.notify.notify_one();
         true
     }
 
@@ -118,6 +120,7 @@ impl ContainerService {
 
 struct Impl {
     sender: mpsc::SyncSender<Vec<u8>>,
+    notify: Arc<Notify>,
     write_shutdown: std::sync::atomic::AtomicBool,
 }
 
@@ -128,8 +131,10 @@ impl Impl {
         mut messages_callback: Pin<&'static mut ffi::MessageCallback>,
     ) -> Result<Self, ContainerError> {
         let (input_sender, input_receiver) = mpsc::sync_channel::<Vec<u8>>(1000);
-        let (mut output_sender, mut output_receiver) = channel();
-        let (mut tokio_input_sender, mut tokio_input_receiver) = channel();
+        let (output_sender, mut output_receiver) = channel();
+        let (mut tokio_input_sender, tokio_input_receiver) = channel();
+        let output_notify = Arc::new(Notify::new());
+        let notifiy_awaiter = output_notify.clone();
 
         let server = Server::connect(address, container_name)?;
         let runtime = Builder::new_current_thread().enable_all().build().unwrap();
@@ -152,7 +157,7 @@ impl Impl {
                 tokio::task::spawn_local(rpc_system);
             });
             local.spawn_local(async move {
-                while true {
+                loop {
                     let mut buf = [0; 8096];
                     let len = output_receiver.read(&mut buf).await;
                     if len.is_err() {
@@ -163,9 +168,12 @@ impl Impl {
                 dbg!("OUTPUT_RECEIVER.RECV() ENDED");
             });
             local.spawn_local(async move {
-                while let Ok(msg) = input_receiver.recv() {
-                    if tokio_input_sender.write_all(&msg).await.is_err() {
-                        break;
+                loop {
+                    notifiy_awaiter.notified().await;
+                    while let Ok(msg) = input_receiver.try_recv() {
+                        if tokio_input_sender.write(&msg).await.is_err() {
+                            break;
+                        }
                     }
                 }
                 dbg!("INPUT_RECEIVER.RECV() ENDED");
@@ -176,6 +184,7 @@ impl Impl {
 
         Ok(Impl {
             sender: input_sender,
+            notify: output_notify,
             write_shutdown: std::sync::atomic::AtomicBool::new(false),
         })
     }

src/workerd/io/container-client.c++

@@ -70,12 +70,11 @@ kj::Promise<size_t> ContainerAsyncStream::tryRead(void* buffer, size_t minBytes,
 
 kj::Promise<void> ContainerAsyncStream::write(kj::ArrayPtr<const kj::byte> buffer) {
   KJ_DBG("WRITE");
-  if (service->write_data(buffer.as<Rust>())) {
-    return kj::READY_NOW;
-  } else {
+  if (!service->write_data(buffer.as<Rust>())) {
     KJ_DBG("WRITE FAILED");
     return KJ_EXCEPTION(DISCONNECTED, "Write failed: stream is disconnected");
   }
+  return kj::READY_NOW;
 }
 
 kj::Promise<void> ContainerAsyncStream::write(

src/workerd/server/server.c++

@@ -2277,28 +2277,24 @@ class Server::WorkerService final: public Service,
         static constexpr uint16_t hibernationEventTypeId = 8;
 
         kj::Maybe<rpc::Container::Client> containerClient = kj::none;
-        kj::Own<io::ContainerAsyncStream> containerStream;
         kj::Own<capnp::TwoPartyClient> containerRpcClient;
 
         KJ_IF_SOME(config, containerOptions) {
-          KJ_IF_SOME(path, service.containerAddress) {
-            KJ_REQUIRE(config.hasName(), "Container name is required");
-            auto container_name = config.getName();
-            containerStream = io::createContainerRpcStream(kj::str(path), kj::str(container_name));
-            containerRpcClient = kj::heap<capnp::TwoPartyClient>(*containerStream);
-            containerClient = containerRpcClient->bootstrap().castAs<rpc::Container>();
-            KJ_DBG("CREATED CONTAINER CLIENT");
-          } else {
-            KJ_FAIL_REQUIRE(
-                "container address needs to be defined in order enable containers on this durable object.");
-          }
+          auto& path = KJ_ASSERT_NONNULL(service.containerAddress,
+              "container address needs to be defined in order enable containers on this durable object.");
+          KJ_REQUIRE(config.hasName(), "Container name is required");
+          auto container_name = config.getName();
+          auto containerStream =
+              io::createContainerRpcStream(kj::str(path), kj::str(container_name));
+          containerRpcClient =
+              kj::heap<capnp::TwoPartyClient>(*containerStream).attach(kj::mv(containerStream));
+          containerClient = containerRpcClient->bootstrap().castAs<rpc::Container>();
         }
 
         auto& actorRef = *actor.emplace(kj::refcounted<Worker::Actor>(*service.worker, getTracker(),
             Worker::Actor::cloneId(id), true, kj::mv(makeActorCache), parent.className,
             kj::mv(makeStorage), kj::mv(loopback), timerChannel, kj::refcounted<ActorObserver>(),
             tryGetManagerRef(), hibernationEventTypeId, kj::mv(containerClient))
-                                            .attach(kj::mv(containerStream))
                                             .attach(kj::mv(containerRpcClient)));
         onBrokenTask = monitorOnBroken(actorRef);
       }