Add FIFO fragmentation backpressure

Summary:
FIFO queues data for later EventBase callbacks, but it ignores downstream `Backpressure`. It never registers for write-ready and never reports the pressure upstream. It also records fragments after the downstream call, so an inline batch completion could run before its tracker entry existed.

This diff makes FIFO handle `Backpressure` like the current SRPT handler and yielding SRPT. FIFO keeps sending in its normal callbacks and limits, registers for write-ready, and reports `Backpressure` upstream until write-ready. Pass-through frames follow the same rule.

When the socket drains, the pipeline wakes each handler that got `Backpressure` once, in the order they got it. So a handler above FIFO can wake first, while FIFO is still blocked. After this diff, FIFO asks for one more round when it unblocks, so that handler can send. It asks only if it answered `Backpressure` itself, for a write that the socket took or that FIFO queued. A writer that the socket blocked waits for the next drain. So a full socket gets one round per drain.

Like yielding SRPT, FIFO now records each fragment before the downstream call. It removes that record if the write returns `Error`. The benchmark handlers remain disconnected from production pipelines. Graceful shutdown is unchanged.

Reviewed By: robertroeser

Differential Revision: D124885417

fbshipit-source-id: 51395eea2fab4f6513eaa4c8226cc00ffd80bdc7

Author: Vitaly Berov

thrift/lib/cpp2/fast_thrift/frame/write/handler/FifoFrameFragmentationHandler.h

@@ -25,6 +25,7 @@
 #include <folly/ExceptionWrapper.h>
 #include <folly/Utility.h>
 
+#include <thrift/lib/cpp2/fast_thrift/channel_pipeline/Backpressure.h>
 #include <thrift/lib/cpp2/fast_thrift/channel_pipeline/Common.h>
 #include <thrift/lib/cpp2/fast_thrift/channel_pipeline/Event.h>
 #include <thrift/lib/cpp2/fast_thrift/channel_pipeline/TypeErasedBox.h>
