feat: listener refcnt

emqx · Oct 19, 2023 · ee1a8a6 · ee1a8a6
1 parent b2fb5fa
commit ee1a8a6
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 24 deletions.
diff --git a/c_src/quicer_config.c b/c_src/quicer_config.c
@@ -760,7 +760,7 @@ getopt3(ErlNifEnv *env,
   ERL_NIF_TERM elevel = argv[2];
 
   void *q_ctx;
-  ERL_NIF_TERM res = ERROR_TUPLE_2(ATOM_CLOSED);
+  ERL_NIF_TERM res = ATOM_ERROR_NOT_FOUND;
 
   if (!enif_is_atom(env, eopt))
     {
@@ -781,32 +781,38 @@ getopt3(ErlNifEnv *env,
     }
   else if (enif_get_resource(env, ctx, ctx_stream_t, &q_ctx))
     {
-      if (get_stream_handle(q_ctx))
+      if (!get_stream_handle(q_ctx))
         {
-          res = get_stream_opt(env, (QuicerStreamCTX *)q_ctx, eopt, elevel);
-          put_stream_handle(q_ctx);
+          goto Exit;
         }
+      res = get_stream_opt(env, (QuicerStreamCTX *)q_ctx, eopt, elevel);
+      put_stream_handle(q_ctx);
     }
   else if (enif_get_resource(env, ctx, ctx_connection_t, &q_ctx))
     {
-      if (get_conn_handle(q_ctx))
+      if (!get_conn_handle(q_ctx))
         {
-          res = get_connection_opt(env, (QuicerConnCTX *)q_ctx, eopt, elevel);
-          put_conn_handle(q_ctx);
+          goto Exit;
         }
+      res = get_connection_opt(env, (QuicerConnCTX *)q_ctx, eopt, elevel);
+      put_conn_handle(q_ctx);
     }
   else if (enif_get_resource(env, ctx, ctx_listener_t, &q_ctx))
     {
-      enif_mutex_lock(((QuicerListenerCTX *)q_ctx)->lock);
+      if (!get_listener_handle(q_ctx))
+        {
+          goto Exit;
+        }
       res = get_listener_opt(env, (QuicerListenerCTX *)q_ctx, eopt, elevel);
-      enif_mutex_unlock(((QuicerListenerCTX *)q_ctx)->lock);
+      put_listener_handle(q_ctx);
     }
   else
     { //@todo support GLOBAL, REGISTRATION and CONFIGURATION
       return ERROR_TUPLE_2(ATOM_BADARG);
     }
-
   return res;
+Exit:
+  return ERROR_TUPLE_2(ATOM_CLOSED);
 }
 
 ERL_NIF_TERM
@@ -873,25 +879,42 @@ setopt4(ErlNifEnv *env, __unused_parm__ int argc, const ERL_NIF_TERM argv[])
     }
   else if (enif_get_resource(env, ctx, ctx_stream_t, &q_ctx))
     {
+      if (!get_stream_handle(q_ctx))
+        {
+          goto Exit;
+        }
       res = set_stream_opt(
           env, (QuicerStreamCTX *)q_ctx, eopt, evalue, elevel);
+      put_stream_handle(q_ctx);
     }
   else if (enif_get_resource(env, ctx, ctx_connection_t, &q_ctx))
     {
+      if (!get_conn_handle(q_ctx))
+        {
+          goto Exit;
+        }
       res = set_connection_opt(
           env, (QuicerConnCTX *)q_ctx, eopt, evalue, elevel);
+      put_conn_handle(q_ctx);
     }
   else if (enif_get_resource(env, ctx, ctx_listener_t, &q_ctx))
     {
+      if (!get_listener_handle(q_ctx))
+        {
+          goto Exit;
+        }
       res = set_listener_opt(
           env, (QuicerListenerCTX *)q_ctx, eopt, evalue, elevel);
+      put_listener_handle(q_ctx);
     }
   else
     { //@todo support GLOBAL, REGISTRATION and CONFIGURATION
       return ERROR_TUPLE_2(ATOM_BADARG);
     }
 
   return res;
+Exit:
+  return ERROR_TUPLE_2(ATOM_CLOSED);
 }
 
 bool

