Skip to content

Commit

Permalink
Merge pull request #2046 from cloudflare/jsnell/hold-strongref-in-wri…
Browse files Browse the repository at this point in the history
…ter-reader

Hold strong refs in certain streams Writer/Reader methods
  • Loading branch information
harrishancock authored Apr 23, 2024
2 parents 35f34fe + adc8f24 commit 03e24eb
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 6 deletions.
13 changes: 10 additions & 3 deletions src/workerd/api/streams/readable.c++
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,6 @@ ReaderImpl::ReaderImpl(ReadableStreamController::Reader& reader) :

ReaderImpl::~ReaderImpl() noexcept(false) {
KJ_IF_SOME(stream, state.tryGet<Attached>()) {
// There's a very good likelihood that this is called during GC or other
// cleanup so we have to make sure that releasing the reader does not also
// trigger resolution of the close promise.
stream->getController().releaseReader(reader, kj::none);
}
}
Expand Down Expand Up @@ -61,6 +58,11 @@ jsg::Promise<void> ReaderImpl::cancel(
KJ_FAIL_ASSERT("this reader was never attached");
}
KJ_CASE_ONEOF(stream, Attached) {
// In some edge cases, this reader is the last thing holding a strong
// reference to the stream. Calling cancel might cause the readers strong
// reference to be cleared, so let's make sure we keep a reference to
// the stream at least until the call to cancel completes.
auto ref = stream.addRef();
return stream->getController().cancel(js, maybeReason);
}
KJ_CASE_ONEOF(r, Released) {
Expand Down Expand Up @@ -141,6 +143,11 @@ void ReaderImpl::releaseLock(jsg::Lock& js) {
KJ_FAIL_ASSERT("this reader was never attached");
}
KJ_CASE_ONEOF(stream, Attached) {
// In some edge cases, this reader is the last thing holding a strong
// reference to the stream. Calling releaseLock might cause the readers strong
// reference to be cleared, so let's make sure we keep a reference to
// the stream at least until the call to releaseLock completes.
auto ref = stream.addRef();
stream->getController().releaseReader(reader, js);
state.init<Released>();
return;
Expand Down
18 changes: 15 additions & 3 deletions src/workerd/api/streams/writable.c++
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,6 @@ WritableStreamDefaultWriter::WritableStreamDefaultWriter()

WritableStreamDefaultWriter::~WritableStreamDefaultWriter() noexcept(false) {
KJ_IF_SOME(stream, state.tryGet<Attached>()) {
// Because this can be called during gc or other cleanup, it is important
// that releasing the writer does not cause the closed promise be resolved
// since that requires v8 heap allocations.
stream->getController().releaseWriter(*this, kj::none);
}
}
Expand All @@ -39,6 +36,11 @@ jsg::Promise<void> WritableStreamDefaultWriter::abort(
KJ_FAIL_ASSERT("this writer was never attached");
}
KJ_CASE_ONEOF(stream, Attached) {
// In some edge cases, this writer is the last thing holding a strong
// reference to the stream. Calling abort can cause the writers strong
// reference to be cleared, so let's make sure we keep a reference to
// the stream at least until the call to abort completes.
auto ref = stream.addRef();
return stream->getController().abort(js, reason);
}
KJ_CASE_ONEOF(r, Released) {
Expand Down Expand Up @@ -68,6 +70,11 @@ jsg::Promise<void> WritableStreamDefaultWriter::close(jsg::Lock& js) {
KJ_FAIL_ASSERT("this writer was never attached");
}
KJ_CASE_ONEOF(stream, Attached) {
// In some edge cases, this writer is the last thing holding a strong
// reference to the stream. Calling close can cause the writers strong
// reference to be cleared, so let's make sure we keep a reference to
// the stream at least until the call to close completes.
auto ref = stream.addRef();
return stream->getController().close(js);
}
KJ_CASE_ONEOF(r, Released) {
Expand Down Expand Up @@ -142,6 +149,11 @@ void WritableStreamDefaultWriter::releaseLock(jsg::Lock& js) {
KJ_FAIL_ASSERT("this writer was never attached");
}
KJ_CASE_ONEOF(stream, Attached) {
// In some edge cases, this writer is the last thing holding a strong
// reference to the stream. Calling releaseWriter can cause the writers
// strong reference to be cleared, so let's make sure we keep a reference
// to the stream at least until the call to releaseLock completes.
auto ref = stream.addRef();
stream->getController().releaseWriter(*this, js);
state.init<Released>();
return;
Expand Down
13 changes: 13 additions & 0 deletions src/workerd/api/tests/streams-test.js
Original file line number Diff line number Diff line change
Expand Up @@ -414,6 +414,19 @@ export const readableStreamFromNoopAsyncGen = {
}
};

export const abortWriterAfterGc = {
async test() {
function getWriter() {
const { writable } = new IdentityTransformStream();
return writable.getWriter();
}

const writer = getWriter();
gc();
await writer.abort();
}
};

export default {
async fetch(request, env) {
strictEqual(request.headers.get('content-length'), '10');
Expand Down

0 comments on commit 03e24eb

Please sign in to comment.