Skip to content

Commit 5e78eaa

Browse files
jasnelljasnell
committed
Use WeakRef for tracking locked streams readers and writers
1 parent 8ed21a8 commit 5e78eaa

File tree

9 files changed

+182
-170
lines changed

9 files changed

+182
-170
lines changed

‎src/workerd/api/streams/common.h

+55-25
Original file line numberDiff line numberDiff line change
@@ -318,16 +318,22 @@ class ReadableStreamController {
318318
// passing along the closed promise that will be used to communicate state to the
319319
// user code.
320320
//
321-
// The Reader will hold a reference to the controller that will be cleared when the reader
322-
// is released or destroyed. The controller is guaranteed to either outlive or detach the
323-
// reader so the ReadableStreamController& reference should remain valid.
321+
// The Reader holds a strong reference to the controller. The Controller will hold a weak
322+
// reference to the reader. It is ok for the reader itself to be freed/garbage collected
323+
// while still being attached to the controller, but not the other way around.
324324
virtual void attach(
325325
ReadableStreamController& controller,
326326
jsg::Promise<void> closedPromise) = 0;
327327

328328
// When a Reader lock is released, the controller will signal to the reader that it has been
329329
// detached.
330330
virtual void detach() = 0;
331+
332+
virtual kj::Own<WeakRef<Reader>> addWeakRef() = 0;
333+
334+
private:
335+
static kj::Badge<Reader> getBadge() { return kj::Badge<Reader>(); }
336+
friend class ReaderImpl;
331337
};
332338

