Skip to content
rachel/workerd
File

Blob: src/workerd/api/streams/readable-source.c++

31.7 KB
1#include "readable-source.h"
2 
3#include "common.h"
4#include "writable-sink.h"
5 
6#include <workerd/api/util.h>
7#include <workerd/io/io-context.h>
8#include <workerd/jsg/jsg.h>
9#include <workerd/util/state-machine.h>
10#include <workerd/util/stream-utils.h>
11#include <workerd/util/string-buffer.h>
12#include <workerd/util/strong-bool.h>
13 
14#include <kj/async-io.h>
15#include <kj/compat/brotli.h>
16#include <kj/compat/gzip.h>
17 
18#include <bit>
19 
20namespace workerd::api::streams {
21 
22namespace {
23// Used to consume and collect all data from a ReadableSource up to a specified
24// limit. Throws if the limit is exceeded before EOF.
25class AllReader final {
26 public:
27  explicit AllReader(ReadableSource& input, size_t limit): input(input), limit(limit) {
28    JSG_REQUIRE(limit > 0, TypeError, "Memory limit exceeded before EOF.");
29    KJ_IF_SOME(length, input.tryGetLength(rpc::StreamEncoding::IDENTITY)) {
30      // Oh hey, we might be able to bail early.
31      JSG_REQUIRE(length <= limit, TypeError, "Memory limit would be exceeded before EOF.");
32    }
33  }
34  KJ_DISALLOW_COPY_AND_MOVE(AllReader);
35 
36  kj::Promise<kj::Array<const kj::byte>> readAllBytes() {
37    co_return co_await read<kj::byte>();
38  }
39 
40  kj::Promise<kj::String> readAllText() {
41    co_return kj::String(co_await read<char>(ReadOption::NULL_TERMINATE));
42  }
43 
44 private:
45  ReadableSource& input;
46  size_t limit;
47 
48  enum class ReadOption {
49    NONE,
50    NULL_TERMINATE,
51  };
52 
53  template <typename T>
54  kj::Promise<kj::Array<T>> read(ReadOption option = ReadOption::NONE) {
55    // Read in chunks and accumulate them. Use an exponential growth strategy
56    // to determine chunk sizes to minimize the number of iterations and
57    // allocations on large streams.
58    kj::Vector<kj::Array<T>> parts;
59    size_t runningTotal = 0;
60    // TODO(later): Make these configurable someday?
61    static constexpr size_t MIN_BUFFER_CHUNK = 1024;
62    static constexpr size_t DEFAULT_BUFFER_CHUNK = 4096;
63    // TODO(later): Consider increasing MAX_BUFFER_CHUNK, maybe up to 1 MB?
64    static constexpr size_t MAX_BUFFER_CHUNK = DEFAULT_BUFFER_CHUNK * 4;
65 
66    // If we know in advance how much data we'll be reading, then we can attempt to optimize the
67    // loop here by setting the value specifically so we are only allocating at most twice. But,
68    // to be safe, let's enforce an upper bound on each allocation even if we do know the total.
69    kj::Maybe<size_t> maybeLength = input.tryGetLength(rpc::StreamEncoding::IDENTITY);
70 
71    size_t amountToRead;
72    KJ_IF_SOME(length, maybeLength) {
73      if (length <= MAX_BUFFER_CHUNK) {
74        amountToRead = kj::min(limit, length);
75      } else {
76        amountToRead = DEFAULT_BUFFER_CHUNK;
77      }
78    } else {
79      amountToRead = MIN_BUFFER_CHUNK;
80    }
81 
82    if (amountToRead != 0) {
83      while (true) {
84        auto bytes = kj::heapArray<T>(amountToRead);
85        size_t amount = co_await input.read(bytes.asBytes(), bytes.size());
86        KJ_DASSERT(amount <= bytes.size());
87        runningTotal += amount;
88        JSG_REQUIRE(runningTotal <= limit, TypeError, "Memory limit exceeded before EOF.");
89 
90        if (amount == bytes.size()) {
91          parts.add(kj::mv(bytes));
92          // Adjust the next allocation size -- double it up to a maximum
93          amountToRead = kj::min(amountToRead * 2, kj::min(MAX_BUFFER_CHUNK, limit - runningTotal));
94        } else {
95          if (amount > 0) {
96            parts.add(bytes.first(amount).attach(kj::mv(bytes)));
97          }
98          break;
99        }
100      }
101    }
102 
103    if (option == ReadOption::NULL_TERMINATE) {
104      auto out = kj::heapArray<T>(runningTotal + 1);
105      out[runningTotal] = '\0';
106      copyInto<T>(out, parts);
107      co_return kj::mv(out);
108    }
109 
110    // As an optimization, if there's only a single part in the list, we can avoid
111    // further copies.
112    if (parts.size() == 1) {
113      co_return kj::mv(parts[0]);
114    }
115 
116    auto out = kj::heapArray<T>(runningTotal);
117    copyInto<T>(out, parts);
118    co_return kj::mv(out);
119  }
120 
121  template <typename T>
122  void copyInto(kj::ArrayPtr<T> out, kj::ArrayPtr<kj::Array<T>> in) {
123    for (auto& part: in) {
124      KJ_DASSERT(part.size() <= out.size());
125      out.first(part.size()).copyFrom(part);
126      out = out.slice(part.size());
127    }
128  }
129};
130 
131// An AsyncInputStream wrapper that translates tee-related kj::Exceptions from read
132// operations into jsg::Exceptions.
133// TODO(later): We might be able to get rid of this and use a KJ exception detail instead.
134class TeeErrorAdapter final: public kj::AsyncInputStream {
135 public:
136  static kj::Own<kj::AsyncInputStream> wrap(kj::Own<kj::AsyncInputStream> inner) {
137    // We make a best effort to avoid double-wrapping.
138    if (dynamic_cast<TeeErrorAdapter*>(inner.get()) == nullptr) {
139      return kj::heap<TeeErrorAdapter>(kj::mv(inner));
140    } else {
141      return kj::mv(inner);
142    }
143  }
144 
145  explicit TeeErrorAdapter(kj::Own<AsyncInputStream> inner): inner(kj::mv(inner)) {}
146 
147  kj::Promise<size_t> tryRead(void* buffer, size_t minBytes, size_t maxBytes) override {
148    return translateErrors([&] { return inner->tryRead(buffer, minBytes, maxBytes); });
149  }
150 
151  kj::Maybe<uint64_t> tryGetLength() override {
152    return inner->tryGetLength();
153  };
154 
155  kj::Promise<uint64_t> pumpTo(kj::AsyncOutputStream& output, uint64_t amount) override {
156    return translateErrors([&] { return inner->pumpTo(output, amount); });
157  }
158 
159  kj::Maybe<kj::Own<kj::AsyncInputStream>> tryTee(uint64_t limit) override {
160    return inner->tryTee(limit);
161  }
162 
163 private:
164  kj::Own<AsyncInputStream> inner;
165 
166  template <typename Func>
167  static auto translateErrors(Func&& f) -> decltype(kj::fwd<Func>(f)()) {
168    try {
169      co_return co_await f();
170    } catch (...) {
171      auto exception = kj::getCaughtExceptionAsKj();
172      KJ_IF_SOME(translated,
173          translateKjException(exception,
174              {
175                {"tee buffer size limit exceeded"_kj,
176                  "ReadableStream.tee() buffer limit exceeded. This error usually occurs "
177                  "when a Request or Response with a large body is cloned, then only one "
178                  "of the clones is read, forcing the Workers runtime to buffer the entire "
179                  "body in memory. To fix this issue, remove unnecessary calls to "
180                  "Request/Response.clone() and ReadableStream.tee(), and always read "
181                  "clones/tees in parallel."_kj},
182              })) {
183        kj::throwFatalException(kj::mv(translated));
184      } else {
185        kj::throwFatalException(kj::mv(exception));
186      }
187    }
188  }
189};
190 
191// A kj::AsyncInputStream implementation that delegates to a provided function
192// to produce data on each read.
193class InputStreamFromProducer final: public kj::AsyncInputStream {
194 public:
195  using Producer = kj::Function<kj::Promise<size_t>(kj::ArrayPtr<kj::byte>, size_t)>;
196  InputStreamFromProducer(Producer producer, kj::Maybe<uint64_t> expectedLength)
197      : producer(kj::mv(producer)),
198        expectedLength(expectedLength) {}
199 
200  KJ_DISALLOW_COPY_AND_MOVE(InputStreamFromProducer);
201 
202  kj::Promise<size_t> tryRead(void* buffer, size_t minBytes, size_t maxBytes) override {
203    KJ_IF_SOME(p, producer) {
204      // If there is an expected length, we won't try to read more than whatever is remaining.
205      maxBytes = kj::min(maxBytes, expectedLength.orDefault(maxBytes));
206      minBytes = kj::min(minBytes, maxBytes);
207      auto amount = co_await p(kj::arrayPtr(static_cast<kj::byte*>(buffer), maxBytes), minBytes);
208      KJ_IF_SOME(length, expectedLength) {
209        KJ_DASSERT(amount <= length, "Producer produced more data than expected.");
210        length -= amount;
211      }
212      if (amount < minBytes) {
213        // The producer is indicating that we're done. Drop the producer.
214        // If the producer did not producer as much data as we expected, that's an error.
215        KJ_IF_SOME(length, expectedLength) {
216          KJ_REQUIRE(length == 0, "jsg.Error: Producer ended stream early.");
217        }
218        producer = kj::none;
219      }
220      co_return amount;
221    } else {
222      co_return 0;  // EOF
223    }
224  }
225 
226  // Returns the expected number of bytes remaining to be read, if known.
227  kj::Maybe<uint64_t> tryGetLength() override {
228    return expectedLength;
229  }
230 
231 private:
232  kj::Maybe<kj::Function<kj::Promise<size_t>(kj::ArrayPtr<kj::byte>, size_t)>> producer;
233  kj::Maybe<uint64_t> expectedLength;
234};
235 
236struct Closed {
237  static constexpr kj::StringPtr NAME KJ_UNUSED = "closed"_kj;
238};
239 
240struct Open {
241  static constexpr kj::StringPtr NAME KJ_UNUSED = "open"_kj;
242  kj::Own<kj::AsyncInputStream> stream;
243};
244 
245// State machine for tracking readable source lifecycle:
246//   Open -> Closed (normal close after EOF or pumpTo)
247//   Open -> kj::Exception (error via cancel() or read failure)
248// Closed is terminal, kj::Exception is implicitly terminal via ErrorState.
249using ReadableSourceState = StateMachine<TerminalStates<Closed>,
250    ErrorState<kj::Exception>,
251    ActiveState<Open>,
252    Open,
253    Closed,
254    kj::Exception>;
255 
256// A base class for ReadableSource implementations that provides default
257// implementations of some methods.
258class ReadableSourceImpl: public ReadableSource {
259 public:
260  ReadableSourceImpl(kj::Own<kj::AsyncInputStream> input,
261      rpc::StreamEncoding encoding = rpc::StreamEncoding::IDENTITY)
262      : state(ReadableSourceState::create<Open>(kj::mv(input))),
263        encoding(encoding) {}
264  ReadableSourceImpl(kj::Exception reason)
265      : state(ReadableSourceState::create<kj::Exception>(kj::mv(reason))),
266        encoding(rpc::StreamEncoding::IDENTITY) {}
267  ReadableSourceImpl()
268      : state(ReadableSourceState::create<Closed>()),
269        encoding(rpc::StreamEncoding::IDENTITY) {}
270  KJ_DISALLOW_COPY_AND_MOVE(ReadableSourceImpl);
271  virtual ~ReadableSourceImpl() noexcept(false) {
272    canceler.cancel(KJ_EXCEPTION(DISCONNECTED, "stream was dropped"));
273  }
274 
275  kj::Promise<size_t> readInner(Open& open, kj::ArrayPtr<kj::byte> buffer, size_t minBytes = 1) {
276    try {
277      auto& stream = setStream(ensureIdentityEncoding(kj::mv(open.stream)));
278      minBytes = kj::max(minBytes, 1u);
279      auto amount = co_await readImpl(stream, buffer, minBytes);
280      if (amount < minBytes) {
281        setClosed();
282      }
283      co_return amount;
284    } catch (...) {
285      handleOperationException();
286    }
287  }
288 
289  kj::Promise<size_t> read(kj::ArrayPtr<kj::byte> buffer, size_t minBytes = 1) override {
290    throwIfErrored();
291    if (state.is<Closed>()) {
292      co_return 0;
293    }
294    auto& open = state.requireActiveUnsafe();
295    KJ_REQUIRE(canceler.isEmpty(), "jsg.Error: Stream is already being read");
296    co_return co_await canceler.wrap(readInner(open, buffer, minBytes));
297    // If the source is dropped while a read is in progress, the canceler will
298    // trigger and abort the read. In such cases, we don't want to wrap this
299    // await in a try catch because it isn't safe to continue using the stream
300    // as it may no longer exist.
301  }
302 
303  kj::Promise<DeferredProxy<void>> pumpTo(
304      WritableSink& output, EndAfterPump end = EndAfterPump::YES) override {
305    // By default, we assume the pump is eligible for deferred proxying.
306    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
307 
308    if (!canceler.isEmpty()) {
309      kj::throwFatalException(KJ_EXCEPTION(FAILED, "jsg.Error: Stream is already being read"));
310    }
311 
312    KJ_IF_SOME(errored, state.tryGetErrorUnsafe()) {
313      output.abort(errored.clone());
314      kj::throwFatalException(errored.clone());
315    }
316 
317    if (state.is<Closed>()) {
318      if (end) {
319        co_await output.end();
320      }
321      co_return;
322    }
323 
324    auto& open = state.requireActiveUnsafe();
325    // Ownership of the underlying inner stream is transferred to the pump operation,
326    // where it will be either fully consumed or errored out. In either case, this
327    // ReadableSource becomes closed and no longer usable once pumpTo() is called.
328    // Critically... it is important that just because the ReadableSource is closed here
329    // does NOT mean that the underlying stream has been fully consumed.
330    auto stream = kj::mv(open.stream);
331    setClosed();
332 
333    if (output.getEncoding() != getEncoding()) {
334      // The target encoding is different from our current encoding.
335      // Let's ensure that our side is in identity encoding. The destination stream will
336      // take care of itself.
337      stream = ensureIdentityEncoding(kj::mv(stream));
338    } else {
339      // Since the encodings match, we can tell the output stream that it doesn't need to
340      // do any of the encoding work since we'll be providing data in the expected encoding.
341      KJ_ASSERT(getEncoding() == output.disownEncodingResponsibility());
342    }
343 
344    // Note that because we are transferring ownership of the stream to the pump operation,
345    // and the pump itself should not rely on the ReadableSource for any state, it is
346    // safe to drop the ReadableSource once the pump operation begins.
347    co_return co_await pumpImpl(kj::mv(stream), output, end);
348  }
349 
350  kj::Maybe<size_t> tryGetLength(rpc::StreamEncoding encoding) override {
351    if (encoding == rpc::StreamEncoding::IDENTITY) {
352      KJ_IF_SOME(open, state.tryGetActiveUnsafe()) {
353        return open.stream->tryGetLength();
354      }
355    }
356    return kj::none;
357  }
358 
359  kj::Promise<kj::Array<const kj::byte>> readAllBytes(size_t limit) override {
360    throwIfErrored();
361    if (state.is<Closed>()) {
362      co_return kj::Array<const kj::byte>();
363    }
364    // Must be active
365    AllReader reader(*this, limit);
366    co_return co_await reader.readAllBytes();
367  }
368 
369  kj::Promise<kj::String> readAllText(size_t limit) override {
370    throwIfErrored();
371    if (state.is<Closed>()) {
372      co_return kj::String();
373    }
374    // Must be active
375    AllReader reader(*this, limit);
376    co_return co_await reader.readAllText();
377  }
378 
379  void cancel(kj::Exception reason) override {
380    canceler.cancel(reason.clone());
381    setErrored(kj::mv(reason));
382  }
383 
384  Tee tee(size_t limit) override {
385    KJ_IF_SOME(errored, state.tryGetErrorUnsafe()) {
386      return Tee{
387        .branch1 = newErroredReadableSource(errored.clone()),
388        .branch2 = newErroredReadableSource(errored.clone()),
389      };
390    }
391 
392    if (state.is<Closed>()) {
393      return Tee{
394        .branch1 = newClosedReadableSource(),
395        .branch2 = newClosedReadableSource(),
396      };
397    }
398 
399    auto& open = state.requireActiveUnsafe();
400    KJ_IF_SOME(result, tryTee(limit)) {
401      setClosed();
402      return kj::mv(result);
403    }
404 
405    auto teeResult = kj::newTee(kj::mv(open.stream), limit);
406    setClosed();
407    return Tee{
408      .branch1 = newReadableSource(wrapTeeBranch(kj::mv(teeResult.branches[0]))),
409      .branch2 = newReadableSource(wrapTeeBranch(kj::mv(teeResult.branches[1]))),
410    };
411  }
412 
413  rpc::StreamEncoding getEncoding() override {
414    return encoding;
415  }
416 
417 protected:
418  // Throws the stored exception if in error state.
419  void throwIfErrored() {
420    KJ_IF_SOME(exception, state.tryGetErrorUnsafe()) {
421      kj::throwFatalException(exception.clone());
422    }
423  }
424 
425  // Handles exceptions from read operations: stores the error and rethrows.
426  [[noreturn]] void handleOperationException() {
427    auto exception = kj::getCaughtExceptionAsKj();
428    setErrored(exception.clone());
429    kj::throwFatalException(kj::mv(exception));
430  }
431 
432  // Implementations really should override this to provide encoding support!
433  virtual kj::Own<kj::AsyncInputStream> ensureIdentityEncoding(
434      kj::Own<kj::AsyncInputStream>&& inner) {
435    // By default, we always use identity encoding so nothing to do here.
436    // It is up to subclasses to override this if they support other encodings.
437    KJ_DASSERT(encoding == rpc::StreamEncoding::IDENTITY);
438    return kj::mv(inner);
439  }
440 
441  // Implementations should override to provide an alternative tee implementation.
442  // This will only be called when the state is known to be not closed or errored.
443  virtual kj::Maybe<Tee> tryTee(size_t limit) {
444    return kj::none;
445  }
446 
447  ReadableSourceState& getState() {
448    return state;
449  }
450 
451  void setClosed() {
452    state.transitionTo<Closed>();
453  }
454 
455  void setErrored(kj::Exception reason) {
456    state.forceTransitionTo<kj::Exception>(kj::mv(reason));
457  }
458 
459  kj::AsyncInputStream& setStream(kj::Own<kj::AsyncInputStream> stream) {
460    auto& inner = *stream;
461    state.getUnsafe<Open>().stream = kj::mv(stream);
462    return inner;
463  }
464 
465  void setEncoding(rpc::StreamEncoding newEncoding) {
466    encoding = newEncoding;
467  }
468 
469 private:
470  ReadableSourceState state;
471  rpc::StreamEncoding encoding;
472  kj::Canceler canceler;
473 
474  struct ReadOption {
475    bool identityEncoding;
476  };
477 
478  // The default pumpTo() implementation which initiates a loop
479  // that reads a chunk from the input stream and writes it to the output
480  // stream until EOF is reached.
481  // The pump is canceled by dropping the returned promise.
482  static kj::Promise<void> pumpImpl(
483      kj::Own<kj::AsyncInputStream> stream, WritableSink& output, EndAfterPump end) {
484    // These are fairly arbitrary but reasonable buffer size choices.
485 
486    // Note: this intentionally contains code that is similar to the
487    // ReadableSourceKjAdapter::pumpToImpl impl in readable-source-adapter.c++.
488    // The optimizations are generally the same but the targets are a bit different
489    // (ReadableStream vs. kj::AsyncInputStream).
490 
491    static constexpr size_t DEFAULT_BUFFER_SIZE = 16384;
492    static constexpr size_t MIN_BUFFER_SIZE = 1024;
493    static constexpr size_t MED_BUFFER_SIZE = MIN_BUFFER_SIZE << 6;
494    static constexpr size_t MAX_BUFFER_SIZE = MIN_BUFFER_SIZE << 7;
495    static constexpr size_t MEDIUM_THRESHOLD = 1048576;
496    static_assert(MIN_BUFFER_SIZE < DEFAULT_BUFFER_SIZE);
497    static_assert(DEFAULT_BUFFER_SIZE < MED_BUFFER_SIZE);
498    static_assert(MED_BUFFER_SIZE < MAX_BUFFER_SIZE);
499    static_assert(MAX_BUFFER_SIZE < MEDIUM_THRESHOLD);
500 
501    // Determine optimal buffer size based on stream length. If the stream does
502    // not report a length, use the default. The logic here is simple: use larger
503    // buffer sizes for larger streams to reduce the number of read/write iterations.
504    // and smaller buffer sizes for smaller streams to reduce memory usage.
505    // If the size is unknown, we defer to a reasonable default.
506    size_t bufferSize = DEFAULT_BUFFER_SIZE;
507    kj::Maybe<uint64_t> maybeRemaining = stream->tryGetLength();
508    KJ_IF_SOME(length, maybeRemaining) {
509      // Streams that advertise their length SHOULD always tell the truth.
510      // But... on the off change they don't, we'll still try to behave
511      // reasonably. At worst we will allocate a backing buffer and
512      // perform a single read. If this proves to be a performance issue,
513      // we can fall back to strictly enforcing the advertised length.
514      if (length <= MEDIUM_THRESHOLD) {
515        // When `length` is below the medium threshold, use
516        // the nearest power of 2 >= length within the range
517        // [MIN_BUFFER_SIZE, MED_BUFFER_SIZE].
518        bufferSize = kj::max(MIN_BUFFER_SIZE, std::bit_ceil(length));
519        bufferSize = kj::min(MED_BUFFER_SIZE, bufferSize);
520      } else {
521        // Otherwise, use the biggest buffer.
522        bufferSize = MAX_BUFFER_SIZE;
523      }
524    }
525 
526    // We use a double-buffering/pipelining strategy here to try to keep both the read
527    // and write operations busy in parallel. While one buffer is being written to the
528    // output, the other buffer is being filled with data from the input stream. It does
529    // mean that we use a bit more memory in the process but should improve throughput on
530    // high-latency streams.
531    int currentReadBuf = 0;
532    kj::SmallArray<kj::byte, 4 * MIN_BUFFER_SIZE> backing(bufferSize * 2);
533    kj::ArrayPtr<kj::byte> buffer[] = {
534      backing.first(bufferSize),
535      backing.slice(bufferSize),
536    };
537 
538    // We will use an adaptive minBytes value to try to optimize read sizes based on
539    // observed stream behavior. We start with a minBytes set to half the buffer size.
540    // As the stream is read, we will adjust minBytes up or down depending on whether
541    // the stream is consistently filling the buffer or not.
542    size_t minBytes = bufferSize >> 1;
543 
544    auto readPromise = readImpl(*stream, buffer[currentReadBuf], minBytes);
545    size_t iterationCount = 0;
546    bool readFailed = false;
547 
548    try {
549      while (true) {
550        // On each iteration, wait for the read to complete...
551        size_t amount;
552        {
553          KJ_ON_SCOPE_FAILURE(readFailed = true);
554          amount = co_await readPromise;
555        }
556        iterationCount++;
557 
558        // If we read less than minBytes, assume EOF.
559        if (amount < minBytes) {
560          // If any bytes were read...
561          if (amount > 0) {
562            // Write our final chunk...
563            co_await output.write(buffer[currentReadBuf].first(amount));
564          }
565          // Then break out of the loop.
566          break;
567        }
568 
569        // Set the write buffer to the one we just filled.
570        auto writeBuf = buffer[currentReadBuf];
571 
572        // Then switch to the other buffer and start the next read.
573        currentReadBuf = 1 - currentReadBuf;
574 
575        // Maybe adjust minBytes based on how much data we read this iteration.
576        if (iterationCount <= 3 || iterationCount % 10 == 0) {
577          if (amount == bufferSize) {
578            // Stream is filling buffer completely... Use smaller minBytes to
579            // increase responsiveness, should produce more reads with less data.
580            if (bufferSize >= 4 * DEFAULT_BUFFER_SIZE) {
581              // For large buffers (≥64KB), be more aggressive about responsiveness.
582              // 25% of a large buffer is still a substantial chunk (e.g., 32KB for 128KB).
583              minBytes = bufferSize >> 2;  // 25%
584            } else {
585              // For smaller buffers, 50% provides better balance, avoiding chunks
586              // that are too small for efficient processing (e.g., keeps 16KB → 8KB).
587              minBytes = bufferSize >> 1;  // 50%
588            }
589          } else {
590            // Stream didn't fill buffer - likely slower or at natural boundary.
591            // Use higher minBytes to accumulate larger chunks and reduce iteration overhead.
592            minBytes = (bufferSize >> 2) + (bufferSize >> 1);  // 75%
593          }
594        }
595 
596        // Start our next read operation.
597        readPromise = readImpl(*stream, buffer[currentReadBuf], minBytes);
598 
599        // Write out the chunk we just read in parallel with the next read.
600        // If the write fails, the exception will propagate and cancel the pump,
601        // including the read operation. If the read fails, it will be picked
602        // up at the start of the next loop iteration.
603        co_await output.write(writeBuf.first(amount));
604      }
605    } catch (...) {
606      auto exception = kj::getCaughtExceptionAsKj();
607      if (readFailed) {
608        output.abort(exception.clone());
609      }
610      kj::throwFatalException(kj::mv(exception));
611    }
612 
613    if (end) {
614      co_await output.end();
615    }
616  }
617 
618  static kj::Promise<size_t> readImpl(
619      kj::AsyncInputStream& inner, kj::ArrayPtr<kj::byte> buffer, size_t minBytes) {
620    KJ_ASSERT(minBytes <= buffer.size());
621    try {
622      // The read() method on AsyncInputStream will throw an exception on short reads,
623      // which is why we're using tryRead() here instead.
624      co_return co_await inner.tryRead(buffer.begin(), minBytes, buffer.size());
625    } catch (...) {
626      auto exception = kj::getCaughtExceptionAsKj();
627      if (exception.getType() == kj::Exception::Type::DISCONNECTED) {
628        // Treat disconnects as EOF.
629        co_return 0;
630      }
631      kj::throwFatalException(kj::mv(exception));
632    }
633  }
634};
635 
636// A ReadableSource wrapper that prevents deferred proxying. This is useful
637// when you expect that the IoContext will need to remain live for the duration
638// of the operations on the stream.
639class NoDeferredProxySource final: public ReadableSourceWrapper {
640 public:
641  NoDeferredProxySource(kj::Own<ReadableSource> inner, IoContext& ioctx)
642      : ReadableSourceWrapper(kj::mv(inner)),
643        ioctx(ioctx) {}
644 
645  kj::Promise<size_t> read(kj::ArrayPtr<kj::byte> buffer, size_t minBytes = 1) override {
646    auto pending = ioctx.registerPendingEvent();
647    co_return co_await getInner().read(buffer, minBytes);
648  }
649 
650  kj::Promise<DeferredProxy<void>> pumpTo(
651      WritableSink& output, EndAfterPump end = EndAfterPump::YES) override {
652    auto pending = ioctx.registerPendingEvent();
653    auto [proxyTask] = co_await getInner().pumpTo(output, end);
654    co_await proxyTask;
655  }
656 
657  Tee tee(size_t limit) override {
658    auto tee = getInner().tee(limit);
659    return Tee{
660      .branch1 = kj::heap<NoDeferredProxySource>(kj::mv(tee.branch1), ioctx),
661      .branch2 = kj::heap<NoDeferredProxySource>(kj::mv(tee.branch2), ioctx),
662    };
663  }
664 
665 private:
666  IoContext& ioctx;
667};
668 
669// A ReadableSource implementation that lazily wraps an innner Gzip or Brotli
670// encoded AsyncInputStream when the first read() is called, or when pumpTo is called,
671// the encoding will be selectively and lazily applied to the inner stream.
672class EncodedAsyncInputStream final: public ReadableSourceImpl {
673 public:
674  EncodedAsyncInputStream(kj::Own<kj::AsyncInputStream> inner, rpc::StreamEncoding encoding)
675      : ReadableSourceImpl(kj::mv(inner), encoding) {}
676 
677  // Read bytes in identity encoding. If the stream is not already in identity encoding, it will be
678  // converted to identity encoding via an appropriate stream wrapper.
679  kj::Promise<size_t> read(kj::ArrayPtr<kj::byte> buffer, size_t minBytes) override {
680    try {
681      co_return co_await ReadableSourceImpl::read(buffer, minBytes);
682    } catch (...) {
683      auto exception = kj::getCaughtExceptionAsKj();
684      KJ_IF_SOME(translated,
685          translateKjException(exception,
686              {
687                {"gzip compressed stream ended prematurely"_kj,
688                  "Gzip compressed stream ended prematurely."_kj},
689                {"gzip decompression failed"_kj, "Gzip decompression failed."},
690                {"brotli state allocation failed"_kj, "Brotli state allocation failed."},
691                {"invalid brotli window size"_kj, "Invalid brotli window size."},
692                {"invalid brotli compression level"_kj, "Invalid brotli compression level."},
693                {"brotli window size too big"_kj, "Brotli window size too big."},
694                {"brotli decompression failed"_kj, "Brotli decompression failed."},
695                {"brotli compression failed"_kj, "Brotli compression failed."},
696                {"brotli compressed stream ended prematurely"_kj,
697                  "Brotli compressed stream ended prematurely."},
698              })) {
699        kj::throwFatalException(kj::mv(translated));
700      } else {
701        kj::throwFatalException(kj::mv(exception));
702      }
703    }
704  }
705 
706  kj::Maybe<Tee> tryTee(size_t limit) override {
707    // Note that if we haven't called read() yet, then the inner stream is still
708    // in its original encoding. If read() has been called, however, then the inner
709    // stream will be wrapped and will be in identity encoding.
710    auto& open = KJ_ASSERT_NONNULL(getState().tryGetActiveUnsafe());
711    auto tee = kj::newTee(kj::mv(open.stream), limit);
712    return Tee{
713      .branch1 =
714          kj::heap<EncodedAsyncInputStream>(wrapTeeBranch(kj::mv(tee.branches[0])), getEncoding()),
715      .branch2 =
716          kj::heap<EncodedAsyncInputStream>(wrapTeeBranch(kj::mv(tee.branches[1])), getEncoding()),
717    };
718  }
719 
720  kj::Own<kj::AsyncInputStream> ensureIdentityEncoding(
721      kj::Own<kj::AsyncInputStream>&& inner) override {
722    auto encoding = getEncoding();
723    if (encoding == rpc::StreamEncoding::IDENTITY) {
724      return kj::mv(inner);
725    }
726    setEncoding(rpc::StreamEncoding::IDENTITY);
727    return wrap(encoding, kj::mv(inner));
728  }
729 
730 private:
731  static kj::Own<kj::AsyncInputStream> wrap(
732      rpc::StreamEncoding encoding, kj::Own<kj::AsyncInputStream> inner) {
733    switch (encoding) {
734      case rpc::StreamEncoding::IDENTITY: {
735        return kj::mv(inner);
736      }
737      case rpc::StreamEncoding::GZIP: {
738        return kj::heap<kj::GzipAsyncInputStream>(*inner).attach(kj::mv(inner));
739      }
740      case rpc::StreamEncoding::BROTLI: {
741        return kj::heap<kj::BrotliAsyncInputStream>(*inner).attach(kj::mv(inner));
742      }
743    }
744    KJ_UNREACHABLE;
745  }
746};
747 
748}  // namespace
749 
750kj::Own<ReadableSource> newReadableSourceFromBytes(
751    kj::ArrayPtr<const kj::byte> bytes, kj::Maybe<kj::Own<void>> maybeBacking) {
752  KJ_IF_SOME(backing, maybeBacking) {
753    return newReadableSource(newMemoryInputStream(bytes, kj::mv(backing)));
754  }
755 
756  auto backing = kj::heapArray<kj::byte>(bytes);
757  auto ptr = backing.asPtr();
758  auto inner = newMemoryInputStream(ptr, kj::heap(kj::mv(backing)));
759  return newReadableSource(kj::mv(inner));
760}
761 
762kj::Own<ReadableSource> newIoContextWrappedReadableSource(
763    IoContext& ioctx, kj::Own<ReadableSource> inner) {
764  return kj::heap<NoDeferredProxySource>(kj::mv(inner), ioctx);
765}
766 
767kj::Own<ReadableSource> newReadableSourceFromProducer(
768    kj::Function<kj::Promise<size_t>(kj::ArrayPtr<kj::byte>, size_t)> producer,
769    kj::Maybe<uint64_t> expectedLength) {
770  return newReadableSource(kj::heap<InputStreamFromProducer>(kj::mv(producer), expectedLength));
771}
772 
773kj::Own<ReadableSource> newClosedReadableSource() {
774  return kj::heap<ReadableSourceImpl>();
775}
776 
777kj::Own<ReadableSource> newErroredReadableSource(kj::Exception exception) {
778  return kj::heap<ReadableSourceImpl>(kj::mv(exception));
779}
780 
781kj::Own<ReadableSource> newReadableSource(kj::Own<kj::AsyncInputStream> inner) {
782  return kj::heap<ReadableSourceImpl>(kj::mv(inner));
783}
784 
785kj::Own<ReadableSource> newEncodedReadableSource(
786    rpc::StreamEncoding encoding, kj::Own<kj::AsyncInputStream> inner) {
787  return kj::heap<EncodedAsyncInputStream>(kj::mv(inner), encoding);
788}
789 
790kj::Own<kj::AsyncInputStream> wrapTeeBranch(kj::Own<kj::AsyncInputStream> branch) {
791  return TeeErrorAdapter::wrap(kj::mv(branch));
792}
793 
794// =======================================================================================
795// MemoryInputStream
796 
797namespace {
798 
799// A ReadableStreamSource backed by in-memory data that does NOT support deferred proxying.
800// This is critical when the backing memory may have V8 heap provenance - if we allowed
801// deferred proxying, the IoContext could complete and V8 GC could free the memory while
802// the deferred pump is still running, causing a use-after-free.
803//
804// TODO(soon): The expectation is that this will be update to implement ReadableSource instead
805// of ReadableStreamSource as we continue the transition.
806class MemoryInputStream final: public ReadableStreamSource {
807 public:
808  MemoryInputStream(kj::ArrayPtr<const kj::byte> bytes, kj::Maybe<kj::Own<void>> backing)
809      : unread(bytes),
810        backing(kj::mv(backing)) {}
811 
812  kj::Promise<size_t> tryRead(void* buffer, size_t minBytes, size_t maxBytes) override {
813    size_t amount = kj::min(maxBytes, unread.size());
814    if (amount > 0) {
815      memcpy(buffer, unread.begin(), amount);
816      unread = unread.slice(amount, unread.size());
817    }
818    return amount;
819  }
820 
821  kj::Maybe<uint64_t> tryGetLength(StreamEncoding encoding) override {
822    if (encoding == StreamEncoding::IDENTITY) {
823      return unread.size();
824    }
825    return kj::none;
826  }
827 
828  kj::Promise<DeferredProxy<void>> pumpTo(WritableStreamSink& output, bool end) override {
829    // Explicitly NOT using KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING here!
830    // The backing memory may be tied to V8 heap (e.g., jsg::BackingStore, Blob data),
831    // so we must complete all I/O before the IoContext can be released.
832    if (unread.size() > 0) {
833      auto data = unread;
834      unread = nullptr;
835      co_await output.write(data);
836    }
837    if (end) {
838      co_await output.end();
839    }
840    co_return;
841  }
842 
843  void cancel(kj::Exception reason) override {
844    // Nothing to do - we're just reading from memory.
845    unread = nullptr;
846  }
847 
848 private:
849  kj::ArrayPtr<const kj::byte> unread;
850  kj::Maybe<kj::Own<void>> backing;
851};
852 
853}  // namespace
854 
855kj::Own<ReadableStreamSource> newMemorySource(
856    kj::ArrayPtr<const kj::byte> bytes, kj::Maybe<kj::Own<void>> maybeBacking) {
857  KJ_IF_SOME(backing, maybeBacking) {
858    return kj::heap<MemoryInputStream>(bytes, kj::mv(backing));
859  }
860  // No backing provided - make a copy of the bytes.
861  auto copy = kj::heapArray<kj::byte>(bytes);
862  auto ptr = copy.asPtr();
863  return kj::heap<MemoryInputStream>(ptr, kj::heap(kj::mv(copy)));
864}
865 
866}  // namespace workerd::api::streams