diff --git a/c_src/quicer_connection.c b/c_src/quicer_connection.c
@@ -977,9 +977,9 @@ sockname1(ErlNifEnv *env, __unused_parm__ int args, const ERL_NIF_TERM argv[])
   if (enif_get_resource(env, argv[0], ctx_connection_t, &q_ctx))
     {
       if (!get_conn_handle((QuicerConnCTX *)q_ctx))
-      {
-        return ERROR_TUPLE_2(ATOM_CLOSED);
-      }
+        {
+          return ERROR_TUPLE_2(ATOM_CLOSED);
+        }
       Handle = (((QuicerConnCTX *)q_ctx))->Connection;
       Param = QUIC_PARAM_CONN_LOCAL_ADDRESS;
       Status = MsQuic->GetParam(Handle, Param, &addrSize, &addr);

diff --git a/c_src/quicer_ctx.c b/c_src/quicer_ctx.c
@@ -342,3 +342,22 @@ get_conn_handle(QuicerConnCTX *c_ctx)
 {
   return CxPlatRefIncrementNonZero(&c_ctx->ref_count, 1);
 }
+
+inline void
+put_listener_handle(QuicerListenerCTX *l_ctx)
+{
+  if (CxPlatRefDecrement(&l_ctx->ref_count) && l_ctx->Listener)
+    {
+      HQUIC Listener = l_ctx->Listener;
+      enif_mutex_lock(l_ctx->lock);
+      l_ctx->Listener = NULL;
+      enif_mutex_unlock(l_ctx->lock);
+      MsQuic->ListenerClose(Listener);
+    }
+}
+
+inline BOOLEAN
+get_listener_handle(QuicerListenerCTX *l_ctx)
+{
+  return CxPlatRefIncrementNonZero(&l_ctx->ref_count, 1);
+}
diff --git a/c_src/quicer_ctx.h b/c_src/quicer_ctx.h
@@ -56,6 +56,8 @@ typedef struct QuicerListenerCTX
   QuicerConfigCTX *config_resource;
   QuicerRegistrationCTX *r_ctx;
   HQUIC Listener;
+  // track lifetime of Connection handle
+  CXPLAT_REF_COUNT ref_count;
   QUICER_ACCEPTOR_QUEUE *acceptor_queue;
   ErlNifPid listenerPid;
   ErlNifMonitor owner_mon;
@@ -177,4 +179,7 @@ BOOLEAN get_stream_handle(QuicerStreamCTX *s_ctx);
 void put_conn_handle(QuicerConnCTX *c_ctx);
 BOOLEAN get_conn_handle(QuicerConnCTX *c_ctx);
 
+void put_listener_handle(QuicerListenerCTX *l_ctx);
+BOOLEAN get_listener_handle(QuicerListenerCTX *l_ctx);
+
 #endif // __QUICER_CTX_H_
diff --git a/c_src/quicer_listener.c b/c_src/quicer_listener.c
@@ -294,6 +294,7 @@ listen2(ErlNifEnv *env, __unused_parm__ int argc, const ERL_NIF_TERM argv[])
       free_certificate(&CredConfig);
       return ERROR_TUPLE_2(ATOM_ERROR_NOT_ENOUGH_MEMORY);
     }
+  CxPlatRefInitialize(&l_ctx->ref_count);
 
   if (is_verify && cacertfile)
     {
@@ -490,20 +491,14 @@ stop_listener1(ErlNifEnv *env,
     {
       return ERROR_TUPLE_2(ATOM_BADARG);
     }
-  enif_mutex_lock(l_ctx->lock);
-  if (!l_ctx->Listener)
+  if (!get_listener_handle(l_ctx))
     {
       ret = ERROR_TUPLE_2(ATOM_CLOSED);
-      goto exit;
+      return ret; // follow otp behaviour?
     }
-  else if (!l_ctx->is_stopped)
-    {
-      l_ctx->is_stopped = TRUE;
-      // void return
-      MsQuic->ListenerStop(l_ctx->Listener);
-    }
-exit:
-  enif_mutex_unlock(l_ctx->lock);
+  l_ctx->is_stopped = TRUE;
+  MsQuic->ListenerStop(l_ctx->Listener);
+  put_listener_handle(l_ctx);
   return ret;
 }