From 69dbaaa015beff0e83cbfdc52363d1a1c9cdb51b Mon Sep 17 00:00:00 2001
From: Tim Harris <tiharr@microsoft.com>
Date: Wed, 30 Sep 2020 11:26:02 +0100
Subject: [PATCH] Add additional test cases to check for leaks in thread pool
 creation / destruction (#5311)

Add additional test cases such as ThreadPoolTest.TestPoolCreation_10Iter to create and destroy thread pools to watch for any memory leaks. Running under Valgrind, these tests should show all of the data allocated being deallocated again.

Two recent issues #5176 and #5292 indicated memory leaks. The test cases help identify whether or not any of the data structures used in the thread pool are being leaked. Currently, on WSL, the only data not being de-allocated in these tests are a small number of nsync waiter objects. This behavior is as expected (the waiter objects should be held on a free list in the nsync library).
---
 onnxruntime/test/platform/threadpool_test.cc | 37 ++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/onnxruntime/test/platform/threadpool_test.cc b/onnxruntime/test/platform/threadpool_test.cc
index fda2cbbc49..34ca40b87e 100644
--- a/onnxruntime/test/platform/threadpool_test.cc
+++ b/onnxruntime/test/platform/threadpool_test.cc
@@ -137,6 +137,30 @@ void TestBurstScheduling(const std::string& name, int num_tasks) {
   }
 }
 
+void TestPoolCreation(const std::string&, int iter) {
+  // Test creating and destroying thread pools.  This can be used with Valgrind to help
+  // check for memory leaks related to the initialization and clean-up code.  For instance
+  //
+  //  valgrind --leak-check=full ./onnxruntime_test_all --gtest_filter=ThreadPoolTest.TestPoolCreation_10Iter
+  //
+  // We create #iter thread pools, and within each of them run a loop of #per_iter steps.
+  std::atomic<std::ptrdiff_t> ctr{0};
+  constexpr std::ptrdiff_t per_iter = 1024;
+  constexpr int num_threads = 4;
+  for (auto i = 0; i < iter; i++) {
+    auto tp = onnxruntime::make_unique<ThreadPool>(&onnxruntime::Env::Default(),
+                                                   onnxruntime::ThreadOptions(),
+                                                   nullptr,
+                                                   num_threads,
+                                                   true);
+    tp->ParallelFor(per_iter, 0.0,
+                    [&](std::ptrdiff_t s, std::ptrdiff_t e) {
+                      ctr += e - s;
+                    });
+  }
+  ASSERT_EQ(ctr, iter * per_iter);
+}
+
 }  // namespace
 
 namespace onnxruntime {
@@ -253,6 +277,19 @@ TEST(ThreadPoolTest, TestBurstScheduling_65536Task) {
   // buffer tasks.
   TestBurstScheduling("TestBurstScheduling_65536Tasks", 65536);
 }
+
+TEST(ThreadPoolTest, TestPoolCreation_1Iter) {
+  TestPoolCreation("TestPoolCreation_1Iter", 1);
+}
+
+TEST(ThreadPoolTest, TestPoolCreation_10Iter) {
+  TestPoolCreation("TestPoolCreation_10Iter", 10);
+}
+
+TEST(ThreadPoolTest, TestPoolCreation_100Iter) {
+  TestPoolCreation("TestPoolCreation_100Iter", 100);
+}
+
 #ifdef _WIN32
 TEST(ThreadPoolTest, TestStackSize) {
   ThreadOptions to;