Add visibility rule for TPEH stream events listeners

Summary:
^
Insert another layer of blocker blocking outside users.

Reviewed By: sazonovkirill

Differential Revision: D75272079

fbshipit-source-id: 53d03fa48f892aa18f3f03c24fa0a20b7bce1a4a

Author: tlj77

third-party/thrift/src/thrift/lib/cpp/ContextStack.cpp

@@ -19,6 +19,7 @@
 
 #include <folly/tracing/StaticTracepoint.h>
 
+#include <thrift/lib/cpp/StreamEventHandler.h>
 #include <thrift/lib/cpp2/async/InterceptorFlags.h>
 #include <thrift/lib/cpp2/detail/EventHandlerRuntime.h>
 #include <thrift/lib/cpp2/server/Cpp2ConnContext.h>
@@ -369,7 +370,9 @@ void ContextStack::onStreamSubscribe() {
 
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onStreamSubscribe(contextAt(i));
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onStreamSubscribe(contextAt(i));
+      }
     }
   }
 }
@@ -382,7 +385,9 @@ void ContextStack::onStreamCredit(uint32_t credits) {
       methodNamePrefixed_);
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onStreamCredit(contextAt(i), credits);
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onStreamCredit(contextAt(i), credits);
+      }
     }
   }
 }
@@ -395,7 +400,9 @@ void ContextStack::onStreamNext() {
       methodNamePrefixed_);
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onStreamNext(contextAt(i));
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onStreamNext(contextAt(i));
+      }
     }
   }
 }
@@ -409,7 +416,9 @@ void ContextStack::onStreamPauseReceive() {
 
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onStreamPauseReceive(contextAt(i));
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onStreamPauseReceive(contextAt(i));
+      }
     }
   }
 }
@@ -423,7 +432,9 @@ void ContextStack::onStreamResumeReceive() {
 
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onStreamResumeReceive(contextAt(i));
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onStreamResumeReceive(contextAt(i));
+      }
     }
   }
 }
@@ -438,7 +449,9 @@ void ContextStack::handleStreamErrorWrapped(
 
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->handleStreamErrorWrapped(contextAt(i), ew);
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->handleStreamErrorWrapped(contextAt(i), ew);
+      }
     }
   }
 }
@@ -452,7 +465,9 @@ void ContextStack::onStreamFinally(details::STREAM_ENDING_TYPES endReason) {
 
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onStreamFinally(contextAt(i), endReason);
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onStreamFinally(contextAt(i), endReason);
+      }
     }
   }
 }
@@ -465,7 +480,9 @@ void ContextStack::onSinkSubscribe() {
       methodNamePrefixed_);
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onSinkSubscribe(contextAt(i));
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onSinkSubscribe(contextAt(i));
+      }
     }
   }
 }
@@ -478,7 +495,9 @@ void ContextStack::onSinkNext() {
       methodNamePrefixed_);
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onSinkNext(contextAt(i));
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onSinkNext(contextAt(i));
+      }
     }
   }
 }
@@ -491,7 +510,9 @@ void ContextStack::onSinkCancel() {
       methodNamePrefixed_);
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onSinkCancel(contextAt(i));
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onSinkCancel(contextAt(i));
+      }
     }
   }
 }
@@ -504,7 +525,9 @@ void ContextStack::onSinkCredit(uint32_t credits) {
       methodNamePrefixed_);
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onSinkCredit(contextAt(i), credits);
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onSinkCredit(contextAt(i), credits);
+      }
     }
   }
 }
@@ -518,7 +541,9 @@ void ContextStack::onSinkFinally(details::SINK_ENDING_TYPES endReason) {
 
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->onSinkFinally(contextAt(i), endReason);
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->onSinkFinally(contextAt(i), endReason);
+      }
     }
   }
 }
@@ -532,7 +557,9 @@ void ContextStack::handleSinkError(const folly::exception_wrapper& ew) {
 
   if (handlers_) {
     for (size_t i = 0; i < handlers_->size(); i++) {
-      (*handlers_)[i]->handleSinkError(contextAt(i), ew);
+      if (auto* streamEventHandler = (*handlers_)[i]->getStreamEventHandler()) {
+        streamEventHandler->handleSinkError(contextAt(i), ew);
+      }
     }
   }
 }

third-party/thrift/src/thrift/lib/cpp/StreamEventHandler.h

