src: support WeakReference in snapshot

Move util::WeakReference to a separate header and implement {de}serialization for it to be snapshotable. PR-URL: #44193 Refs: #44014 Refs: #37476 Reviewed-By: Matteo Collina <matteo.collina@gmail.com> Reviewed-By: Chengzhong Wu <legendecas@gmail.com>
nodejs · joyeecheung · May 2, 2022 · Jul 11, 2022 · Jun 21, 2022 · Aug 9, 2022
commit 4670f7cd1f48d9e57a449c4b466d99f6d1691e9c
diff --git a/node.gyp b/node.gyp
@@ -643,6 +643,7 @@
         'src/node_stat_watcher.h',
         'src/node_union_bytes.h',
         'src/node_url.h',
+        'src/node_util.h',
         'src/node_version.h',
         'src/node_v8.h',
         'src/node_v8_platform-inl.h',

diff --git a/src/node_snapshotable.cc b/src/node_snapshotable.cc
@@ -16,6 +16,7 @@
 #include "node_metadata.h"
 #include "node_process.h"
 #include "node_snapshot_builder.h"
+#include "node_util.h"
 #include "node_v8.h"
 #include "node_v8_platform-inl.h"
 

diff --git a/src/node_snapshotable.h b/src/node_snapshotable.h
@@ -18,7 +18,8 @@ class ExternalReferenceRegistry;
   V(fs_binding_data, fs::BindingData)                                          \
   V(v8_binding_data, v8_utils::BindingData)                                    \
   V(blob_binding_data, BlobBindingData)                                        \
