Create Borrowing/Owning reference types

packages/react-native-nitro-modules/cpp/core/HybridObject.cpp

@@ -76,8 +76,8 @@ jsi::Value HybridObject::get(facebook::jsi::Runtime& runtime, const facebook::js
 
   if (functionCache.contains(name)) [[likely]] {
     // cache hit - let's see if the function is still alive..
-    std::shared_ptr<jsi::Function> function = functionCache[name].lock();
-    if (function != nullptr) [[likely]] {
+    OwningReference<jsi::Function> function = functionCache[name];
+    if (function) [[likely]] {
       // function is still alive, we can use it.
       return jsi::Value(runtime, *function);
     }
@@ -97,10 +97,10 @@ jsi::Value HybridObject::get(facebook::jsi::Runtime& runtime, const facebook::js
     jsi::Function function = jsi::Function::createFromHostFunction(runtime, jsi::PropNameID::forUtf8(runtime, name),
                                                                    hybridFunction.parameterCount, hybridFunction.function);
     // throw it into the cache
-    auto globalFunction = runtimeCache->makeGlobal(std::move(function));
+    OwningReference<jsi::Function> globalFunction = runtimeCache->makeGlobal(std::move(function));
     functionCache[name] = globalFunction;
     // copy the reference & return it to JS
-    return jsi::Value(runtime, *globalFunction.lock());
+    return jsi::Value(runtime, *globalFunction);
   }
 
   if (name == "toString") {

packages/react-native-nitro-modules/cpp/core/HybridObject.hpp

@@ -5,6 +5,7 @@
 #pragma once
 
 #include "JSIConverter.hpp"
+#include "OwningReference.hpp"
 #include "NitroLogger.hpp"
 #include <functional>
 #include <jsi/jsi.h>
@@ -70,7 +71,7 @@ class HybridObject : public jsi::HostObject, public std::enable_shared_from_this
   std::unordered_map<std::string, HybridFunction> _methods;
   std::unordered_map<std::string, jsi::HostFunctionType> _getters;
   std::unordered_map<std::string, jsi::HostFunctionType> _setters;
-  std::unordered_map<jsi::Runtime*, std::unordered_map<std::string, std::weak_ptr<jsi::Function>>> _functionCache;
+  std::unordered_map<jsi::Runtime*, std::unordered_map<std::string, OwningReference<jsi::Function>>> _functionCache;
 
 private:
   inline void ensureInitialized(facebook::jsi::Runtime& runtime);

packages/react-native-nitro-modules/cpp/jsi/FunctionCache.cpp

@@ -15,6 +15,17 @@ static constexpr auto CACHE_PROP_NAME = "__nitroModulesFunctionCache";
 
 FunctionCache::FunctionCache(jsi::Runtime* runtime): _runtime(runtime) {}
 
+FunctionCache::~FunctionCache() {
+  // TODO: Do we need to throw Mutexes on OwningReference so there are no race conditions with nullchecks and then pointer accessses while we delete the pointer?
+  for (auto& func : _cache) {
+    OwningReference<jsi::Function> owning = func.lock();
+    if (owning) {
+      // Destroy all functions that we might still have in cache, some callbacks and Promises may now become invalid.
+      owning.destroy();
+    }
+  }
+}
+
 std::weak_ptr<FunctionCache> FunctionCache::getOrCreateCache(jsi::Runtime &runtime) {
   if (_globalCache.contains(&runtime)) {
     // Fast path: get weak_ptr to FunctionCache from our global list.
@@ -36,10 +47,10 @@ std::weak_ptr<FunctionCache> FunctionCache::getOrCreateCache(jsi::Runtime &runti
   return nativeState;
 }
 
-std::weak_ptr<jsi::Function> FunctionCache::makeGlobal(jsi::Function&& function) {
-  auto shared = std::make_shared<jsi::Function>(std::move(function));
-  _cache.push_back(shared);
-  return std::weak_ptr(shared);
+OwningReference<jsi::Function> FunctionCache::makeGlobal(jsi::Function&& function) {
+  auto owning = OwningReference<jsi::Function>(new jsi::Function(std::move(function)));
+  _cache.push_back(owning.weak());
+  return owning;
 }
 
 }

packages/react-native-nitro-modules/cpp/jsi/FunctionCache.hpp

@@ -11,6 +11,8 @@
 #include <memory>
 #include <vector>
 #include <unordered_map>
+#include "OwningReference.hpp"
+#include "BorrowingReference.hpp"
 
 namespace margelo {
 
@@ -22,6 +24,7 @@ using namespace facebook;
 class FunctionCache final: public jsi::NativeState {
 public:
   explicit FunctionCache(jsi::Runtime* runtime);
+  ~FunctionCache();
 
 public:
   /**
@@ -43,11 +46,11 @@ class FunctionCache final: public jsi::NativeState {
    Do not hold the returned `shared_ptr` in memory, only use it in the calling function's scope.
    Note: By design, this is not thread-safe, the returned `weak_ptr` must only be locked on the same thread as it was created on.
    */
-  std::weak_ptr<jsi::Function> makeGlobal(jsi::Function&& function);
+  OwningReference<jsi::Function> makeGlobal(jsi::Function&& function);
 
 private:
   jsi::Runtime* _runtime;
-  std::vector<std::shared_ptr<jsi::Function>> _cache;
+  std::vector<BorrowingReference<jsi::Function>> _cache;
 
 private:
   static std::unordered_map<jsi::Runtime*, std::weak_ptr<FunctionCache>> _globalCache;

packages/react-native-nitro-modules/cpp/jsi/JSIConverter.hpp

@@ -210,28 +210,26 @@ template <typename ReturnType, typename... Args> struct JSIConverter<std::functi
     // Make function global - it'll be managed by the Runtime's memory, and we only have a weak_ref to it.
     auto cache = FunctionCache::getOrCreateCache(runtime).lock();
     jsi::Function function = arg.getObject(runtime).getFunction(runtime);
-    auto sharedFunction = cache->makeGlobal(std::move(function));
+    OwningReference<jsi::Function> sharedFunction = cache->makeGlobal(std::move(function));
 
     // Create a C++ function that can be called by the consumer.
     // This will call the jsi::Function if it is still alive.
     return [&runtime, sharedFunction](Args... args) -> ReturnType {
       if constexpr (std::is_same_v<ReturnType, void>) {
         // it is a void function (returns undefined)
-        auto function = sharedFunction.lock();
-        if (!function) {
+        if (!sharedFunction) {
           // runtime has already been deleted. since this returns void, we can just ignore it being deleted.
           return;
         }
-        function->call(runtime, JSIConverter<std::decay_t<Args>>::toJSI(runtime, args)...);
+        sharedFunction->call(runtime, JSIConverter<std::decay_t<Args>>::toJSI(runtime, args)...);
         return;
       } else {
         // it returns a custom type, parse it from the JSI value.
-        auto function = sharedFunction.lock();
-        if (!function) {
+        if (!sharedFunction) {
           // runtime has already been deleted. since we expect a return value here, we need to throw.
           throw std::runtime_error("Cannot call the given Function - the Runtime has already been destroyed!");
         }
-        jsi::Value result = function->call(runtime, JSIConverter<std::decay_t<Args>>::toJSI(runtime, args)...);
+        jsi::Value result = sharedFunction->call(runtime, JSIConverter<std::decay_t<Args>>::toJSI(runtime, args)...);
         return JSIConverter<ReturnType>::fromJSI(runtime, std::move(result));
       }
     };

packages/react-native-nitro-modules/cpp/jsi/Promise.cpp

@@ -37,22 +37,20 @@ jsi::Value Promise::createPromise(jsi::Runtime& runtime, RunPromise run) {
 }
 
 void Promise::resolve(jsi::Runtime& runtime, jsi::Value&& result) {
-  auto resolver = _resolver.lock();
-  if (resolver == nullptr) {
+  if (!_resolver) {
     Logger::log(TAG, "Promise resolver function has already been deleted! Ignoring call..");
     return;
   }
-  resolver->call(runtime, std::move(result));
+  _resolver->call(runtime, std::move(result));
 }
 
 void Promise::reject(jsi::Runtime& runtime, std::string message) {
-  auto rejecter = _rejecter.lock();
-  if (rejecter == nullptr) {
+  if (!_rejecter) {
     Logger::log(TAG, "Promise rejecter function has already been deleted! Ignoring call..");
     return;
   }
   jsi::JSError error(runtime, message);
-  rejecter->call(runtime, error.value());
+  _rejecter->call(runtime, error.value());
 }
 
 } // namespace margelo

packages/react-native-nitro-modules/cpp/jsi/Promise.hpp

@@ -41,8 +41,8 @@ class Promise final {
 
 private:
   std::weak_ptr<FunctionCache> _functionCache;
-  std::weak_ptr<jsi::Function> _resolver;
-  std::weak_ptr<jsi::Function> _rejecter;
+  OwningReference<jsi::Function> _resolver;
+  OwningReference<jsi::Function> _rejecter;
   static constexpr auto TAG = "Promise";
 
 public:

packages/react-native-nitro-modules/cpp/utils/BorrowingReference.cpp

@@ -0,0 +1,34 @@
+//
+//  BorrowingReference.cpp
+//  NitroModules
+//
+//  Created by Marc Rousavy on 21.06.24.
+//
+
+#include "BorrowingReference.hpp"
+#include "OwningReference.hpp"
+
+namespace margelo {
+
+template<typename T>
+BorrowingReference<T>::BorrowingReference(const OwningReference<T>& ref) {
+  _value = ref._value;
+  _isDeleted = ref._isDeleted;
+  _strongRefCount = ref._strongRefCount;
+  _weakRefCount = ref._weakRefCount;
+  (*_weakRefCount)++;
+}
+
+template<typename T>
+OwningReference<T> BorrowingReference<T>::lock() {
+  if (*_isDeleted) {
+    // return nullptr
+    return OwningReference<T>();
+  }
+
+  return OwningReference(*this);
+}
+
+template class BorrowingReference<int>;
+
+}

packages/react-native-nitro-modules/cpp/utils/BorrowingReference.hpp

@@ -0,0 +1,97 @@
+//
+//  BorrowingReference.hpp
+//  NitroModules
+//
+//  Created by Marc Rousavy on 21.06.24.
+//
+
+#pragma once
+
+namespace margelo {
+
+// forward-declaration to avoid duplicate symbols
+template<typename T>
+class OwningReference;
+
+template<typename T>
+class BorrowingReference {
+private:
+  explicit BorrowingReference(const OwningReference<T>& ref);
+
+public:
+  BorrowingReference(): _value(nullptr), _isDeleted(nullptr), _strongRefCount(nullptr), _weakRefCount(nullptr) { }
+
+  BorrowingReference(const BorrowingReference& ref):
+    _value(ref._value),
+    _isDeleted(ref._isDeleted),
+    _strongRefCount(ref._strongRefCount),
+    _weakRefCount(ref._weakRefCount) {
+      // increment ref count after copy
+      (*_strongRefCount)++;
+  }
+
+  BorrowingReference(BorrowingReference&& ref):
+    _value(ref._value),
+    _isDeleted(ref._isDeleted),
+    _strongRefCount(ref._strongRefCount),
+    _weakRefCount(ref._weakRefCount) {
+      ref._value = nullptr;
+      ref._isDeleted = nullptr;
+      ref._strongRefCount = nullptr;
+      ref._weakRefCount = nullptr;
+  }
+
+  BorrowingReference& operator=(const BorrowingReference& ref) {
+    if (this == &ref) return *this;
+
+    if (_weakRefCount != nullptr) {
+      // destroy previous pointer
+      (*_weakRefCount)--;
+      maybeDestroy();
+    }
+
+    _value = ref._value;
+    _isDeleted = ref._isDeleted;
+    _strongRefCount = ref._strongRefCount;
+    _weakRefCount = ref._weakRefCount;
+    if (_weakRefCount != nullptr) {
+      (*_weakRefCount)++;
+    }
+  }
+
+  ~BorrowingReference() {
+    if (_weakRefCount == nullptr) {
+      // we are just a dangling nullptr.
+      return;
+    }
+
+    (*_weakRefCount)--;
+    maybeDestroy();
+  }
+
+  /**
+   Try to lock the borrowing reference to an owning reference, or `nullptr` if it has already been deleted.
+   */
+  OwningReference<T> lock();
+
+private:
+  void maybeDestroy() {
+    if (*_strongRefCount < 0 && *_weakRefCount < 0) {
+      // free the full memory if there are no more references at all
+      if (!(*_isDeleted)) {
+        delete _value;
+      }
+      delete _isDeleted;
+      delete _strongRefCount;
+      delete _weakRefCount;
+    }
+  }
+
+private:
+  T* _value;
+  bool* _isDeleted;
+  int* _strongRefCount;
+  int* _weakRefCount;
+};
+
+} // namespace margelo