@@ -0,0 +1,57 @@
+/*
+ * 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 <folly/ExceptionWrapper.h>
+
+namespace apache::thrift {
+namespace details {
+enum class STREAM_ENDING_TYPES {
+  COMPLETE = 0,
+  ERROR = 1,
+  CANCEL = 2,
+};
+
+enum class SINK_ENDING_TYPES {
+  COMPLETE = 0,
+  COMPLETE_WITH_ERROR = 1,
+  ERROR = 2,
+};
+} // namespace details
+
+// EXPERIMENTAL: DO NOT USE WITHOUT TALKING TO THRIFT TEAM
+class StreamEventHandler {
+ public:
+  virtual ~StreamEventHandler() {}
+  virtual void onStreamSubscribe(void*) {}
+  virtual void onStreamNext(void*) {}
+  virtual void onStreamCredit(void*, uint32_t) {}
+  virtual void onStreamPauseReceive(void*) {}
+  virtual void onStreamResumeReceive(void*) {}
+  virtual void handleStreamErrorWrapped(
+      void*, const folly::exception_wrapper&) {}
+  virtual void onStreamFinally(void*, details::STREAM_ENDING_TYPES) {}
+
+  virtual void onSinkSubscribe(void*) {}
+  virtual void onSinkNext(void*) {}
+  virtual void onSinkCancel(void*) {}
+  virtual void onSinkCredit(void*, uint32_t) {}
+  virtual void handleSinkError(void*, const folly::exception_wrapper&) {}
+  virtual void onSinkFinally(void*, details::SINK_ENDING_TYPES) {}
+};
+
+} // namespace apache::thrift

third-party/thrift/src/thrift/lib/cpp/TProcessorEventHandler.h

@@ -22,23 +22,11 @@
 #include <folly/Demangle.h>
 #include <folly/ExceptionWrapper.h>
 #include <thrift/lib/cpp/SerializedMessage.h>
+#include <thrift/lib/cpp/StreamEventHandler.h>
 #include <thrift/lib/cpp/server/TConnectionContext.h>
 #include <thrift/lib/cpp/transport/THeader.h>
 
 namespace apache::thrift {
-namespace details {
-enum class STREAM_ENDING_TYPES {
-  COMPLETE = 0,
-  ERROR = 1,
-  CANCEL = 2,
-};
-
-enum class SINK_ENDING_TYPES {
-  COMPLETE = 0,
-  COMPLETE_WITH_ERROR = 1,
-  ERROR = 2,
-};
-} // namespace details
 
 using server::TConnectionContext;
 
@@ -197,28 +185,18 @@ class TProcessorEventHandler {
     return userException(ctx, fn_name, type, what);
   }
 
-  // Experimental: callbacks for stream events, please do not use without asking
-  // help from Thrift team.
-  virtual void onStreamSubscribe(void*) {}
-  virtual void onStreamNext(void*) {}
-  virtual void onStreamCredit(void*, uint32_t) {}
-  virtual void onStreamPauseReceive(void*) {}
-  virtual void onStreamResumeReceive(void*) {}
-  virtual void handleStreamErrorWrapped(
-      void*, const folly::exception_wrapper&) {}
-  virtual void onStreamFinally(void*, details::STREAM_ENDING_TYPES) {}
-
-  // Experimental: callbacks for sink events, please do not use without asking
-  // help from Thrift team.
-  virtual void onSinkSubscribe(void*) {}
-  virtual void onSinkNext(void*) {}
-  virtual void onSinkCancel(void*) {}
-  virtual void onSinkCredit(void*, uint32_t) {}
-  virtual void handleSinkError(void*, const folly::exception_wrapper&) {}
-  virtual void onSinkFinally(void*, details::SINK_ENDING_TYPES) {}
+  StreamEventHandler* getStreamEventHandler() {
+    return streamEventHandler_.get();
+  }
 
  protected:
-  TProcessorEventHandler() {}
+  TProcessorEventHandler() : streamEventHandler_(nullptr) {}
+  explicit TProcessorEventHandler(
+      std::unique_ptr<StreamEventHandler> streamEventHandler)
+      : streamEventHandler_(std::move(streamEventHandler)) {}
+
+ private:
+  std::unique_ptr<StreamEventHandler> streamEventHandler_;
 };
 
 /**