pw_rpc: Guard the on_client_stream_end callback

Bug: b/234876851 Change-Id: Ia6588afe588923d64c8c4a27be4e761bed8a111b Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/126050 Pigweed-Auto-Submit: Wyatt Hepler <[email protected]> Reviewed-by: Alexei Frolov <[email protected]> Commit-Queue: Auto-Submit <[email protected]>
google · Jan 4, 2023 · db5bc5d · db5bc5d
1 parent ca6e845
commit db5bc5d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/pw_rpc/call_test.cc b/pw_rpc/call_test.cc
@@ -244,7 +244,7 @@ TEST_F(ServerReaderWriterTest, Move_MovesCallbacks) {
   reader_writer_.set_on_next([&calls](ConstByteSpan) { calls += 1; });
 
 #if PW_RPC_CLIENT_STREAM_END_CALLBACK
-  reader_writer.set_on_client_stream_end([&calls]() { calls += 1; });
+  reader_writer_.set_on_client_stream_end([&calls]() { calls += 1; });
 #endif  // PW_RPC_CLIENT_STREAM_END_CALLBACK
 
   FakeServerReaderWriter destination(std::move(reader_writer_));

diff --git a/pw_rpc/public/pw_rpc/internal/server_call.h b/pw_rpc/public/pw_rpc/internal/server_call.h
@@ -25,13 +25,15 @@ class ServerCall : public Call {
  public:
   void HandleClientStreamEnd() PW_UNLOCK_FUNCTION(rpc_lock()) {
     MarkClientStreamCompleted();
-    // TODO(b/234876851): Ensure on_client_stream_end_ is properly guarded.
-    rpc_lock().unlock();
 
 #if PW_RPC_CLIENT_STREAM_END_CALLBACK
-    if (on_client_stream_end_) {
-      on_client_stream_end_();
+    auto on_client_stream_end_local = std::move(on_client_stream_end_);
+    rpc_lock().unlock();
+    if (on_client_stream_end_local) {
+      on_client_stream_end_local();
     }
+#else
+    rpc_lock().unlock();
 #endif  // PW_RPC_CLIENT_STREAM_END_CALLBACK
   }
 
@@ -63,22 +65,23 @@ class ServerCall : public Call {
   // disabled with a helpful static_assert message.
   template <typename UnusedType = void>
   void set_on_client_stream_end(
-      [[maybe_unused]] Function<void()>&& on_client_stream_end) {
-    // TODO(b/234876851): Ensure on_client_stream_end_ is properly guarded.
+      [[maybe_unused]] Function<void()>&& on_client_stream_end)
+      PW_LOCKS_EXCLUDED(rpc_lock()) {
     static_assert(
         cfg::kClientStreamEndCallbackEnabled<UnusedType>,
         "The client stream end callback is disabled, so "
         "set_on_client_stream_end cannot be called. To enable the client end "
         "callback, set PW_RPC_CLIENT_STREAM_END_CALLBACK to 1.");
 #if PW_RPC_CLIENT_STREAM_END_CALLBACK
+    LockGuard lock(rpc_lock());
     on_client_stream_end_ = std::move(on_client_stream_end);
 #endif  // PW_RPC_CLIENT_STREAM_END_CALLBACK
   }
 
  private:
 #if PW_RPC_CLIENT_STREAM_END_CALLBACK
   // Called when a client stream completes.
-  Function<void()> on_client_stream_end_;
+  Function<void()> on_client_stream_end_ PW_GUARDED_BY(rpc_lock());
 #endif  // PW_RPC_CLIENT_STREAM_END_CALLBACK
 };