@@ -48,6 +49,12 @@ namespace apache::thrift::fast_thrift::frame::write::handler {
  * turns as there are live streams when it starts, and each turn sends one
  * fragment. The callback starts no new turn after its budget reaches zero. The
  * turn that crosses the budget still sends its whole fragment.
+ * Backpressure does not stop a callback or delay the next one. After
+ * downstream returns Backpressure, `onWrite()` returns Backpressure for each
+ * accepted frame until write-ready. If downstream took one of those frames
+ * without Backpressure, or the handler queued it, the handler asks for another
+ * write-ready round. A writer above that already ran in the round can then
+ * write.
  *
  * A frame that cannot be fragmented goes downstream at once. It can reach the
  * peer before a queued data frame of its stream. `CANCEL` and `ERROR` first
@@ -58,6 +65,8 @@ namespace apache::thrift::fast_thrift::frame::write::handler {
 template <FragmentCompletionTracker Tracker = NoOpFragmentCompletionTracker>
 class FifoFrameFragmentationHandlerT : private folly::NonCopyableNonMovable {
  public:
+  channel_pipeline::WriteReadyHook writeReadyHook_;
+
   explicit FifoFrameFragmentationHandlerT(
       QuantumFragmentationHandlerConfig config = {}) noexcept
       : state_(config) {}
@@ -75,7 +84,8 @@ class FifoFrameFragmentationHandlerT : private folly::NonCopyableNonMovable {
   }
 
   template <typename Context>
-  void handlerRemoved(Context& /*ctx*/) noexcept {
+  void handlerRemoved(Context& ctx) noexcept {
+    cancelAwaitWriteReady(ctx);
     state_.detach();
   }
 
@@ -94,17 +104,30 @@ class FifoFrameFragmentationHandlerT : private folly::NonCopyableNonMovable {
       const uint32_t streamId = frame.streamId;
       return passThrough(ctx, std::move(msg), streamId);
     }
-    if (state_.enqueue(msg.take<ComposedFrame>())) {
-      return channel_pipeline::Result::Success;
+    if (!state_.enqueue(msg.take<ComposedFrame>())) {
+      return channel_pipeline::Result::Error;
     }
-    return channel_pipeline::Result::Error;
+    return acceptedResult();
   }
 
   template <typename Context>
-  void onWriteReady(Context& /*ctx*/) noexcept {}
+  void onWriteReady(Context& ctx) noexcept {
+    const bool blockedWriterAbove = blockedWriterAbove_;
+    cancelAwaitWriteReady(ctx);
+    // A round wakes each handler once. A writer above may have run earlier in
+    // this round and got Backpressure only from this handler. That writer needs
+    // another round now. A writer that got Backpressure from downstream waits
+    // for the next write-ready of downstream. If it got a round here too, a
+    // writer that writes in every round would keep getting rounds while the
+    // transport stays full.
+    if (blockedWriterAbove) {
+      ctx.pipeline()->onWriteReady();
+    }
+  }
 
   template <typename Context>
-  void onPipelineInactive(Context& /*ctx*/) noexcept {
+  void onPipelineInactive(Context& ctx) noexcept {
+    cancelAwaitWriteReady(ctx);
     state_.clearPending();
   }
 
@@ -134,6 +157,8 @@ class FifoFrameFragmentationHandlerT : private folly::NonCopyableNonMovable {
     return state_.isScheduled();
   }
 
+  [[nodiscard]] bool isBackpressured() const noexcept { return backpressured_; }
+
   Tracker& tracker() noexcept { return tracker_; }
 
   using PublishedEvents = typename Tracker::PublishedEvents;
@@ -145,22 +170,52 @@ class FifoFrameFragmentationHandlerT : private folly::NonCopyableNonMovable {
   }
 
  private:
+  template <typename Context>
+  void awaitWriteReady(Context& ctx) noexcept {
+    backpressured_ = true;
+    ctx.awaitWriteReady();
+  }
+
+  template <typename Context>
+  void cancelAwaitWriteReady(Context& ctx) noexcept {
+    if (backpressured_) {
+      ctx.cancelAwaitWriteReady();
+      backpressured_ = false;
+    }
+    blockedWriterAbove_ = false;
+  }
+
+  channel_pipeline::Result acceptedResult() noexcept {
+    if (!backpressured_) {
+      return channel_pipeline::Result::Success;
+    }
+    blockedWriterAbove_ = true;
+    return channel_pipeline::Result::Backpressure;
+  }
+
   template <typename Context>
   [[nodiscard]] channel_pipeline::Result passThrough(
       Context& ctx,
       channel_pipeline::TypeErasedBox&& msg,
       uint32_t streamId) noexcept {
+    tracker_.onFragment(streamId, true);
     const channel_pipeline::Result result = ctx.fireWrite(std::move(msg));
-    if (result != channel_pipeline::Result::Error) {
-      tracker_.onFragment(streamId, true);
+    if (result == channel_pipeline::Result::Error) {
+      tracker_.onFragmentAborted();
+      return result;
+    }
+    if (result == channel_pipeline::Result::Backpressure) {
+      awaitWriteReady(ctx);
+      return result;
     }
-    return result;
+    return acceptedResult();
   }
 
   template <typename Context>
   void flushAndPropagateErrors(Context& ctx) noexcept {
     if (flushOnePass(ctx) == channel_pipeline::Result::Error) {
       state_.stopAfterError();
+      cancelAwaitWriteReady(ctx);
       ctx.fireException(
           folly::make_exception_wrapper<std::runtime_error>(
               "FifoFrameFragmentationHandler: downstream write failed"));
@@ -177,19 +232,27 @@ class FifoFrameFragmentationHandlerT : private folly::NonCopyableNonMovable {
       if (std::optional<FifoStreamQueue::ScheduledFragment> scheduled =
               state_.nextFragment();
           scheduled.has_value()) {
-        if (const channel_pipeline::Result result = ctx.fireWrite(
-                channel_pipeline::erase_and_box(
-                    std::move(scheduled->fragment.outFrame)));
-            result == channel_pipeline::Result::Error) {
-          return result;
-        }
         tracker_.onFragment(
             scheduled->streamId, scheduled->fragment.currentFrameDone);
+        const channel_pipeline::Result result = ctx.fireWrite(
+            channel_pipeline::erase_and_box(
+                std::move(scheduled->fragment.outFrame)));
+        if (result == channel_pipeline::Result::Error) {
+          tracker_.onFragmentAborted();
+          return result;
+        }
+        if (result == channel_pipeline::Result::Backpressure) {
+          awaitWriteReady(ctx);
+        }
       }
     }
     return channel_pipeline::Result::Success;
   }
 
+  bool backpressured_{false};
+  // Set when onWrite() answers Backpressure for a frame that downstream took or
+  // the handler queued. The writer above then waits for this handler alone.
+  bool blockedWriterAbove_{false};
   [[no_unique_address]] Tracker tracker_{};
   FifoFrameFragmentationHandlerState state_;
 };

thrift/lib/cpp2/fast_thrift/frame/write/handler/FragmentCompletionTracker.h

@@ -81,11 +81,9 @@ concept FrameWriteCompleteEventFactory =
 /**
  * Maps downstream batch completions to original frames.
  *
- * `onFragment()` appends {streamId, isLastFragment} to a FIFO. A caller that
- * records after the downstream write calls it only for a non-Error result. A
- * caller that records before the write calls `onFragmentAborted()` if the
- * result is Error. The caller must abort before another `onFragment()` call
- * because the abort removes the newest record.
+ * `onFragment()` appends {streamId, isLastFragment} to a FIFO.
+ * `onFragmentAborted()` removes the newest record, and does nothing when the
+ * FIFO is empty.
  *
  * Each `BatchWriteComplete` pops `frameCount` records from the front and fires
  * FrameWriteComplete{streamId} for each record whose `isLastFragment` is true.