[MachineLearning.Train] Change internal logic to resolve a memory leak

There were issues with heap use after free and heap overflow occurring (TNINE-4152, 4153, 4154). In order to solve this problem, modified the internal logic of the API. Signed-off-by: SeoHyungjun <[email protected]>
Samsung · Oct 16, 2024 · f92be9d · f92be9d
1 parent 4bc6308
commit f92be9d
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 4 deletions.
diff --git a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Dataset.cs b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Dataset.cs
@@ -37,13 +37,19 @@ public class Dataset: IDisposable
         private IntPtr handle = IntPtr.Zero;
         private bool disposed = false;
 
+        /// if false, model will be destroy dataset handle
+        private bool hasOwnership = true;
+
         /// <summary>
         ///  Constructs the dataset.
         /// </summary>
         /// <since_tizen> 10 </since_tizen>
         public Dataset()
         {
             NNTrainerError ret = Interop.Dataset.Create(out handle);
+            if (ret != NNTrainerError.None) {
+                handle = IntPtr.Zero;
+            }
             NNTrainer.CheckException(ret, "Failed to create dataset instance");
             Log.Info(NNTrainer.Tag, "Create Dataset");
         }
@@ -82,6 +88,14 @@ protected virtual void Dispose(bool disposing)
             {
                 // release managed object
             }
+
+            disposed = true;
+
+            if (!hasOwnership){
+                Log.Info(NNTrainer.Tag, "Cannot destroy dataset already added in a Model. Model will destroy this dataset");
+                return;
+            }
+
             // release unmanaged object
             if (handle != IntPtr.Zero)
             {
@@ -92,7 +106,6 @@ protected virtual void Dispose(bool disposing)
 
                 handle = IntPtr.Zero;
             }
-            disposed = true;
         }
 
         /// <summary>
@@ -124,6 +137,11 @@ internal IntPtr GetHandle()
             return handle;
         }
 
+        internal void RemoveOwnership()
+        {
+            this.hasOwnership = false;
+        }
+
         /// <summary>
         /// Sets the neural network dataset property.
         /// </summary>

diff --git a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Layer.cs b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Layer.cs
@@ -49,6 +49,9 @@ public class Layer: IDisposable
         public Layer(NNTrainerLayerType type)
         {
             NNTrainerError ret = Interop.Layer.Create(out handle, type);
+            if (ret != NNTrainerError.None) {
+                handle = IntPtr.Zero;
+            }
             NNTrainer.CheckException(ret, "Failed to create model instance");
             Log.Info(NNTrainer.Tag, $"Create layer with type:{type}");
         }
@@ -98,7 +101,7 @@ protected virtual void Dispose(bool disposing)
             disposed = true;
 
             if (!hasOwnership){
-                Log.Error(NNTrainer.Tag, "Cannot destroy layer already added in a Model. Model will destroy this layer");
+                Log.Info(NNTrainer.Tag, "Cannot destroy layer already added in a Model. Model will destroy this layer");
                 return;
             }
 

diff --git a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Model.cs b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Model.cs
@@ -65,6 +65,9 @@ public Model(string modelConf)
                 NNTrainer.CheckException(NNTrainerError.InvalidParameter, "modelConf is null");
 
             NNTrainerError ret = Interop.Model.ConstructWithConf(modelConf, out handle);
+            if (ret != NNTrainerError.None) {
+                handle = IntPtr.Zero;
+            }
             NNTrainer.CheckException(ret, "Failed to create model instance with modelConf");
             Log.Info(NNTrainer.Tag, "Created Model with Conf path: "+ modelConf);
         }
@@ -281,7 +284,7 @@ public void AddLayer(Layer layer)
         public Layer GetLayer(string layerName)
         {
             IntPtr layerHandle = IntPtr.Zero;
-             if (string.IsNullOrEmpty(layerName))
+            if (string.IsNullOrEmpty(layerName))
                 NNTrainer.CheckException(NNTrainerError.InvalidParameter, "layerName is null");
 
             NNTrainerError ret = Interop.Model.GetLayer(handle, layerName, out layerHandle);
@@ -311,6 +314,7 @@ public void SetOptimizer(Optimizer optimizer)
 
             NNTrainerError ret = Interop.Model.SetOptimizer(handle, optimizer.GetHandle());
             NNTrainer.CheckException(ret, "Failed to set optimizer");
+            optimizer.RemoveOwnership();
         }
 
         /// <summary>
@@ -335,6 +339,7 @@ public void SetDataset(Dataset dataset)
 
             NNTrainerError ret = Interop.Model.SetDataset(handle, dataset.GetHandle());
             NNTrainer.CheckException(ret, "Failed to set dataset");
+            dataset.RemoveOwnership();
         }
 
         internal static TensorsInfo CreateTensorsInfoFormHandle(IntPtr handle)

diff --git a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Optimizer.cs b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Optimizer.cs
@@ -37,6 +37,9 @@ public class Optimizer: IDisposable
         private IntPtr handle = IntPtr.Zero;
         private bool disposed = false;
 
+        /// if false, model will be destroy optimizer handle
+        private bool hasOwnership = true;
+
         /// <summary>
         /// Creates a neural network optimizer.
         /// </summary>
@@ -45,6 +48,9 @@ public class Optimizer: IDisposable
         public Optimizer(NNTrainerOptimizerType type)
         {
             NNTrainerError ret = Interop.Optimizer.Create(out handle, type);
+            if (ret != NNTrainerError.None) {
+                handle = IntPtr.Zero;
+            }
             NNTrainer.CheckException(ret, "Failed to create optimizer instance");
             Log.Info(NNTrainer.Tag, $"Create optimizer with type:{type}");
         }
@@ -83,6 +89,14 @@ protected virtual void Dispose(bool disposing)
             {
                 // release managed object
             }
+
+            disposed = true;
+
+            if (!hasOwnership){
+                Log.Info(NNTrainer.Tag, "Cannot destroy optimizer already added in a Model. Model will destroy this optimizer");
+                return;
+            }
+
             // release unmanaged object
             if (handle != IntPtr.Zero)
             {
@@ -93,7 +107,6 @@ protected virtual void Dispose(bool disposing)
 
                 handle = IntPtr.Zero;
             }
-            disposed = true;
         }
 
         /// <summary>
@@ -122,5 +135,10 @@ internal IntPtr GetHandle()
         {
             return handle;
         }
+
+        internal void RemoveOwnership()
+        {
+            this.hasOwnership = false;
+        }
     } 
 }