Reject interaction requests if interaction has been loadshedded

Summary: Reject requests if the interaction is already marked as loadshedded by the InteractionOverloadPolicy. Reviewed By: robertroeser, sazonovkirill Differential Revision: D66111259 fbshipit-source-id: b11123c04e6c70cfff4b7b6f20bc51a5ce5d1a13
facebook · Dec 19, 2024 · 957ea84 · 957ea84
1 parent 897ab79
commit 957ea84
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/third-party/thrift/src/thrift/lib/cpp2/transport/rocket/server/ThriftRocketServerHandler.cpp b/third-party/thrift/src/thrift/lib/cpp2/transport/rocket/server/ThriftRocketServerHandler.cpp
@@ -602,6 +602,22 @@ void ThriftRocketServerHandler::handleRequestCommon(
       ? (*metadata.frameworkMetadata())->clone()
       : nullptr;
 
+  if (interactionIdOpt &&
+      THRIFT_FLAG(enable_interaction_overload_protection_server)) {
+    if (auto* interaction =
+            apache::thrift::detail::Cpp2ConnContextInternalAPI(connContext_)
+                .findTile(*interactionIdOpt)) {
+      if (auto* overloadPolicy =
+              apache::thrift::detail::TileInternalAPI(*interaction)
+                  .getOverloadPolicy();
+          overloadPolicy && !overloadPolicy->allowNewRequest()) {
+        handleInteractionLoadshedded(makeActiveRequest(
+            std::move(metadata), std::move(debugPayload), std::move(reqCtx)));
+        return;
+      }
+    }
+  }
+
   // A request should not be active until the overload checking is done.
   auto request = makeRequest(
       std::move(metadata), std::move(debugPayload), std::move(reqCtx));
@@ -857,6 +873,19 @@ void ThriftRocketServerHandler::handleServerShutdown(
       kQueueOverloadedErrorCode);
 }
 
+void ThriftRocketServerHandler::handleInteractionLoadshedded(
+    ThriftRequestCoreUniquePtr request) {
+  if (auto* observer = serverConfigs_->getObserver()) {
+    observer->taskKilled();
+  }
+
+  request->sendErrorWrapped(
+      folly::make_exception_wrapper<TApplicationException>(
+          TApplicationException::LOADSHEDDING,
+          "Interaction already loadshedded"),
+      kInteractionLoadsheddedErrorCode);
+}
+
 void ThriftRocketServerHandler::handleInjectedFault(
     ThriftRequestCoreUniquePtr request, InjectedFault fault) {
   switch (fault) {

diff --git a/third-party/thrift/src/thrift/lib/cpp2/transport/rocket/server/ThriftRocketServerHandler.h b/third-party/thrift/src/thrift/lib/cpp2/transport/rocket/server/ThriftRocketServerHandler.h
@@ -160,6 +160,8 @@ class ThriftRocketServerHandler : public RocketServerHandler {
       ThriftRequestCoreUniquePtr request, std::string&& reason);
   FOLLY_NOINLINE void handleServerNotReady(ThriftRequestCoreUniquePtr request);
   FOLLY_NOINLINE void handleServerShutdown(ThriftRequestCoreUniquePtr request);
+  FOLLY_NOINLINE void handleInteractionLoadshedded(
+      ThriftRequestCoreUniquePtr request);
 
   enum class InjectedFault { ERROR, DROP, DISCONNECT };
   FOLLY_NOINLINE void handleInjectedFault(