Add EndableAsyncOutputStream with end() method

Author: jasnell

src/workerd/api/streams/writable-sink-test.c++

@@ -3,6 +3,7 @@
 #include <workerd/jsg/jsg-test.h>
 #include <workerd/tests/test-fixture.h>
 #include <workerd/util/own-util.h>
+#include <workerd/util/stream-utils.h>
 
 #include <kj/async-io.h>
 #include <kj/compat/gzip.h>
@@ -132,6 +133,32 @@ class MemoryAsyncOutputStream final: public kj::AsyncOutputStream {
   kj::Vector<kj::byte> data;
 };
 
+struct MockEndable final: public EndableAsyncOutputStream {
+  bool isEnded = false;
+  kj::Vector<kj::byte> data;
+
+  kj::Promise<void> write(kj::ArrayPtr<const kj::byte> buffer) override {
+    data.addAll(buffer);
+    co_return;
+  }
+
+  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const kj::byte>> pieces) override {
+    for (auto piece: pieces) {
+      data.addAll(piece);
+    }
+    co_return;
+  }
+
+  kj::Promise<void> whenWriteDisconnected() override {
+    return kj::NEVER_DONE;
+  }
+
+  kj::Promise<void> end() override {
+    isEnded = true;
+    co_return;
+  }
+};
+
 // ======================================================================================
 // Core WritableStreamSink Interface Tests
 
@@ -517,5 +544,23 @@ KJ_TEST("IoContext aware wrapper") {
   KJ_ASSERT(inner.data == "some data"_kjb);
 }
 
+// ======================================================================================
+// EndableAsyncOutputStream Tests
+
+KJ_TEST("EndableAsyncOutputStream") {
+  TestFixture fixture;
+  MockEndable inner;
+  auto fakeOwn = kj::Own<MockEndable>(&inner, kj::NullDisposer::instance);
+  auto sink = newWritableStreamSink(kj::mv(fakeOwn));
+
+  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
+    co_await sink->write("some data"_kjb);
+    co_await sink->end();
+  });
+
+  KJ_ASSERT(inner.data == "some data"_kjb);
+  KJ_ASSERT(inner.isEnded);
+}
+
 }  // namespace
 }  // namespace workerd::api::streams

src/workerd/api/streams/writable-sink.c++

@@ -3,6 +3,7 @@
 #include <workerd/io/io-context.h>
 #include <workerd/util/stream-utils.h>
 
+#include <capnp/compat/byte-stream.h>
 #include <kj/async-io.h>
 #include <kj/compat/brotli.h>
 #include <kj/compat/gzip.h>
@@ -139,6 +140,11 @@ class WritableStreamSinkImpl: public WritableStreamSink {
   virtual kj::Promise<void> flush(kj::AsyncOutputStream& output) {
     // When using the default implementation, we assume IDENTITY encoding.
     KJ_ASSERT(encoding == rpc::StreamEncoding::IDENTITY);
+    if (auto endable = dynamic_cast<EndableAsyncOutputStream*>(&output)) {
+      co_await endable->end();
+    } else if (auto endable = dynamic_cast<capnp::ExplicitEndOutputStream*>(&output)) {
+      co_await endable->end();
+    }
     // By default there's nothing to flush.
     co_return;
   }
@@ -199,6 +205,10 @@ class EncodedAsyncOutputStream final: public WritableStreamSinkImpl {
       co_await gzip->end();
     } else if (auto br = dynamic_cast<kj::BrotliAsyncOutputStream*>(&output)) {
       co_await br->end();
+    } else if (auto endable = dynamic_cast<EndableAsyncOutputStream*>(&output)) {
+      co_await endable->end();
+    } else if (auto endable = dynamic_cast<capnp::ExplicitEndOutputStream*>(&output)) {
+      co_await endable->end();
     }
     // By default there's nothing to flush.
   }

src/workerd/util/stream-utils.h

@@ -33,6 +33,16 @@ class NeuterableIoStream: public kj::AsyncIoStream {
   virtual void neuter(kj::Exception ex) = 0;
 };
 
+// Until kj::AsyncOutputStream has an end() method of its own... We
+// provide this subclass that adds it.
+class EndableAsyncOutputStream: public kj::AsyncOutputStream {
+ public:
+  // By default, end() is a no-op. Subclasses may override.
+  virtual kj::Promise<void> end() {
+    co_return;
+  }
+};
+
 kj::Own<NeuterableInputStream> newNeuterableInputStream(kj::AsyncInputStream&);
 kj::Own<NeuterableIoStream> newNeuterableIoStream(kj::AsyncIoStream&);