Fix RunTimePoolInfo leak of shared memory regions.

SamplePreparedModel and StepExecutor call RunTimePoolInfo::set(), which
calls mmap(), but there is no corresponding call to munmap().  The fix
is to add a new method runTimePoolInfo::release(), and call it from the
appropriate points.

Without this fix, even if we don't run out of virtual address space,
we can still run up against a kernel limit in the number of mapped regions
("/proc/sys/vm/max_map_count").

ALSO: Add a missing check for HIDL transport error (isOk()).

Test: nn/runtime/test

Bug: 69685100
Change-Id: If6ac1396e89cd738738a6ae3397450bd1d80f506

7 files changed, 98 insertions, 33 deletions

diff --git a/nn/common/CpuExecutor.cpp b/nn/common/CpuExecutor.cpp
index 9c6df76e6..ce6efd6eb 100644
--- a/nn/common/CpuExecutor.cpp
+++ b/nn/common/CpuExecutor.cpp
@@ -29,6 +29,10 @@ namespace nn {
 // TODO: short term, make share memory mapping and updating a utility function.
 // TODO: long term, implement mmap_fd as a hidl IMemory service.
 bool RunTimePoolInfo::set(const hidl_memory& hidlMemory) {
+    if (buffer != nullptr) {
+        release();
+    }
+
     this->hidlMemory = hidlMemory;
     auto memType = hidlMemory.name();
     if (memType == "ashmem") {
@@ -52,16 +56,40 @@ bool RunTimePoolInfo::set(const hidl_memory& hidlMemory) {
                                         hidlMemory.handle()->data[3]);
         buffer = static_cast<uint8_t*>(mmap(nullptr, size, prot, MAP_SHARED, fd, offset));
         if (buffer == MAP_FAILED) {
-            LOG(ERROR) << "Can't mmap the file descriptor.";
+            LOG(ERROR) << "RunTimePoolInfo::set(): Can't mmap the file descriptor.";
             return false;
         }
         return true;
     } else {
-        LOG(ERROR) << "unsupported hidl_memory type";
+        LOG(ERROR) << "RunTimePoolInfo::set(): unsupported hidl_memory type";
         return false;
     }
 }
 
+void RunTimePoolInfo::release() {
+    if (buffer == nullptr) {
+        // Never initialized, or already released.
+        LOG(ERROR) << "RunTimePoolInfo::release(): incorrect RunTimePoolInfo resource management";
+    }
+
+    auto memType = hidlMemory.name();
+    if (memType == "ashmem") {
+        // nothing to do
+    } else if (memType == "mmap_fd") {
+        size_t size = hidlMemory.size();
+        if (munmap(buffer, size)) {
+            LOG(ERROR) << "RunTimePoolInfo::release(): Can't munmap";
+        }
+    } else if (memType == "") {
+        // Represents a POINTER argument; nothing to do
+    } else {
+        LOG(ERROR) << "RunTimePoolInfo::release(): unsupported hidl_memory type";
+    }
+    memory = nullptr;
+    hidlMemory = hidl_memory();
+    buffer = nullptr;
+}
+
 // Making sure the output data are correctly updated after execution.
 bool RunTimePoolInfo::update() {
     auto memType = hidlMemory.name();
@@ -86,12 +114,21 @@ bool setRunTimePoolInfosFromHidlMemories(std::vector<RunTimePoolInfo>* poolInfos
         auto& poolInfo = (*poolInfos)[i];
         if (!poolInfo.set(pools[i])) {
             LOG(ERROR) << "Could not map pool";
+            for (size_t j = 0; j < i; j++) {
+                (*poolInfos)[j].release();
+            }
             return false;
         }
     }
     return true;
 }
 
+void releaseRunTimePoolInfos(std::vector<RunTimePoolInfo>* poolInfos) {
+    for (auto& info : *poolInfos) {
+        info.release();
+    }
+}
+
 // Updates the RunTimeOperandInfo with the newly calculated shape.
 // Allocate the buffer if we need to.
 static bool setInfoAndAllocateIfNeeded(RunTimeOperandInfo* info, const Shape& shape) {
diff --git a/nn/common/include/CpuExecutor.h b/nn/common/include/CpuExecutor.h
index b765efc7c..1ddcfe1e6 100644
--- a/nn/common/include/CpuExecutor.h
+++ b/nn/common/include/CpuExecutor.h
@@ -63,15 +63,21 @@ struct RunTimeOperandInfo {
 struct RunTimePoolInfo {
     sp<IMemory> memory;
     hidl_memory hidlMemory;
-    uint8_t* buffer;
+    uint8_t* buffer = nullptr;
 
     bool set(const hidl_memory& hidlMemory);
+
+    // Release resources owned by this info (acquired by set()).
+    void release();
+
     bool update();
 };
 
 bool setRunTimePoolInfosFromHidlMemories(std::vector<RunTimePoolInfo>* poolInfos,
                                          const hidl_vec<hidl_memory>& pools);
 
+void releaseRunTimePoolInfos(std::vector<RunTimePoolInfo>* poolInfos);
+
 // This class is used to execute a model on the CPU.
 class CpuExecutor {
 public:
diff --git a/nn/driver/sample/SampleDriver.cpp b/nn/driver/sample/SampleDriver.cpp
index 664607fd4..b686cfcb5 100644
--- a/nn/driver/sample/SampleDriver.cpp
+++ b/nn/driver/sample/SampleDriver.cpp
@@ -86,6 +86,7 @@ void SamplePreparedModel::asyncExecute(const Request& request,
     CpuExecutor executor;
     int n = executor.run(mModel, request, mPoolInfos, requestPoolInfos);
     VLOG(DRIVER) << "executor.run returned " << n;
+    releaseRunTimePoolInfos(&requestPoolInfos);
     ErrorStatus executionStatus =
             n == ANEURALNETWORKS_NO_ERROR ? ErrorStatus::NONE : ErrorStatus::GENERAL_FAILURE;
     Return<void> returned = callback->notify(executionStatus);
@@ -113,6 +114,10 @@ Return<ErrorStatus> SamplePreparedModel::execute(const Request& request,
     return ErrorStatus::NONE;
 }
 
+SamplePreparedModel::~SamplePreparedModel() {
+    releaseRunTimePoolInfos(&mPoolInfos);
+}
+
 } // namespace sample_driver
 } // namespace nn
 } // namespace android
diff --git a/nn/driver/sample/SampleDriver.h b/nn/driver/sample/SampleDriver.h
index 7e95c952b..3b2e97356 100644
--- a/nn/driver/sample/SampleDriver.h
+++ b/nn/driver/sample/SampleDriver.h
@@ -52,7 +52,7 @@ public:
     SamplePreparedModel(const Model& model)
           : // Make a copy of the model, as we need to preserve it.
             mModel(model) {}
-    ~SamplePreparedModel() override {}
+    ~SamplePreparedModel() override;
     bool initialize();
     Return<ErrorStatus> execute(const Request& request,
                                 const sp<IExecutionCallback>& callback) override;
diff --git a/nn/runtime/ExecutionBuilder.cpp b/nn/runtime/ExecutionBuilder.cpp
index f87b7b4b8..9d041c256 100644
--- a/nn/runtime/ExecutionBuilder.cpp
+++ b/nn/runtime/ExecutionBuilder.cpp
@@ -595,7 +595,8 @@ int StepExecutor::startComputeOnDevice(sp<ExecutionCallback>* synchronizationCal
     // maybe the HIDL infrastructure handles this magically? At worst,
     // it seems like this is a small memory leak, if the Callback stays
     // alive forever.
-    if (mPreparedModel->execute(request, executionCallback) != ErrorStatus::NONE) {
+    Return<ErrorStatus> executeStatus = mPreparedModel->execute(request, executionCallback);
+    if (!executeStatus.isOk() || executeStatus != ErrorStatus::NONE) {
         VLOG(EXECUTION) << "**Execute failed**";
         return ANEURALNETWORKS_OP_FAILED;
     }
@@ -603,8 +604,8 @@ int StepExecutor::startComputeOnDevice(sp<ExecutionCallback>* synchronizationCal
     // TODO: Remove this synchronization point when the block of code below is
     // removed.
     executionCallback->wait();
-    Return<ErrorStatus> executionStatus = executionCallback->getStatus();
-    if (!executionStatus.isOk() || executionStatus != ErrorStatus::NONE) {
+    Return<ErrorStatus> callbackStatus = executionCallback->getStatus();
+    if (!callbackStatus.isOk() || callbackStatus != ErrorStatus::NONE) {
         VLOG(EXECUTION) << "**Execute async failed**";
         return ANEURALNETWORKS_OP_FAILED;
     }
@@ -636,6 +637,8 @@ static void asyncStartComputeOnCpu(const Model& model, const Request& request,
                                    const sp<IExecutionCallback>& executionCallback) {
     CpuExecutor executor;
     int err = executor.run(model, request, modelPoolInfos, requestPoolInfos);
+    releaseRunTimePoolInfos(const_cast<std::vector<RunTimePoolInfo>*>(&modelPoolInfos));
+    releaseRunTimePoolInfos(const_cast<std::vector<RunTimePoolInfo>*>(&requestPoolInfos));
     ErrorStatus status = err == ANEURALNETWORKS_NO_ERROR ?
             ErrorStatus::NONE : ErrorStatus::GENERAL_FAILURE;
     executionCallback->notify(status);
@@ -665,6 +668,10 @@ int StepExecutor::startComputeOnCpu(sp<ExecutionCallback>* synchronizationCallba
     for (uint32_t i = 0; i < count; i++) {
         const Memory* mem = mMemories[i];
         if (!requestPoolInfos[i].set(mem->getHidlMemory())) {
+            for (uint32_t j = 0; j < i; j++) {
+                requestPoolInfos[j].release();
+            }
+            releaseRunTimePoolInfos(&modelPoolInfos);
             return ANEURALNETWORKS_UNMAPPABLE;
         }
     }
diff --git a/nn/runtime/Memory.cpp b/nn/runtime/Memory.cpp
index beac661f5..cbaff17e7 100644
--- a/nn/runtime/Memory.cpp
+++ b/nn/runtime/Memory.cpp
@@ -116,7 +116,7 @@ int MemoryFd::getPointer(uint8_t** buffer) const {
     size_t offset = getSizeFromInts(mHandle->data[2], mHandle->data[3]);
     void* data = mmap(nullptr, mHidlMemory.size(), prot, MAP_SHARED, fd, offset);
     if (data == MAP_FAILED) {
-        LOG(ERROR) << "Can't mmap the file descriptor.";
+        LOG(ERROR) << "MemoryFd::getPointer(): Can't mmap the file descriptor.";
         return ANEURALNETWORKS_UNMAPPABLE;
     } else {
         mMapping = *buffer = static_cast<uint8_t*>(data);
diff --git a/nn/runtime/test/TestMemory.cpp b/nn/runtime/test/TestMemory.cpp
index ef11c2ba1..f63f4fa93 100644
--- a/nn/runtime/test/TestMemory.cpp
+++ b/nn/runtime/test/TestMemory.cpp
@@ -52,11 +52,23 @@ protected:
     const Matrix3x4 expected3 = {{121.f, 232.f, 343.f, 454.f},
                                  {526.f, 628.f, 730.f, 832.f},
                                  {940.f, 1042.f, 1144.f, 1246.f}};
-    const Matrix3x4 expected3b = {{22.f, 34.f, 46.f, 58.f},
-                                  {31.f, 34.f, 37.f, 40.f},
-                                  {49.f, 52.f, 55.f, 58.f}};
 };
 
+// Tests to ensure that various kinds of memory leaks do not occur.
+class MemoryLeakTest : public MemoryTest {
+protected:
+    void SetUp() override;
+
+    int mMaxMapCount = 0;
+};
+
+void MemoryLeakTest::SetUp() {
+    std::ifstream maxMapCountStream("/proc/sys/vm/max_map_count");
+    if (maxMapCountStream) {
+        maxMapCountStream >> mMaxMapCount;
+    }
+}
+
 // Check that the values are the same. This works only if dealing with integer
 // value, otherwise we should accept values that are similar if not exact.
 int CompareMatrices(const Matrix3x4& expected, const Matrix3x4& actual) {
@@ -73,8 +85,14 @@ int CompareMatrices(const Matrix3x4& expected, const Matrix3x4& actual) {
     return errors;
 }
 
+// As well as serving as a functional test for ASharedMemory, also
+// serves as a regression test for http://b/69685100 "RunTimePoolInfo
+// leaks shared memory regions".
+//
 // TODO: test non-zero offset.
-TEST_F(MemoryTest, TestASharedMemory) {
+TEST_F(MemoryLeakTest, TestASharedMemory) {
+    ASSERT_GT(mMaxMapCount, 0);
+
     // Layout where to place matrix2 and matrix3 in the memory we'll allocate.
     // We have gaps to test that we don't assume contiguity.
     constexpr uint32_t offsetForMatrix2 = 20;
@@ -135,13 +153,18 @@ TEST_F(MemoryTest, TestASharedMemory) {
     WrapperCompilation compilation2(&model);
     ASSERT_EQ(compilation2.finish(), WrapperResult::NO_ERROR);
 
-    WrapperExecution execution2(&compilation2);
-    ASSERT_EQ(execution2.setInputFromMemory(0, &input, offsetForMatrix1, sizeof(Matrix3x4)),
-              WrapperResult::NO_ERROR);
-    ASSERT_EQ(execution2.setOutputFromMemory(0, &actual, offsetForActual, sizeof(Matrix3x4)),
-              WrapperResult::NO_ERROR);
-    ASSERT_EQ(execution2.compute(), WrapperResult::NO_ERROR);
-    ASSERT_EQ(CompareMatrices(expected3, *reinterpret_cast<Matrix3x4*>(outputData + offsetForActual)), 0);
+    for (int i = 0, e = mMaxMapCount + 10; i < e; i++) {
+        SCOPED_TRACE(i);
+        WrapperExecution execution2(&compilation2);
+        ASSERT_EQ(execution2.setInputFromMemory(0, &input, offsetForMatrix1, sizeof(Matrix3x4)),
+                  WrapperResult::NO_ERROR);
+        ASSERT_EQ(execution2.setOutputFromMemory(0, &actual, offsetForActual, sizeof(Matrix3x4)),
+                  WrapperResult::NO_ERROR);
+        ASSERT_EQ(execution2.compute(), WrapperResult::NO_ERROR);
+        ASSERT_EQ(CompareMatrices(expected3,
+                                  *reinterpret_cast<Matrix3x4*>(outputData + offsetForActual)), 0);
+    }
+
     close(weightsFd);
     close(inputFd);
     close(outputFd);
@@ -199,20 +222,6 @@ TEST_F(MemoryTest, TestFd) {
 }
 
 // Regression test for http://b/69621433 "MemoryFd leaks shared memory regions".
-class MemoryLeakTest : public ::testing::Test {
-protected:
-    void SetUp() override;
-
-    int mMaxMapCount = 0;
-};
-
-void MemoryLeakTest::SetUp() {
-    std::ifstream maxMapCountStream("/proc/sys/vm/max_map_count");
-    if (maxMapCountStream) {
-        maxMapCountStream >> mMaxMapCount;
-    }
-}
-
 TEST_F(MemoryLeakTest, IterativelyGetPointer) {
     ASSERT_GT(mMaxMapCount, 0);
 
@@ -252,6 +261,7 @@ TEST_F(MemoryLeakTest, IterativelyGetPointer) {
     close(fd);
 }
 
+// Regression test for http://b/69621433 "MemoryFd leaks shared memory regions".
 TEST_F(MemoryLeakTest, IterativelyInstantiate) {
     ASSERT_GT(mMaxMapCount, 0);