From 5424d5feeac70e9c62c935617755ab5879de502f Mon Sep 17 00:00:00 2001
From: liufeng27 <liufeng27@baidu.com>
Date: Sat, 7 Nov 2020 15:13:01 +0800
Subject: [PATCH] fix memleak

---
 examples/load_model/CMakeLists.txt |  4 +++-
 examples/load_model/main.cpp       |  5 +++++
 include/cppflow/defer.h            | 25 +++++++++++++++++++++++++
 include/cppflow/model.h            |  7 +++++++
 include/cppflow/tensor.h           |  9 +++++++--
 5 files changed, 47 insertions(+), 3 deletions(-)
 create mode 100644 include/cppflow/defer.h

diff --git a/examples/load_model/CMakeLists.txt b/examples/load_model/CMakeLists.txt
index 64662c9..430da95 100644
--- a/examples/load_model/CMakeLists.txt
+++ b/examples/load_model/CMakeLists.txt
@@ -4,7 +4,9 @@ project(example)
 find_library(TENSORFLOW_LIB tensorflow HINT $ENV{HOME}/libtensorflow2/lib)
 
 set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-omit-frame-pointer -fsanitize=address")
+set(CMAKE_LINKER_FLAGS "${CMAKE_LINKER_FLAGS} -lasan")
 
 add_executable(example main.cpp)
 target_include_directories(example PRIVATE ../../include $ENV{HOME}/libtensorflow2/include)
-target_link_libraries (example "${TENSORFLOW_LIB}")
\ No newline at end of file
+target_link_libraries (example "${TENSORFLOW_LIB}")
diff --git a/examples/load_model/main.cpp b/examples/load_model/main.cpp
index 9111db1..58a8f95 100644
--- a/examples/load_model/main.cpp
+++ b/examples/load_model/main.cpp
@@ -12,5 +12,10 @@ int main() {
     
     std::cout << output << std::endl;
 
+    auto values = output.get_data<float>();
+
+    for (auto v : values) {
+        std::cout << v << std::endl;
+    }
     return 0;
 }
diff --git a/include/cppflow/defer.h b/include/cppflow/defer.h
new file mode 100644
index 0000000..b127ba3
--- /dev/null
+++ b/include/cppflow/defer.h
@@ -0,0 +1,25 @@
+#pragma once
+#include <functional>
+
+namespace cppflow {
+
+class defer {
+public:
+    typedef std::function<void ()> Func;
+
+    explicit defer(const Func& func) : _func(func) {}
+    ~defer() {
+        _func();
+    }
+
+    defer(const defer&) = delete;
+    defer(defer&&) = delete;
+    defer& operator=(const defer&) = delete;
+    void* operator new (size_t) = delete;
+    void operator delete (void*) = delete;
+
+private:
+    Func _func;
+};
+
+} // namespace cppflow
diff --git a/include/cppflow/model.h b/include/cppflow/model.h
index 54de644..b413855 100644
--- a/include/cppflow/model.h
+++ b/include/cppflow/model.h
@@ -12,6 +12,7 @@
 #include <vector>
 
 #include "context.h"
+#include "defer.h"
 #include "tensor.h"
 
 namespace cppflow {
@@ -84,6 +85,12 @@ namespace cppflow {
 
         std::vector<TF_Output> inp_ops(inputs.size());
         std::vector<TF_Tensor*> inp_val(inputs.size());
+
+        defer d([&inp_val]{
+            for (auto* tf_tensor : inp_val) {
+                TF_DeleteTensor(tf_tensor);
+            }
+        });
         for (int i=0; i<inputs.size(); i++) {
 
             // Operations
diff --git a/include/cppflow/tensor.h b/include/cppflow/tensor.h
index 62a0bbd..6de73b9 100644
--- a/include/cppflow/tensor.h
+++ b/include/cppflow/tensor.h
@@ -167,7 +167,7 @@ namespace cppflow {
 
         // EXECUTE
         int n = 1;
-        TFE_TensorHandle* res[1];
+        TFE_TensorHandle* res[1] = { nullptr };
         TFE_Execute(op, res, &n, context::get_status());
         status_check(context::get_status());
         TFE_DeleteOp(op);
@@ -175,6 +175,8 @@ namespace cppflow {
         tensor r;
         r.tf_tensor = { TFE_TensorHandleResolve(res[0], context::get_status()), TF_DeleteTensor};
         status_check(context::get_status());
+        TFE_DeleteTensorHandle(res[0]);
+
         r.tfe_handle = {TFE_NewTensorHandle(r.tf_tensor.get(), context::get_status()), TFE_DeleteTensorHandle};
         status_check(context::get_status());
 
@@ -205,7 +207,10 @@ namespace cppflow {
 
         // Convert to correct type
         const auto T_data = static_cast<T*>(raw_data);
-        return std::vector<T>(T_data, T_data + size);
+        std::vector<T> r(T_data, T_data + size);
+        TF_DeleteTensor(res_tensor);
+
+        return r;
     }
 
     datatype tensor::dtype() const {