Carry the per-connection and per-request internal fields

Summary: Gives `ThriftConnContext` and `ThriftRequestContext` the `getInternalFields<T>()` slot the classic contexts carry, over fast_thrift's own `TypeErasedValue<128>` and pluggable functions rather than a dependency on `Cpp2ConnContext`.

Reviewed By: robertroeser

Differential Revision: D116839752

fbshipit-source-id: e7f5da7ae639c2b7a3e20c8b275adc18e323820c

Author: Ankit Kumar

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

@@ -27,10 +27,22 @@
 #include <folly/SocketAddress.h>
 #include <folly/io/async/AsyncTransportCertificate.h>
 
+#include <folly/CppAttributes.h>
+
 #include <thrift/lib/cpp2/fast_thrift/rocket/common/TypeErasedPtr.h>
+#include <thrift/lib/cpp2/util/TypeErasedValue.h>
 
 namespace apache::thrift::fast_thrift::thrift {
 
+namespace detail {
+
+// The classic server's internal-fields slot, which fast_thrift points at
+// rather than owning: the object lives inline in whichever context owns it, so
+// two contexts can only reach one object by one of them holding a pointer.
+using InternalFieldsT = apache::thrift::util::TypeErasedValue<128>;
+
+} // namespace detail
+
 // Per-connection context. Lives for the duration of one accepted connection.
 //
 // Refcount is non-atomic (boost::thread_unsafe_counter), which avoids an
@@ -75,11 +87,41 @@ class ThriftConnContext : public boost::intrusive_ref_counter<
   }
 
   // Opaque per-connection slot. The deleter runs at connection close.
+  //
+  // There is one, so it has one owner. Anything shared between handlers, or
+  // published from a handler to the service, belongs in the connection's
+  // ExtensionStateStore instead — that is keyed by type and so cannot be
+  // claimed twice.
   void setUserData(rocket::TypeErasedPtr userData) noexcept {
     userData_ = std::move(userData);
   }
   void* getUserData() const noexcept { return userData_.get(); }
 
+  // Points this context at the per-connection fields the security layer keeps
+  // for this connection. Non-owning: the fields live on the Cpp2ConnContext
+  // built for the connection, which outlives this call, and nothing is copied
+  // — both contexts hand out the one object.
+  void setInternalFields(detail::InternalFieldsT* fields) noexcept {
+    internalFields_ = fields;
+  }
+
+  // The security layer's per-connection fields, or null while nothing holds
+  // any — a server with no security layer never fills the slot, unlike the
+  // classic server, where the fields always exist.
+  //
+  // Unchecked: only valid for the `T` the owner constructed.
+  template <class T>
+  T* FOLLY_NULLABLE getInternalFields() noexcept {
+    return hasInternalFields() ? &internalFields_->value_unchecked<T>()
+                               : nullptr;
+  }
+
+  template <class T>
+  const T* FOLLY_NULLABLE getInternalFields() const noexcept {
+    return hasInternalFields() ? &internalFields_->value_unchecked<T>()
+                               : nullptr;
+  }
+
   void setPeerAddress(folly::SocketAddress addr) noexcept {
     peerAddress_ = std::move(addr);
   }
@@ -92,10 +134,15 @@ class ThriftConnContext : public boost::intrusive_ref_counter<
   }
 
  private:
+  bool hasInternalFields() const noexcept {
+    return internalFields_ != nullptr && internalFields_->has_value();
+  }
+
   folly::SocketAddress peerAddress_{};
   std::string securityProtocol_;
   std::shared_ptr<const folly::AsyncTransportCertificate> peerCertificate_;
   rocket::TypeErasedPtr userData_{};
+  detail::InternalFieldsT* internalFields_{nullptr};
 };
 
 } // namespace apache::thrift::fast_thrift::thrift

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

@@ -24,6 +24,8 @@
 
 #include <folly/container/F14Map.h>
 
+#include <folly/CppAttributes.h>
+
 #include <thrift/lib/cpp2/fast_thrift/thrift/server/common/context/ThriftConnContext.h>
 #include <thrift/lib/thrift/gen-cpp2/RpcMetadata_types.h>
 
