From 9b80a9503be190421837516c9801a7878451cab6 Mon Sep 17 00:00:00 2001
From: Craig Tiller <ctiller@google.com>
Date: Fri, 2 Feb 2024 18:16:00 +0000
Subject: [PATCH] OnServerTrailingMetadata

---
 src/core/lib/transport/call_filters.h    | 51 ++++++++++++++++++++++++
 test/core/transport/call_filters_test.cc | 25 ++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/src/core/lib/transport/call_filters.h b/src/core/lib/transport/call_filters.h
index 6366b248c46..9c1b1a1069f 100644
--- a/src/core/lib/transport/call_filters.h
+++ b/src/core/lib/transport/call_filters.h
@@ -1137,6 +1137,42 @@ class PipeState {
   bool started_ = false;
 };
 
+template <typename Fn>
+class ServerTrailingMetadataInterceptor {
+ public:
+  class Call {
+   public:
+    static const NoInterceptor OnClientInitialMetadata;
+    static const NoInterceptor OnServerInitialMetadata;
+    static const NoInterceptor OnClientToServerMessage;
+    static const NoInterceptor OnServerToClientMessage;
+    static const NoInterceptor OnFinalize;
+    void OnServerTrailingMetadata(ServerMetadata& md,
+                                  ServerTrailingMetadataInterceptor* filter) {
+      filter->fn_(md);
+    }
+  };
+
+  explicit ServerTrailingMetadataInterceptor(Fn fn) : fn_(std::move(fn)) {}
+
+ private:
+  GPR_NO_UNIQUE_ADDRESS Fn fn_;
+};
+template <typename Fn>
+const NoInterceptor
+    ServerTrailingMetadataInterceptor<Fn>::Call::OnClientInitialMetadata;
+template <typename Fn>
+const NoInterceptor
+    ServerTrailingMetadataInterceptor<Fn>::Call::OnServerInitialMetadata;
+template <typename Fn>
+const NoInterceptor
+    ServerTrailingMetadataInterceptor<Fn>::Call::OnClientToServerMessage;
+template <typename Fn>
+const NoInterceptor
+    ServerTrailingMetadataInterceptor<Fn>::Call::OnServerToClientMessage;
+template <typename Fn>
+const NoInterceptor ServerTrailingMetadataInterceptor<Fn>::Call::OnFinalize;
+
 }  // namespace filters_detail
 
 // Execution environment for a stack of filters.
@@ -1144,6 +1180,7 @@ class PipeState {
 class CallFilters {
  public:
   class StackBuilder;
+  class StackTestSpouse;
 
   // A stack is an opaque, immutable type that contains the data necessary to
   // execute a call through a given set of filters.
@@ -1157,6 +1194,7 @@ class CallFilters {
    private:
     friend class CallFilters;
     friend class StackBuilder;
+    friend class StackTestSpouse;
     explicit Stack(filters_detail::StackData data) : data_(std::move(data)) {}
     const filters_detail::StackData data_;
   };
@@ -1187,6 +1225,19 @@ class CallFilters {
       AddOwnedObject([](void* p) { static_cast<T*>(p)->Unref(); }, p.release());
     }
 
+    template <typename T>
+    void AddOwnedObject(std::unique_ptr<T> p) {
+      AddOwnedObject([](void* p) { delete static_cast<T*>(p); }, p.release());
+    }
+
+    template <typename Fn>
+    void AddOnServerTrailingMetadata(Fn fn) {
+      auto filter = std::make_unique<
+          filters_detail::ServerTrailingMetadataInterceptor<Fn>>(std::move(fn));
+      Add(filter.get());
+      AddOwnedObject(std::move(filter));
+    }
+
     RefCountedPtr<Stack> Build();
 
    private:
diff --git a/test/core/transport/call_filters_test.cc b/test/core/transport/call_filters_test.cc
index 0a88a9e242f..08df49c0057 100644
--- a/test/core/transport/call_filters_test.cc
+++ b/test/core/transport/call_filters_test.cc
@@ -997,6 +997,31 @@ TEST(StackDataTest,
 
 }  // namespace filters_detail
 
+///////////////////////////////////////////////////////////////////////////////
+// StackBuilder
+
+class CallFilters::StackTestSpouse {
+ public:
+  static const filters_detail::StackData& StackDataFrom(const Stack& stack) {
+    return stack.data_;
+  }
+};
+
+TEST(StackBuilderTest, AddOnServerTrailingMetadata) {
+  CallFilters::StackBuilder b;
+  b.AddOnServerTrailingMetadata(
+      [x = std::make_unique<int>(42)](ServerMetadata&) { EXPECT_EQ(*x, 42); });
+  auto stack = b.Build();
+  const auto& data = CallFilters::StackTestSpouse().StackDataFrom(*stack);
+  ASSERT_EQ(data.server_trailing_metadata.ops.size(), 1u);
+  ASSERT_EQ(data.client_initial_metadata.ops.size(), 0u);
+  ASSERT_EQ(data.client_to_server_messages.ops.size(), 0u);
+  ASSERT_EQ(data.server_to_client_messages.ops.size(), 0u);
+  ASSERT_EQ(data.server_initial_metadata.ops.size(), 0u);
+  EXPECT_EQ(data.server_trailing_metadata.ops[0].call_offset, 0);
+  EXPECT_NE(data.server_trailing_metadata.ops[0].channel_data, nullptr);
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // OperationExecutor