Implement per connection storage for extensions via pipeline state

Summary: Extensions declaring `using ConnState = T` are handed the connection's single `T`, resolved once from a type-keyed store owned by the connection and threaded through the pipeline handler factories.

Reviewed By: robertroeser

Differential Revision: D116684106

fbshipit-source-id: c3a93bda4cf47db13697ea87b7db2d48d0523f96

Author: Ankit Kumar

third-party/thrift/src/thrift/lib/cpp2/fast_thrift/thrift/server/ThriftServerConnectionFactory.cpp

@@ -330,7 +330,7 @@ ThriftServerConnection ThriftServerConnectionFactory::buildConnectionImpl(
   // order — the first sits closest to the head, the last immediately above
   // the tail adapter. Each factory constructs a fresh per-connection instance.
   for (const auto& factory : config_.thriftPipelineHandlerFactories) {
-    thriftPipelineBuilder.addErasedHandler(factory());
+    thriftPipelineBuilder.addErasedHandler(factory(conn.extensionStates));
   }
   // Last before the tail, and deliberately after the embedder handlers: this
   // terminates the connection-lifecycle messages, so everything that might

third-party/thrift/src/thrift/lib/cpp2/fast_thrift/thrift/server/common/ExtensionStateStore.h

@@ -0,0 +1,95 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <memory>
+#include <typeindex>
+#include <typeinfo>
+#include <utility>
+
+#include <folly/container/F14Map.h>
+
+#include <thrift/lib/cpp2/fast_thrift/rocket/common/TypeErasedPtr.h>
+
+namespace apache::thrift::fast_thrift::thrift {
+
+/**
+ * Per-connection state shared between extensions, keyed by state type.
+ *
+ * One extension instance is constructed per connection, so an extension that
+ * needs nothing from its peers keeps its state as a member. This store exists
+ * for the other case: several cooperating extensions on the same connection
+ * that must read and write one object — an identity resolved by one and
+ * authorized against by the next.
+ *
+ * A slot is keyed by the state type itself, so two extensions reach the same
+ * object exactly when they name the same type, and unrelated extensions cannot
+ * collide however they are composed. The object is heap-owned, so a reference
+ * handed out at construction stays valid for the connection even as later
+ * slots are added.
+ *
+ * Not thread-safe, and needs not be: the store, the pipeline it belongs to, and
+ * every extension reading it live on one connection's EventBase.
+ */
+class ExtensionStateStore {
+ public:
+  ExtensionStateStore() = default;
+
+  ExtensionStateStore(const ExtensionStateStore&) = delete;
+  ExtensionStateStore& operator=(const ExtensionStateStore&) = delete;
+  ExtensionStateStore(ExtensionStateStore&&) noexcept = default;
+  ExtensionStateStore& operator=(ExtensionStateStore&&) noexcept = default;
+
+  /**
+   * The connection's `T`, constructing it from `args` on the first call for
+   * that type. Later calls return the same object and ignore `args`.
+   *
+   * The extension path never passes `args`: an extension's ConnState is
+   * default-constructed, because the extension that happens to be spliced in
+   * first is not a meaningful choice of constructor arguments. `args` is for
+   * direct users of the store, which today means tests.
+   */
+  template <typename T, typename... Args>
+  T& getOrCreate(Args&&... args) {
+    const std::type_index key{typeid(T)};
+    if (auto it = slots_.find(key); it != slots_.end()) {
+      // Sound because the slot is keyed by typeid(T): only a T was ever
+      // stored under this key.
+      return *static_cast<T*>(it->second.get());
+    }
+    auto owned = std::make_unique<T>(std::forward<Args>(args)...);
+    T* state = owned.get();
+    slots_.emplace(key, rocket::from_unique_ptr(std::move(owned)));
+    return *state;
+  }
+
+  /**
+   * Whether a slot for `T` has been created. For tests and diagnostics; the
+   * message path uses getOrCreate.
+   */
+  template <typename T>
+  bool contains() const noexcept {
+    return slots_.contains(std::type_index{typeid(T)});
+  }
+
+  std::size_t size() const noexcept { return slots_.size(); }
+
+ private:
+  folly::F14FastMap<std::type_index, rocket::TypeErasedPtr> slots_;
+};
+
+} // namespace apache::thrift::fast_thrift::thrift

third-party/thrift/src/thrift/lib/cpp2/fast_thrift/thrift/server/common/ThriftServerConnection.h

@@ -30,6 +30,7 @@
 #include <thrift/lib/cpp2/fast_thrift/thrift/server/adapter/ThriftServerAppAdapter.h>
 #include <thrift/lib/cpp2/fast_thrift/thrift/server/adapter/ThriftServerCompositeAppAdapter.h>
 #include <thrift/lib/cpp2/fast_thrift/thrift/server/adapter/ThriftServerTransportAdapter.h>
+#include <thrift/lib/cpp2/fast_thrift/thrift/server/common/ExtensionStateStore.h>
 #include <thrift/lib/cpp2/fast_thrift/thrift/server/common/context/ThriftConnContext.h>
 
 namespace apache::thrift::fast_thrift::thrift::server {
@@ -77,6 +78,12 @@ struct ThriftServerConnection {
   // Buffer allocator used by the thrift pipeline.
   channel_pipeline::SimpleBufferAllocator thriftAllocator;
 
+  // Per-connection state shared between this connection's extensions. Handlers
+  // hold references into it, so it must outlive thriftPipeline — as a value
+  // member it is destroyed after the dtor body, which is where the pipeline
+  // goes.
+  ExtensionStateStore extensionStates;
+
   // Head of the thrift pipeline. Owns the rocket connection (transport
   // handler, app adapter, rocket pipeline) via its rocketConnection()
   // member.

third-party/thrift/src/thrift/lib/cpp2/fast_thrift/thrift/server/common/test/ExtensionStateStoreTest.cpp

@@ -0,0 +1,136 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Unit test for ExtensionStateStore — the per-connection slot table that lets
+// cooperating extensions share one object. What it must guarantee: the same
+// type reaches the same object, different types never alias, and a reference
+// handed out early survives later slots being added.
+
+#include <cstddef>
+#include <string>
+#include <utility>
+
+#include <gtest/gtest.h>
+
+#include <thrift/lib/cpp2/fast_thrift/thrift/server/common/ExtensionStateStore.h>
+
+namespace apache::thrift::fast_thrift::thrift {
+
+namespace {
+
+struct Identity {
+  std::string name;
+};
+
+struct Counters {
+  int requests{0};
+};
+
+struct Tracked {
+  explicit Tracked(int& live) : live_(live) { ++live_; }
+  ~Tracked() { --live_; }
+  Tracked(const Tracked&) = delete;
+  Tracked& operator=(const Tracked&) = delete;
+
+  int& live_;
+};
+
+// A family of distinct state types, so a test can fill the store with slots
+// that actually grow it — resolving one type repeatedly would reuse the single
+// slot it already has.
+template <size_t N>
+struct Slot {
+  size_t value{N};
+};
+
+template <size_t... Ns>
+void createSlots(ExtensionStateStore& store, std::index_sequence<Ns...>) {
+  (store.getOrCreate<Slot<Ns>>(), ...);
+}
+
+} // namespace
+
+// The property the whole mechanism rests on: two lookups of the same type are
+// the same object, so one extension's write is another's read.
+TEST(ExtensionStateStoreTest, SameTypeResolvesToTheSameObject) {
+  ExtensionStateStore store;
+
+  store.getOrCreate<Identity>().name = "svc:caller";
+
+  EXPECT_EQ(store.getOrCreate<Identity>().name, "svc:caller");
+  EXPECT_EQ(&store.getOrCreate<Identity>(), &store.getOrCreate<Identity>());
+}
+
+// Distinct types are distinct slots — the collision the type-keyed store
+// exists to make impossible.
+TEST(ExtensionStateStoreTest, DistinctTypesGetDistinctSlots) {
+  ExtensionStateStore store;
+
+  store.getOrCreate<Identity>().name = "svc:caller";
+  store.getOrCreate<Counters>().requests = 3;
+
+  EXPECT_EQ(store.getOrCreate<Identity>().name, "svc:caller");
+  EXPECT_EQ(store.getOrCreate<Counters>().requests, 3);
+  EXPECT_EQ(store.size(), 2);
+}
+
+// Extensions take their reference at construction and keep it for the
+// connection, so a slot added later must not move an earlier one. The slots
+// added here are all distinct types, which is what makes the table grow and
+// rehash — repeating one type would reuse its existing slot and prove nothing.
+TEST(ExtensionStateStoreTest, EarlierReferencesSurviveLaterSlots) {
+  ExtensionStateStore store;
+
+  Identity& identity = store.getOrCreate<Identity>();
+  identity.name = "svc:caller";
+  constexpr size_t kLaterSlots = 64;
+  createSlots(store, std::make_index_sequence<kLaterSlots>{});
+
+  EXPECT_EQ(store.size(), kLaterSlots + 1);
+  EXPECT_EQ(identity.name, "svc:caller");
+  EXPECT_EQ(&identity, &store.getOrCreate<Identity>());
+}
+
+// Only the first call constructs: a second extension naming the type joins the
+// existing object rather than resetting what the first one wrote.
+TEST(ExtensionStateStoreTest, ArgumentsApplyOnlyToTheFirstCall) {
+  ExtensionStateStore store;
+
+  EXPECT_EQ(store.getOrCreate<Identity>(Identity{"first"}).name, "first");
+  EXPECT_EQ(store.getOrCreate<Identity>(Identity{"second"}).name, "first");
+}
+
+TEST(ExtensionStateStoreTest, ContainsReportsOnlyCreatedSlots) {
+  ExtensionStateStore store;
+
+  EXPECT_FALSE(store.contains<Identity>());
+  store.getOrCreate<Identity>();
+  EXPECT_TRUE(store.contains<Identity>());
+  EXPECT_FALSE(store.contains<Counters>());
+}
+
+// The store owns its slots, so a connection's state goes away with it.
+TEST(ExtensionStateStoreTest, SlotsAreDestroyedWithTheStore) {
+  int live = 0;
+  {
+    ExtensionStateStore store;
+    store.getOrCreate<Tracked>(live);
+    EXPECT_EQ(live, 1);
+  }
+  EXPECT_EQ(live, 0);
+}
+
+} // namespace apache::thrift::fast_thrift::thrift

third-party/thrift/src/thrift/lib/cpp2/fast_thrift/thrift/server/extension/ThriftExtension.h

@@ -307,4 +307,22 @@ template <typename H>
 concept HasResponseCallback =
     HasResponseViewCallback<H> || HasResponseMutatorCallback<H>;
 
+/**
+ * True iff H shares per-connection state with its peer extensions, by defining
+ *
+ *   using ConnState = SomeType;
+ *
+ * The adapter then resolves the connection's `ConnState` from the
+ * ExtensionStateStore and passes it as H's first constructor argument, ahead of
+ * the arguments given to addThriftExtension. Extensions naming the same type
+ * share one object per connection; an extension that declares none is
+ * constructed from its own arguments alone.
+ *
+ * The state is default-constructed — whichever extension names it first brings
+ * it into being, so no one extension's arguments could construct it. A
+ * ConnState that is not default-constructible is a compile error.
+ */
+template <typename H>
+concept HasConnState = requires { typename H::ConnState; };
+
 } // namespace apache::thrift::fast_thrift::thrift

third-party/thrift/src/thrift/lib/cpp2/fast_thrift/thrift/server/extension/ThriftExtensionPipelineHandler.h

@@ -115,7 +115,9 @@ const std::string& extensionName() {
  *
  * Lifetime: one instance per connection (constructed by the pipeline handler
  * factory); H is owned by value and its constructor arguments are forwarded
- * from addThriftExtension.
+ * from addThriftExtension. An H declaring a ConnState additionally receives
+ * that connection's shared state as its first constructor argument, resolved
+ * once here so the extension never handles the store itself.
  */
 template <typename H>
 class ThriftExtensionPipelineHandler {
@@ -132,8 +134,9 @@ class ThriftExtensionPipelineHandler {
 
  public:
   template <typename... Args>
-  explicit ThriftExtensionPipelineHandler(Args&&... args)
-      : handler_(std::forward<Args>(args)...) {}
+  explicit ThriftExtensionPipelineHandler(
+      ExtensionStateStore& store, Args&&... args)
+      : handler_(makeHandler(store, std::forward<Args>(args)...)) {}
 
   channel_pipeline::Result onRead(
       ThriftPipelineHandlerContext& ctx,
@@ -353,6 +356,25 @@ class ThriftExtensionPipelineHandler {
   void handlerRemoved(ThriftPipelineHandlerContext&) noexcept {}
 
  private:
+  // Returned as a prvalue so H is built directly into handler_ — H needs no
+  // move constructor, and its shared state is resolved exactly once per
+  // connection, before any callback can run.
+  template <typename... Args>
+  static H makeHandler(ExtensionStateStore& store, Args&&... args) {
+    if constexpr (HasConnState<H>) {
+      static_assert(
+          std::is_default_constructible_v<typename H::ConnState>,
+          "H::ConnState must be default-constructible: the connection's state "
+          "is created by whichever extension names it first, so there is no "
+          "one extension whose arguments could construct it.");
+      return H(
+          store.getOrCreate<typename H::ConnState>(),
+          std::forward<Args>(args)...);
+    } else {
+      return H(std::forward<Args>(args)...);
+    }
+  }
+
   H handler_;
   // Latched from the setup messages so the payload-less ConnectionClosed can
   // still present the connection. Non-owning: the connection-context handler
@@ -363,4 +385,35 @@ class ThriftExtensionPipelineHandler {
   bool established_{false};
 };
 
+/**
+ * Build a factory that splices extension H into a connection's thrift pipeline.
+ *
+ * The peer of makeThriftPipelineHandlerFactory for the extension API: it wraps
+ * H in the adapter and threads the connection's ExtensionStateStore into it, so
+ * an H declaring a ConnState is handed that state at construction. `args` are
+ * copied and forwarded to every per-connection instance, after the state.
+ */
+template <typename H, typename... Args>
+ThriftPipelineHandlerFactory makeThriftExtensionHandlerFactory(
+    channel_pipeline::HandlerId id, Args... args) {
+  using Adapter = ThriftExtensionPipelineHandler<H>;
+  // The adapter is framework-owned, so this can only fire if the adapter itself
+  // stops being a pipeline handler — the same named diagnostic the native path
+  // gets from makeThriftPipelineHandlerFactory, rather than a failure deep
+  // inside makeHandlerNode.
+  static_assert(
+      channel_pipeline::InboundHandler<Adapter, ThriftPipelineHandlerContext> ||
+          channel_pipeline::
+              OutboundHandler<Adapter, ThriftPipelineHandlerContext> ||
+          channel_pipeline::
+              DuplexHandler<Adapter, ThriftPipelineHandlerContext>,
+      "ThriftExtensionPipelineHandler<H> must satisfy the Inbound, Outbound, or "
+      "Duplex handler concept over ThriftPipelineHandlerContext");
+  return [id, args...](ExtensionStateStore& store) {
+    return channel_pipeline::detail::
+        makeHandlerNode<Adapter, ThriftServerEventType>(
+            id, std::make_unique<Adapter>(store, args...));
+  };
+}
+
 } // namespace apache::thrift::fast_thrift::thrift::server