Merge pull request #5461 from cloudflare/jasnell/replace-body-stream-take-2

Author: jasnell

src/workerd/api/http.c++

@@ -715,10 +715,21 @@ Body::ExtractedBody Body::extractBody(jsg::Lock& js, Initializer init) {
     }
   }
 
-  auto bodyStream = kj::heap<BodyBufferInputStream>(buffer.clone(js));
+  auto buf = buffer.clone(js);
 
-  return {js.alloc<ReadableStream>(IoContext::current(), kj::mv(bodyStream)), kj::mv(buffer),
-    kj::mv(contentType)};
+  if (util::Autogate::isEnabled(util::AutogateKey::BODY_BUFFER_INPUT_STREAM_REPLACEMENT)) {
+    auto memStream = newMemoryInputStream(buf.view, kj::heap(kj::mv(buf.ownBytes)));
+    auto rs = newSystemStream(kj::mv(memStream), StreamEncoding::IDENTITY);
+
+    return {js.alloc<ReadableStream>(IoContext::current(), kj::mv(rs)), kj::mv(buffer),
+      kj::mv(contentType)};
+  } else {
+    // TODO(cleanup): Remove once the Autogate is removed.
+    auto bodyStream = kj::heap<BodyBufferInputStream>(kj::mv(buf));
+
+    return {js.alloc<ReadableStream>(IoContext::current(), kj::mv(bodyStream)), kj::mv(buffer),
+      kj::mv(contentType)};
+  }
 }
 
 Body::Body(jsg::Lock& js, kj::Maybe<ExtractedBody> init, Headers& headers)
@@ -765,8 +776,15 @@ void Body::rewindBody(jsg::Lock& js) {
 
   KJ_IF_SOME(i, impl) {
     auto bufferCopy = KJ_ASSERT_NONNULL(i.buffer).clone(js);
-    auto bodyStream = kj::heap<BodyBufferInputStream>(kj::mv(bufferCopy));
-    i.stream = js.alloc<ReadableStream>(IoContext::current(), kj::mv(bodyStream));
+    if (util::Autogate::isEnabled(util::AutogateKey::BODY_BUFFER_INPUT_STREAM_REPLACEMENT)) {
+      auto memStream = newMemoryInputStream(bufferCopy.view, kj::heap(kj::mv(bufferCopy.ownBytes)));
+      auto rs = newSystemStream(kj::mv(memStream), StreamEncoding::IDENTITY);
+      i.stream = js.alloc<ReadableStream>(IoContext::current(), kj::mv(rs));
+    } else {
+      // TODO(cleanup): Remove once the Autogate is removed.
+      auto bodyStream = kj::heap<BodyBufferInputStream>(kj::mv(bufferCopy));
+      i.stream = js.alloc<ReadableStream>(IoContext::current(), kj::mv(bodyStream));
+    }
   }
 }
 

src/workerd/util/autogate.c++

@@ -25,6 +25,8 @@ kj::StringPtr KJ_STRINGIFY(AutogateKey key) {
       return "streaming-tail-worker"_kj;
     case AutogateKey::TAIL_STREAM_REFACTOR:
       return "tail-stream-refactor"_kj;
+    case AutogateKey::BODY_BUFFER_INPUT_STREAM_REPLACEMENT:
+      return "body-buffer-input-stream-replacement"_kj;
     case AutogateKey::NumOfKeys:
       KJ_FAIL_ASSERT("NumOfKeys should not be used in getName");
   }

src/workerd/util/autogate.h

@@ -20,6 +20,8 @@ enum class AutogateKey {
   STREAMING_TAIL_WORKER,
   // Enable refactor used to consolidate the different tail worker stream implementations.
   TAIL_STREAM_REFACTOR,
+  // Enable the BodyBufferInputStream replacement
+  BODY_BUFFER_INPUT_STREAM_REPLACEMENT,
   NumOfKeys  // Reserved for iteration.
 };
 

src/workerd/util/stream-utils.c++

@@ -37,9 +37,9 @@ class MemoryInputStream final: public kj::AsyncInputStream {
   kj::Promise<size_t> tryRead(void* buffer, size_t minBytes, size_t maxBytes) override {
     auto ptr = kj::arrayPtr<kj::byte>(static_cast<kj::byte*>(buffer), maxBytes);
     size_t toRead = kj::min(data.size(), ptr.size());
-    KJ_DEFER(advance(toRead));
     if (toRead == 0) return toRead;
     ptr.first(toRead).copyFrom(data.first(toRead));
+    data = data.slice(toRead);
     return toRead;
   }
 
@@ -51,8 +51,11 @@ class MemoryInputStream final: public kj::AsyncInputStream {
     // An optimized pumpTo... we know we have all the data right here. We can
     // just write it all at once up to `amount`.
     uint64_t toRead = kj::min(data.size(), amount);
-    KJ_DEFER(advance(toRead));
+    if (toRead == 0) {
+      co_return toRead;
+    }
     co_await output.write(data.first(toRead));
+    data = data.slice(toRead);
     co_return toRead;
   }
 
@@ -74,14 +77,6 @@ class MemoryInputStream final: public kj::AsyncInputStream {
  private:
   kj::ArrayPtr<const kj::byte> data;
   kj::Maybe<kj::Rc<OwnedBacking>> ownedBacking;
-
-  void advance(size_t amount) {
-    data = data.slice(amount);
-    // If we've consumed all the data, drop our reference to the backing storage eagerly.
-    if (data.size() == 0) {
-      ownedBacking = kj::none;
-    }
-  }
 };
 
 class NeuterableInputStreamImpl final: public NeuterableInputStream {