Skip to content
rachel/workerd
File

Blob: src/workerd/api/streams/identity-transform-stream.c++

14.2 KB
1#include "identity-transform-stream.h"
2 
3#include "common.h"
4 
5#include <workerd/util/autogate.h>
6#include <workerd/util/state-machine.h>
7 
8namespace workerd::api {
9 
10namespace {
11// An implementation of ReadableStreamSource and WritableStreamSink which communicates read and
12// write requests via a StateMachine.
13//
14// This class is also used as the implementation of FixedLengthStream, in which case `limit` is
15// non-nullptr.
16 
17struct Idle {
18  static constexpr kj::StringPtr NAME KJ_UNUSED = "idle"_kj;
19};
20 
21struct ReadRequest {
22  static constexpr kj::StringPtr NAME KJ_UNUSED = "read-request"_kj;
23  kj::ArrayPtr<kj::byte> bytes;
24  // WARNING: `bytes` may be invalid if fulfiller->isWaiting() returns false! (This indicates the
25  //   read was canceled.)
26  kj::Own<kj::PromiseFulfiller<size_t>> fulfiller;
27};
28 
29struct WriteRequest {
30  static constexpr kj::StringPtr NAME KJ_UNUSED = "write-request"_kj;
31  kj::ArrayPtr<const kj::byte> bytes;
32  kj::Own<kj::PromiseFulfiller<void>> fulfiller;
33};
34 
35struct Closed {
36  static constexpr kj::StringPtr NAME KJ_UNUSED = "closed"_kj;
37};
38 
39// State machine for IdentityTransformStream:
40//   Idle -> ReadRequest (read arrives when no write pending)
41//   Idle -> WriteRequest (write arrives when no read pending)
42//   Idle -> Closed (empty write = close)
43//   ReadRequest -> Idle (write fulfills read completely)
44//   WriteRequest -> Idle (read fulfills write completely)
45//   ReadRequest -> Closed (empty write closes while read pending)
46//   Any -> kj::Exception (cancel/abort)
47//   Closed -> kj::Exception (abort can force-transition a closed stream to error)
48// Closed is terminal, kj::Exception is implicitly terminal via ErrorState.
49// abort() uses forceTransitionTo to allow the exceptional Closed -> Exception transition.
50using IdentityTransformState = StateMachine<TerminalStates<Closed>,
51    ErrorState<kj::Exception>,
52    Idle,
53    ReadRequest,
54    WriteRequest,
55    Closed,
56    kj::Exception>;
57 
58class IdentityTransformStreamImpl final: public kj::Refcounted,
59                                         public ReadableStreamSource,
60                                         public WritableStreamSink {
61 public:
62  // The limit is the maximum number of bytes that can be fed through the stream.
63  // If kj::none, there is no limit.
64  explicit IdentityTransformStreamImpl(kj::Maybe<uint64_t> limit = kj::none)
65      : limit(limit),
66        state(IdentityTransformState::create<Idle>()) {}
67 
68  ~IdentityTransformStreamImpl() noexcept(false) {
69    // Due to the different natures of JS and C++ disposal, there is no point in enforcing the limit
70    // for a FixedLengthStream here.
71    //
72    // 1. Creating but not using a `new FixedLengthStream(n)` should not be an error, and ought not
73    //    to logspam us.
74    // 2. Chances are high that by the time this object gets destroyed, it's too late to tell the
75    //    user about the failure.
76  }
77 
78  // ReadableStreamSource implementation -------------------------------------------------
79 
80  kj::Promise<size_t> tryRead(void* buffer, size_t minBytes, size_t maxBytes) override {
81    size_t total = 0;
82    while (total < minBytes) {
83      // TODO(perf): tryReadInternal was written assuming minBytes would always be 1 but we've now
84      // introduced an API for user to specify a larger minBytes. For now, this is implemented as a
85      // naive loop dispatching to the 1 byte version but would be better to bake it deeper into
86      // the implementation where it can be more efficient.
87      auto amount = co_await tryReadInternal(buffer, maxBytes);
88      KJ_ASSERT(amount <= maxBytes);
89      if (amount == 0) {
90        // EOF.
91        break;
92      }
93 
94      total += amount;
95      buffer = reinterpret_cast<char*>(buffer) + amount;
96      maxBytes -= amount;
97    }
98 
99    co_return total;
100  }
101 
102  kj::Promise<size_t> tryReadInternal(void* buffer, size_t maxBytes) {
103    auto promise = readHelper(kj::arrayPtr(static_cast<kj::byte*>(buffer), maxBytes));
104 
105    KJ_IF_SOME(l, limit) {
106      promise = promise.then([this, &l = l](size_t amount) -> kj::Promise<size_t> {
107        if (amount > l) {
108          auto exception = JSG_KJ_EXCEPTION(
109              FAILED, TypeError, "Attempt to write too many bytes through a FixedLengthStream.");
110          cancel(exception.clone());
111          return kj::mv(exception);
112        } else if (amount == 0 && l != 0) {
113          auto exception = JSG_KJ_EXCEPTION(FAILED, TypeError,
114              "FixedLengthStream did not see all expected bytes before close().");
115          cancel(exception.clone());
116          return kj::mv(exception);
117        }
118        l -= amount;
119        return amount;
120      });
121    }
122 
123    return promise;
124  }
125 
126  kj::Promise<DeferredProxy<void>> pumpTo(WritableStreamSink& output, bool end) override {
127#ifdef KJ_NO_RTTI
128    // Yes, I'm paranoid.
129    static_assert(!KJ_NO_RTTI, "Need RTTI for correctness");
130#endif
131 
132    // HACK: If `output` is another TransformStream, we don't allow pumping to it, in order to
133    //   guarantee that we can't create cycles.
134    JSG_REQUIRE(!isIdentityTransformStream(output), TypeError,
135        "Inter-TransformStream ReadableStream.pipeTo() is not implemented.");
136 
137    return ReadableStreamSource::pumpTo(output, end);
138  }
139 
140  kj::Maybe<uint64_t> tryGetLength(StreamEncoding encoding) override {
141    if (encoding == StreamEncoding::IDENTITY) {
142      return limit;
143    } else {
144      return kj::none;
145    }
146  }
147 
148  void cancel(kj::Exception reason) override {
149    // Already errored - nothing to do.
150    if (state.isErrored()) return;
151 
152    // Already closed by writable side - nothing to do.
153    if (state.is<Closed>()) return;
154 
155    KJ_IF_SOME(request, state.tryGetUnsafe<ReadRequest>()) {
156      request.fulfiller->fulfill(static_cast<size_t>(0));
157    } else KJ_IF_SOME(request, state.tryGetUnsafe<WriteRequest>()) {
158      request.fulfiller->reject(reason.clone());
159    }
160    // Idle state is fine, just transition to error.
161 
162    state.forceTransitionTo<kj::Exception>(kj::mv(reason));
163 
164    // TODO(conform): Proactively put WritableStream into Errored state.
165  }
166 
167  // WritableStreamSink implementation ---------------------------------------------------
168 
169  kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override {
170    if (buffer == nullptr) {
171      return kj::READY_NOW;
172    }
173    return writeHelper(buffer);
174  }
175 
176  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const kj::byte>> pieces) override {
177    KJ_UNIMPLEMENTED("IdentityTransformStreamImpl piecewise write() not currently supported");
178    // TODO(soon): This will be called by TeeBranch::pumpTo(). We disallow that anyway, since we
179    //   disallow inter-TransformStream pumping.
180  }
181 
182  kj::Promise<void> end() override {
183    // If we're already closed, there's nothing else we need to do here.
184    if (state.is<Closed>()) return kj::READY_NOW;
185 
186    return writeHelper(kj::ArrayPtr<const kj::byte>());
187  }
188 
189  void abort(kj::Exception reason) override {
190    // Already errored - nothing to do.
191    if (state.isErrored()) return;
192 
193    KJ_IF_SOME(request, state.tryGetUnsafe<ReadRequest>()) {
194      request.fulfiller->reject(reason.clone());
195    } else KJ_IF_SOME(request, state.tryGetUnsafe<WriteRequest>()) {
196      // If the fulfiller is not waiting, the write promise was already
197      // canceled and no one is waiting on it.
198      KJ_ASSERT(!request.fulfiller->isWaiting(),
199          "abort() is supposed to wait for any pending write() to finish");
200    }
201    // Idle and Closed states are fine, just transition to error.
202    // (Closed can transition to error via abort)
203 
204    state.forceTransitionTo<kj::Exception>(kj::mv(reason));
205 
206    // TODO(conform): Proactively put ReadableStream into Errored state.
207  }
208 
209 private:
210  kj::Promise<size_t> readHelper(kj::ArrayPtr<kj::byte> bytes) {
211    // Handle error state first.
212    KJ_IF_SOME(exception, state.tryGetErrorUnsafe()) {
213      return exception.clone();
214    }
215 
216    // Handle closed state.
217    if (state.is<Closed>()) {
218      return static_cast<size_t>(0);
219    }
220 
221    // Check for already in-flight read.
222    if (state.is<ReadRequest>()) {
223      KJ_FAIL_ASSERT("read operation already in flight");
224    }
225 
226    // Check for pending write request.
227    KJ_IF_SOME(request, state.tryGetUnsafe<WriteRequest>()) {
228      if (bytes.size() >= request.bytes.size()) {
229        // The write buffer will entirely fit into our read buffer; fulfill both requests.
230        memmove(bytes.begin(), request.bytes.begin(), request.bytes.size());
231        auto result = request.bytes.size();
232        request.fulfiller->fulfill();
233 
234        // Switch to idle state.
235        state.transitionTo<Idle>();
236 
237        return result;
238      }
239 
240      // The write buffer won't quite fit into our read buffer; fulfill only the read request.
241      memmove(bytes.begin(), request.bytes.begin(), bytes.size());
242      request.bytes = request.bytes.slice(bytes.size(), request.bytes.size());
243      return bytes.size();
244    }
245 
246    // Must be idle - no outstanding write request, switch to ReadRequest state.
247    KJ_ASSERT(state.is<Idle>());
248    auto paf = kj::newPromiseAndFulfiller<size_t>();
249    state.transitionTo<ReadRequest>(bytes, kj::mv(paf.fulfiller));
250    return kj::mv(paf.promise);
251  }
252 
253  kj::Promise<void> writeHelper(kj::ArrayPtr<const kj::byte> bytes) {
254    // Handle error state first.
255    KJ_IF_SOME(exception, state.tryGetErrorUnsafe()) {
256      return exception.clone();
257    }
258 
259    // Handle closed state.
260    if (state.is<Closed>()) {
261      KJ_FAIL_ASSERT("close operation already in flight");
262    }
263 
264    // Check for already in-flight write.
265    if (state.is<WriteRequest>()) {
266      KJ_FAIL_ASSERT("write operation already in flight");
267    }
268 
269    // Check for pending read request.
270    KJ_IF_SOME(request, state.tryGetUnsafe<ReadRequest>()) {
271      if (!request.fulfiller->isWaiting()) {
272        // Oops, the request was canceled. Currently, this happen in particular when pumping a
273        // response body to the client, and the client disconnects, cancelling the pump. In this
274        // specific case, we want to propagate the error back to the write end of the transform
275        // stream. In theory, though, there could be other cases where propagation is incorrect.
276        //
277        // TODO(cleanup): This cancellation should probably be handled at a higher level, e.g.
278        //   in pumpTo(), but I need a quick fix.
279        state.forceTransitionTo<kj::Exception>(KJ_EXCEPTION(DISCONNECTED, "reader canceled"));
280 
281        // I was going to use a `goto` but Harris choked on his bagel. Recursion it is.
282        return writeHelper(bytes);
283      }
284 
285      if (bytes.size() == 0) {
286        // This is a close operation.
287        request.fulfiller->fulfill(static_cast<size_t>(0));
288        state.transitionTo<Closed>();
289        return kj::READY_NOW;
290      }
291 
292      KJ_ASSERT(request.bytes.size() > 0);
293 
294      if (request.bytes.size() >= bytes.size()) {
295        // Our write buffer will entirely fit into the read buffer; fulfill both requests.
296        memmove(request.bytes.begin(), bytes.begin(), bytes.size());
297        request.fulfiller->fulfill(bytes.size());
298        state.transitionTo<Idle>();
299        return kj::READY_NOW;
300      }
301 
302      // Our write buffer won't quite fit into the read buffer; fulfill only the read request.
303      memmove(request.bytes.begin(), bytes.begin(), request.bytes.size());
304      bytes = bytes.slice(request.bytes.size(), bytes.size());
305      request.fulfiller->fulfill(request.bytes.size());
306 
307      auto paf = kj::newPromiseAndFulfiller<void>();
308      state.transitionTo<WriteRequest>(bytes, kj::mv(paf.fulfiller));
309      return kj::mv(paf.promise);
310    }
311 
312    // Must be idle.
313    KJ_ASSERT(state.is<Idle>());
314    if (bytes.size() == 0) {
315      // This is a close operation.
316      state.transitionTo<Closed>();
317      return kj::READY_NOW;
318    }
319 
320    auto paf = kj::newPromiseAndFulfiller<void>();
321    state.transitionTo<WriteRequest>(bytes, kj::mv(paf.fulfiller));
322    return kj::mv(paf.promise);
323  }
324 
325  kj::Maybe<uint64_t> limit;
326  IdentityTransformState state;
327};
328 
329struct Pair {
330  kj::Own<ReadableStreamSource> readable;
331  kj::Own<WritableStreamSink> writable;
332};
333Pair newIdentityPair(kj::Maybe<uint64_t> expectedLength = kj::none) {
334  auto readableSide = kj::refcounted<IdentityTransformStreamImpl>(kj::mv(expectedLength));
335  auto writableSide = kj::addRef(*readableSide);
336  return Pair{.readable = kj::mv(readableSide), .writable = kj::mv(writableSide)};
337}
338}  // namespace
339 
340jsg::Ref<IdentityTransformStream> IdentityTransformStream::constructor(
341    jsg::Lock& js, jsg::Optional<IdentityTransformStream::QueuingStrategy> maybeQueuingStrategy) {
342 
343  auto& ioContext = IoContext::current();
344  auto pipe = newIdentityPipe();
345 
346  kj::Maybe<uint64_t> maybeHighWaterMark = kj::none;
347  KJ_IF_SOME(queuingStrategy, maybeQueuingStrategy) {
348    maybeHighWaterMark = queuingStrategy.highWaterMark;
349  }
350  return js.alloc<IdentityTransformStream>(js.alloc<ReadableStream>(ioContext, kj::mv(pipe.in)),
351      js.alloc<WritableStream>(ioContext, kj::mv(pipe.out),
352          ioContext.getMetrics().tryCreateWritableByteStreamObserver(), maybeHighWaterMark));
353}
354 
355jsg::Ref<FixedLengthStream> FixedLengthStream::constructor(jsg::Lock& js,
356    uint64_t expectedLength,
357    jsg::Optional<IdentityTransformStream::QueuingStrategy> maybeQueuingStrategy) {
358  constexpr uint64_t MAX_SAFE_INTEGER = (1ull << 53) - 1;
359 
360  JSG_REQUIRE(expectedLength <= MAX_SAFE_INTEGER, TypeError,
361      "FixedLengthStream requires an integer expected length less than 2^53.");
362 
363  auto& ioContext = IoContext::current();
364  auto pipe = newIdentityPipe(expectedLength);
365 
366  kj::Maybe<uint64_t> maybeHighWaterMark = kj::none;
367  // For a FixedLengthStream we do not want a highWaterMark higher than the expectedLength.
368  KJ_IF_SOME(queuingStrategy, maybeQueuingStrategy) {
369    maybeHighWaterMark = queuingStrategy.highWaterMark.map(
370        [&](uint64_t highWaterMark) { return kj::min(expectedLength, highWaterMark); });
371  }
372 
373  return js.alloc<FixedLengthStream>(js.alloc<ReadableStream>(ioContext, kj::mv(pipe.in)),
374      js.alloc<WritableStream>(ioContext, kj::mv(pipe.out),
375          ioContext.getMetrics().tryCreateWritableByteStreamObserver(), maybeHighWaterMark));
376}
377 
378OneWayPipe newIdentityPipe(kj::Maybe<uint64_t> expectedLength) {
379  auto pair = newIdentityPair(kj::mv(expectedLength));
380  return OneWayPipe{.in = kj::mv(pair.readable), .out = kj::mv(pair.writable)};
381}
382 
383bool isIdentityTransformStream(WritableStreamSink& sink) {
384  return kj::dynamicDowncastIfAvailable<IdentityTransformStreamImpl>(sink) != kj::none;
385}
386 
387}  // namespace workerd::api