Fix memory leak: Avoid disposing python stream in finalizer thread (#…

…1142)
kendryte · Dec 13, 2023 · 980f977 · 980f977
1 parent c4f7b5f
commit 980f977
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/Native/include/nncase/compiler.h b/src/Native/include/nncase/compiler.h
@@ -109,6 +109,7 @@ typedef struct {
     clr_object_handle_t (*calibration_dataset_provider_create)(
         clr_object_handle_t dataset, size_t samplesCount,
         clr_object_handle_t fn_params);
+    void (*handle_dispose)(clr_object_handle_t handle);
     void (*handle_free)(clr_object_handle_t handle);
     clr_object_handle_t (*compile_options_create)();
     void (*compile_options_set_input_file)(clr_object_handle_t compile_options,
@@ -478,6 +479,8 @@ class cstream : public clr_object_base {
     cstream(const nncase_stream_mt_t *mt, void *handle) {
         obj_ = nncase_clr_api()->stream_create(mt, handle);
     }
+
+    ~cstream() { nncase_clr_api()->handle_dispose(obj_.get()); }
 };
 
 class compile_options : public clr_object_base {

diff --git a/src/Nncase.Compiler/Interop/CApi.cs b/src/Nncase.Compiler/Interop/CApi.cs
@@ -41,6 +41,7 @@ public unsafe struct CApiMT
     public delegate* unmanaged<IntPtr, nuint, IntPtr> ArrayGetItemPtr;
     public delegate* unmanaged<IntPtr, nuint> ArrayGetLengthPtr;
     public delegate* unmanaged<IntPtr, nuint, IntPtr, IntPtr> CalibrationDatasetProviderCreatePtr;
+    public delegate* unmanaged<IntPtr, void> ClrHandleDisposePtr;
     public delegate* unmanaged<IntPtr, void> ClrHandleFreePtr;
     public delegate* unmanaged<IntPtr> CompileOptionsCreatePtr;
     public delegate* unmanaged<IntPtr, byte*, nuint, void> CompileOptionsSetInputFilePtr;
@@ -112,6 +113,7 @@ public static void Initialize(CApiMT* mt)
         mt->ArrayGetItemPtr = &ArrayGetItem;
         mt->ArrayGetLengthPtr = &ArrayGetLength;
         mt->CalibrationDatasetProviderCreatePtr = &CalibrationDatasetProviderCreate;
+        mt->ClrHandleDisposePtr = &ClrHandleDispose;
         mt->ClrHandleFreePtr = &ClrHandleFree;
         mt->CompileOptionsCreatePtr = &CompileOptionsCreate;
         mt->CompileOptionsSetInputFilePtr = &CompileOptionsSetInputFile;
@@ -225,6 +227,12 @@ private static IntPtr CalibrationDatasetProviderCreate(IntPtr datasetHandle, nui
         return GCHandle.ToIntPtr(GCHandle.Alloc(new CCalibrationDatasetProvider(samples, (int)samplesCount)));
     }
 
+    [UnmanagedCallersOnly]
+    private static void ClrHandleDispose(IntPtr handle)
+    {
+        Get<IDisposable>(handle).Dispose();
+    }
+
     [UnmanagedCallersOnly]
     private static void ClrHandleFree(IntPtr handle)
     {