Fix Python Shm Client Leak (#614)

* Fix Python Shm Client Leak

src/python/library/tritonclient/utils/shared_memory/__init__.py

@@ -284,9 +284,6 @@ def destroy_shared_memory_region(shm_handle):
     SharedMemoryException
         If unable to unlink the shared memory region.
     """
-
-    _raise_if_error(c_int(_cshm_shared_memory_region_destroy(shm_handle)))
-
     shm_fd = c_int()
     offset = c_uint64()
     byte_size = c_uint64()
@@ -304,8 +301,13 @@ def destroy_shared_memory_region(shm_handle):
             )
         )
     )
+    # It is safer to remove the shared memory key from the list before
+    # deleting the shared memory region because if the deletion should
+    # fail, a re-attempt could result in a segfault. Secondarily, if we
+    # fail to delete a region, we should not report it back to the user
+    # as a valid memory region.
     mapped_shm_regions.remove(shm_key.value.decode("utf-8"))
-
+    _raise_if_error(c_int(_cshm_shared_memory_region_destroy(shm_handle)))
     return
 
 

src/python/library/tritonclient/utils/shared_memory/shared_memory.cc

@@ -33,6 +33,7 @@
 
 #include <cstring>
 #include <iostream>
+#include <memory>
 
 #include "shared_memory_handle.h"
 
@@ -132,8 +133,8 @@ GetSharedMemoryHandleInfo(
 int
 SharedMemoryRegionDestroy(void* shm_handle)
 {
-  SharedMemoryHandle* handle =
-      reinterpret_cast<SharedMemoryHandle*>(shm_handle);
+  std::unique_ptr<SharedMemoryHandle> handle(
+      reinterpret_cast<SharedMemoryHandle*>(shm_handle));
   void* shm_addr = reinterpret_cast<char*>(handle->base_addr_);
   int status = munmap(shm_addr, handle->byte_size_);
   if (status == -1) {
@@ -144,7 +145,6 @@ SharedMemoryRegionDestroy(void* shm_handle)
   if (shm_fd == -1) {
     return -5;
   }
-
   return 0;
 }