-  V(process_binding_data, process::BindingData)
+  V(process_binding_data, process::BindingData)                                \
+  V(util_weak_reference, util::WeakReference)
 
 enum class EmbedderObjectType : uint8_t {
 #define V(PropertyName, NativeType) k_##PropertyName,

diff --git a/src/node_util.cc b/src/node_util.cc
@@ -1,3 +1,4 @@
+#include "node_util.h"
 #include "base_object-inl.h"
 #include "node_errors.h"
 #include "node_external_reference.h"
@@ -15,7 +16,7 @@ using v8::Context;
 using v8::External;
 using v8::FunctionCallbackInfo;
 using v8::FunctionTemplate;
-using v8::Global;
+using v8::HandleScope;
 using v8::IndexFilter;
 using v8::Integer;
 using v8::Isolate;
@@ -207,52 +208,106 @@ void ArrayBufferViewHasBuffer(const FunctionCallbackInfo<Value>& args) {
   args.GetReturnValue().Set(args[0].As<ArrayBufferView>()->HasBuffer());
 }
 
-class WeakReference : public BaseObject {
- public:
-  WeakReference(Environment* env, Local<Object> object, Local<Object> target)
-    : BaseObject(env, object) {
-    MakeWeak();
+WeakReference::WeakReference(Environment* env,
+                             Local<Object> object,
+                             Local<Object> target)
+    : WeakReference(env, object, target, 0) {}
+
+WeakReference::WeakReference(Environment* env,
+                             Local<Object> object,
+                             Local<Object> target,
+                             uint64_t reference_count)
+    : SnapshotableObject(env, object, type_int),
+      reference_count_(reference_count) {
+  MakeWeak();
+  if (!target.IsEmpty()) {
     target_.Reset(env->isolate(), target);
-    target_.SetWeak();
+    if (reference_count_ == 0) {
+      target_.SetWeak();
+    }
   }
+}
 
-  static void New(const FunctionCallbackInfo<Value>& args) {
-    Environment* env = Environment::GetCurrent(args);
-    CHECK(args.IsConstructCall());
-    CHECK(args[0]->IsObject());
-    new WeakReference(env, args.This(), args[0].As<Object>());
+bool WeakReference::PrepareForSerialization(Local<Context> context,
+                                            v8::SnapshotCreator* creator) {
+  if (target_.IsEmpty()) {
+    target_index_ = 0;
+    return true;
   }
 
-  static void Get(const FunctionCallbackInfo<Value>& args) {
-    WeakReference* weak_ref = Unwrap<WeakReference>(args.Holder());
-    Isolate* isolate = args.GetIsolate();
-    if (!weak_ref->target_.IsEmpty())
-      args.GetReturnValue().Set(weak_ref->target_.Get(isolate));
-  }
+  // Users can still hold strong references to target in addition to the
+  // reference that we manage here, and they could expect that the referenced
+  // object remains the same as long as that external strong reference
+  // is alive. Since we have no way to know if there is any other reference
+  // keeping the target alive, the best we can do to maintain consistency is to
+  // simply save a reference to the target in the snapshot (effectively making
+  // it strong) during serialization, and restore it during deserialization.
+  // If there's no known counted reference from our side, we'll make the
+  // reference here weak upon deserialization so that it can be GC'ed if users
+  // do not hold additional references to it.
+  Local<Object> target = target_.Get(context->GetIsolate());
+  target_index_ = creator->AddData(context, target);
+  DCHECK_NE(target_index_, 0);
+  target_.Reset();
+  return true;
+}
 
-  static void IncRef(const FunctionCallbackInfo<Value>& args) {
-    WeakReference* weak_ref = Unwrap<WeakReference>(args.Holder());
-    weak_ref->reference_count_++;
-    if (weak_ref->target_.IsEmpty()) return;
-    if (weak_ref->reference_count_ == 1) weak_ref->target_.ClearWeak();
-  }
+InternalFieldInfoBase* WeakReference::Serialize(int index) {
+  DCHECK_EQ(index, BaseObject::kEmbedderType);
+  InternalFieldInfo* info =
+      InternalFieldInfoBase::New<InternalFieldInfo>(type());
+  info->target = target_index_;
+  info->reference_count = reference_count_;
+  return info;
+}
 
-  static void DecRef(const FunctionCallbackInfo<Value>& args) {
-    WeakReference* weak_ref = Unwrap<WeakReference>(args.Holder());
-    CHECK_GE(weak_ref->reference_count_, 1);
-    weak_ref->reference_count_--;
-    if (weak_ref->target_.IsEmpty()) return;
-    if (weak_ref->reference_count_ == 0) weak_ref->target_.SetWeak();
+void WeakReference::Deserialize(Local<Context> context,
+                                Local<Object> holder,
+                                int index,
+                                InternalFieldInfoBase* info) {
+  DCHECK_EQ(index, BaseObject::kEmbedderType);
+  HandleScope scope(context->GetIsolate());
+
+  InternalFieldInfo* weak_info = reinterpret_cast<InternalFieldInfo*>(info);
+  Local<Object> target;
+  if (weak_info->target != 0) {
+    target = context->GetDataFromSnapshotOnce<Object>(weak_info->target)
+                 .ToLocalChecked();
   }
+  new WeakReference(Environment::GetCurrent(context),
+                    holder,
+                    target,
+                    weak_info->reference_count);
+}
+
+void WeakReference::New(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  CHECK(args.IsConstructCall());
+  CHECK(args[0]->IsObject());
+  new WeakReference(env, args.This(), args[0].As<Object>());
+}
+
+void WeakReference::Get(const FunctionCallbackInfo<Value>& args) {
+  WeakReference* weak_ref = Unwrap<WeakReference>(args.Holder());
+  Isolate* isolate = args.GetIsolate();
+  if (!weak_ref->target_.IsEmpty())
+    args.GetReturnValue().Set(weak_ref->target_.Get(isolate));
+}
 
-  SET_MEMORY_INFO_NAME(WeakReference)
-  SET_SELF_SIZE(WeakReference)
-  SET_NO_MEMORY_INFO()
+void WeakReference::IncRef(const FunctionCallbackInfo<Value>& args) {
+  WeakReference* weak_ref = Unwrap<WeakReference>(args.Holder());
+  weak_ref->reference_count_++;
+  if (weak_ref->target_.IsEmpty()) return;
+  if (weak_ref->reference_count_ == 1) weak_ref->target_.ClearWeak();
+}
 
- private:
-  Global<Object> target_;
-  uint64_t reference_count_ = 0;
-};
+void WeakReference::DecRef(const FunctionCallbackInfo<Value>& args) {
+  WeakReference* weak_ref = Unwrap<WeakReference>(args.Holder());
+  CHECK_GE(weak_ref->reference_count_, 1);
+  weak_ref->reference_count_--;
+  if (weak_ref->target_.IsEmpty()) return;
+  if (weak_ref->reference_count_ == 0) weak_ref->target_.SetWeak();
+}
 
 static void GuessHandleType(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);

diff --git a/src/node_util.h b/src/node_util.h
@@ -0,0 +1,54 @@
+
+#ifndef SRC_NODE_UTIL_H_
+#define SRC_NODE_UTIL_H_
+
+#if defined(NODE_WANT_INTERNALS) && NODE_WANT_INTERNALS
+#include "base_object.h"
+#include "node_snapshotable.h"
+#include "v8.h"
+
+namespace node {
+namespace util {
+
+class WeakReference : public SnapshotableObject {
+ public:
+  SERIALIZABLE_OBJECT_METHODS();
+
+  static constexpr FastStringKey type_name{"node::util::WeakReference"};
+  static constexpr EmbedderObjectType type_int =
+      EmbedderObjectType::k_util_weak_reference;
+
+  WeakReference(Environment* env,
+                v8::Local<v8::Object> object,
+                v8::Local<v8::Object> target);
+  static void New(const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void Get(const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void IncRef(const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void DecRef(const v8::FunctionCallbackInfo<v8::Value>& args);
+
+  SET_MEMORY_INFO_NAME(WeakReference)
+  SET_SELF_SIZE(WeakReference)
+  SET_NO_MEMORY_INFO()
+
+  struct InternalFieldInfo : public node::InternalFieldInfoBase {
+    SnapshotIndex target;
+    uint64_t reference_count;
+  };
+
+ private:
+  WeakReference(Environment* env,
+                v8::Local<v8::Object> object,
+                v8::Local<v8::Object> target,
+                uint64_t reference_count);
+  v8::Global<v8::Object> target_;
+  uint64_t reference_count_ = 0;
+
+  SnapshotIndex target_index_ = 0;  // 0 means target_ is not snapshotted
+};
+
+}  // namespace util
+}  // namespace node
+
+#endif  // defined(NODE_WANT_INTERNALS) && NODE_WANT_INTERNALS
+
+#endif  // SRC_NODE_UTIL_H_
diff --git a/src/util.cc b/src/util.cc
@@ -27,6 +27,7 @@
 #include "node_buffer.h"
 #include "node_errors.h"
 #include "node_internals.h"
+#include "node_util.h"
 #include "string_bytes.h"
 #include "uv.h"
 

diff --git a/test/fixtures/snapshot/weak-reference-gc.js b/test/fixtures/snapshot/weak-reference-gc.js
@@ -0,0 +1,20 @@
+'use strict';
+
+const { internalBinding } = require('internal/test/binding');
+const { WeakReference } = internalBinding('util');
+const {
+  setDeserializeMainFunction
+} = require('v8').startupSnapshot
+const assert = require('assert');
+
+let obj = { hello: 'world' };
+const ref = new WeakReference(obj);
+
+setDeserializeMainFunction(() => {
+  obj = null;
+  globalThis.gc();
+
+  setImmediate(() => {
+    assert.strictEqual(ref.get(), undefined);
+  });
+});
diff --git a/test/fixtures/snapshot/weak-reference.js b/test/fixtures/snapshot/weak-reference.js
@@ -0,0 +1,15 @@
+'use strict';
+
+const { internalBinding } = require('internal/test/binding');
+const { WeakReference } = internalBinding('util');
+const {
+  setDeserializeMainFunction
+} = require('v8').startupSnapshot
+const assert = require('assert');
+
+let obj = { hello: 'world' };
+const ref = new WeakReference(obj);
+
+setDeserializeMainFunction(() => {
+  assert.strictEqual(ref.get(), obj);
+});
diff --git a/test/parallel/test-snapshot-weak-reference.js b/test/parallel/test-snapshot-weak-reference.js
@@ -0,0 +1,60 @@
+'use strict';
+
+// This tests that weak references work across serialization.
+
+require('../common');
+const assert = require('assert');
+const { spawnSync } = require('child_process');
+const tmpdir = require('../common/tmpdir');
+const fixtures = require('../common/fixtures');
+const path = require('path');
+const fs = require('fs');
+
+tmpdir.refresh();
+const blobPath = path.join(tmpdir.path, 'snapshot.blob');
+
+function runTest(entry) {
+  console.log('running test with', entry);
+  {
+    const child = spawnSync(process.execPath, [
+      '--expose-internals',
+      '--expose-gc',
+      '--snapshot-blob',
+      blobPath,
+      '--build-snapshot',
+      entry,
+    ], {
+      cwd: tmpdir.path
+    });
+    if (child.status !== 0) {
+      console.log(child.stderr.toString());
+      console.log(child.stdout.toString());
+      assert.strictEqual(child.status, 0);
+    }
+    const stats = fs.statSync(path.join(tmpdir.path, 'snapshot.blob'));
+    assert(stats.isFile());
+  }
+
+  {
+    const child = spawnSync(process.execPath, [
+      '--expose-internals',
+      '--expose-gc',
+      '--snapshot-blob',
+      blobPath,
+    ], {
+      cwd: tmpdir.path,
+      env: {
+        ...process.env,
+      }
+    });
+
+    const stdout = child.stdout.toString().trim();
+    const stderr = child.stderr.toString().trim();
+    console.log(`[stdout]:\n${stdout}\n`);
+    console.log(`[stderr]:\n${stderr}\n`);
+    assert.strictEqual(child.status, 0);
+  }
+}
+
+runTest(fixtures.path('snapshot', 'weak-reference.js'));
+runTest(fixtures.path('snapshot', 'weak-reference-gc.js'));