ByteFlow: backpressure draft

GitOrigin-RevId: 09adce82dd88fcf84f41e525e45b07da03acc0f6

6 files changed, 67 insertions, 61 deletions

diff --git a/tdnet/td/net/HttpChunkedByteFlow.cpp b/tdnet/td/net/HttpChunkedByteFlow.cpp
index 918fef0d2..d0d462138 100644
--- a/tdnet/td/net/HttpChunkedByteFlow.cpp
+++ b/tdnet/td/net/HttpChunkedByteFlow.cpp
@@ -14,25 +14,26 @@
 
 namespace td {
 
-void HttpChunkedByteFlow::loop() {
-  bool was_updated = false;
-  size_t need_size;
-  while (true) {
+bool HttpChunkedByteFlow::loop() {
+  bool result = false;
+  do {
     if (state_ == State::ReadChunkLength) {
       bool ok = find_boundary(input_->clone(), "\r\n", len_);
       if (len_ > 10) {
-        return finish(Status::Error(PSLICE() << "Too long length in chunked "
-                                             << input_->cut_head(len_).move_as_buffer_slice().as_slice()));
+        finish(Status::Error(PSLICE() << "Too long length in chunked "
+                                      << input_->cut_head(len_).move_as_buffer_slice().as_slice()));
+        return false;
       }
       if (!ok) {
-        need_size = input_->size() + 1;
+        set_need_size(input_->size() + 1);
         break;
       }
       auto s_len = input_->cut_head(len_).move_as_buffer_slice();
       input_->advance(2);
       len_ = hex_to_integer<size_t>(s_len.as_slice());
       if (len_ > MAX_CHUNK_SIZE) {
-        return finish(Status::Error(PSLICE() << "Invalid chunk size " << tag("size", len_)));
+        finish(Status::Error(PSLICE() << "Invalid chunk size " << tag("size", len_)));
+        return false;
       }
       save_len_ = len_;
       state_ = State::ReadChunkContent;
@@ -40,21 +41,23 @@ void HttpChunkedByteFlow::loop() {
 
     auto size = input_->size();
     auto ready = min(len_, size);
-    need_size = min(MIN_UPDATE_SIZE, len_ + 2);
+    auto need_size = min(MIN_UPDATE_SIZE, len_ + 2);
     if (size < need_size) {
+      set_need_size(need_size);
       break;
     }
     total_size_ += ready;
     uncommited_size_ += ready;
     if (total_size_ > MAX_SIZE) {
-      return finish(Status::Error(PSLICE() << "Too big query " << tag("size", input_->size())));
+      finish(Status::Error(PSLICE() << "Too big query " << tag("size", input_->size())));
+      return false;
     }
 
     output_.append(input_->cut_head(ready));
+    result = true;
     len_ -= ready;
     if (uncommited_size_ >= MIN_UPDATE_SIZE) {
       uncommited_size_ = 0;
-      was_updated = true;
     }
 
     if (len_ == 0) {
@@ -65,19 +68,17 @@ void HttpChunkedByteFlow::loop() {
       input_->advance(2);
       total_size_ += 2;
       if (save_len_ == 0) {
-        return finish(Status::OK());
+        finish(Status::OK());
+        return false;
       }
       state_ = State::ReadChunkLength;
       len_ = 0;
     }
+  } while (0);
+  if (!is_input_active_ && !result) {
+    finish(Status::Error("Unexpected end of stream"));
   }
-  if (was_updated) {
-    on_output_updated();
-  }
-  if (!is_input_active_) {
-    return finish(Status::Error("Unexpected end of stream"));
-  }
-  set_need_size(need_size);
+  return result;
 }
 
 }  // namespace td
diff --git a/tdnet/td/net/HttpChunkedByteFlow.h b/tdnet/td/net/HttpChunkedByteFlow.h
index ba3c05d5e..184fff104 100644
--- a/tdnet/td/net/HttpChunkedByteFlow.h
+++ b/tdnet/td/net/HttpChunkedByteFlow.h
@@ -12,7 +12,7 @@ namespace td {
 
 class HttpChunkedByteFlow final : public ByteFlowBase {
  public:
-  void loop() override;
+  bool loop() override;
 
  private:
   static constexpr int MAX_CHUNK_SIZE = 15 << 20;  // some reasonable limit
diff --git a/tdnet/td/net/HttpContentLengthByteFlow.cpp b/tdnet/td/net/HttpContentLengthByteFlow.cpp
index cfd31d670..f7d15c81a 100644
--- a/tdnet/td/net/HttpContentLengthByteFlow.cpp
+++ b/tdnet/td/net/HttpContentLengthByteFlow.cpp
@@ -10,7 +10,7 @@
 
 namespace td {
 
-void HttpContentLengthByteFlow::loop() {
+bool HttpContentLengthByteFlow::loop() {
   auto ready_size = input_->size();
   if (ready_size > len_) {
     ready_size = len_;
@@ -18,17 +18,19 @@ void HttpContentLengthByteFlow::loop() {
   auto need_size = min(MIN_UPDATE_SIZE, len_);
   if (ready_size < need_size) {
     set_need_size(need_size);
-    return;
+    return false;
   }
   output_.append(input_->cut_head(ready_size));
   len_ -= ready_size;
   if (len_ == 0) {
-    return finish(Status::OK());
+    finish(Status::OK());
+    return false;
   }
   if (!is_input_active_) {
-    return finish(Status::Error("Unexpected end of stream"));
+    finish(Status::Error("Unexpected end of stream"));
+    return false;
   }
-  on_output_updated();
+  return true;
 }
 
 }  // namespace td
diff --git a/tdnet/td/net/HttpContentLengthByteFlow.h b/tdnet/td/net/HttpContentLengthByteFlow.h
index 1c4129c73..38e5a70d7 100644
--- a/tdnet/td/net/HttpContentLengthByteFlow.h
+++ b/tdnet/td/net/HttpContentLengthByteFlow.h
@@ -15,7 +15,7 @@ class HttpContentLengthByteFlow final : public ByteFlowBase {
   HttpContentLengthByteFlow() = default;
   explicit HttpContentLengthByteFlow(size_t len) : len_(len) {
   }
-  void loop() override;
+  bool loop() override;
 
  private:
   static constexpr size_t MIN_UPDATE_SIZE = 1 << 14;
diff --git a/tdnet/td/net/HttpReader.cpp b/tdnet/td/net/HttpReader.cpp
index 4804cd60a..508f80131 100644
--- a/tdnet/td/net/HttpReader.cpp
+++ b/tdnet/td/net/HttpReader.cpp
@@ -68,6 +68,7 @@ Result<size_t> HttpReader::read_next(HttpQuery *query) {
   size_t need_size = input_->size() + 1;
   while (true) {
     if (state_ != State::ReadHeaders) {
+      gzip_flow_.wakeup();
       flow_source_.wakeup();
       if (flow_sink_.is_ready() && flow_sink_.status().is_error()) {
         if (!temp_file_.empty()) {
@@ -108,7 +109,11 @@ Result<size_t> HttpReader::read_next(HttpQuery *query) {
         if (content_encoding_.empty()) {
         } else if (content_encoding_ == "gzip" || content_encoding_ == "deflate") {
           gzip_flow_ = GzipByteFlow(Gzip::Mode::Decode);
-          gzip_flow_.set_max_output_size(MAX_CONTENT_SIZE);
+          GzipByteFlow::Options options;
+          options.write_watermark.low = 0;
+          options.write_watermark.hight = max_post_size_ + 10;
+          gzip_flow_.set_options(options);
+          //gzip_flow_.set_max_output_size(MAX_CONTENT_SIZE);
           *source >> gzip_flow_;
           source = &gzip_flow_;
         } else {
@@ -170,6 +175,10 @@ Result<size_t> HttpReader::read_next(HttpQuery *query) {
       case State::ReadContent: {
         if (content_->size() > max_post_size_) {
           state_ = State::ReadContentToFile;
+          GzipByteFlow::Options options;
+          options.write_watermark.low = 4 << 20;
+          options.write_watermark.hight = 8 << 20;
+          gzip_flow_.set_options(options);
           continue;
         }
         if (flow_sink_.is_ready()) {
@@ -191,14 +200,19 @@ Result<size_t> HttpReader::read_next(HttpQuery *query) {
         }
 
         auto size = content_->size();
-        if (size) {
+        bool restart = false;
+        if (size > (1 << 20) || flow_sink_.is_ready()) {
           TRY_STATUS(save_file_part(content_->cut_head(size).move_as_buffer_slice()));
+          restart = true;
         }
         if (flow_sink_.is_ready()) {
           query_->files_.emplace_back("file", "", content_type_.str(), file_size_, temp_file_name_);
           close_temp_file();
           break;
         }
+        if (restart) {
+          continue;
+        }
 
         return need_size;
       }
diff --git a/tdnet/td/net/SslStream.cpp b/tdnet/td/net/SslStream.cpp
index 81d5db497..92d68e363 100644
--- a/tdnet/td/net/SslStream.cpp
+++ b/tdnet/td/net/SslStream.cpp
@@ -387,24 +387,19 @@ class SslStreamImpl {
    public:
     explicit SslReadByteFlow(SslStreamImpl *stream) : stream_(stream) {
     }
-    void loop() override {
-      bool was_append = false;
-      while (true) {
-        auto to_read = output_.prepare_append();
-        auto r_size = stream_->read(to_read);
-        if (r_size.is_error()) {
-          return finish(r_size.move_as_error());
-        }
-        auto size = r_size.move_as_ok();
-        if (size == 0) {
-          break;
-        }
-        output_.confirm_append(size);
-        was_append = true;
+    bool loop() override {
+      auto to_read = output_.prepare_append();
+      auto r_size = stream_->read(to_read);
+      if (r_size.is_error()) {
+        finish(r_size.move_as_error());
+        return false;
       }
-      if (was_append) {
-        on_output_updated();
+      auto size = r_size.move_as_ok();
+      if (size == 0) {
+        return false;
       }
+      output_.confirm_append(size);
+      return true;
     }
 
     size_t read(MutableSlice data) {
@@ -419,34 +414,28 @@ class SslStreamImpl {
    public:
     explicit SslWriteByteFlow(SslStreamImpl *stream) : stream_(stream) {
     }
-    void loop() override {
-      while (!input_->empty()) {
-        auto to_write = input_->prepare_read();
-        auto r_size = stream_->write(to_write);
-        if (r_size.is_error()) {
-          return finish(r_size.move_as_error());
-        }
-        auto size = r_size.move_as_ok();
-        if (size == 0) {
-          break;
-        }
-        input_->confirm_read(size);
+    bool loop() override {
+      auto to_write = input_->prepare_read();
+      auto r_size = stream_->write(to_write);
+      if (r_size.is_error()) {
+        finish(r_size.move_as_error());
+        return false;
       }
-      if (output_updated_) {
-        output_updated_ = false;
-        on_output_updated();
+      auto size = r_size.move_as_ok();
+      if (size == 0) {
+        return false;
       }
+      input_->confirm_read(size);
+      return true;
     }
 
     size_t write(Slice data) {
       output_.append(data);
-      output_updated_ = true;
       return data.size();
     }
 
    private:
     SslStreamImpl *stream_;
-    bool output_updated_{false};
   };
 
   SslReadByteFlow read_flow_{this};