Add JSI APIs setRuntimeData and getRuntimeData (#50197)

Summary:
Pull Request resolved: #50197

Adds `setRuntimeData` and `getRuntimeData` JSI APIs. This provides a
convenient way for users to store some custom data associated with an
UUID.

For the default implementation of this feature, store the runtime data
in the global map that is shared between all VMs. This is done to keep
JSI lightweight and stateless, instead of adding data members.

Changelog: [Internal]

Reviewed By: neildhar

Differential Revision: D71579532

fbshipit-source-id: 553e28fbf80c93e268d475860197a00d2c5bacf7

Author: tsaichien

packages/react-native/ReactCommon/jsi/jsi/decorator.h

@@ -393,6 +393,17 @@ class RuntimeDecorator : public Base, private jsi::Instrumentation {
     return plain_.callAsConstructor(f, args, count);
   };
 
+  void setRuntimeDataImpl(
+      const UUID& uuid,
+      const void* data,
+      void (*deleter)(const void* data)) override {
+    return plain_.setRuntimeDataImpl(uuid, data, deleter);
+  }
+
+  const void* getRuntimeDataImpl(const UUID& uuid) override {
+    return plain_.getRuntimeDataImpl(uuid);
+  }
+
   // Private data for managing scopes.
   Runtime::ScopeState* pushScope() override {
     return plain_.pushScope();
@@ -957,6 +968,19 @@ class WithRuntimeDecorator : public RuntimeDecorator<Plain, Base> {
     RD::setExternalMemoryPressure(obj, amount);
   };
 
+  void setRuntimeDataImpl(
+      const UUID& uuid,
+      const void* data,
+      void (*deleter)(const void* data)) override {
+    Around around{with_};
+    RD::setRuntimeDataImpl(uuid, data, deleter);
+  }
+
+  const void* getRuntimeDataImpl(const UUID& uuid) override {
+    Around around{with_};
+    return RD::getRuntimeDataImpl(uuid);
+  }
+
  private:
   // Wrap an RAII type around With& to guarantee after always happens.
   struct Around {

packages/react-native/ReactCommon/jsi/jsi/jsi-inl.h

@@ -84,6 +84,20 @@ inline const Runtime::PointerValue* Runtime::getPointerValue(
   return value.data_.pointer.ptr_;
 }
 
+inline void Runtime::setRuntimeData(
+    const UUID& uuid,
+    const std::shared_ptr<void>& data) {
+  auto* dataPtr = new std::shared_ptr<void>(data);
+  setRuntimeDataImpl(uuid, dataPtr, [](const void* data) {
+    delete (const std::shared_ptr<void>*)data;
+  });
+}
+
+inline std::shared_ptr<void> Runtime::getRuntimeData(const UUID& uuid) {
+  auto* data = (const std::shared_ptr<void>*)getRuntimeDataImpl(uuid);
+  return data ? *data : nullptr;
+}
+
 Value Object::getPrototype(Runtime& runtime) const {
   return runtime.getPrototypeOf(*this);
 }

packages/react-native/ReactCommon/jsi/jsi/jsi.cpp

@@ -8,6 +8,8 @@
 #include <cassert>
 #include <cmath>
 #include <cstdlib>
+#include <map>
+#include <mutex>
 #include <stdexcept>
 
 #include <jsi/instrumentation.h>
@@ -18,6 +20,63 @@ namespace jsi {
 
 namespace {
 
+/// A global map used to store custom runtime data for VMs that do not provide
+/// their own default implementation of setRuntimeData and getRuntimeData.
+struct RuntimeDataGlobal {
+  /// Mutex protecting the Runtime data map
+  std::mutex mutex_{};
+  /// Maps a runtime pointer to a map of its custom data. At destruction of the
+  /// runtime, its entry will be removed from the global map.
+  std::unordered_map<
+      Runtime*,
+      std::unordered_map<
+          UUID,
+          std::pair<const void*, void (*)(const void* data)>,
+          UUID::Hash>>
+      dataMap_;
+};
+
+RuntimeDataGlobal& getRuntimeDataGlobal() {
+  static RuntimeDataGlobal runtimeData{};
+  return runtimeData;
+}
+
+/// A host object that, when destructed, will remove the runtime's custom data
+/// entry from the global map of custom data.
+class RemoveRuntimeDataHostObject : public jsi::HostObject {
+ public:
+  explicit RemoveRuntimeDataHostObject(Runtime* runtime) : runtime_(runtime) {}
+
+  RemoveRuntimeDataHostObject(const RemoveRuntimeDataHostObject&) = default;
+  RemoveRuntimeDataHostObject(RemoveRuntimeDataHostObject&&) = default;
+  RemoveRuntimeDataHostObject& operator=(const RemoveRuntimeDataHostObject&) =
+      default;
+  RemoveRuntimeDataHostObject& operator=(RemoveRuntimeDataHostObject&&) =
+      default;
+
+  ~RemoveRuntimeDataHostObject() override {
+    auto& runtimeDataGlobal = getRuntimeDataGlobal();
+    std::lock_guard<std::mutex> lock(runtimeDataGlobal.mutex_);
+    auto runtimeMapIt = runtimeDataGlobal.dataMap_.find(runtime_);
+    // We install the RemoveRuntimeDataHostObject only when the first custom
+    // data for the runtime is added, and only this object is responsible for
+    // clearing runtime data. Thus, we should always be able to find the data
+    // entry.
+    assert(
+        runtimeMapIt != runtimeDataGlobal.dataMap_.end() &&
+        "Custom runtime data not found for this runtime");
+
+    for (auto [_, entry] : runtimeMapIt->second) {
+      auto* deleter = entry.second;
+      deleter(entry.first);
+    }
+    runtimeDataGlobal.dataMap_.erase(runtime_);
+  }
+
+ private:
+  Runtime* runtime_;
+};
+
 // This is used for generating short exception strings.
 std::string kindToString(const Value& v, Runtime* rt = nullptr) {
   if (v.isUndefined()) {
@@ -353,6 +412,62 @@ Object Runtime::createObjectWithPrototype(const Value& prototype) {
   return createFn.call(*this, prototype).asObject(*this);
 }
 
+void Runtime::setRuntimeDataImpl(
+    const UUID& uuid,
+    const void* data,
+    void (*deleter)(const void* data)) {
+  auto& runtimeDataGlobal = getRuntimeDataGlobal();
+  std::lock_guard<std::mutex> lock(runtimeDataGlobal.mutex_);
+  if (auto it = runtimeDataGlobal.dataMap_.find(this);
+      it != runtimeDataGlobal.dataMap_.end()) {
+    auto& map = it->second;
+    if (auto entryIt = map.find(uuid); entryIt != map.end()) {
+      // Free the old data
+      auto oldData = entryIt->second.first;
+      auto oldDataDeleter = entryIt->second.second;
+      oldDataDeleter(oldData);
+    }
+    map[uuid] = {data, deleter};
+    return;
+  }
+  // No custom data entry exist for this runtime in the global map, so create
+  // one.
+  runtimeDataGlobal.dataMap_[this][uuid] = {data, deleter};
+
+  // The first time data is added for this runtime is added to the map, install
+  // a host object on the global object of the runtime. This host object is used
+  // to release the runtime's entry from the global custom data map when the
+  // runtime is destroyed.
+  // Also, try to protect the host object by making it non-configurable,
+  // non-enumerable, and non-writable. These JSI operations are purposely
+  // performed after runtime-specific data map is added and the host object is
+  // created to prevent data leaks if any operations fail.
+  Object ho = Object::createFromHostObject(
+      *this, std::make_shared<RemoveRuntimeDataHostObject>(this));
+  global().setProperty(*this, "_jsiRuntimeDataCleanUp", ho);
+  auto definePropertyFn = global()
+                              .getPropertyAsObject(*this, "Object")
+                              .getPropertyAsFunction(*this, "defineProperty");
+  auto desc = Object(*this);
+  desc.setProperty(*this, "configurable", Value(false));
+  desc.setProperty(*this, "enumerable", Value(false));
+  desc.setProperty(*this, "writable", Value(false));
+  definePropertyFn.call(*this, global(), "_jsiRuntimeDataCleanUp", desc);
+}
+
+const void* Runtime::getRuntimeDataImpl(const UUID& uuid) {
+  auto& runtimeDataGlobal = getRuntimeDataGlobal();
+  std::lock_guard<std::mutex> lock(runtimeDataGlobal.mutex_);
+  if (auto runtimeMapIt = runtimeDataGlobal.dataMap_.find(this);
+      runtimeMapIt != runtimeDataGlobal.dataMap_.end()) {
+    if (auto customDataIt = runtimeMapIt->second.find(uuid);
+        customDataIt != runtimeMapIt->second.end()) {
+      return customDataIt->second.first;
+    }
+  }
+  return nullptr;
+}
+
 Pointer& Pointer::operator=(Pointer&& other) noexcept {
   if (ptr_) {
     ptr_->invalidate();

packages/react-native/ReactCommon/jsi/jsi/jsi.h

@@ -349,6 +349,16 @@ class JSI_EXPORT Runtime {
   /// which returns no metrics.
   virtual Instrumentation& instrumentation();
 
+  /// Stores the pointer \p data with the \p uuid in the runtime. This can be
+  /// used to store some custom data within the runtime. When the runtime is
+  /// destroyed, or if an entry at an existing key is overwritten, the runtime
+  /// will release its ownership of the held object.
+  void setRuntimeData(const UUID& uuid, const std::shared_ptr<void>& data);
+
+  /// Returns the data associated with the \p uuid in the runtime. If there's no
+  /// data associated with the uuid, return a null pointer.
+  std::shared_ptr<void> getRuntimeData(const UUID& uuid);
+
  protected:
   friend class Pointer;
   friend class PropNameID;
@@ -364,6 +374,19 @@ class JSI_EXPORT Runtime {
   friend class Scope;
   friend class JSError;
 
+  /// Stores the pointer \p data with the \p uuid in the runtime. This can be
+  /// used to store some custom data within the runtime. When the runtime is
+  /// destroyed, or if an entry at an existing key is overwritten, the runtime
+  /// will release its ownership by calling \p deleter.
+  virtual void setRuntimeDataImpl(
+      const UUID& uuid,
+      const void* data,
+      void (*deleter)(const void* data));
+
+  /// Returns the data associated with the \p uuid in the runtime. If there's no
+  /// data associated with the uuid, return a null pointer.
+  virtual const void* getRuntimeDataImpl(const UUID& uuid);
+
   // Potential optimization: avoid the cloneFoo() virtual dispatch,
   // and instead just fix the number of fields, and copy them, since
   // in practice they are trivially copyable.  Sufficient use of

packages/react-native/ReactCommon/jsi/jsi/test/testlib.cpp

@@ -1773,6 +1773,94 @@ TEST_P(JSITest, ObjectCreateWithPrototype) {
   EXPECT_TRUE(child.getPrototype(rd).isNull());
 }
 
+TEST_P(JSITest, SetRuntimeData) {
+  class RD : public RuntimeDecorator<Runtime, Runtime> {
+   public:
+    explicit RD(Runtime& rt) : RuntimeDecorator(rt) {}
+
+    void setRuntimeDataImpl(
+        const UUID& uuid,
+        const void* data,
+        void (*deleter)(const void* data)) override {
+      Runtime::setRuntimeDataImpl(uuid, data, deleter);
+    }
+
+    const void* getRuntimeDataImpl(const UUID& uuid) override {
+      return Runtime::getRuntimeDataImpl(uuid);
+    }
+  };
+
+  RD rd1 = RD(rt);
+  UUID uuid1{0xe67ab3d6, 0x09a0, 0x11f0, 0xa641, 0x325096b39f47};
+  auto str = std::make_shared<std::string>("hello world");
+  rd1.setRuntimeData(uuid1, str);
+
+  UUID uuid2{0xa12f99fc, 0x09a2, 0x11f0, 0x84de, 0x325096b39f47};
+  auto obj1 = std::make_shared<Object>(rd1);
+  rd1.setRuntimeData(uuid2, obj1);
+
+  auto storedStr =
+      std::static_pointer_cast<std::string>(rd1.getRuntimeData(uuid1));
+  auto storedObj = std::static_pointer_cast<Object>(rd1.getRuntimeData(uuid2));
+  EXPECT_EQ(storedStr, str);
+  EXPECT_EQ(storedObj, obj1);
+
+  // Override the existing value at uuid1
+  auto weakOldStr = std::weak_ptr<std::string>(str);
+  str = std::make_shared<std::string>("goodbye world");
+  rd1.setRuntimeData(uuid1, str);
+  storedStr = std::static_pointer_cast<std::string>(rd1.getRuntimeData(uuid1));
+  EXPECT_EQ(str, storedStr);
+  // Verify that the old data was not held on after it was overwritten.
+  EXPECT_EQ(weakOldStr.use_count(), 0);
+
+  auto rt2 = factory();
+  RD* rd2 = new RD(*rt2);
+  UUID uuid3{0x16f55892, 0x1034, 0x11f0, 0x8f65, 0x325096b39f47};
+  auto obj2 = std::make_shared<Object>(*rd2);
+  rd2->setRuntimeData(uuid3, obj2);
+
+  auto storedObj2 =
+      std::static_pointer_cast<Object>(rd2->getRuntimeData(uuid3));
+  EXPECT_EQ(storedObj2, obj2);
+
+  // UUID1 is for some data in runtime rd1, not rd2
+  EXPECT_FALSE(rd2->getRuntimeData(uuid1));
+
+  // Verify that when runtime is deleted, its runtime data map gets removed from
+  // the global map. So nothing should be holding on to the stored data.
+  auto weakObj2 = std::weak_ptr<Object>(obj2);
+  obj2.reset();
+  storedObj2.reset();
+  delete rd2;
+  rt2.reset();
+  EXPECT_EQ(weakObj2.use_count(), 0);
+
+  // Only the second runtime was destroyed, so custom data from the first
+  // runtime should remain unaffected.
+  storedStr = std::static_pointer_cast<std::string>(rd1.getRuntimeData(uuid1));
+  EXPECT_EQ(storedStr, str);
+
+  // Overwrite Object.defineProperty, which is called in the default
+  // implementation of setRuntimeDataImpl, to test that even if the JSI
+  // operations on this secret property fail, we are still able to properly
+  // clean up the custom data.
+  auto rt3 = factory();
+  RD* rd3 = new RD(*rt3);
+  UUID uuid4{0xa5682986, 0x1edc, 0x11f0, 0xa4fa, 0x325096b39f47};
+  rd3->global()
+      .getPropertyAsObject(*rd3, "Object")
+      .setProperty(*rd3, "defineProperty", Value(false));
+
+  auto obj3 = std::make_shared<Object>(*rd3);
+  auto weakObj3 = std::weak_ptr<Object>(obj3);
+  EXPECT_THROW(rd3->setRuntimeData(uuid4, obj3), JSIException);
+  obj3.reset();
+  delete rd3;
+  rt3.reset();
+  EXPECT_EQ(weakObj3.use_count(), 0);
+}
+
 INSTANTIATE_TEST_CASE_P(
     Runtimes,
     JSITest,