chore(🧹): remove RNFRuntimeCache (#278)

Author: wcandillon

packages/webgpu/android/CMakeLists.txt

@@ -46,8 +46,7 @@ add_library(${PACKAGE_NAME} SHARED
     ../cpp/rnwgpu/RNWebGPUManager.cpp
     ../cpp/jsi/RNFPromise.cpp
     ../cpp/jsi/RNFHybridObject.cpp
-    ../cpp/jsi/RNFRuntimeCache.cpp
-    ../cpp/jsi/RNFWorkletRuntimeRegistry.cpp
+    ../cpp/jsi/RNFRuntimeState.cpp
     ../cpp/rnwgpu/async/AsyncRunner.cpp
     ../cpp/rnwgpu/async/AsyncTaskHandle.cpp
     ../cpp/rnwgpu/async/JSIMicrotaskDispatcher.cpp

packages/webgpu/cpp/jsi/RNFHybridObject.cpp

@@ -129,18 +129,22 @@ void HybridObject::ensureInitialized(facebook::jsi::Runtime& runtime) {
   if (!_didLoadMethods) {
     [[unlikely]];
     _creationRuntime = &runtime;
+    _runtimeState = rnwgpu::RNFRuntimeState::get(runtime);
     // lazy-load all exposed methods
     loadHybridMethods();
     _didLoadMethods = true;
   }
 }
 
 bool HybridObject::isRuntimeAlive() {
-  if (_creationRuntime == nullptr) {
-    [[unlikely]];
-    return false;
+  return !_runtimeState.expired();
+}
+
+facebook::jsi::Runtime* HybridObject::getCreationRuntime() const {
+  if (_runtimeState.expired()) {
+    return nullptr;
   }
-  return RNFWorkletRuntimeRegistry::isRuntimeAlive(_creationRuntime);
+  return _creationRuntime;
 }
 
 } // namespace margelo

packages/webgpu/cpp/jsi/RNFHybridObject.h

@@ -5,7 +5,7 @@
 #pragma once
 
 #include "WGPULogger.h"
-#include "RNFWorkletRuntimeRegistry.h"
+#include "RNFRuntimeState.h"
 #include <functional>
 #include <jsi/jsi.h>
 #include <memory>
@@ -91,8 +91,9 @@ class HybridObject : public jsi::HostObject, public std::enable_shared_from_this
 
 protected:
   const char* _name = TAG;
-  // Store a pointer to the runtime. Needed for checking if the runtime is still active, see WorkletRuntimeRegistry.
+  // Store a pointer to the runtime for convenience; ensure it is only used while the runtime state is alive.
   jsi::Runtime* _creationRuntime = nullptr;
+  std::weak_ptr<rnwgpu::RNFRuntimeState> _runtimeState;
 
 private:
   inline void ensureInitialized(facebook::jsi::Runtime& runtime);
@@ -128,6 +129,10 @@ class HybridObject : public jsi::HostObject, public std::enable_shared_from_this
     };
   }
 