@@ -83,11 +85,39 @@ class ThriftRequestContext {
     return checksumAlgorithm_;
   }
 
+  // Points this context at the per-request fields the security layer keeps for
+  // this request. Non-owning, as on the connection context; the fields live on
+  // the Cpp2RequestContext built for the request, which outlives it.
+  void setInternalFields(detail::InternalFieldsT* fields) noexcept {
+    internalFields_ = fields;
+  }
+
+  // The security layer's per-request fields, or null while nothing holds any.
+  // The connection's own fields are on getConnectionContext().
+  //
+  // Unchecked: only valid for the `T` the owner constructed.
+  template <class T>
+  T* FOLLY_NULLABLE getInternalFields() noexcept {
+    return hasInternalFields() ? &internalFields_->value_unchecked<T>()
+                               : nullptr;
+  }
+
+  template <class T>
+  const T* FOLLY_NULLABLE getInternalFields() const noexcept {
+    return hasInternalFields() ? &internalFields_->value_unchecked<T>()
+                               : nullptr;
+  }
+
  private:
+  bool hasInternalFields() const noexcept {
+    return internalFields_ != nullptr && internalFields_->has_value();
+  }
+
   boost::intrusive_ptr<ThriftConnContext> connContext_;
   HeaderMap headers_;
   apache::thrift::ChecksumAlgorithm checksumAlgorithm_{
       apache::thrift::ChecksumAlgorithm::NONE};
+  detail::InternalFieldsT* internalFields_{nullptr};
 };
 
 } // namespace apache::thrift::fast_thrift::thrift

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

@@ -18,12 +18,20 @@
 
 #include <gtest/gtest.h>
 
