Fixes for concurrency in WS streams

hpnkv · hpnkv · commit b761511bea8c · 2025-05-28T13:36:12.000+01:00
diff --git a/src/eglt/actions/action.h b/src/eglt/actions/action.h
@@ -333,10 +333,9 @@ class Action : public std::enable_shared_from_this<Action> {
     AsyncNode* node = node_map_->Get(GetInputId(name));
     if (stream_ != nullptr &&
         bind_stream.value_or(bind_streams_on_inputs_default_)) {
-      absl::flat_hash_map<std::string_view,
-                          std::shared_ptr<EvergreenWireStream>>
+      absl::flat_hash_map<std::string, std::shared_ptr<EvergreenWireStream>>
           peers;
-      peers.insert({stream_->GetId(), stream_});
+      peers.insert({std::string(stream_->GetId()), stream_});
       node->BindPeers(std::move(peers));
       nodes_with_bound_streams_.insert(node);
     }
@@ -375,10 +374,9 @@ class Action : public std::enable_shared_from_this<Action> {
     AsyncNode* node = node_map_->Get(GetOutputId(name));
     if (stream_ != nullptr &&
         bind_stream.value_or(bind_streams_on_outputs_default_)) {
-      absl::flat_hash_map<std::string_view,
-                          std::shared_ptr<EvergreenWireStream>>
+      absl::flat_hash_map<std::string, std::shared_ptr<EvergreenWireStream>>
           peers;
-      peers.insert({stream_->GetId(), stream_});
+      peers.insert({std::string(stream_->GetId()), stream_});
       node->BindPeers(std::move(peers));
       nodes_with_bound_streams_.insert(node);
     }
diff --git a/src/eglt/concurrency/concurrency.h b/src/eglt/concurrency/concurrency.h
@@ -142,9 +142,7 @@ class ABSL_LOCKABLE ABSL_ATTRIBUTE_WARN_UNUSED ExclusiveAccessGuard {
     mutex_->Unlock();
   }
 
-  void StartBlockingPendingOperation() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
-    ++pending_operations_;
-  }
+  void StartBlockingPendingOperation() { ++pending_operations_; }
 
   void FinishPendingOperation() ABSL_LOCKS_EXCLUDED(mutex_)
       ABSL_EXCLUSIVE_LOCK_FUNCTION(mutex_) {
@@ -153,7 +151,7 @@ class ABSL_LOCKABLE ABSL_ATTRIBUTE_WARN_UNUSED ExclusiveAccessGuard {
     cv_->SignalAll();
   }
 
-  void FinishBlockingPendingOperation() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+  void FinishBlockingPendingOperation() {
     --pending_operations_;
     cv_->SignalAll();
   }
@@ -163,8 +161,8 @@ class ABSL_LOCKABLE ABSL_ATTRIBUTE_WARN_UNUSED ExclusiveAccessGuard {
   friend class PreventExclusiveAccess;
 
   Mutex* absl_nonnull const mutex_;
-  CondVar* absl_nonnull const cv_ ABSL_GUARDED_BY(mutex_);
-  int pending_operations_ ABSL_GUARDED_BY(mutex_) = 0;
+  CondVar* absl_nonnull const cv_;
+  int pending_operations_ = 0;
 };
 
 class ABSL_SCOPED_LOCKABLE ABSL_ATTRIBUTE_WARN_UNUSED EnsureExclusiveAccess {
@@ -242,14 +240,34 @@ inline Case OnCancel() {
   return impl::OnCancel();
 }
 
-inline int Select(const CaseArray& cases) {
+inline int Select(const CaseArray& cases) noexcept {
   return impl::Select(cases);
 }
 
-inline int SelectUntil(const absl::Time deadline, const CaseArray& cases) {
+inline int SelectWithScopedUnlock(Mutex* absl_nonnull mu,
+                                  const CaseArray& cases)
+    ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu) {
+  mu->Unlock();
+  const int selected = concurrency::Select(cases);
+  mu->Lock();
+  return selected;
+}
+
+inline int SelectUntil(const absl::Time deadline,
+                       const CaseArray& cases) noexcept {
   return impl::SelectUntil(deadline, cases);
 }
 
+inline int SelectWithScopedUnlockUntil(Mutex* absl_nonnull mu,
+                                       const absl::Time deadline,
+                                       const CaseArray& cases)
+    ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu) {
+  mu->Unlock();
+  const int selected = concurrency::SelectUntil(deadline, cases);
+  mu->Lock();
+  return selected;
+}
+
 inline void Detach(std::unique_ptr<Fiber> fiber) {
   impl::Detach(std::move(fiber));
 }
diff --git a/src/eglt/nodes/async_node.h b/src/eglt/nodes/async_node.h
@@ -50,15 +50,11 @@ class AsyncNode {
   AsyncNode& operator=(AsyncNode& other) = delete;
   AsyncNode& operator=(AsyncNode&& other) noexcept;
 
-  ~AsyncNode() {
-    concurrency::MutexLock lock(&mutex_);
-    finalized_ = true;
-    concurrency::EnsureExclusiveAccess guard(&finalization_guard_);
-  }
+  ~AsyncNode() { concurrency::MutexLock lock(&mutex_); }
 
-  void BindPeers(absl::flat_hash_map<std::string_view,
-                                     std::shared_ptr<EvergreenWireStream>>
-                     peers) {
+  void BindPeers(
+      absl::flat_hash_map<std::string, std::shared_ptr<EvergreenWireStream>>
+          peers) {
     concurrency::MutexLock lock(&mutex_);
     peers_ = std::move(peers);
   }
@@ -163,13 +159,10 @@ class AsyncNode {
 
   mutable concurrency::Mutex mutex_;
   mutable concurrency::CondVar cv_ ABSL_GUARDED_BY(mutex_);
-  concurrency::ExclusiveAccessGuard finalization_guard_ ABSL_GUARDED_BY(mutex_){
-      &mutex_, &cv_};
-  bool finalized_ ABSL_GUARDED_BY(mutex_) = false;
   std::unique_ptr<ChunkStoreReader> default_reader_ ABSL_GUARDED_BY(mutex_);
   std::unique_ptr<ChunkStoreWriter> default_writer_ ABSL_GUARDED_BY(mutex_);
-  absl::flat_hash_map<std::string_view, std::shared_ptr<EvergreenWireStream>>
-      peers_ ABSL_GUARDED_BY(mutex_);
+  absl::flat_hash_map<std::string, std::shared_ptr<EvergreenWireStream>> peers_
+      ABSL_GUARDED_BY(mutex_);
 };
 
 template <>
@@ -187,6 +180,9 @@ inline auto AsyncNode::Put<Chunk>(Chunk value, int seq_id, bool final)
 template <>
 inline auto AsyncNode::Put(NodeFragment value, int seq_id, bool final)
     -> absl::Status {
+  if (seq_id == -1) {
+    seq_id = value.seq;
+  }
   return PutFragment(std::move(value), seq_id);
 }
 
diff --git a/src/eglt/sdk/serving/websockets.h b/src/eglt/sdk/serving/websockets.h
@@ -7,6 +7,7 @@
 #define BOOST_ASIO_NO_DEPRECATED
 
 #include <boost/asio/ip/tcp.hpp>
+#include <boost/asio/strand.hpp>
 #include <boost/asio/thread_pool.hpp>
 #include <boost/beast/core.hpp>
 #include <boost/beast/websocket.hpp>
@@ -87,41 +88,52 @@ class WebsocketEvergreenWireStream final : public EvergreenWireStream {
   }
 
   absl::Status Send(SessionMessage message) override {
-
+    auto message_bytes = cppack::Pack(std::move(message));
     concurrency::MutexLock lock(&mutex_);
     if (!status_.ok()) {
       return status_;
     }
 
+    if (concurrency::Cancelled()) {
+      return absl::CancelledError("Send cancelled");
+    }
+
+    while (send_pending_) {
+      cv_.Wait(&mutex_);
+    }
     send_pending_ = true;
-    mutex_.Unlock();
-    boost::system::error_code error;
     stream_.binary(true);
-    RunInAsioContext(
-        [this, &error, &message]() {
-          auto message_bytes = cppack::Pack(std::move(message));
-          stream_.write(asio::buffer(message_bytes), error);
-        },
-        {concurrency::OnCancel()});
+
+    boost::system::error_code error;
+    concurrency::PermanentEvent write_done;
+    stream_.async_write(asio::buffer(message_bytes),
+                        [&error, &write_done, this](
+                            const boost::system::error_code& ec, std::size_t) {
+                          concurrency::MutexLock lock(&mutex_);
+                          error = ec;
+                          write_done.Notify();
+                          send_pending_ = false;
+                          cv_.SignalAll();
+                        });
+    mutex_.Unlock();
+    concurrency::Select({write_done.OnEvent(), concurrency::OnCancel()});
     mutex_.Lock();
-    send_pending_ = false;
-    last_send_status_ = absl::OkStatus();
 
+    absl::Status status = absl::OkStatus();
     if (concurrency::Cancelled()) {
-      stream_.next_layer().shutdown(asio::socket_base::shutdown_send);
+      stream_.next_layer().shutdown(asio::socket_base::shutdown_send, error);
       DLOG(INFO) << absl::StrFormat("WESt %s Send cancelled", id_);
-      last_send_status_ = absl::CancelledError("Send cancelled");
+      status = absl::CancelledError("Send cancelled");
     }
 
+    // send_pending_ = false;
+    // cv_.SignalAll();
     if (error) {
-      last_send_status_ = absl::InternalError(error.message());
-      DLOG(INFO) << absl::StrFormat("WESt %s Send failed: %v", id_,
-                                    last_send_status_);
+      status = absl::InternalError(error.message());
+      DLOG(INFO) << absl::StrFormat("WESt %s Send failed: %v", id_, status);
     }
 
-    cv_.SignalAll();
-
-    return last_send_status_;
+    return status;
   }
 
   std::optional<SessionMessage> Receive() override {
@@ -135,6 +147,14 @@ class WebsocketEvergreenWireStream final : public EvergreenWireStream {
         return std::nullopt;
       }
 
+      if (concurrency::Cancelled()) {
+        return std::nullopt;
+      }
+
+      while (recv_pending_) {
+        cv_.Wait(&mutex_);
+      }
+
       recv_pending_ = true;
       mutex_.Unlock();
       boost::system::error_code error;
@@ -144,11 +164,23 @@ class WebsocketEvergreenWireStream final : public EvergreenWireStream {
             stream_.read(dynamic_buffer, error);
           },
           {concurrency::OnCancel()});
+
+      if (concurrency::Cancelled()) {
+        stream_.next_layer().shutdown(asio::socket_base::shutdown_receive,
+                                      error);
+        if (error && error != asio::error::not_connected) {
+          LOG(ERROR) << absl::StrFormat(
+              "WESt %s Cannot shut down receive on socket: %v", id_,
+              error.message());
+        }
+      }
+
       mutex_.Lock();
       recv_pending_ = false;
 
+      cv_.SignalAll();
+
       if (concurrency::Cancelled()) {
-        stream_.next_layer().shutdown(asio::socket_base::shutdown_receive);
         DLOG(INFO) << absl::StrFormat("WESt %s Receive cancelled", id_);
         return std::nullopt;
       }
@@ -158,8 +190,6 @@ class WebsocketEvergreenWireStream final : public EvergreenWireStream {
                                       error.message());
         return std::nullopt;
       }
-
-      cv_.SignalAll();
     }
 
     if (auto unpacked = cppack::Unpack<SessionMessage>(buffer); unpacked.ok()) {
@@ -185,6 +215,7 @@ class WebsocketEvergreenWireStream final : public EvergreenWireStream {
               beast::http::field::server,
               "Action Engine / Evergreen Light 0.1.0 WebsocketEvergreenServer");
         }));
+    stream_.write_buffer_bytes(16);
 
     boost::system::error_code error;
 
@@ -214,7 +245,7 @@ class WebsocketEvergreenWireStream final : public EvergreenWireStream {
       status_ = absl::InternalError(error.message());
     }
     status_ = absl::CancelledError("Cancelled");
-    while (send_pending_ || recv_pending_) {
+    while (send_pending_) {
       cv_.Wait(&mutex_);
     }
     stream_.next_layer().wait(tcp::socket::wait_error, error);
@@ -235,7 +266,7 @@ class WebsocketEvergreenWireStream final : public EvergreenWireStream {
     }
   }
 
-  absl::Status GetStatus() const override { return last_send_status_; }
+  absl::Status GetStatus() const override { return status_; }
 
   [[nodiscard]] std::string_view GetId() const override { return id_; }
 
@@ -255,7 +286,6 @@ class WebsocketEvergreenWireStream final : public EvergreenWireStream {
   std::string id_;
 
   absl::Status status_;
-  absl::Status last_send_status_;
 
   mutable concurrency::Mutex mutex_;
   mutable concurrency::CondVar cv_ ABSL_GUARDED_BY(mutex_);
@@ -317,7 +347,8 @@ class WebsocketEvergreenServer {
     concurrency::MutexLock lock(&mutex_);
     main_loop_ = concurrency::NewTree({}, [this]() {
       while (!concurrency::Cancelled()) {
-        tcp::socket socket{*GetDefaultAsioExecutionContext()};
+        tcp::socket socket{
+            asio::make_strand(*GetDefaultAsioExecutionContext())};
 
         DLOG(INFO) << "WES waiting for connection.";
         boost::system::error_code error;
@@ -499,6 +530,8 @@ MakeWebsocketEvergreenWireStream(std::string_view address = "127.0.0.1",
     return absl::CancelledError("Cancelled");
   }
 
+  ws_stream.write_buffer_bytes(16);
+
   if (prepare_stream) {
     if (auto status = std::move(prepare_stream)(&ws_stream); !status.ok()) {
       return status;
diff --git a/src/thread_on_boost/thread_on_boost/fiber.cc b/src/thread_on_boost/thread_on_boost/fiber.cc
@@ -155,6 +155,7 @@ void Fiber::Body() {
     // MarkFinished returns whether the fiber was detached when finished.
     // Detached fibers are self-joining.
     InternalJoin();
+    context_.detach();
     delete this;
   }
 }
@@ -222,9 +223,8 @@ void Fiber::Join() {
     CHECK(parent_ == current_fiber) << "Join() called from non-parent fiber";
   }
 
-  context_ = nullptr;
-
   InternalJoin();
+  context_->join();
 }
 
 // Update *this to a FINISHED state.  Preparing it to be Join()-ed (and

Original file line number	Diff line number	Diff line change
`@@ -155,6 +155,7 @@ void Fiber::Body() {`
`155`	`155`	`// MarkFinished returns whether the fiber was detached when finished.`
`156`	`156`	`// Detached fibers are self-joining.`
`157`	`157`	`InternalJoin();`
	`158`	`+ context_.detach();`
`158`	`159`	`delete this;`
`159`	`160`	`}`
`160`	`161`	`}`
`@@ -222,9 +223,8 @@ void Fiber::Join() {`
`222`	`223`	`CHECK(parent_ == current_fiber) << "Join() called from non-parent fiber";`
`223`	`224`	`}`
`224`	`225`
`225`		`- context_ = nullptr;`
`226`		`-`
`227`	`226`	`InternalJoin();`
	`227`	`+ context_->join();`
`228`	`228`	`}`
`229`	`229`
`230`	`230`	`// Update *this to a FINISHED state. Preparing it to be Join()-ed (and`