Expose InterceptorFrameworkMetadata to ServiceInterceptor

Summary: Allow ServiceInterceptor to access InterceptorFrameworkMetadata.

Reviewed By: praihan

Differential Revision: D70279672

fbshipit-source-id: 5479b57d49382eb2605718dd7345126b221b7e8e

third-party/thrift/src/thrift/lib/cpp2/async/AsyncProcessor.cpp

@@ -1167,7 +1167,8 @@ HandlerCallbackBase::processServiceInterceptorsOnRequest(
         arguments,
         serviceName_,
         definingServiceName_,
-        methodName_};
+        methodName_,
+        reqCtx_->getInterceptorFrameworkMetadata()};
     try {
       co_await serviceInterceptorsInfo[i].interceptor->internal_onRequest(
           std::move(connectionInfo), std::move(requestInfo));

third-party/thrift/src/thrift/lib/cpp2/async/InterceptorFrameworkMetadata.cpp

@@ -33,4 +33,11 @@ THRIFT_PLUGGABLE_FUNC_REGISTER(
   return nullptr;
 }
 
+THRIFT_PLUGGABLE_FUNC_REGISTER(
+    InterceptorFrameworkMetadataStorage,
+    deserializeFrameworkMetadata,
+    const folly::IOBuf& /* unused */) {
+  return InterceptorFrameworkMetadataStorage{};
+}
+
 } // namespace apache::thrift::detail

third-party/thrift/src/thrift/lib/cpp2/async/InterceptorFrameworkMetadata.h

@@ -46,6 +46,11 @@ THRIFT_PLUGGABLE_FUNC_DECLARE(
     serializeFrameworkMetadata,
     InterceptorFrameworkMetadataStorage&& storage);
 
+THRIFT_PLUGGABLE_FUNC_DECLARE(
+    InterceptorFrameworkMetadataStorage,
+    deserializeFrameworkMetadata,
+    const folly::IOBuf& buf);
+
 } // namespace detail
 
 } // namespace apache::thrift

third-party/thrift/src/thrift/lib/cpp2/server/Cpp2ConnContext.h

@@ -575,6 +575,7 @@ class Cpp2ConnContextInternalAPI {
  private:
   Cpp2ConnContext& connContext_;
 };
+class Cpp2RequestContextUnsafeAPI;
 } // namespace detail
 
 // Request-specific context
@@ -801,10 +802,20 @@ class Cpp2RequestContext : public apache::thrift::server::TConnectionContext {
     return &serviceInterceptorsStorage_.onRequest[index];
   }
 
+  const InterceptorFrameworkMetadataStorage* getInterceptorFrameworkMetadata() {
+    return &serviceInterceptorsStorage_.frameworkMetadata;
+  }
+
  protected:
   apache::thrift::server::TServerObserver::CallTimestamps timestamps_;
 
  private:
+  void initializeInterceptorFrameworkMetadata(
+      const folly::IOBuf& interceptorFrameworkMetadata) {
+    serviceInterceptorsStorage_.frameworkMetadata =
+        detail::deserializeFrameworkMetadata(interceptorFrameworkMetadata);
+  }
+
   Cpp2ConnContext* ctx_;
   folly::erased_unique_ptr requestData_{nullptr, nullptr};
   std::chrono::milliseconds requestTimeout_{0};
@@ -822,8 +833,30 @@ class Cpp2RequestContext : public apache::thrift::server::TConnectionContext {
       serviceInterceptorsStorage_;
   detail::RequestInternalFieldsT internalFields_;
   size_t wiredRequestBytes_{0};
+
+  friend class detail::Cpp2RequestContextUnsafeAPI;
 };
 
+namespace detail {
+class Cpp2RequestContextUnsafeAPI {
+ public:
+  explicit Cpp2RequestContextUnsafeAPI(Cpp2RequestContext& requestContext)
+      : requestContext_(requestContext) {}
+
+  void initializeInterceptorFrameworkMetadata(
+      const folly::IOBuf& interceptorFrameworkMetadata) {
+    if (!THRIFT_FLAG(enable_interceptor_framework_metadata)) {
+      return;
+    }
+    requestContext_.initializeInterceptorFrameworkMetadata(
+        interceptorFrameworkMetadata);
+  }
+
+ private:
+  Cpp2RequestContext& requestContext_;
+};
+} // namespace detail
+
 } // namespace apache::thrift
 
 #endif // #ifndef THRIFT_ASYNC_CPP2CONNCONTEXT_H_

third-party/thrift/src/thrift/lib/cpp2/server/ServiceInterceptorBase.h

@@ -75,6 +75,13 @@ class ServiceInterceptorBase {
      * in the format `{interaction_name}.{method_name}`.
      */
     const char* methodName = nullptr;
+    /**
+     * This is a pointer to the deserialized frameworkMetadata buffer sent as
+     * part of the request. InterceptorFrameworkMetadataStorage may be empty
+     * even if the pointer here is not null - it should always be checked with
+     * has_value()
+     */
+    const InterceptorFrameworkMetadataStorage* frameworkMetadata = nullptr;
   };
   virtual folly::coro::Task<void> internal_onRequest(
       ConnectionInfo, RequestInfo) = 0;

third-party/thrift/src/thrift/lib/cpp2/server/ServiceInterceptorStorage.h

@@ -17,9 +17,8 @@
 #pragma once
 
 #include <cstddef>
-#include <optional>
-#include <tuple>
 
+#include <thrift/lib/cpp2/async/InterceptorFrameworkMetadata.h>
 #include <thrift/lib/cpp2/util/AllocationColocator.h>
 #include <thrift/lib/cpp2/util/TypeErasedTupleRef.h>
 #include <thrift/lib/cpp2/util/TypeErasedValue.h>
@@ -33,6 +32,7 @@ struct ServiceInterceptorRequestStorageContext {
   std::size_t count = 0;
   util::AllocationColocator<>::ArrayPtr<ServiceInterceptorOnRequestStorage>
       onRequest = nullptr;
+  InterceptorFrameworkMetadataStorage frameworkMetadata = {};
 };
 
 using ServiceInterceptorOnConnectionStorage =

third-party/thrift/src/thrift/lib/cpp2/transport/core/ThriftRequest.cpp

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <thrift/lib/cpp2/async/InterceptorFrameworkMetadata.h>
 #include <thrift/lib/cpp2/server/CPUConcurrencyController.h>
 #include <thrift/lib/cpp2/server/LoggingEvent.h>
 #include <thrift/lib/cpp2/transport/core/ThriftRequest.h>
@@ -90,11 +91,17 @@ ThriftRequestCore::ThriftRequestCore(
     header_.setCallPriority(static_cast<concurrency::PRIORITY>(*priority));
   }
   auto otherMetadata = metadata.otherMetadata_ref();
+
   // When processing ThriftFrameworkMetadata, the header takes priority.
   if (!otherMetadata ||
       !detail::handleFrameworkMetadataHeader(*otherMetadata)) {
     if (auto frameworkMetadata = metadata.frameworkMetadata_ref()) {
       DCHECK(*frameworkMetadata && !(**frameworkMetadata).empty());
+      // This path handles setting InterceptorFrameworkMetata for
+      // ServiceInterceptors
+      detail::Cpp2RequestContextUnsafeAPI(reqContext_)
+          .initializeInterceptorFrameworkMetadata(**frameworkMetadata);
+      // Handling path using pluggable function
       detail::handleFrameworkMetadata(std::move(*frameworkMetadata));
     }
   }