fix: SerializedProto's memory leaks (#461)

homuler · Feb 22, 2022 · ac6316d · ac6316d
1 parent 0dafc8a
commit ac6316d
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 16 deletions.
diff --git a/Packages/com.github.homuler.mediapipe/Runtime/Scripts/External/SerializedProto.cs b/Packages/com.github.homuler.mediapipe/Runtime/Scripts/External/SerializedProto.cs
@@ -12,20 +12,20 @@
 namespace Mediapipe
 {
   [StructLayout(LayoutKind.Sequential)]
-  internal struct SerializedProto
+  internal readonly struct SerializedProto
   {
-    public IntPtr str;
-    public int length;
+    private readonly IntPtr _str;
+    private readonly int _length;
 
     public void Dispose()
     {
-      UnsafeNativeMethods.delete_array__PKc(str);
+      UnsafeNativeMethods.delete_array__PKc(_str);
     }
 
     public T Deserialize<T>(pb::MessageParser<T> parser) where T : pb::IMessage<T>
     {
-      var bytes = new byte[length];
-      Marshal.Copy(str, bytes, 0, bytes.Length);
+      var bytes = new byte[_length];
+      Marshal.Copy(_str, bytes, 0, bytes.Length);
       return parser.ParseFrom(bytes);
     }
   }

diff --git a/Packages/com.github.homuler.mediapipe/Runtime/Scripts/External/SerializedProtoVector.cs b/Packages/com.github.homuler.mediapipe/Runtime/Scripts/External/SerializedProtoVector.cs
@@ -13,25 +13,25 @@
 namespace Mediapipe
 {
   [StructLayout(LayoutKind.Sequential)]
-  internal struct SerializedProtoVector
+  internal readonly struct SerializedProtoVector
   {
-    public IntPtr data;
-    public int size;
+    private readonly IntPtr _data;
+    private readonly int _size;
 
     public void Dispose()
     {
-      UnsafeNativeMethods.mp_api_SerializedProtoArray__delete(data);
+      UnsafeNativeMethods.mp_api_SerializedProtoArray__delete(_data, _size);
     }
 
     public List<T> Deserialize<T>(pb::MessageParser<T> parser) where T : pb::IMessage<T>
     {
-      var protos = new List<T>(size);
+      var protos = new List<T>(_size);
 
       unsafe
       {
-        var protoPtr = (SerializedProto*)data;
+        var protoPtr = (SerializedProto*)_data;
 
-        for (var i = 0; i < size; i++)
+        for (var i = 0; i < _size; i++)
         {
           var serializedProto = Marshal.PtrToStructure<SerializedProto>((IntPtr)protoPtr++);
           protos.Add(serializedProto.Deserialize(parser));

diff --git a/...ithub.homuler.mediapipe/Runtime/Scripts/PInvoke/NativeMethods/External/Protobuf_Unsafe.cs b/...ithub.homuler.mediapipe/Runtime/Scripts/PInvoke/NativeMethods/External/Protobuf_Unsafe.cs
@@ -16,7 +16,7 @@ internal static partial class UnsafeNativeMethods
         [MarshalAs(UnmanagedType.FunctionPtr)] Protobuf.ProtobufLogHandler logHandler);
 
     [DllImport(MediaPipeLibrary, ExactSpelling = true)]
-    public static extern void mp_api_SerializedProtoArray__delete(IntPtr serializedProtoVectorData);
+    public static extern void mp_api_SerializedProtoArray__delete(IntPtr serializedProtoVectorData, int size);
 
     #region MessageProto
     [DllImport(MediaPipeLibrary, ExactSpelling = true)]

diff --git a/mediapipe_api/external/protobuf.cc b/mediapipe_api/external/protobuf.cc
@@ -27,4 +27,10 @@ MpReturnCode google_protobuf__SetLogHandler__PF(LogHandler* handler) {
   CATCH_EXCEPTION
 }
 
-void mp_api_SerializedProtoArray__delete(mp_api::SerializedProto* serialized_proto_vector_data) { delete[] serialized_proto_vector_data; }
+void mp_api_SerializedProtoArray__delete(mp_api::SerializedProto* serialized_proto_vector_data, int size) {
+  auto serialized_proto = serialized_proto_vector_data;
+  for (auto i = 0; i < size; ++i) {
+    delete (serialized_proto++)->str;
+  }
+  delete[] serialized_proto_vector_data;
+}
diff --git a/mediapipe_api/external/protobuf.h b/mediapipe_api/external/protobuf.h
@@ -63,7 +63,7 @@ typedef void LogHandler(int level, const char* filename, int line, const char* m
 
 MP_CAPI(MpReturnCode) google_protobuf__SetLogHandler__PF(LogHandler* handler);
 
-MP_CAPI(void) mp_api_SerializedProtoArray__delete(mp_api::SerializedProto* serialized_proto_vector_data);
+MP_CAPI(void) mp_api_SerializedProtoArray__delete(mp_api::SerializedProto* serialized_proto_vector_data, int size);
 
 }  // extern "C"