Fixup destruction order in streams (#5906)

Author: jasnell

src/workerd/api/streams/queue-test.c++

@@ -1817,5 +1817,130 @@ KJ_TEST("ValueQueue draining read with default maxRead (unlimited)") {
 
 #pragma endregion Draining Read maxRead Tests
 
+#pragma region Queue/Consumer Destruction Order Tests
+
+// These tests verify that destroying the queue before its consumers doesn't crash.
+// This can happen during isolate teardown when the destruction order isn't guaranteed
+// to follow the ownership hierarchy.
+// These will typically only catch in builds with AddressSanitizer enabled.
+
+KJ_TEST("ValueQueue destroyed before consumer doesn't crash") {
+  preamble([](jsg::Lock& js) {
+    // Heap-allocate the queue so we can control its destruction order
+    auto queue = kj::heap<ValueQueue>(2);
+
+    // Create a consumer attached to the queue
+    auto consumer = kj::heap<ValueQueue::Consumer>(*queue);
+
+    // Push some data to make sure the consumer has state
+    queue->push(js, getEntry(js, 4));
+    KJ_ASSERT(consumer->size() == 4);
+
+    // Now destroy the queue FIRST - this simulates the production scenario
+    // where wrapper cleanup destroys the controller (and its queue) before
+    // the consumer that holds a reference to it.
+    queue = nullptr;
+
+    // When the consumer is destroyed (here, or when going out of scope),
+    // its destructor calls queue.removeConsumer(this).
+    // Without the fix, this is a use-after-free.
+    // With the fix, the consumer knows the queue is gone and skips the call.
+    consumer = nullptr;
+
+    // If we get here without crashing, the test passes
+  });
+}
+
+KJ_TEST("ValueQueue destroyed before multiple consumers doesn't crash") {
+  preamble([](jsg::Lock& js) {
+    auto queue = kj::heap<ValueQueue>(2);
+
+    auto consumer1 = kj::heap<ValueQueue::Consumer>(*queue);
+    auto consumer2 = kj::heap<ValueQueue::Consumer>(*queue);
+
+    queue->push(js, getEntry(js, 4));
+    KJ_ASSERT(consumer1->size() == 4);
+    KJ_ASSERT(consumer2->size() == 4);
+
+    // Destroy queue before consumers
+    queue = nullptr;
+
+    // Both consumers should handle destruction gracefully
+    consumer1 = nullptr;
+    consumer2 = nullptr;
+  });
+}
+
+KJ_TEST("ByteQueue destroyed before consumer doesn't crash") {
+  preamble([](jsg::Lock& js) {
+    auto queue = kj::heap<ByteQueue>(2);
+    auto consumer = kj::heap<ByteQueue::Consumer>(*queue);
+
+    auto store = jsg::BackingStore::alloc(js, 4);
+    store.asArrayPtr().fill('a');
+    queue->push(js, kj::rc<ByteQueue::Entry>(jsg::BufferSource(js, kj::mv(store))));
+    KJ_ASSERT(consumer->size() == 4);
+
+    // Destroy queue before consumer
+    queue = nullptr;
+    consumer = nullptr;
+  });
+}
+
+KJ_TEST("ValueQueue destroyed with pending read requests doesn't crash") {
+  preamble([](jsg::Lock& js) {
+    auto queue = kj::heap<ValueQueue>(2);
+    auto consumer = kj::heap<ValueQueue::Consumer>(*queue);
+
+    // Queue a read request (no data pushed, so it will be pending)
+    auto prp = js.newPromiseAndResolver<ReadResult>();
+    consumer->read(js, ValueQueue::ReadRequest{.resolver = kj::mv(prp.resolver)});
+
+    KJ_ASSERT(consumer->hasReadRequests());
+
+    // Destroy queue while there are pending reads
+    queue = nullptr;
+
+    // Consumer destruction should handle this gracefully
+    consumer = nullptr;
+
+    js.runMicrotasks();
+  });
+}
+
+KJ_TEST("ValueQueue close then destroy before consumer doesn't crash") {
+  preamble([](jsg::Lock& js) {
+    auto queue = kj::heap<ValueQueue>(2);
+    auto consumer = kj::heap<ValueQueue::Consumer>(*queue);
+
+    // Close the queue first
+    queue->close(js);
+
+    // Then destroy it
+    queue = nullptr;
+
+    // Consumer should still handle destruction gracefully
+    consumer = nullptr;
+  });
+}
+
+KJ_TEST("ValueQueue error then destroy before consumer doesn't crash") {
+  preamble([](jsg::Lock& js) {
+    auto queue = kj::heap<ValueQueue>(2);
+    auto consumer = kj::heap<ValueQueue::Consumer>(*queue);
+
+    // Error the queue first
+    queue->error(js, js.v8Ref(js.v8Error("boom"_kj)));
+
+    // Then destroy it
+    queue = nullptr;
+
+    // Consumer should still handle destruction gracefully
+    consumer = nullptr;
+  });
+}
+
+#pragma endregion Queue / Consumer Destruction Order Tests
+
 }  // namespace
 }  // namespace workerd::api

