From 2958f7eacedc2e498ee232ba9c76e54239cf829d Mon Sep 17 00:00:00 2001
From: Stanley Tsang <stanley.tsang@amd.com>
Date: Fri, 13 Nov 2020 10:32:42 -0700
Subject: [PATCH] Fixing IPC handle leak (#302)

---
 src/clique/CliqueManager.cc |  6 +++-
 src/clique/HandleCache.h    | 72 ++++++++++++++++++++++++++++++++-----
 2 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/src/clique/CliqueManager.cc b/src/clique/CliqueManager.cc
index fef2a91599..dde0cf0085 100644
--- a/src/clique/CliqueManager.cc
+++ b/src/clique/CliqueManager.cc
@@ -82,7 +82,11 @@ void CliqueManager::CleanUp()
   {
     // Release caches
     if (m_ipcHandleSendCache) delete m_ipcHandleSendCache;
-    if (m_ipcHandleSendCache) delete m_ipcHandleRecvCache;
+    if (m_ipcHandleRecvCache)
+    {
+        m_ipcHandleRecvCache->close();
+        delete m_ipcHandleRecvCache;
+    }
 
     // Close shared memory
     m_shmHandles.Close();
diff --git a/src/clique/HandleCache.h b/src/clique/HandleCache.h
index dc479e00e8..245b78f80d 100644
--- a/src/clique/HandleCache.h
+++ b/src/clique/HandleCache.h
@@ -103,14 +103,16 @@ public:
 
         return std::pair<iterator, bool>(it, inserted);
     }
-
+    ncclResult_t close();
 private:
-    void pop()
-    {
-        typename LRUCache::iterator it = m_cache.find(m_lruHistory.front());
-        m_cache.erase(it);
-        m_lruHistory.pop_front();
-    }
+    // tag for dispatch
+    template<class U>
+        struct CloseTag{};
+
+    hipError_t CloseIfPointer(CloseTag<hipIpcMemHandle_t> tag, iterator it);
+    hipError_t CloseIfPointer(CloseTag<void*> tag, iterator it);
+
+    void pop();
 
     void updateHistory(const iterator& it)
     {
@@ -133,8 +135,60 @@ auto hipIpcMemHandleEqual = [](const hipIpcMemHandle_t& l, const hipIpcMemHandle
     return memcmp(l.reserved, r.reserved, sizeof(l.reserved)) == 0;
 };
 
-//typedef llvm::DenseMap<uint64_t, hipIpcMemHandle_t> SendCache;
-//typedef llvm::DenseMap<hipIpcMemHandle_t, void*, decltype(&HandleHash), decltype(HandleEqual)> RecvCache;
+template <
+    class Key,
+    class Value,
+    class Hash,
+    class KeyEqual,
+    class Allocator
+>
+ncclResult_t NcclIpcHandleCache<Key, Value, Hash, KeyEqual, Allocator>::close()
+{
+    for (auto it = m_cache.begin(); it != m_cache.end(); it ++)
+    {
+        CUDACHECK(CloseIfPointer(CloseTag<Value>{}, it));
+    }
+    return ncclSuccess;
+}
+
+template <
+    class Key,
+    class Value,
+    class Hash,
+    class KeyEqual,
+    class Allocator
+>
+void NcclIpcHandleCache<Key, Value, Hash, KeyEqual, Allocator>::pop()
+{
+    typename LRUCache::iterator it = m_cache.find(m_lruHistory.front());
+    CloseIfPointer(CloseTag<Value>{}, it);
+    m_cache.erase(it);
+    m_lruHistory.pop_front();
+}
+
+template <
+    class Key,
+    class Value,
+    class Hash,
+    class KeyEqual,
+    class Allocator
+>
+hipError_t NcclIpcHandleCache<Key, Value, Hash, KeyEqual, Allocator>::CloseIfPointer(CloseTag<void*> tag, iterator it)
+{
+    return hipIpcCloseMemHandle(it->second.first);
+}
+
+template <
+    class Key,
+    class Value,
+    class Hash,
+    class KeyEqual,
+    class Allocator
+>
+hipError_t NcclIpcHandleCache<Key, Value, Hash, KeyEqual, Allocator>::CloseIfPointer(CloseTag<hipIpcMemHandle_t> tag, iterator it)
+{
+    return hipSuccess;
+}
 
 typedef NcclIpcHandleCache<uint64_t, hipIpcMemHandle_t, std::hash<uint64_t>, std::equal_to<uint64_t>, std::allocator< std::pair<const uint64_t, std::pair<hipIpcMemHandle_t, std::list<uint64_t>::iterator>>>> NcclIpcHandleSendCache;
 typedef NcclIpcHandleCache<hipIpcMemHandle_t, void*, decltype(&hipIpcMemHandleHash), decltype(hipIpcMemHandleEqual), std::allocator< std::pair<const hipIpcMemHandle_t, std::pair<void*, std::list<hipIpcMemHandle_t>::iterator>>>> NcclIpcHandleRecvCache;