From feaeecdcd8aa9d13883074c6ac3fa6c5b430963d Mon Sep 17 00:00:00 2001
From: Jason Lee <jasonlee@lanl.gov>
Date: Mon, 1 Jul 2024 13:53:04 -0600
Subject: [PATCH] remove provider handle refcount

---
 include/dpusm/provider.h |  1 -
 src/provider.c           | 23 ++++++++++-------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/include/dpusm/provider.h b/include/dpusm/provider.h
index 324ecde..1e1c1fd 100644
--- a/include/dpusm/provider.h
+++ b/include/dpusm/provider.h
@@ -13,7 +13,6 @@ typedef struct dpusm_provider_handle {
     struct module *module;
     dpusm_pc_t capabilities; /* constant set of capabilities */
     const dpusm_pf_t *funcs; /* reference to a struct */
-    atomic_t refs;           /* how many users are holding this provider */
     struct list_head list;
     struct dpusm_provider_handle *self;
 } dpusm_ph_t;
diff --git a/src/provider.c b/src/provider.c
index 904ba0e..609bc20 100644
--- a/src/provider.c
+++ b/src/provider.c
@@ -227,7 +227,6 @@ dpusmph_init(struct module *module, const dpusm_pf_t *funcs)
         dpusmph->module = module;
         dpusmph->funcs = funcs;
         dpusmph->self = dpusmph;
-        atomic_set(&dpusmph->refs, 0);
     }
 
     return dpusmph;
@@ -236,6 +235,7 @@ dpusmph_init(struct module *module, const dpusm_pf_t *funcs)
 /* add a new provider */
 int
 dpusm_provider_register(dpusm_t *dpusm, struct module *module, const dpusm_pf_t *funcs) {
+    printk("refcount %d\n", module_refcount(module));
     const int rc = dpusm_provider_sane_at_load(funcs);
     if (rc != DPUSM_OK) {
         static const size_t max =
@@ -263,7 +263,6 @@ dpusm_provider_register(dpusm_t *dpusm, struct module *module, const dpusm_pf_t
 
         dpusm_mem_free(buf, size);
 
-        /* module_put(module); */
         return -EINVAL;
     }
 
@@ -274,14 +273,12 @@ dpusm_provider_register(dpusm_t *dpusm, struct module *module, const dpusm_pf_t
         printk("%s: DPUSM Provider with the name \"%s\" (%p) already exists. %zu providers registered.\n",
                __func__, module_name(module), *found, dpusm->count);
         dpusm_provider_write_unlock(dpusm);
-        /* module_put(module); */
         return -EEXIST;
     }
 
     dpusm_ph_t *provider = dpusmph_init(module, funcs);
     if (!provider) {
         dpusm_provider_write_unlock(dpusm);
-        /* module_put(module); */
         return -ECANCELED;
     }
 
@@ -304,7 +301,7 @@ dpusm_provider_unregister_handle(dpusm_t *dpusm, dpusm_ph_t **provider) {
     }
 
     int rc = 0;
-    const int refs = atomic_read(&(*provider)->refs);
+    const int refs = module_refcount((*provider)->module) - 1;
     if (refs) {
         printk("%s: Unregistering provider \"%s\" with %d references remaining.\n",
                __func__, module_name((*provider)->module), refs);
@@ -358,11 +355,10 @@ dpusm_provider_get(dpusm_t *dpusm, const char *name) {
             return NULL;
         }
 
-        atomic_inc(&(*provider)->refs);
         atomic_inc(&dpusm->active);
 
         printk("%s: User has been given a handle to \"%s\" (%p) (now %d users).\n",
-               __func__, name, *provider, atomic_read(&(*provider)->refs));
+               __func__, name, *provider, module_refcount((*provider)->module) - 1);
 
         if ((*provider)->funcs->at_connect) {
             (*provider)->funcs->at_connect();
@@ -385,14 +381,15 @@ dpusm_provider_put(dpusm_t *dpusm, void *handle) {
         return DPUSM_ERROR;
     }
 
-    if (!atomic_read(&(*provider)->refs)) {
+    struct module *module = (*provider)->module;
+
+    if (module_refcount(module) == 1) {
         printk("%s Error: Cannot decrement provider \"%s\" user count already at 0.\n",
-               __func__, module_name((*provider)->module));
+               __func__, module_name(module));
         return DPUSM_ERROR;
     }
 
-    module_put((*provider)->module);
-    atomic_dec(&(*provider)->refs);
+    module_put(module);
     atomic_dec(&dpusm->active);
 
     if ((*provider)->funcs) { /* provider might have been invalidated */
@@ -402,7 +399,7 @@ dpusm_provider_put(dpusm_t *dpusm, void *handle) {
     }
 
     printk("%s: User has returned a handle to \"%s\" (%p) (now %d users).\n",
-           __func__, module_name((*provider)->module), *provider, atomic_read(&(*provider)->refs));
+           __func__, module_name(module), *provider, module_refcount(module) - 1);
     return DPUSM_OK;
 }
 
@@ -423,7 +420,7 @@ void dpusm_provider_invalidate(dpusm_t *dpusm, const char *name) {
         (*provider)->funcs = NULL;
         memset(&(*provider)->capabilities, 0, sizeof((*provider)->capabilities));
         printk("%s: Provider \"%s\" has been invalidated with %d users active.\n",
-               __func__, name, atomic_read(&(*provider)->refs));
+               __func__, name, module_refcount((*provider)->module) - 1);
         /* not decrementing module reference count here - provider is still registered */
     }
     else {