333339
struct ByobOptions {
@@ -479,12 +485,12 @@ class ReadableStreamController {
479485

480486
// Locks this controller to the given reader, returning true if the lock was successful, or false
481487
// if the controller was already locked.
482-
virtual bool lockReader(jsg::Lock& js, Reader& reader) = 0;
488+
virtual bool lockReader(jsg::Lock& js, kj::Own<WeakRef<Reader>> reader) = 0;
483489

484490
// Removes the lock and releases the reader from this controller.
485491
// maybeJs will be nullptr when the isolate lock is not available.
486492
// If maybeJs is set, the reader's closed promise will be resolved.
487-
virtual void releaseReader(Reader& reader, kj::Maybe<jsg::Lock&> maybeJs) = 0;
493+
virtual void releaseReader(jsg::Lock& js, Reader& reader) = 0;
488494

489495
virtual kj::Maybe<PipeController&> tryPipeLock(jsg::Ref<WritableStream> destination) = 0;
490496

@@ -574,8 +580,9 @@ class WritableStreamController {
574580
// passing along the closed and ready promises that will be used to communicate state to the
575581
// user code.
576582
//
577-
// The controller is guaranteed to either outlive the Writer or will detach the Writer so the
578-
// WritableStreamController& reference should always remain valid.
583+
// The Writer holds a strong reference to the controller. The Controller will hold a weak
584+
// reference to the writer. It is ok for the writer itself to be freed/garbage collected
585+
// while still being attached to the controller, but not the other way around.
579586
virtual void attach(
580587
WritableStreamController& controller,
581588
jsg::Promise<void> closedPromise,
@@ -588,6 +595,12 @@ class WritableStreamController {
588595
// The ready promise can be replaced whenever backpressure is signaled by the underlying
589596
// controller.
590597
virtual void replaceReadyPromise(jsg::Promise<void> readyPromise) = 0;
598+
599+
virtual kj::Own<WeakRef<Writer>> addWeakRef() = 0;
600+
601+
private:
602+
static kj::Badge<Writer> getBadge() { return kj::Badge<Writer>(); }
603+
friend class WritableStreamDefaultWriter;
591604
};
592605

593606
struct PendingAbort {
@@ -672,12 +685,10 @@ class WritableStreamController {
672685

673686
// Locks this controller to the given writer, returning true if the lock was successful, or false
674687
// if the controller was already locked.
675-
virtual bool lockWriter(jsg::Lock& js, Writer& writer) = 0;
688+
virtual bool lockWriter(jsg::Lock& js, kj::Own<WeakRef<Writer>>) = 0;
676689

677690
// Removes the lock and releases the writer from this controller.
678-
// maybeJs will be nullptr when the isolate lock is not available.
679-
// If maybeJs is set, the writer's closed and ready promises will be resolved.
680-
virtual void releaseWriter(Writer& writer, kj::Maybe<jsg::Lock&> maybeJs) = 0;
691+
virtual void releaseWriter(jsg::Lock& js, Writer& writer) = 0;
681692

682693
virtual kj::Maybe<v8::Local<v8::Value>> isErroring(jsg::Lock& js) = 0;
683694

@@ -710,28 +721,37 @@ struct Locked {};
710721

711722
// When a reader is locked to a ReadableStream, a ReaderLock instance
712723
// is used internally to represent the locked state in the ReadableStreamController.
724+
// ReaderLocked maintains a weak referene to the actual Reader instance. It's ok
725+
// for the Reader to be garbage collected while the ReadableStream is still alive but
726+
// not vis versa. The Reader holds a strong reference to the ReadableStream only while
727+
// it is attached.
713728
class ReaderLocked {
714729
public:
715730
ReaderLocked(
716-
ReadableStreamController::Reader& reader,
731+
kj::Own<WeakRef<ReadableStreamController::Reader>> reader,
717732
jsg::Promise<void>::Resolver closedFulfiller,
718733
kj::Maybe<IoOwn<kj::Canceler>> canceler = kj::none)
719-
: reader(reader),
734+
: reader(kj::mv(reader)),
720735
closedFulfiller(kj::mv(closedFulfiller)),
721736
canceler(kj::mv(canceler)) {}
722737

723738
ReaderLocked(ReaderLocked&&) = default;
724739
~ReaderLocked() noexcept(false) {
725-
KJ_IF_SOME(r, reader) { r.detach(); }
740+
if (reader.get() != nullptr) {
741+
reader->runIfAlive([](auto& r) {
742+
r.detach();
743+
});
744+
}
726745
}
727746
KJ_DISALLOW_COPY(ReaderLocked);
728747

729748
void visitForGc(jsg::GcVisitor& visitor) {
730749
visitor.visit(closedFulfiller);
731750
}
732751

733-
ReadableStreamController::Reader& getReader() {
734-
return KJ_ASSERT_NONNULL(reader);
752+
kj::Maybe<ReadableStreamController::Reader&> getReader() {
753+
if (reader.get() == nullptr) return kj::none;
754+
return reader->tryGet();
735755
}
736756

737757
kj::Maybe<jsg::Promise<void>::Resolver>& getClosedFulfiller() {
@@ -743,34 +763,43 @@ class ReaderLocked {
743763
}
744764

745765
private:
746-
kj::Maybe<ReadableStreamController::Reader&> reader;
766+
kj::Own<WeakRef<ReadableStreamController::Reader>> reader;
747767
kj::Maybe<jsg::Promise<void>::Resolver> closedFulfiller;
748768
kj::Maybe<IoOwn<kj::Canceler>> canceler;
749769
};
750770

751771
// When a writer is locked to a WritableStream, a WriterLock instance
752772
// is used internally to represent the locked state in the WritableStreamController.
773+
// WriterLocked maintains a weak reference to the actual Writer instance. It's ok
774+
// for the Writer to be garbage collected while the WritableStream is still alive but
775+
// not vis versa. The Writer holds a strong reference to the WritableStream only while
776+
// it is attached.
753777
class WriterLocked {
754778
public:
755779
WriterLocked(
756-
WritableStreamController::Writer& writer,
780+
kj::Own<WeakRef<WritableStreamController::Writer>> writer,
757781
jsg::Promise<void>::Resolver closedFulfiller,
758782
kj::Maybe<jsg::Promise<void>::Resolver> readyFulfiller = kj::none)
759-
: writer(writer),
783+
: writer(kj::mv(writer)),
760784
closedFulfiller(kj::mv(closedFulfiller)),
761785
readyFulfiller(kj::mv(readyFulfiller)) {}
762786

763787
WriterLocked(WriterLocked&&) = default;
764788
~WriterLocked() noexcept(false) {
765-
KJ_IF_SOME(w, writer) { w.detach(); }
789+
if (writer.get() != nullptr) {
790+
writer->runIfAlive([&](auto& w) {
791+
w.detach();
792+
});
793+
}
766794
}
767795

768796
void visitForGc(jsg::GcVisitor& visitor) {
769797
visitor.visit(closedFulfiller, readyFulfiller);
770798
}
771799

772-
WritableStreamController::Writer& getWriter() {
773-
return KJ_ASSERT_NONNULL(writer);
800+
kj::Maybe<WritableStreamController::Writer&> getWriter() {
801+
if (writer.get() == nullptr) return kj::none;
802+
return writer->tryGet();
774803
}
775804

776805
kj::Maybe<jsg::Promise<void>::Resolver>& getClosedFulfiller() {
@@ -782,14 +811,15 @@ class WriterLocked {
782811
}
783812

784813
void setReadyFulfiller(jsg::PromiseResolverPair<void>& pair) {
785-
KJ_IF_SOME(w, writer) {
814+
if (writer.get() == nullptr) return;
815+
writer->runIfAlive([&](auto& w) {
786816
readyFulfiller = kj::mv(pair.resolver);
787817
w.replaceReadyPromise(kj::mv(pair.promise));
788-
}
818+
});
789819
}
790820

791821
private:
792-
kj::Maybe<WritableStreamController::Writer&> writer;
822+
kj::Own<WeakRef<WritableStreamController::Writer>> writer;
793823
kj::Maybe<jsg::Promise<void>::Resolver> closedFulfiller;
794824
kj::Maybe<jsg::Promise<void>::Resolver> readyFulfiller;
795825
};

‎src/workerd/api/streams/internal.c++

+35-54
Original file line numberDiff line numberDiff line change
@@ -471,12 +471,7 @@ kj::Maybe<kj::Promise<DeferredProxy<void>>> WritableStreamSink::tryPumpFrom(
471471

472472
// =======================================================================================
473473

474-
ReadableStreamInternalController::~ReadableStreamInternalController() noexcept(false) {
475-
KJ_IF_SOME(locked, readState.tryGet<ReaderLocked>()) {
476-
auto lock = kj::mv(locked);
477-
readState.init<Unlocked>();
478-
}
479-
}
474+
ReadableStreamInternalController::~ReadableStreamInternalController() noexcept(false) {}
480475

481476
jsg::Ref<ReadableStream> ReadableStreamInternalController::addRef() {
482477
return KJ_ASSERT_NONNULL(owner).addRef();
@@ -761,16 +756,17 @@ kj::Maybe<kj::Own<ReadableStreamSource>> ReadableStreamInternalController::remov
761756
KJ_UNREACHABLE;
762757
}
763758

764-
bool ReadableStreamInternalController::lockReader(jsg::Lock& js, Reader& reader) {
759+
bool ReadableStreamInternalController::lockReader(jsg::Lock& js, kj::Own<WeakRef<Reader>> reader) {
765760
if (isLockedToReader()) {
766761
return false;
767762
}
768763

769764
auto prp = js.newPromiseAndResolver<void>();
770765
prp.promise.markAsHandled(js);
771766

772-
auto lock = ReaderLocked(reader, kj::mv(prp.resolver),
767+
auto lock = ReaderLocked(kj::mv(reader), kj::mv(prp.resolver),
773768
IoContext::current().addObject(kj::heap<kj::Canceler>()));
769+
// Take care not to access reader directly after this point. Use the lock.
774770

775771
KJ_SWITCH_ONEOF(state) {
776772
KJ_CASE_ONEOF(closed, StreamStates::Closed) {
@@ -785,42 +781,30 @@ bool ReadableStreamInternalController::lockReader(jsg::Lock& js, Reader& reader)
785781
}
786782

787783
readState = kj::mv(lock);
788-
reader.attach(*this, kj::mv(prp.promise));
784+
785+
auto& inner = KJ_ASSERT_NONNULL(readState.get<ReaderLocked>().getReader());
786+
inner.attach(*this, kj::mv(prp.promise));
789787
return true;
790788
}
791789

792-
void ReadableStreamInternalController::releaseReader(
793-
Reader& reader,
794-
kj::Maybe<jsg::Lock&> maybeJs) {
790+
void ReadableStreamInternalController::releaseReader(jsg::Lock& js, Reader& reader) {
795791
KJ_IF_SOME(locked, readState.tryGet<ReaderLocked>()) {
796-
KJ_ASSERT(&locked.getReader() == &reader);
797-
KJ_IF_SOME(js, maybeJs) {
798-
JSG_REQUIRE(KJ_ASSERT_NONNULL(locked.getCanceler())->isEmpty(), TypeError,
799-
"Cannot call releaseLock() on a reader with outstanding read promises.");
800-
maybeRejectPromise<void>(js,
801-
locked.getClosedFulfiller(),
802-
js.v8TypeError("This ReadableStream reader has been released."_kj));
792+
KJ_IF_SOME(r, locked.getReader()) {
793+
KJ_ASSERT(&r == &reader);
803794
}
804-
auto lock = kj::mv(locked);
805-
806-
// When maybeJs is nullptr, that means releaseReader was called when the reader is
807-
// being deconstructed and not as the result of explicitly calling releaseLock. In
808-
// that case, we don't want to change the lock state itself because we do not have
809-
// an isolate lock. Moving the lock above will free the lock state while keeping the
810-
// ReadableStream marked as locked.
811-
if (maybeJs != kj::none) {
812-
readState.template init<Unlocked>();
813-
}
814-
}
815-
}
795+
JSG_REQUIRE(KJ_ASSERT_NONNULL(locked.getCanceler())->isEmpty(), TypeError,
796+
"Cannot call releaseLock() on a reader with outstanding read promises.");
797+
maybeRejectPromise<void>(js,
798+
locked.getClosedFulfiller(),
799+
js.v8TypeError("This ReadableStream reader has been released."_kj));
816800

817-
WritableStreamInternalController::~WritableStreamInternalController() noexcept(false) {
818-
KJ_IF_SOME(locked, writeState.tryGet<WriterLocked>()) {
819801
auto lock = kj::mv(locked);
820-
writeState.init<Unlocked>();
802+
readState.template init<Unlocked>();
821803
}
822804
}
823805

806+
WritableStreamInternalController::~WritableStreamInternalController() noexcept(false) {}
807+
824808
jsg::Ref<WritableStream> WritableStreamInternalController::addRef() {
825809
return KJ_ASSERT_NONNULL(owner).addRef();
826810
}
@@ -1246,7 +1230,7 @@ kj::Maybe<int> WritableStreamInternalController::getDesiredSize() {
12461230
KJ_UNREACHABLE;
12471231
}
12481232

1249-
bool WritableStreamInternalController::lockWriter(jsg::Lock& js, Writer& writer) {
1233+
bool WritableStreamInternalController::lockWriter(jsg::Lock& js, kj::Own<WeakRef<Writer>> writer) {
12501234
if (isLockedToWriter()) {
12511235
return false;
12521236
}
@@ -1257,7 +1241,8 @@ bool WritableStreamInternalController::lockWriter(jsg::Lock& js, Writer& writer)
12571241
auto readyPrp = js.newPromiseAndResolver<void>();
12581242
readyPrp.promise.markAsHandled(js);
12591243

1260-
auto lock = WriterLocked(writer, kj::mv(closedPrp.resolver), kj::mv(readyPrp.resolver));
1244+
auto lock = WriterLocked(kj::mv(writer), kj::mv(closedPrp.resolver), kj::mv(readyPrp.resolver));
1245+
// Careful not to access writer directly after this point. Access is through the lock.
12611246

12621247
KJ_SWITCH_ONEOF(state) {
12631248
KJ_CASE_ONEOF(closed, StreamStates::Closed) {
@@ -1274,30 +1259,26 @@ bool WritableStreamInternalController::lockWriter(jsg::Lock& js, Writer& writer)
12741259
}
12751260

12761261
writeState = kj::mv(lock);
1277-
writer.attach(*this, kj::mv(closedPrp.promise), kj::mv(readyPrp.promise));
1262+
1263+
auto& inner = KJ_ASSERT_NONNULL(writeState.get<WriterLocked>().getWriter());
1264+
inner.attach(*this, kj::mv(closedPrp.promise), kj::mv(readyPrp.promise));
1265+
12781266
return true;
12791267
}
12801268

1281-
void WritableStreamInternalController::releaseWriter(
1282-
Writer& writer,
1283-
kj::Maybe<jsg::Lock&> maybeJs) {
1269+
void WritableStreamInternalController::releaseWriter(jsg::Lock& js, Writer& writer) {
12841270
KJ_IF_SOME(locked, writeState.tryGet<WriterLocked>()) {
1285-
KJ_ASSERT(&locked.getWriter() == &writer);
1286-
KJ_IF_SOME(js, maybeJs) {
1287-
maybeRejectPromise<void>(js,
1288-
locked.getClosedFulfiller(),
1289-
js.v8TypeError("This WritableStream writer has been released."_kj));
1271+
KJ_IF_SOME(w, locked.getWriter()) {
1272+
// Just an extra verification.
1273+
KJ_ASSERT(&w == &writer);
12901274
}
1291-
auto lock = kj::mv(locked);
12921275

1293-
// When maybeJs is nullptr, that means releaseWriter was called when the writer is
1294-
// being deconstructed and not as the result of explicitly calling releaseLock and
1295-
// we do not have an isolate lock. In that case, we don't want to change the lock
1296-
// state itself. Moving the lock above will free the lock state while keeping the
1297-
// WritableStream marked as locked.
1298-
if (maybeJs != kj::none) {
1299-
writeState.template init<Unlocked>();
1300-
}
1276+
maybeRejectPromise<void>(js,
1277+
locked.getClosedFulfiller(),
1278+
js.v8TypeError("This WritableStream writer has been released."_kj));
1279+
1280+
auto lock = kj::mv(locked);
1281+
writeState.template init<Unlocked>();
13011282
}
13021283
}
13031284

‎src/workerd/api/streams/internal.h

+4-5
Original file line numberDiff line numberDiff line change
@@ -81,10 +81,9 @@ class ReadableStreamInternalController: public ReadableStreamController {
8181

8282
bool isLockedToReader() const override { return !readState.is<Unlocked>(); }
8383

84-
bool lockReader(jsg::Lock& js, Reader& reader) override;
84+
bool lockReader(jsg::Lock& js, kj::Own<WeakRef<Reader>> reader) override;
8585

86-
void releaseReader(Reader& reader, kj::Maybe<jsg::Lock&> maybeJs) override;
87-
// See the comment for releaseReader in common.h for details on the use of maybeJs
86+
void releaseReader(jsg::Lock& js, Reader& reader) override;
8887

8988
kj::Maybe<PipeController&> tryPipeLock(jsg::Ref<WritableStream> destination) override;
9089

@@ -201,9 +200,9 @@ class WritableStreamInternalController: public WritableStreamController {
201200

202201
bool isLockedToWriter() const override { return !writeState.is<Unlocked>(); }
203202

204-
bool lockWriter(jsg::Lock& js, Writer& writer) override;
203+
bool lockWriter(jsg::Lock& js, kj::Own<WeakRef<Writer>> writer) override;
205204

206-
void releaseWriter(Writer& writer, kj::Maybe<jsg::Lock&> maybeJs) override;
205+
void releaseWriter(jsg::Lock& js, Writer& writer) override;
207206
// See the comment for releaseWriter in common.h for details on the use of maybeJs
208207

209208
kj::Maybe<v8::Local<v8::Value>> isErroring(jsg::Lock& js) override {

0 commit comments

Comments
 (0)