Add reference counter of users of TRACKER

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

src/provider/provider_tracking.c

@@ -158,13 +158,20 @@ static umf_result_t trackingInitialize(void *params, void **ret) {
 
     *provider = *((umf_tracking_memory_provider_t *)params);
     *ret = provider;
+    umfMemoryTrackerIncRef();
     return UMF_RESULT_SUCCESS;
 }
 
 static void trackingFinalize(void *provider) {
     umf_tracking_memory_provider_t *p =
         (umf_tracking_memory_provider_t *)provider;
 
+    if (umfMemoryTrackerDecRef()) {
+        // it is not the last user of the TRACKER
+        free(provider);
+        return;
+    }
+
     uintptr_t rkey;
     void *rvalue;
     size_t n_items = 0;

src/provider/provider_tracking.h

@@ -39,6 +39,11 @@ void umfTrackingMemoryProviderGetUpstreamProvider(
     umf_memory_provider_handle_t hTrackingProvider,
     umf_memory_provider_handle_t *hUpstream);
 
+void umfMemoryTrackerIncRef(void);
+
+// returns 0 only if n_ref_TRACKER==0 (at exit), !=0 otherwise
+int umfMemoryTrackerDecRef(void);
+
 #ifdef __cplusplus
 }
 #endif

src/provider/provider_tracking_linux.c

@@ -13,6 +13,7 @@
 #include <stddef.h>
 
 static critnib *TRACKER = NULL;
+static int n_ref_TRACKER = 0;
 
 void __attribute__((constructor)) createLibTracker(void) {
     TRACKER = critnib_new();
@@ -28,3 +29,12 @@ void umfTrackingMemoryProviderInit(void) {
 umf_memory_tracker_handle_t umfMemoryTrackerGet(void) {
     return (umf_memory_tracker_handle_t)TRACKER;
 }
+
+void umfMemoryTrackerIncRef(void) {
+    __atomic_fetch_add(&n_ref_TRACKER, 1, __ATOMIC_SEQ_CST);
+}
+
+// returns 0 only if n_ref_TRACKER==0 (at exit), !=0 otherwise
+int umfMemoryTrackerDecRef(void) {
+    return (__atomic_fetch_sub(&n_ref_TRACKER, 1, __ATOMIC_SEQ_CST) != 1);
+}

src/provider/provider_tracking_windows.c

@@ -11,8 +11,10 @@
 #include "provider_tracking.h"
 
 #include <windows.h>
+#include <winnt.h>
 
 static critnib *TRACKER = NULL;
+static long n_ref_TRACKER = 0;
 
 #if defined(UMF_SHARED_LIBRARY)
 BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {
@@ -47,3 +49,12 @@ void umfTrackingMemoryProviderInit(void) {
 umf_memory_tracker_handle_t umfMemoryTrackerGet(void) {
     return (umf_memory_tracker_handle_t)TRACKER;
 }
+
+void umfMemoryTrackerIncRef(void) {
+    InterlockedExchangeAdd((volatile long *)&n_ref_TRACKER, 1);
+}
+
+// returns 0 only if n_ref_TRACKER==0 (at exit), !=0 otherwise
+int umfMemoryTrackerDecRef(void) {
+    return (InterlockedExchangeAdd(&n_ref_TRACKER, -1) != 1);
+}