src/workerd/api/streams/queue.c++

@@ -68,6 +68,9 @@ ValueQueue::Consumer::Consumer(
     QueueImpl& impl, kj::Maybe<ConsumerImpl::StateListener&> stateListener)
     : impl(impl, stateListener) {}
 
+ValueQueue::Consumer::Consumer(kj::Maybe<ConsumerImpl::StateListener&> stateListener)
+    : impl(stateListener) {}
+
 void ValueQueue::Consumer::cancel(jsg::Lock& js, jsg::Optional<v8::Local<v8::Value>> maybeReason) {
   impl.cancel(js, maybeReason);
 }
@@ -102,7 +105,14 @@ size_t ValueQueue::Consumer::size() {
 
 kj::Own<ValueQueue::Consumer> ValueQueue::Consumer::clone(
     jsg::Lock& js, kj::Maybe<ConsumerImpl::StateListener&> stateListener) {
-  auto consumer = kj::heap<Consumer>(impl.queue, stateListener);
+  // If the queue was destroyed (e.g., stream was closed), we can still clone
+  // the consumer - the cloneTo() will copy the closed/errored state.
+  kj::Own<Consumer> consumer;
+  KJ_IF_SOME(q, impl.queue) {
+    consumer = kj::heap<Consumer>(q, stateListener);
+  } else {
+    consumer = kj::heap<Consumer>(stateListener);
+  }
   impl.cloneTo(js, consumer->impl);
   return kj::mv(consumer);
 }
@@ -160,7 +170,7 @@ jsg::Promise<DrainingReadResult> ValueQueue::Consumer::drainingRead(jsg::Lock& j
   }
 
   auto& ready = impl.state.requireActiveUnsafe();
-  ConsumerImpl::UpdateBackpressureScope scope(impl.queue);
+  ConsumerImpl::UpdateBackpressureScope scope(impl);
 
   // Mark that we're doing a draining read. This allows onConsumerWantsData()
   // to use forcePull() which bypasses backpressure checks. The flag is cleared
@@ -320,7 +330,7 @@ size_t ValueQueue::size() const {
 }
 
 void ValueQueue::handlePush(
-    jsg::Lock& js, ConsumerImpl::Ready& state, QueueImpl& queue, kj::Rc<Entry> entry) {
+    jsg::Lock& js, ConsumerImpl::Ready& state, kj::Maybe<QueueImpl&> queue, kj::Rc<Entry> entry) {
   // If there are no pending reads, just add the entry to the buffer and return, adjusting
   // the size of the queue in the process.
   if (state.readRequests.empty()) {
@@ -339,7 +349,7 @@ void ValueQueue::handlePush(
 void ValueQueue::handleRead(jsg::Lock& js,
     ConsumerImpl::Ready& state,
     ConsumerImpl& consumer,
-    QueueImpl& queue,
+    kj::Maybe<QueueImpl&> queue,
     ReadRequest request) {
   // If there are no pending read requests and there is data in the buffer,
   // we will try to fulfill the read request immediately.
@@ -392,8 +402,10 @@ void ValueQueue::handleRead(jsg::Lock& js,
   }
 }
 
-bool ValueQueue::handleMaybeClose(
-    jsg::Lock& js, ConsumerImpl::Ready& state, ConsumerImpl& consumer, QueueImpl& queue) {
+bool ValueQueue::handleMaybeClose(jsg::Lock& js,
+    ConsumerImpl::Ready& state,
+    ConsumerImpl& consumer,
+    kj::Maybe<QueueImpl&> queue) {
   // If the value queue is not yet empty we have to keep waiting for more reads to consume it.
   // Return false to indicate that we cannot close yet.
   return false;
@@ -518,6 +530,9 @@ ByteQueue::Consumer::Consumer(
     QueueImpl& impl, kj::Maybe<ConsumerImpl::StateListener&> stateListener)
     : impl(impl, stateListener) {}
 
+ByteQueue::Consumer::Consumer(kj::Maybe<ConsumerImpl::StateListener&> stateListener)
+    : impl(stateListener) {}
+
 void ByteQueue::Consumer::cancel(jsg::Lock& js, jsg::Optional<v8::Local<v8::Value>> maybeReason) {
   impl.cancel(js, maybeReason);
 }
@@ -552,7 +567,14 @@ size_t ByteQueue::Consumer::size() const {
 
 kj::Own<ByteQueue::Consumer> ByteQueue::Consumer::clone(
     jsg::Lock& js, kj::Maybe<ConsumerImpl::StateListener&> stateListener) {
-  auto consumer = kj::heap<Consumer>(impl.queue, stateListener);
+  // If the queue was destroyed (e.g., stream was closed), we can still clone
+  // the consumer - the cloneTo() will copy the closed/errored state.
+  kj::Own<Consumer> consumer;
+  KJ_IF_SOME(q, impl.queue) {
+    consumer = kj::heap<Consumer>(q, stateListener);
+  } else {
+    consumer = kj::heap<Consumer>(stateListener);
+  }
   impl.cloneTo(js, consumer->impl);
   return kj::mv(consumer);
 }
@@ -582,7 +604,7 @@ jsg::Promise<DrainingReadResult> ByteQueue::Consumer::drainingRead(jsg::Lock& js
   }
 
   auto& ready = impl.state.requireActiveUnsafe();
-  ConsumerImpl::UpdateBackpressureScope scope(impl.queue);
+  ConsumerImpl::UpdateBackpressureScope scope(impl);
 
   // Mark that we're doing a draining read. This allows onConsumerWantsData()
   // to use forcePull() which bypasses backpressure checks. The flag is cleared
@@ -898,8 +920,10 @@ size_t ByteQueue::size() const {
   return impl.size();
 }
 
-void ByteQueue::handlePush(
-    jsg::Lock& js, ConsumerImpl::Ready& state, QueueImpl& queue, kj::Rc<Entry> newEntry) {
+void ByteQueue::handlePush(jsg::Lock& js,
+    ConsumerImpl::Ready& state,
+    kj::Maybe<QueueImpl&> queue,
+    kj::Rc<Entry> newEntry) {
   const auto bufferData = [&](size_t offset) {
     state.queueTotalSize += newEntry->getSize() - offset;
     state.buffer.emplace_back(QueueEntry{
@@ -1046,7 +1070,7 @@ void ByteQueue::handlePush(
 void ByteQueue::handleRead(jsg::Lock& js,
     ConsumerImpl::Ready& state,
     ConsumerImpl& consumer,
-    QueueImpl& queue,
+    kj::Maybe<QueueImpl&> queue,
     ReadRequest request) {
   const auto pendingRead = [&]() {
     bool isByob = request.pullInto.type == ReadRequest::Type::BYOB;
@@ -1055,11 +1079,14 @@ void ByteQueue::handleRead(jsg::Lock& js,
       // Because ReadRequest is movable, and because the ByobRequest captures
       // a reference to the ReadRequest, we wait until after it is added to
       // state.readRequests to create the associated ByobRequest.
-      // If the queue state is nullptr here, it means the queue has already
-      // been closed.
-      KJ_IF_SOME(queueState, queue.getState()) {
-        queueState.pendingByobReadRequests.push_back(
-            state.readRequests.back()->makeByobReadRequest(consumer, queue));
+      // If the queue is none, the consumer was cloned from a closed stream
+      // and we can't create a ByobRequest. If the queue state is none,
+      // the queue has already been closed.
+      KJ_IF_SOME(q, queue) {
+        KJ_IF_SOME(queueState, q.getState()) {
+          queueState.pendingByobReadRequests.push_back(
+              state.readRequests.back()->makeByobReadRequest(consumer, q));
+        }
       }
     }
     KJ_IF_SOME(listener, consumer.stateListener) {
@@ -1178,8 +1205,10 @@ void ByteQueue::handleRead(jsg::Lock& js,
   }
 }
 
-bool ByteQueue::handleMaybeClose(
-    jsg::Lock& js, ConsumerImpl::Ready& state, ConsumerImpl& consumer, QueueImpl& queue) {
+bool ByteQueue::handleMaybeClose(jsg::Lock& js,
+    ConsumerImpl::Ready& state,
+    ConsumerImpl& consumer,
+    kj::Maybe<QueueImpl&> queue) {
   // This is called when we know that we are closing and we still have data in
   // the queue. We want to see if we can drain as much of it into pending reads
   // as possible. If we're able to drain all of it, then yay! We can go ahead and