Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement memory tracking for Blob/File/Streams #1608

Merged
merged 4 commits into from
Feb 6, 2024
Merged

Implement memory tracking for Blob/File/Streams #1608

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
515d0df
Implement memory tracking for Blob/File
jasnell Feb 1, 2024
55a4e5b
Implement heap snapshot memory tracking for WritableStream
jasnell Feb 1, 2024
add1018
Implement heap snapshot memory tracking for ReadableStream
jasnell Feb 2, 2024
13de7bd
Implement heap snapshot memory tracking for TransformStream
jasnell Feb 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions src/workerd/api/blob.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,18 @@ class Blob: public jsg::Object {
JSG_METHOD(stream);
}

void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
KJ_SWITCH_ONEOF(ownData) {
KJ_CASE_ONEOF(data, kj::Array<byte>) {
tracker.trackField("ownData", data);
}
KJ_CASE_ONEOF(data, jsg::Ref<Blob>) {
tracker.trackField("ownData", data);
}
}
Comment on lines +64 to +71
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be convenient if there were a way to call trackField with a kj::OneOf<S, T, ...> where there is an impl for trackField for each of them. Perhaps it's not worth the hassle to save a little bit of boilerplate.

But maybe you could make a concept Trackable with T::trackField(tracker, name, value) and then define tracker.trackField<Trackable T> to call T::trackField. Then maybe you can define OneOf<Trackable S, Trackable T, Trackable U>::trackField or something?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd considered adding that but I think I prefer separating it out like this in order to give a bit more control over how things are tracked. In some cases, it won't need to be tracked at all, in other cases it make use trackField(...) for one and trackFieldWithSize(...) for others.

tracker.trackField("type", type);
}

private:
kj::OneOf<kj::Array<byte>, jsg::Ref<Blob>> ownData;
kj::ArrayPtr<const byte> data;
Expand Down Expand Up @@ -106,6 +118,10 @@ class File: public Blob {
}
}

void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
tracker.trackField("name", name);
}

private:
kj::String name;
double lastModified;
Expand Down
26 changes: 26 additions & 0 deletions src/workerd/api/streams/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -523,6 +523,10 @@ class ReadableStreamController {
// Used by sockets to signal that the ReadableStream shouldn't allow reads due to pending
// closure.
virtual void setPendingClosure() = 0;

virtual kj::StringPtr jsgGetMemoryName() const = 0;
virtual size_t jsgGetMemorySelfSize() const = 0;
virtual void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const = 0;
};

kj::Own<ReadableStreamController> newReadableStreamJsController();
Expand Down Expand Up @@ -624,6 +628,12 @@ class WritableStreamController {
}

static kj::Maybe<PendingAbort> dequeue(kj::Maybe<PendingAbort>& maybePendingAbort);

JSG_MEMORY_INFO(PendingAbort) {
tracker.trackField("resolver", resolver);
tracker.trackField("promise", promise);
tracker.trackField("reason", reason);
}
};

virtual ~WritableStreamController() noexcept(false) {}
Expand Down Expand Up @@ -696,6 +706,11 @@ class WritableStreamController {
// Used by sockets to signal that the WritableStream shouldn't allow writes due to pending
// closure.
virtual void setPendingClosure() = 0;

// For menmory tracking
virtual kj::StringPtr jsgGetMemoryName() const = 0;
virtual size_t jsgGetMemorySelfSize() const = 0;
virtual void jsgGetMemoryInfo(jsg::MemoryTracker& info) const = 0;
};

kj::Own<WritableStreamController> newWritableStreamJsController();
Expand Down Expand Up @@ -748,6 +763,12 @@ class ReaderLocked {
canceler = kj::none;
}

JSG_MEMORY_INFO(ReaderLocked) {
tracker.trackField("closedFulfiller", closedFulfiller);
tracker.trackFieldWithSize("IoOwn<kj::Canceler>",
sizeof(IoOwn<kj::Canceler>));
}

