Add ConsumerCallback to ServerSinkBridge

Reviewed By: evanjzou

Differential Revision: D89158038

fbshipit-source-id: 26a968bc5d8809c6391c0cb8db5269f59c1f8213

Author: Dan Deng

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

@@ -79,6 +79,8 @@ class ServerSinkBridge : public TwoWayBridge<
   // start() should be called on the CPU thread
   folly::coro::Task<void> start();
 
+  uint64_t getBufferSize() const { return consumer_.bufferSize; }
+
   //
   // SinkServerCallback method
   //
@@ -103,6 +105,11 @@ class ServerSinkBridge : public TwoWayBridge<
   // end of TwoWayBridge methods
   //
 
+  using TwoWayBridge::serverClose;
+  using TwoWayBridge::serverGetMessages;
+  using TwoWayBridge::serverPush;
+  using TwoWayBridge::serverWait;
+
  private:
   ServerSinkBridge(
       SinkConsumerImpl&& sinkConsumer,

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

@@ -92,6 +92,38 @@ void ServerSinkFactory::setContextStack(ContextStack::UniquePtr contextStack) {
   contextStack_ = std::move(contextStack);
 }
 
+ServerSinkFactory::ServerSinkFactory(
+    ConsumerCallback* consumerCallback,
+    uint64_t bufferSize,
+    std::chrono::milliseconds chunkTimeout)
+    : bufferSize_{bufferSize}, chunkTimeout_{chunkTimeout} {
+  startFunction_ = [consumerCallback](
+                       uint64_t bufferSize,
+                       std::chrono::milliseconds chunkTimeout,
+                       folly::EventBase* evb,
+                       TilePtr&& interaction,
+                       ContextStack::UniquePtr contextStack,
+                       FirstResponsePayload&& firstResponsePayload,
+                       SinkClientCallback* clientCallback) {
+    DCHECK(evb->isInEventBaseThread());
+    SinkConsumerImpl sinkConsumer;
+    sinkConsumer.bufferSize = bufferSize;
+    sinkConsumer.chunkTimeout = chunkTimeout;
+    sinkConsumer.interaction = std::move(interaction);
+    sinkConsumer.contextStack = std::move(contextStack);
+
+    auto sink =
+        new ServerSinkBridge(std::move(sinkConsumer), *evb, clientCallback);
+    auto sinkPtr = sink->copy();
+
+    std::ignore = clientCallback->onFirstResponse(
+        std::move(firstResponsePayload), evb, sink);
+
+    sink->serverPush(bufferSize);
+    consumerCallback->provideSink(std::move(sinkPtr));
+  };
+}
+
 } // namespace apache::thrift::detail
 
 #endif

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

@@ -36,6 +36,13 @@ class ServerSinkFactory {
       FirstResponsePayload&&,
       SinkClientCallback*)>;
 
+  class ConsumerCallback {
+   public:
+    virtual void provideSink(ServerSinkBridge::Ptr sink) = 0;
+    ConsumerCallback() = default;
+    virtual ~ConsumerCallback() = default;
+  };
+
   explicit ServerSinkFactory();
 
   explicit ServerSinkFactory(StartFunction&& fn);
@@ -46,6 +53,11 @@ class ServerSinkFactory {
       uint64_t bufferSize,
       std::chrono::milliseconds timeout);
 
+  explicit ServerSinkFactory(
+      ConsumerCallback* consumerCallback,
+      uint64_t bufferSize,
+      std::chrono::milliseconds timeout);
+
   void start(
       FirstResponsePayload&& payload,
       SinkClientCallback* callback,