ByteFlow: backpressure draft

GitOrigin-RevId: 09adce82dd88fcf84f41e525e45b07da03acc0f6

1 files changed, 96 insertions, 8 deletions

diff --git a/tdutils/td/utils/ByteFlow.h b/tdutils/td/utils/ByteFlow.h
index d46d1386d..d68857ef3 100644
--- a/tdutils/td/utils/ByteFlow.h
+++ b/tdutils/td/utils/ByteFlow.h
@@ -9,6 +9,7 @@
 #include "td/utils/buffer.h"
 #include "td/utils/common.h"
 #include "td/utils/Status.h"
+#include <limits>
 
 namespace td {
 
@@ -19,6 +20,8 @@ class ByteFlowInterface {
   virtual void set_parent(ByteFlowInterface &other) = 0;
   virtual void set_input(ChainBufferReader *input) = 0;
   virtual size_t get_need_size() = 0;
+  virtual size_t get_read_size() = 0;
+  virtual size_t get_write_size() = 0;
   ByteFlowInterface() = default;
   ByteFlowInterface(const ByteFlowInterface &) = delete;
   ByteFlowInterface &operator=(const ByteFlowInterface &) = delete;
@@ -45,32 +48,92 @@ class ByteFlowBaseCommon : public ByteFlowInterface {
       return;
     }
     input_->sync_with_writer();
+
     if (waiting_flag_) {
       if (!is_input_active_) {
         finish(Status::OK());
       }
       return;
     }
-    if (is_input_active_) {
-      if (need_size_ != 0 && input_->size() < need_size_) {
-        return;
+    while (true) {
+      if (stop_flag_) {
+        break;
+      }
+
+      // update can_read
+      if (is_input_active_) {
+        auto read_size = get_read_size();
+        if (read_size < min(need_size_, options_.read_watermark.low)) {
+          can_read = false;
+        }
+        if (read_size >= max(need_size_, options_.read_watermark.hight)) {
+          can_read = true;
+        }
+
+      } else {
+        //Alway can read when input is closed
+        can_read = true;
+      }
+
+      // update can_write
+      {
+        auto write_size = get_write_size();
+        if (write_size > options_.write_watermark.hight) {
+          can_write = false;
+        }
+        if (write_size <= options_.write_watermark.low) {
+          can_write = true;
+        }
+      }
+
+      if (!can_read || !can_write) {
+        break;
+      }
+      need_size_ = 0;
+
+      if (!loop()) {
+        if (need_size_ <= get_read_size()) {
+          need_size_ = get_read_size() + 1;
+        }
       }
     }
-    need_size_ = 0;
-    loop();
+    on_output_updated();
   }
 
   size_t get_need_size() final {
     return need_size_;
   }
+  size_t get_read_size() override {
+    input_->sync_with_writer();
+    return input_->size();
+  }
+  size_t get_write_size() override {
+    CHECK(parent_);
+    return parent_->get_read_size();
+  }
+
+  struct Watermark {
+    size_t low{std::numeric_limits<size_t>::max()};
+    size_t hight{0};
+  };
+  struct Options {
+    Watermark write_watermark;
+    Watermark read_watermark;
+  };
+  void set_options(Options options) {
+    options_ = options;
+  }
 
-  virtual void loop() = 0;
+  virtual bool loop() = 0;
 
  protected:
   bool waiting_flag_ = false;
   ChainBufferReader *input_ = nullptr;
   bool is_input_active_ = true;
   size_t need_size_ = 0;
+  bool can_read{true};
+  bool can_write{true};
+  Options options_;
   void finish(Status status) {
     stop_flag_ = true;
     need_size_ = 0;
@@ -114,7 +177,7 @@ class ByteFlowBase : public ByteFlowBaseCommon {
     parent_ = &other;
     parent_->set_input(&output_reader_);
   }
-  void loop() override = 0;
+  bool loop() override = 0;
 
   // ChainBufferWriter &get_output() {
   // return output_;
@@ -137,7 +200,7 @@ class ByteFlowInplaceBase : public ByteFlowBaseCommon {
     parent_ = &other;
     parent_->set_input(&output_);
   }
-  void loop() override = 0;
+  bool loop() override = 0;
 
   ChainBufferReader &get_output() {
     return output_;
@@ -195,6 +258,14 @@ class ByteFlowSource : public ByteFlowInterface {
     }
     return parent_->get_need_size();
   }
+  size_t get_read_size() final {
+    UNREACHABLE();
+    return 0;
+  }
+  size_t get_write_size() final {
+    CHECK(parent_);
+    return parent_->get_read_size();
+  }
 
  private:
   ChainBufferReader *buffer_ = nullptr;
@@ -223,6 +294,14 @@ class ByteFlowSink : public ByteFlowInterface {
     UNREACHABLE();
     return 0;
   }
+  size_t get_read_size() final {
+    buffer_->sync_with_writer();
+    return buffer_->size();
+  }
+  size_t get_write_size() final {
+    UNREACHABLE();
+    return 0;
+  }
   bool is_ready() {
     return !active_;
   }
@@ -270,6 +349,15 @@ class ByteFlowMoveSink : public ByteFlowInterface {
     UNREACHABLE();
     return 0;
   }
+  size_t get_read_size() final {
+    input_->sync_with_writer();
+    //TODO: must be input_->size() + output_->size()
+    return input_->size();
+  }
+  size_t get_write_size() final {
+    UNREACHABLE();
+    return 0;
+  }
   void set_output(ChainBufferWriter *output) {
     CHECK(!output_);
     output_ = output;