private:
kj::Maybe<ReadableStreamController::Reader&> reader;
kj::Maybe<jsg::Promise<void>::Resolver> closedFulfiller;
Expand Down Expand Up @@ -800,6 +821,11 @@ class WriterLocked {
readyFulfiller = kj::none;
}

JSG_MEMORY_INFO(WriterLocked) {
tracker.trackField("closedFulfiller", closedFulfiller);
tracker.trackField("readyFulfiller", readyFulfiller);
}

private:
kj::Maybe<WritableStreamController::Writer&> writer;
kj::Maybe<jsg::Promise<void>::Resolver> closedFulfiller;
Expand Down
78 changes: 78 additions & 0 deletions src/workerd/api/streams/internal.c++
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2287,4 +2287,82 @@ kj::Own<WritableStreamController> newWritableStreamInternalController(
kj::mv(maybeClosureWaitable));
}

kj::StringPtr WritableStreamInternalController::jsgGetMemoryName() const {
return "WritableStreamInternalController"_kjc;
}

size_t WritableStreamInternalController::jsgGetMemorySelfSize() const {
return sizeof(WritableStreamInternalController);
}
void WritableStreamInternalController::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
KJ_SWITCH_ONEOF(state) {
KJ_CASE_ONEOF(closed, StreamStates::Closed) {}
KJ_CASE_ONEOF(errored, StreamStates::Errored) {
tracker.trackField("error", errored);
}
KJ_CASE_ONEOF(_, Writable) {
// Ideally we'd be able to track the size of any pending writes held in the sink's
// queue but since it is behind an IoOwn and we won't be holding the IoContext here,
// we can't.
tracker.trackFieldWithSize("IoOwn<WritableStreamSink>",
sizeof(IoOwn<WritableStreamSink>));
}
}
KJ_IF_SOME(writerLocked, writeState.tryGet<WriterLocked>()) {
tracker.trackField("writerLocked", writerLocked);
}
tracker.trackField("pendingAbort", maybePendingAbort);
tracker.trackField("maybeClosureWaitable", maybeClosureWaitable);

for (auto& event : queue) {
tracker.trackField("event", event);
}
}

kj::StringPtr ReadableStreamInternalController::PipeLocked::jsgGetMemoryName() const {
return "ReadableStreamInternalController::PipeLocked"_kjc;
}
size_t ReadableStreamInternalController::PipeLocked::jsgGetMemorySelfSize() const {
return sizeof(PipeLocked);
}

void ReadableStreamInternalController::PipeLocked::jsgGetMemoryInfo(
jsg::MemoryTracker& tracker) const {
tracker.trackField("ref", ref);
}

kj::StringPtr ReadableStreamInternalController::jsgGetMemoryName() const {
return "ReadableStreamInternalController"_kjc;
}

size_t ReadableStreamInternalController::jsgGetMemorySelfSize() const {
return sizeof(ReadableStreamInternalController);
}

void ReadableStreamInternalController::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
KJ_SWITCH_ONEOF(state) {
KJ_CASE_ONEOF(closed, StreamStates::Closed) {}
KJ_CASE_ONEOF(error, StreamStates::Errored) {
tracker.trackField("error", error);
}
KJ_CASE_ONEOF(readable, Readable) {
// Ideally we'd be able to track the size of any pending reads held in the source's
// queue but since it is behind an IoOwn and we won't be holding the IoContext here,
// we can't.
tracker.trackFieldWithSize("IoOwn<ReadableStreamSource>",
sizeof(IoOwn<ReadableStreamSource>));
}
}
KJ_SWITCH_ONEOF(readState) {
KJ_CASE_ONEOF(unlocked, Unlocked) {}
KJ_CASE_ONEOF(locked, Locked) {}
KJ_CASE_ONEOF(pipeLocked, PipeLocked) {
tracker.trackField("pipeLocked", pipeLocked);
}
KJ_CASE_ONEOF(readerLocked, ReaderLocked) {
tracker.trackField("readerLocked", readerLocked);
}
}
}

} // namespace workerd::api
43 changes: 43 additions & 0 deletions src/workerd/api/streams/internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,10 @@ class ReadableStreamInternalController: public ReadableStreamController {
isPendingClosure = true;
}

