BiDi works up until transport

Summary: Get bidi to send the response from transport, granted that part is incomplete

Reviewed By: sazonovkirill

Differential Revision: D83386024

fbshipit-source-id: f3de4fca723d2adf91b6ea03637f85c0bf217cd9

Author: evanjzou

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

@@ -142,12 +142,8 @@ class BiDiStreamReplyInfo {
         crc32c_{std::move(crc32c)} {}
 
   void operator()() noexcept {
-    // TODO(ezou) stubbed
-    req_->sendErrorWrapped(
-        folly::make_exception_wrapper<TApplicationException>(
-            TApplicationException::INTERNAL_ERROR,
-            "TODO(ezou) BiDiStreamReplyInfo not implemented"),
-        kUnknownErrorCode);
+    req_->sendBiDiReply(
+        std::move(response_), std::move(serverBiDiStreamFactory_), crc32c_);
   }
 
  private:

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

@@ -26,6 +26,7 @@
 #include <thrift/lib/cpp/Thrift.h>
 #include <thrift/lib/cpp/server/TServerObserver.h>
 #include <thrift/lib/cpp2/async/MessageChannel.h>
+#include <thrift/lib/cpp2/async/ServerBiDiStreamFactory.h>
 #include <thrift/lib/cpp2/async/ServerStream.h>
 #include <thrift/lib/cpp2/async/Sink.h>
 #include <thrift/lib/cpp2/async/StreamCallbacks.h>
@@ -147,6 +148,13 @@ class ResponseChannelRequest {
   }
 #endif
 
+  virtual void sendBiDiReply(
+      ResponsePayload&&,
+      detail::ServerBiDiStreamFactory&&,
+      folly::Optional<uint32_t> = folly::none) {
+    throw std::logic_error("unimplemented");
+  }
+
   FOLLY_NODISCARD static bool sendBiDiReply(
       ResponseChannelRequest::UniquePtr request,
       folly::EventBase* eb,

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

@@ -27,6 +27,11 @@ void ServerBiDiStreamFactory::setInteraction(TilePtr&& interaction) {
   interaction_ = std::move(interaction);
 }
 
+bool ServerBiDiStreamFactory::valid() const {
+  // TODO(ezou) this is a stub, revisit this.
+  return startFunction_ != nullptr;
+}
+
 void ServerBiDiStreamFactory::start(
     FirstResponsePayload&& payload,
     BiDiClientCallback* clientCallback,

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

@@ -50,16 +50,21 @@ class ServerBiDiStreamFactory {
                          std::shared_ptr<ContextStack>,
                          TilePtr&&,
                          FirstResponsePayload&&,
-                         BiDiClientCallback*,
+                         BiDiClientCallback* clientCallback,
                          folly::EventBase*) mutable -> void {
-      LOG(FATAL) << "Not implemented";
+      clientCallback->onFirstResponseError(
+          folly::make_exception_wrapper<rocket::RocketException>(
+              rocket::ErrorCode::APPLICATION_ERROR,
+              "TODO(ezou) in Transport not implemented"));
     };
   }
 
   void setContextStack(std::shared_ptr<ContextStack> contextStack);
 
   void setInteraction(TilePtr&& interaction);
 
+  bool valid() const;
+
   void start(
       FirstResponsePayload&& firstResponsePayload,
       BiDiClientCallback* clientCallback,

third-party/thrift/src/thrift/lib/cpp2/async/tests/BiDiServiceE2ETest.cpp

@@ -79,7 +79,7 @@ CO_TEST_F(BiDiServiceTest, BiDiNoResponse) {
     co_await client->co_echo();
     CO_FAIL() << "No error - error expected";
   } catch (TApplicationException e) {
-    EXPECT_EQ(e.getMessage(), "TODO(ezou) BiDiStreamReplyInfo not implemented");
+    EXPECT_EQ(e.getMessage(), "Unexpected error frame type: 513");
   }
 }
 
@@ -89,7 +89,7 @@ CO_TEST_F(BiDiServiceTest, BiDiWithResponse) {
     co_await client->co_echoWithResponse("Test");
     CO_FAIL() << "No error - error expected";
   } catch (TApplicationException e) {
-    EXPECT_EQ(e.getMessage(), "TODO(ezou) BiDiStreamReplyInfo not implemented");
+    EXPECT_EQ(e.getMessage(), "Unexpected error frame type: 513");
   }
 }
 

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

@@ -250,6 +250,30 @@ class ThriftRequestCore : public ResponseChannelRequest {
   }
 #endif
 
+  void sendBiDiReply(
+      ResponsePayload&& response,
+      detail::ServerBiDiStreamFactory&& bidiStreamFactory,
+      folly::Optional<uint32_t> crc32c) final {
+    if (tryCancel()) {
+      cancelTimeout();
+      auto metadata = makeResponseRpcMetadata(
+          header_.extractAllWriteHeaders(),
+          header_.extractProxiedPayloadMetadata(),
+          header_.getChecksum());
+      if (crc32c) {
+        metadata.crc32c() = *crc32c;
+      }
+      sendReplyInternal(
+          std::move(metadata),
+          std::move(response).buffer(),
+          std::move(bidiStreamFactory));
+
+      if (auto* observer = serverConfigs_.getObserver()) {
+        observer->sentReply();
+      }
+    }
+  }
+
   bool sendBiDiReply(
       ResponsePayload&& response,
       BiDiServerCallbackPtr callback,
@@ -393,6 +417,13 @@ class ThriftRequestCore : public ResponseChannelRequest {
   }
 #endif
 
+  virtual void sendBiDiThriftResponse(
+      ResponseRpcMetadata&&,
+      std::unique_ptr<folly::IOBuf>,
+      detail::ServerBiDiStreamFactory&&) noexcept {
+    LOG(FATAL) << "sendBiDiThriftResponse not implemented";
+  }
+
   virtual bool sendBiDiThriftResponse(
       ResponseRpcMetadata&&,
       std::unique_ptr<folly::IOBuf>,
@@ -508,6 +539,18 @@ class ThriftRequestCore : public ResponseChannelRequest {
   }
 #endif
 
+  void sendReplyInternal(
+      ResponseRpcMetadata&& metadata,
+      std::unique_ptr<folly::IOBuf> buf,
+      detail::ServerBiDiStreamFactory&& bidiStreamFactory) {
+    if (checkResponseSize(*buf)) {
+      sendBiDiThriftResponse(
+          std::move(metadata), std::move(buf), std::move(bidiStreamFactory));
+    } else {
+      sendResponseTooBigEx();
+    }
+  }
+
   bool sendReplyInternal(
       ResponseRpcMetadata&& metadata,
       std::unique_ptr<folly::IOBuf> buf,

third-party/thrift/src/thrift/lib/cpp2/transport/rocket/server/RocketThriftRequests.cpp

@@ -881,6 +881,36 @@ void ThriftServerRequestBiDi::sendSerializedError(
           FirstResponsePayload(std::move(exbuf), std::move(metadata))));
 }
 
+void ThriftServerRequestBiDi::sendBiDiThriftResponse(
+    ResponseRpcMetadata&& metadata,
+    std::unique_ptr<folly::IOBuf> data,
+    apache::thrift::detail::ServerBiDiStreamFactory&&
+        bidiStreamFactory) noexcept {
+  if (!bidiStreamFactory.valid()) {
+    sendSerializedError(std::move(metadata), std::move(data));
+    return;
+  }
+
+  if (auto responseRpcError = processFirstResponse(
+          metadata, data, getProtoId(), version_, getCompressionConfig())) {
+    auto ex = makeRocketException(
+        *responseRpcError, *context_.connection().getPayloadSerializer());
+    handleStreamError(std::move(ex), clientCallback_);
+    return;
+  }
+
+  context_.unsetMarkRequestComplete();
+  // Missing some configuration here
+  // clientCallback_->setProtoId(getProtoId());
+  // clientCallback_->setChunkTimeout(sinkConsumer.chunkTimeout);
+  auto payload = apache::thrift::FirstResponsePayload{
+      std::move(data), std::move(metadata)};
+  payload.fds =
+      std::move(getRequestContext()->getHeader()->fds.dcheckToSendOrEmpty());
+  std::move(bidiStreamFactory)
+      .start(std::move(payload), clientCallback_, getEventBase());
+}
+
 bool ThriftServerRequestBiDi::sendBiDiThriftResponse(
     ResponseRpcMetadata&& metadata,
     std::unique_ptr<folly::IOBuf> data,

third-party/thrift/src/thrift/lib/cpp2/transport/rocket/server/RocketThriftRequests.h

@@ -353,6 +353,11 @@ class ThriftServerRequestBiDi final : public RocketThriftRequest {
   void sendSerializedError(
       ResponseRpcMetadata&&, std::unique_ptr<folly::IOBuf>) noexcept override;
 
+  void sendBiDiThriftResponse(
+      ResponseRpcMetadata&&,
+      std::unique_ptr<folly::IOBuf>,
+      apache::thrift::detail::ServerBiDiStreamFactory&&) noexcept override;
+
   bool sendBiDiThriftResponse(
       ResponseRpcMetadata&&,
       std::unique_ptr<folly::IOBuf>,