+protected:
+  facebook::jsi::Runtime* getCreationRuntime() const;
+  std::weak_ptr<rnwgpu::RNFRuntimeState> getRuntimeStateWeak() const { return _runtimeState; }
+
 protected:
   template <typename Derived, typename ReturnType, typename... Args>
   void registerHybridMethod(std::string name, ReturnType (Derived::*method)(Args...), Derived* derivedInstance, bool override = false) {

packages/webgpu/cpp/jsi/RNFJSIConverter.h

@@ -22,7 +22,6 @@
 #include "RNFEnumMapper.h"
 #include "RNFJSIHelper.h"
 #include "RNFPromise.h"
-#include "RNFWorkletRuntimeRegistry.h"
 
 #include "Unions.h"
 

packages/webgpu/cpp/jsi/RNFJSIHelper.h

@@ -4,7 +4,7 @@
 
 #pragma once
 
-#include "RNFWorkletRuntimeRegistry.h"
+#include "RNFRuntimeState.h"
 #include <jsi/jsi.h>
 #include <memory>
 #include <utility>
@@ -20,8 +20,10 @@ class JSIHelper {
    * Every jsi::Function you intend to share or hold should be wrapped using this function.
    */
   static std::shared_ptr<jsi::Function> createSharedJsiFunction(jsi::Runtime& runtime, jsi::Function&& function) {
-    std::shared_ptr<jsi::Function> sharedFunction(new jsi::Function(std::move(function)), [&runtime](jsi::Function* ptr) {
-      if (RNFWorkletRuntimeRegistry::isRuntimeAlive(&runtime)) {
+    auto runtimeState = rnwgpu::RNFRuntimeState::get(runtime);
+    std::weak_ptr<rnwgpu::RNFRuntimeState> runtimeStateWeak = runtimeState;
+    std::shared_ptr<jsi::Function> sharedFunction(new jsi::Function(std::move(function)), [runtimeStateWeak](jsi::Function* ptr) {
+      if (!runtimeStateWeak.expired()) {
         // Only delete the jsi::Function when the runtime it created is still alive.
         // Otherwise leak memory. We do this on purpose, as sometimes we would keep
         // references to JSI objects past the lifetime of its runtime (e.g.,

packages/webgpu/cpp/jsi/RNFRuntimeCache.cpp

@@ -0,0 +1,18 @@
+#include "RNFRuntimeState.h"
+
+namespace rnwgpu {
+
+const facebook::jsi::UUID RNFRuntimeState::kRuntimeStateKey = facebook::jsi::UUID();
+
+std::shared_ptr<RNFRuntimeState> RNFRuntimeState::get(facebook::jsi::Runtime& runtime) {
+  auto existing = runtime.getRuntimeData(kRuntimeStateKey);
+  if (existing) {
+    return std::static_pointer_cast<RNFRuntimeState>(existing);
+  }
+
+  auto state = std::shared_ptr<RNFRuntimeState>(new RNFRuntimeState());
+  runtime.setRuntimeData(kRuntimeStateKey, state);
+  return state;
+}
+
+} // namespace rnwgpu

packages/webgpu/cpp/jsi/RNFRuntimeCache.h

@@ -0,0 +1,106 @@
+#pragma once
+
+#include <jsi/jsi.h>
+
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace rnwgpu {
+
+namespace jsi = facebook::jsi;
+
+/**
+ * Runtime state management using Hermes' runtime.setRuntimeData API.
+ * This replaces the old RuntimeLifecycleMonitor/RuntimeAwareCache system.
+ */
+class RNFRuntimeState {
+public:
+  // UUID key for storing our runtime state
+  static const jsi::UUID kRuntimeStateKey;
+
+  /**
+   * Get or create the runtime state for the given runtime
+   */
+  static std::shared_ptr<RNFRuntimeState> get(jsi::Runtime& runtime);
+
+  /**
+   * Template cache that can store any type T per object pointer
+   */
+  template <typename T>
+  class ObjectCache {
+  public:
+    std::shared_ptr<T> getOrCreate(void* object) {
+      std::lock_guard<std::mutex> lock(mutex_);
+      auto it = cache_.find(object);
+      if (it != cache_.end()) {
+        return it->second;
+      }
+
+      auto value = std::make_shared<T>();
+      cache_[object] = value;
+      return value;
+    }
+
+    void remove(void* object) {
+      std::lock_guard<std::mutex> lock(mutex_);
+      cache_.erase(object);
+    }
+
+    void clear() {
+      std::lock_guard<std::mutex> lock(mutex_);
+      cache_.clear();
+    }
+
+  private:
+    std::mutex mutex_;
+    std::unordered_map<void*, std::shared_ptr<T>> cache_;
+  };
+
+  /**
+   * Get or create a cache for a specific type T
+   */
+  template <typename T>
+  std::shared_ptr<ObjectCache<T>> getCache() {
+    std::lock_guard<std::mutex> lock(mutex_);
+
+    // Use type_info as key for the cache type
+    const std::type_info& typeId = typeid(T);
+    auto it = typeCaches_.find(&typeId);
+
+    if (it != typeCaches_.end()) {
+      return std::static_pointer_cast<ObjectCache<T>>(it->second);
+    }
+
+    auto cache = std::make_shared<ObjectCache<T>>();
+    typeCaches_[&typeId] = cache;
+    return cache;
+  }
+
+private:
+  RNFRuntimeState() = default;
+
+  std::mutex mutex_;
+  // Map from type_info to cache instance
+  std::unordered_map<const std::type_info*, std::shared_ptr<void>> typeCaches_;
+};
+
+/**
+ * Template helper for runtime-aware caching compatible with the old API
+ * This provides a migration path from RuntimeAwareCache
+ */
+template <typename T>
+class RuntimeAwareCache {
+public:
+  T& get(jsi::Runtime& rt) {
+    auto state = RNFRuntimeState::get(rt);
+    auto cache = state->getCache<T>();
+
+    // For compatibility, we use the runtime pointer as the object key
+    auto ptr = cache->getOrCreate(&rt);
+    return *ptr;
+  }
+};
+
+} // namespace rnwgpu