kj::StringPtr jsgGetMemoryName() const override;
size_t jsgGetMemorySelfSize() const override;
void jsgGetMemoryInfo(jsg::MemoryTracker& info) const override;

private:
void doCancel(jsg::Lock& js, jsg::Optional<v8::Local<v8::Value>> reason);
void doClose(jsg::Lock& js);
Expand Down Expand Up @@ -136,6 +140,10 @@ class ReadableStreamInternalController: public ReadableStreamController {

void visitForGc(jsg::GcVisitor& visitor) { visitor.visit(ref); }

kj::StringPtr jsgGetMemoryName() const;
size_t jsgGetMemorySelfSize() const;
void jsgGetMemoryInfo(jsg::MemoryTracker& info) const;

private:
ReadableStreamInternalController& inner;
jsg::Ref<WritableStream> ref;
Expand Down Expand Up @@ -224,6 +232,10 @@ class WritableStreamInternalController: public WritableStreamController {
void setPendingClosure() override {
isPendingClosure = true;
}

kj::StringPtr jsgGetMemoryName() const override;
size_t jsgGetMemorySelfSize() const override;
void jsgGetMemoryInfo(jsg::MemoryTracker& info) const override;
private:

struct AbortOptions {
Expand Down Expand Up @@ -283,12 +295,25 @@ class WritableStreamInternalController: public WritableStreamController {
kj::Maybe<jsg::Promise<void>::Resolver> promise;
std::shared_ptr<v8::BackingStore> ownBytes;
kj::ArrayPtr<const kj::byte> bytes;

JSG_MEMORY_INFO(Write) {
tracker.trackField("resolver", promise);
if (ownBytes != nullptr) {
tracker.trackFieldWithSize("backing", ownBytes->ByteLength());
}
}
};
struct Close {
kj::Maybe<jsg::Promise<void>::Resolver> promise;
JSG_MEMORY_INFO(Close) {
tracker.trackField("promise", promise);
}
};
struct Flush {
kj::Maybe<jsg::Promise<void>::Resolver> promise;
JSG_MEMORY_INFO(Flush) {
tracker.trackField("promise", promise);
}
};
struct Pipe {
WritableStreamInternalController& parent;
Expand All @@ -302,10 +327,28 @@ class WritableStreamInternalController: public WritableStreamController {
bool checkSignal(jsg::Lock& js);
jsg::Promise<void> pipeLoop(jsg::Lock& js);
jsg::Promise<void> write(v8::Local<v8::Value> value);

JSG_MEMORY_INFO(Pipe) {
tracker.trackField("resolver", promise);
tracker.trackField("signal", maybeSignal);
}
};
struct WriteEvent {
kj::Maybe<IoOwn<kj::Promise<void>>> outputLock; // must wait for this before actually writing
kj::OneOf<Write, Pipe, Close, Flush> event;

JSG_MEMORY_INFO(WriteEvent) {
if (outputLock != kj::none) {
tracker.trackFieldWithSize("outputLock",
sizeof(IoOwn<kj::Promise<void>>));
}
KJ_SWITCH_ONEOF(event) {
KJ_CASE_ONEOF(w, Write) { tracker.trackField("inner", w); }
KJ_CASE_ONEOF(p, Pipe) { tracker.trackField("inner", p); }
KJ_CASE_ONEOF(c, Close) { tracker.trackField("inner", c); }
KJ_CASE_ONEOF(f, Flush) { tracker.trackField("inner", f); }
}
}
};

std::deque<WriteEvent> queue;
Expand Down
Loading
Loading