1
#include "memory.h"
2
 
3
#include <v8-array-buffer.h>
4
 
5
#include <kj/one-of.h>
6
 
7
namespace workerd::jsg {
8
 
9
class MemoryRetainerNode final: public v8::EmbedderGraph::Node {
10
 public:
11
  static constexpr auto PREFIX = "workerd /";
12
 
13
  const char* Name() override {
14
    return name_.cStr();
15
  }
16
 
17
  const char* NamePrefix() override {
18
    return PREFIX;
19
  }
20
 
21
  size_t SizeInBytes() override {
22
    return size_;
23
  }
24
 
25
  bool IsRootNode() override {
26
    KJ_SWITCH_ONEOF(isRootNode) {
27
      KJ_CASE_ONEOF(b, bool) {
28
        return b;
29
      }
30
      KJ_CASE_ONEOF(fn, kj::Function<bool()>) {
31
        return fn();
32
      }
33
    }
34
    KJ_UNREACHABLE;
35
  }
36
 
37
  v8::EmbedderGraph::Node::Detachedness GetDetachedness() override {
38
    return detachedness_;
39
  }
40
 
41
  inline Node* JSWrapperNode() {
42
    return wrapper_node_;
43
  }
44
 
45
  KJ_DISALLOW_COPY_AND_MOVE(MemoryRetainerNode);
46
  ~MemoryRetainerNode() noexcept(true) {}
47
 
48
 private:
49
  static v8::EmbedderGraph::Node::Detachedness fromDetachedState(MemoryInfoDetachedState state) {
50
    switch (state) {
51
      case MemoryInfoDetachedState::UNKNOWN:
52
        return v8::EmbedderGraph::Node::Detachedness::kUnknown;
53
      case MemoryInfoDetachedState::ATTACHED:
54
        return v8::EmbedderGraph::Node::Detachedness::kAttached;
55
      case MemoryInfoDetachedState::DETACHED:
56
        return v8::EmbedderGraph::Node::Detachedness::kDetached;
57
    }
58
    KJ_UNREACHABLE;
59
  }
60
 
61
  static v8::EmbedderGraph::Node* maybeWrapperNode(
62
      MemoryTracker* tracker, v8::Local<v8::Object> obj) {
63
    if (!obj.IsEmpty()) return tracker->graph_->V8Node(obj.As<v8::Value>());
64
    return nullptr;
65
  }
66
 
67
  MemoryRetainerNode(MemoryTracker* tracker,
68
      const void* retainer,
69
      const kj::StringPtr name,
70
      const size_t size,
71
      v8::Local<v8::Object> obj,
72
      kj::Maybe<kj::Function<bool()>> checkIsRootNode,
73
      MemoryInfoDetachedState detachedness)
74
      : name_(name),
75
        size_(size),
76
        wrapper_node_(maybeWrapperNode(tracker, obj)),
77
        detachedness_(fromDetachedState(detachedness)) {
78
    KJ_IF_SOME(fn, checkIsRootNode) {
79
      isRootNode = kj::mv(fn);
80
    }
81
  }
82
 
83
  MemoryRetainerNode(
84
      MemoryTracker* tracker, kj::StringPtr name, size_t size, bool isRootNode = false)
85
      : name_(name),
86
        size_(size),
87
        isRootNode(isRootNode) {}
88
 
89
  kj::StringPtr name_;
90
  size_t size_ = 0;
91
  v8::EmbedderGraph::Node* wrapper_node_ = nullptr;
92
 
93
  kj::OneOf<bool, kj::Function<bool()>> isRootNode = false;
94
 
95
  v8::EmbedderGraph::Node::Detachedness detachedness_ =
96
      v8::EmbedderGraph::Node::Detachedness::kUnknown;
97
 
98
  friend class MemoryTracker;
99
};
100
 
101
namespace {
102
kj::Maybe<MemoryRetainerNode&> getCurrentNode(const std::stack<MemoryRetainerNode*> stack) {
103
  if (stack.empty()) return kj::none;
104
  return *stack.top();
105
}
106
}  // namespace
107
 
108
MemoryTracker::MemoryTracker(v8::Isolate* isolate, v8::EmbedderGraph* graph)
109
    : isolate_(isolate),
110
      graph_(graph) {}
111
 
112
MemoryRetainerNode* MemoryTracker::addNode(const void* retainer,
113
    const kj::StringPtr name,
114
    const size_t size,
115
    v8::Local<v8::Object> obj,
116
    kj::Maybe<kj::Function<bool()>> checkIsRootNode,
117
    MemoryInfoDetachedState detachedness,
118
    kj::Maybe<kj::StringPtr> edgeName) {
119
  KJ_IF_SOME(found, seen_.find(retainer)) {
120
    return found;
121
  }
122
 
123
  MemoryRetainerNode* n = new MemoryRetainerNode(
124
      this, retainer, name, size, obj, kj::mv(checkIsRootNode), detachedness);
125
  graph_->AddNode(std::unique_ptr<v8::EmbedderGraph::Node>(n));
126
  seen_.insert(retainer, n);
127
 
128
  KJ_IF_SOME(currentNode, getCurrentNode(nodeStack_)) {
129
    KJ_IF_SOME(name, edgeName) {
130
      graph_->AddEdge(&currentNode, n, name.cStr());
131
    } else {
132
      graph_->AddEdge(&currentNode, n, nullptr);
133
    }
134
  }
135
 
136
  if (n->JSWrapperNode() != nullptr) {
137
    graph_->AddEdge(n, n->JSWrapperNode(), "native_to_javascript");
138
    graph_->AddEdge(n->JSWrapperNode(), n, "javascript_to_native");
139
  }
140
 
141
  return n;
142
}
143
 
144
MemoryRetainerNode* MemoryTracker::addNode(
145
    kj::StringPtr nodeName, size_t size, kj::Maybe<kj::StringPtr> edgeName) {
146
  MemoryRetainerNode* n = new MemoryRetainerNode(this, nodeName, size);
147
  graph_->AddNode(std::unique_ptr<v8::EmbedderGraph::Node>(n));
148
 
149
  KJ_IF_SOME(currentNode, getCurrentNode(nodeStack_)) {
150
    KJ_IF_SOME(name, edgeName) {
151
      graph_->AddEdge(&currentNode, n, name.cStr());
152
    } else {
153
      graph_->AddEdge(&currentNode, n, nullptr);
154
    }
155
  }
156
 
157
  return n;
158
}
159
 
160
MemoryRetainerNode* MemoryTracker::pushNode(
161
    kj::StringPtr nodeName, size_t size, kj::Maybe<kj::StringPtr> edgeName) {
162
  MemoryRetainerNode* n = addNode(nodeName, size, edgeName);
163
  nodeStack_.push(n);
164
  return n;
165
}
166
 
167
void MemoryTracker::decCurrentNodeSize(size_t size) {
168
  KJ_IF_SOME(currentNode, getCurrentNode(nodeStack_)) {
169
    currentNode.size_ -= size;
170
  }
171
}
172
 
173
void MemoryTracker::addEdge(v8::EmbedderGraph::Node* node, kj::StringPtr edgeName) {
174
  KJ_IF_SOME(currentNode, getCurrentNode(nodeStack_)) {
175
    graph_->AddEdge(&currentNode, node, edgeName.cStr());
176
  } else {
177
    graph_->AddEdge(nullptr, node, edgeName.cStr());
178
  }
179
}
180
 
181
void MemoryTracker::addEdge(MemoryRetainerNode* node, kj::StringPtr edgeName) {
182
  KJ_IF_SOME(currentNode, getCurrentNode(nodeStack_)) {
183
    graph_->AddEdge(&currentNode, node, edgeName.cStr());
184
  } else {
185
    graph_->AddEdge(nullptr, node, edgeName.cStr());
186
  }
187
}
188
 
189
// ======================================================================================
190
 
191
HeapSnapshotActivity::HeapSnapshotActivity(Callback callback): callback(kj::mv(callback)) {}
192
 
193
v8::ActivityControl::ControlOption HeapSnapshotActivity::ReportProgressValue(
194
    uint32_t done, uint32_t total) {
195
  return callback(done, total) ? ControlOption::kContinue : ControlOption::kAbort;
196
}
197
 
198
HeapSnapshotWriter::HeapSnapshotWriter(Callback callback, size_t chunkSize)
199
    : callback(kj::mv(callback)),
200
      chunkSize(chunkSize) {}
201
 
202
void HeapSnapshotWriter::EndOfStream() {
203
  callback(kj::none);
204
}
205
 
206
int HeapSnapshotWriter::GetChunkSize() {
207
  return chunkSize;
208
}
209
 
210
v8::OutputStream::WriteResult HeapSnapshotWriter::WriteAsciiChunk(char* data, int size) {
211
  return callback(kj::ArrayPtr<char>(data, size)) ? v8::OutputStream::WriteResult::kContinue
212
                                                  : v8::OutputStream::WriteResult::kAbort;
213
}
214
 
215
void MemoryTracker::trackField(
216
    kj::StringPtr edgeName, const v8::BackingStore* value, kj::Maybe<kj::StringPtr> nodeName) {
217
  trackFieldWithSize(edgeName, value->ByteLength(), "BackingStore"_kjc);
218
}
219
 
220
}  // namespace workerd::jsg