feat: implement dispose() for deterministic native resource release

cpp/HybridTfliteModel.cpp

@@ -28,11 +28,28 @@ HybridTfliteModel::HybridTfliteModel(std::shared_ptr<TfLiteInterpreter> interpre
   }
 }
 
+void HybridTfliteModel::dispose() {
+  // The lock waits out an in-flight inference on another thread - freeing the
+  // interpreter under a running TfLiteInterpreterInvoke would be a native crash.
+  std::lock_guard<std::mutex> lock(_lifecycleMutex);
+  if (_interpreter == nullptr) {
+    return; // already disposed
+  }
+  // We hold the only reference to the interpreter, so this runs its deleter
+  // now: TfLiteInterpreterDelete, then the delegates. The model bytes are
+  // released once nobody else (e.g. JS) references them either.
+  _interpreter.reset();
+  _modelData.reset();
+  _outputBuffers.clear();
+}
+
 std::vector<TensorflowModelDelegate> HybridTfliteModel::getDelegates() {
   return _delegates;
 }
 
 std::vector<Tensor> HybridTfliteModel::getInputs() {
+  std::lock_guard<std::mutex> lock(_lifecycleMutex);
+  throwIfDisposed();
   int count = TfLiteInterpreterGetInputTensorCount(_interpreter.get());
   std::vector<Tensor> tensors;
   tensors.reserve(count);
@@ -55,6 +72,8 @@ std::vector<Tensor> HybridTfliteModel::getInputs() {
 }
 
 std::vector<Tensor> HybridTfliteModel::getOutputs() {
+  std::lock_guard<std::mutex> lock(_lifecycleMutex);
+  throwIfDisposed();
   int count = TfLiteInterpreterGetOutputTensorCount(_interpreter.get());
   std::vector<Tensor> tensors;
   tensors.reserve(count);
@@ -148,19 +167,31 @@ void HybridTfliteModel::invoke() {
 
 std::vector<std::shared_ptr<ArrayBuffer>>
 HybridTfliteModel::runSync(const std::vector<std::shared_ptr<ArrayBuffer>>& input) {
+  // Held for the whole inference so dispose() can never free the interpreter
+  // mid-invoke. The disposed-throw is a catchable JS error on any runtime.
+  std::lock_guard<std::mutex> lock(_lifecycleMutex);
+  throwIfDisposed();
   copyInputBuffers(input);
   invoke();
   return copyOutputBuffers();
 }
 
 std::shared_ptr<Promise<std::vector<std::shared_ptr<ArrayBuffer>>>>
 HybridTfliteModel::run(const std::vector<std::shared_ptr<ArrayBuffer>>& input) {
-  // Copy input buffers on caller (JS) thread first — input ArrayBuffers are
-  // non-owning JS buffers that may be GC'd if we access them async.
-  copyInputBuffers(input);
+  {
+    // Copy input buffers on caller (JS) thread first — input ArrayBuffers are
+    // non-owning JS buffers that may be GC'd if we access them async.
+    std::lock_guard<std::mutex> lock(_lifecycleMutex);
+    throwIfDisposed();
+    copyInputBuffers(input);
+  }
   std::shared_ptr<HybridTfliteModel> sharedThis = shared_cast<HybridTfliteModel>();
   return Promise<std::vector<std::shared_ptr<ArrayBuffer>>>::async(
       [sharedThis]() -> std::vector<std::shared_ptr<ArrayBuffer>> {
+        // Re-acquire on the async thread: dispose() may have landed between
+        // the input copy above and this lambda running.
+        std::lock_guard<std::mutex> lock(sharedThis->_lifecycleMutex);
+        sharedThis->throwIfDisposed();
         sharedThis->invoke();
         return sharedThis->copyOutputBuffers();
       });

cpp/HybridTfliteModel.hpp

@@ -2,6 +2,8 @@
 
 #include "HybridTfliteModelSpec.hpp"
 #include <memory>
+#include <mutex>
+#include <stdexcept>
 #include <string>
 #include <unordered_map>
 
@@ -22,6 +24,17 @@ class HybridTfliteModel : public HybridTfliteModelSpec {
                              std::vector<TensorflowModelDelegate> delegates);
   ~HybridTfliteModel() override = default;
 
+  /**
+   * Free the interpreter, its delegates and our reference to the model bytes
+   * NOW, without waiting for every runtime's GC to drop its reference (worklet
+   * runtimes may not GC for a long time, especially while backgrounded).
+   * Thread-safe: blocks until an in-flight inference on another thread has
+   * completed. Every later runSync/run/getInputs/getOutputs call throws a
+   * catchable JS error. Idempotent. Called by Nitro when JS invokes
+   * `model.dispose()`.
+   */
+  void dispose() override;
+
   // Properties (from HybridTfliteModelSpec)
   std::vector<TensorflowModelDelegate> getDelegates() override;
   std::vector<Tensor> getInputs() override;
@@ -39,11 +52,23 @@ class HybridTfliteModel : public HybridTfliteModelSpec {
   std::vector<std::shared_ptr<ArrayBuffer>> copyOutputBuffers();
   std::shared_ptr<ArrayBuffer> getOutputBufferForTensor(const TfLiteTensor* tensor);
 
+  // Caller must hold _lifecycleMutex. A disposed model has released its
+  // interpreter; std::runtime_error surfaces as a catchable JS error.
+  void throwIfDisposed() const {
+    if (_interpreter == nullptr) {
+      throw std::runtime_error("TFLite: Model was disposed!");
+    }
+  }
+
 private:
   std::shared_ptr<TfLiteInterpreter> _interpreter;
   std::vector<TensorflowModelDelegate> _delegates;
   std::shared_ptr<ArrayBuffer> _modelData;
   std::unordered_map<std::string, std::shared_ptr<ArrayBuffer>> _outputBuffers;
+
+  // Serializes inference against dispose(). Uncontended in normal operation
+  // (one lock per inference, ~ns vs ~ms inference cost).
+  std::mutex _lifecycleMutex;
 };
 
 } // namespace margelo::nitro::tflite

cpp/HybridTfliteModule.cpp

@@ -2,11 +2,18 @@
 #include "TfliteHelpers.hpp"
 
 #include <memory>
+#include <utility>
+#include <vector>
 
 #if defined(ANDROID)
 #include <tflite/c/c_api.h>
+#include <tflite/delegates/gpu/delegate.h>
+#include <tflite/delegates/nnapi/nnapi_delegate_c_api.h>
 #elif defined(__APPLE__)
 #include <TensorFlowLiteC/TensorFlowLiteC.h>
+#if FAST_TFLITE_ENABLE_CORE_ML
+#include <TensorFlowLiteCCoreML/TensorFlowLiteCCoreML.h>
+#endif
 #else
 #error "Invalid Platform!"
 #endif
@@ -32,6 +39,37 @@ TfLiteDelegate* getDelegate(TensorflowModelDelegate delegateType) {
                            "\"!");
 }
 
+/**
+ * TFLite's C API does not transfer delegate ownership to the interpreter: the
+ * caller must keep a delegate alive for the interpreter's lifetime and free it
+ * afterwards with the delegate's own delete function.
+ */
+struct DelegateDeleter {
+  TensorflowModelDelegate delegateType;
+
+  void operator()(TfLiteDelegate* delegate) const {
+    switch (delegateType) {
+#if defined(__APPLE__) && FAST_TFLITE_ENABLE_CORE_ML
+      case TensorflowModelDelegate::CORE_ML:
+        TfLiteCoreMlDelegateDelete(delegate);
+        return;
+#endif
+#if defined(ANDROID)
+      case TensorflowModelDelegate::ANDROID_GPU:
+        TfLiteGpuDelegateV2Delete(delegate);
+        return;
+      case TensorflowModelDelegate::NNAPI:
+        TfLiteNnapiDelegateDelete(delegate);
+        return;
+#endif
+      default:
+        // getDelegate() throws for every other type on this platform.
+        return;
+    }
+  }
+};
+using OwnedDelegate = std::unique_ptr<TfLiteDelegate, DelegateDeleter>;
+
 std::shared_ptr<HybridTfliteModelSpec>
 HybridTfliteModule::createModel(const std::shared_ptr<ArrayBuffer>& modelData,
                                 const std::vector<TensorflowModelDelegate>& delegates) {
@@ -50,20 +88,39 @@ HybridTfliteModule::createModel(const std::shared_ptr<ArrayBuffer>& modelData,
 
   // Add all hardware accelerated delegates (e.g. GPU, NPU, ...)
   // if any. The default CPU delegate will always be available.
+  std::vector<TensorflowModelDelegate> effectiveDelegates;
+  std::vector<OwnedDelegate> ownedDelegates;
+  effectiveDelegates.reserve(delegates.size());
+  ownedDelegates.reserve(delegates.size());
   for (const TensorflowModelDelegate& delegateType : delegates) {
-    TfLiteDelegate* delegate = getDelegate(delegateType);
-    TfLiteInterpreterOptionsAddDelegate(options.get(), delegate);
+    OwnedDelegate delegate(getDelegate(delegateType), DelegateDeleter{delegateType});
+    if (delegate == nullptr) {
+      // e.g. CoreML on devices without a Neural Engine — fall back to CPU
+      // instead of registering a null delegate with the interpreter.
+      continue;
+    }
+    TfLiteInterpreterOptionsAddDelegate(options.get(), delegate.get());
+    effectiveDelegates.push_back(delegateType);
+    ownedDelegates.push_back(std::move(delegate));
   }
 
   TfLiteInterpreter* rawInterpreter = TfLiteInterpreterCreate(model.get(), options.get());
   if (rawInterpreter == nullptr) {
+    // `ownedDelegates` frees the delegates on unwind.
     throw std::runtime_error("Failed to create TFLite interpreter!");
   }
+  // The delegates travel with the interpreter and are freed right after it,
+  // so they can never be deleted while the interpreter still uses them.
   const std::shared_ptr<TfLiteInterpreter> interpreter(
-      rawInterpreter, [modelData](TfLiteInterpreter* value) { TfLiteInterpreterDelete(value); });
+      rawInterpreter,
+      [modelData, ownedDelegates = std::move(ownedDelegates)](TfLiteInterpreter* value) mutable {
+        TfLiteInterpreterDelete(value);
+        ownedDelegates.clear();
+      });
 
-  // Wrap in HybridTfliteModel — stores shared_ptr<ArrayBuffer> to keep model data bytes alive
-  return std::make_shared<HybridTfliteModel>(interpreter, modelData, delegates);
+  // Wrap in HybridTfliteModel — stores shared_ptr<ArrayBuffer> to keep model data bytes alive.
+  // Only the delegates that were actually registered are reported via `getDelegates()`.
+  return std::make_shared<HybridTfliteModel>(interpreter, modelData, effectiveDelegates);
 }
 
 } // namespace margelo::nitro::tflite