+#include <utility>
+
 #include <folly/SocketAddress.h>
 
 #include <thrift/lib/cpp2/fast_thrift/rocket/common/TypeErasedPtr.h>
 
 namespace apache::thrift::fast_thrift::thrift {
 
+// Stands in for the security layer's per-connection fields. Owned outside the
+// context, as the real ones are — they live on the Cpp2ConnContext.
+struct TestConnFields {
+  int value{0};
+};
+
 TEST(ThriftConnContextTest, DefaultsAreEmpty) {
   boost::intrusive_ptr<ThriftConnContext> ctx{new ThriftConnContext()};
   EXPECT_TRUE(ctx->getPeerAddress().empty());
@@ -55,6 +63,53 @@ TEST(ThriftConnContextTest, IntrusivePtrSharesOwnership) {
   EXPECT_EQ(b->use_count(), 1);
 }
 
+// Nothing to read until something puts fields there: a server with no
+// security layer has none, and asking must not invent one. An installed but
+// empty slot reads the same way — there is no object to hand back.
+TEST(ThriftConnContextTest, InternalFieldsAreNullUntilFilled) {
+  detail::InternalFieldsT empty;
+  boost::intrusive_ptr<ThriftConnContext> ctx{new ThriftConnContext()};
+  EXPECT_EQ(ctx->getInternalFields<TestConnFields>(), nullptr);
+  EXPECT_EQ(std::as_const(*ctx).getInternalFields<TestConnFields>(), nullptr);
+
+  ctx->setInternalFields(&empty);
+  EXPECT_EQ(ctx->getInternalFields<TestConnFields>(), nullptr);
+  EXPECT_EQ(std::as_const(*ctx).getInternalFields<TestConnFields>(), nullptr);
+}
+
+// The point of the pointer: the context hands out the object its owner holds,
+// so a write through the owner is the write this context reads.
+TEST(ThriftConnContextTest, InternalFieldsAliasTheOwnersObject) {
+  detail::InternalFieldsT owned{std::in_place_type<TestConnFields>};
+  boost::intrusive_ptr<ThriftConnContext> ctx{new ThriftConnContext()};
+  ctx->setInternalFields(&owned);
+
+  owned.value<TestConnFields>().value = 7;
+  ASSERT_NE(ctx->getInternalFields<TestConnFields>(), nullptr);
+  EXPECT_EQ(ctx->getInternalFields<TestConnFields>()->value, 7);
+  EXPECT_EQ(
+      ctx->getInternalFields<TestConnFields>(), &owned.value<TestConnFields>());
+
+  ctx->getInternalFields<TestConnFields>()->value = 9;
+  EXPECT_EQ(owned.value<TestConnFields>().value, 9);
+}
+
+// Two connections point at their own owner's fields, never at each other's.
+TEST(ThriftConnContextTest, InternalFieldsAreNotSharedBetweenConnections) {
+  detail::InternalFieldsT firstFields{std::in_place_type<TestConnFields>};
+  detail::InternalFieldsT secondFields{std::in_place_type<TestConnFields>};
+  boost::intrusive_ptr<ThriftConnContext> first{new ThriftConnContext()};
+  boost::intrusive_ptr<ThriftConnContext> second{new ThriftConnContext()};
+  first->setInternalFields(&firstFields);
+  second->setInternalFields(&secondFields);
+
+  first->getInternalFields<TestConnFields>()->value = 1;
+  second->getInternalFields<TestConnFields>()->value = 2;
+
+  EXPECT_EQ(first->getInternalFields<TestConnFields>()->value, 1);
+  EXPECT_EQ(second->getInternalFields<TestConnFields>()->value, 2);
+}
+
 TEST(ThriftConnContextTest, IntrusivePtrDeletesOnLastReference) {
   // Probe destruction via the user-data deleter — it runs from
   // ~ThriftConnContext, which fires when the refcount drops to zero.

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

@@ -22,6 +22,10 @@
 
 namespace apache::thrift::fast_thrift::thrift {
 
+struct TestRequestFields {
+  int value{0};
+};
+
 TEST(ThriftRequestContextTest, DefaultConstructedHasNoConnContext) {
   ThriftRequestContext rc;
   EXPECT_EQ(rc.getConnectionContext(), nullptr);
@@ -37,6 +41,41 @@ TEST(ThriftRequestContextTest, SetConnectionContextStoresIt) {
   EXPECT_EQ(rc.getConnectionContext()->getSecurityProtocol(), "TLS1.3");
 }
 
+// An installed but empty slot reads the same as no slot at all — there is no
+// object to hand back.
+TEST(ThriftRequestContextTest, InternalFieldsAreNullUntilFilled) {
+  detail::InternalFieldsT empty;
+  ThriftRequestContext rc;
+  EXPECT_EQ(rc.getInternalFields<TestRequestFields>(), nullptr);
+
+  rc.setInternalFields(&empty);
+  EXPECT_EQ(rc.getInternalFields<TestRequestFields>(), nullptr);
+}
+
+// Request fields belong to their request: two requests on one connection point
+// at their own, and neither at the connection's.
+TEST(ThriftRequestContextTest, InternalFieldsArePerRequest) {
+  detail::InternalFieldsT firstFields{std::in_place_type<TestRequestFields>};
+  detail::InternalFieldsT secondFields{std::in_place_type<TestRequestFields>};
+  boost::intrusive_ptr<ThriftConnContext> conn{new ThriftConnContext()};
+
+  ThriftRequestContext first;
+  first.setConnectionContext(conn);
+  first.setInternalFields(&firstFields);
+  ThriftRequestContext second;
+  second.setConnectionContext(conn);
+  second.setInternalFields(&secondFields);
+
+  first.getInternalFields<TestRequestFields>()->value = 1;
+  second.getInternalFields<TestRequestFields>()->value = 2;
+
+  EXPECT_EQ(firstFields.value<TestRequestFields>().value, 1);
+  EXPECT_EQ(secondFields.value<TestRequestFields>().value, 2);
+  EXPECT_EQ(
+      first.getConnectionContext()->getInternalFields<TestRequestFields>(),
+      nullptr);
+}
+
 TEST(ThriftRequestContextTest, KeepsConnContextAliveAfterLocalReset) {
   boost::intrusive_ptr<ThriftConnContext> conn{new ThriftConnContext()};
   auto* raw = conn.get();