pw_rpc: Release acquired buffer on ServerWriter close

This updates the BaseServerWriter to release its payload buffer when it
is closed.

Change-Id: I8808d618bfc1930fc04f869a47a81a53de1d490e
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/23420
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Pigweed-Auto-Submit: Alexei Frolov <frolv@google.com>
Reviewed-by: Prashanth Swaminathan <prashanthsw@google.com>
Reviewed-by: Wyatt Hepler <hepler@google.com>

2 files changed, 27 insertions, 1 deletions

diff --git a/pw_rpc/base_server_writer.cc b/pw_rpc/base_server_writer.cc
index 8c57ee1ab..1d9ea4d06 100644
--- a/pw_rpc/base_server_writer.cc
+++ b/pw_rpc/base_server_writer.cc
@@ -50,6 +50,12 @@ void BaseServerWriter::Finish(Status status) {
     return;
   }
 
+  // If the ServerWriter implementer or user forgets to release an acquired
+  // buffer before finishing, release it here.
+  if (!response_.empty()) {
+    ReleasePayloadBuffer();
+  }
+
   Close();
 
   // Send a control packet indicating that the stream (and RPC) has terminated.
@@ -66,7 +72,11 @@ std::span<std::byte> BaseServerWriter::AcquirePayloadBuffer() {
     return {};
   }
 
-  response_ = call_.channel().AcquireBuffer();
+  // Only allow having one active buffer at a time.
+  if (response_.empty()) {
+    response_ = call_.channel().AcquireBuffer();
+  }
+
   return response_.payload(ResponsePacket());
 }
 
diff --git a/pw_rpc/base_server_writer_test.cc b/pw_rpc/base_server_writer_test.cc
index c090fb1d1..d5767df29 100644
--- a/pw_rpc/base_server_writer_test.cc
+++ b/pw_rpc/base_server_writer_test.cc
@@ -70,6 +70,9 @@ class FakeServerWriter : public BaseServerWriter {
                 std::min(buffer.size(), response.size()));
     return ReleasePayloadBuffer(buffer.first(response.size()));
   }
+
+  ByteSpan PayloadBuffer() { return AcquirePayloadBuffer(); }
+  const Channel::OutputBuffer& output_buffer() { return buffer(); }
 };
 
 TEST(ServerWriter, DefaultConstruct_Closed) {
@@ -131,6 +134,19 @@ TEST(ServerWriter, Close) {
   EXPECT_FALSE(writer.open());
 }
 
+TEST(ServerWriter, Close_ReleasesBuffer) {
+  ServerContextForTest<TestService> context(TestService::method.method());
+  FakeServerWriter writer(context.get());
+
+  ASSERT_TRUE(writer.open());
+  auto buffer = writer.PayloadBuffer();
+  buffer[0] = std::byte{0};
+  EXPECT_FALSE(writer.output_buffer().empty());
+  writer.Finish();
+  EXPECT_FALSE(writer.open());
+  EXPECT_TRUE(writer.output_buffer().empty());
+}
+
 TEST(ServerWriter, Open_SendsPacketWithPayload) {
   ServerContextForTest<TestService> context(TestService::method.method());
   FakeServerWriter writer(context.get());