Aio reap noreschedule

src/core/aio.c

@@ -9,6 +9,7 @@
 //
 
 #include "core/nng_impl.h"
+#include "core/taskq.h"
 #include <string.h>
 
 struct nni_aio_expire_q {
@@ -18,6 +19,7 @@ struct nni_aio_expire_q {
 	nni_thr  eq_thr;
 	nni_time eq_next; // next expiration
 	bool     eq_exit;
+	bool     eq_stop;
 };
 
 static nni_aio_expire_q **nni_aio_expire_q_list;
@@ -343,7 +345,7 @@ nni_aio_begin(nni_aio *aio)
 	aio->a_cancel_fn = NULL;
 
 	// We should not reschedule anything at this point.
-	if (aio->a_stop) {
+	if (aio->a_stop || eq->eq_stop) {
 		aio->a_result    = NNG_ECANCELED;
 		aio->a_cancel_fn = NULL;
 		aio->a_expire    = NNI_TIME_NEVER;
@@ -380,7 +382,7 @@ nni_aio_schedule(nni_aio *aio, nni_aio_cancel_fn cancel, void *data)
 	}
 
 	nni_mtx_lock(&eq->eq_mtx);
-	if (aio->a_stop) {
+	if (aio->a_stop || eq->eq_stop) {
 		nni_task_abort(&aio->a_task);
 		nni_mtx_unlock(&eq->eq_mtx);
 		return (NNG_ECLOSED);
@@ -595,16 +597,15 @@ nni_aio_expire_loop(void *arg)
 			nni_mtx_unlock(mtx);
 			return;
 		}
-		if (now < next) {
-			// Early wake up (just to reschedule), no need to
-			// rescan the list.  This is an optimization.
+		if (now < next && !(q->eq_stop && aio != NULL)) {
+			// nothing to do!
 			nni_cv_until(cv, next);
 			continue;
 		}
 		q->eq_next = NNI_TIME_NEVER;
 		exp_idx    = 0;
 		while (aio != NULL) {
-			if ((aio->a_expire < now) &&
+			if ((q->eq_stop || aio->a_expire < now) &&
 			    (exp_idx < NNI_EXPIRE_BATCH)) {
 				nni_aio *nxt;
 
@@ -627,7 +628,14 @@ nni_aio_expire_loop(void *arg)
 
 		for (uint32_t i = 0; i < exp_idx; i++) {
 			aio = expires[i];
-			rv  = aio->a_expire_ok ? 0 : NNG_ETIMEDOUT;
+			if (q->eq_stop) {
+				rv = NNG_ECANCELED;
+			} else if (aio->a_expire_ok) {
+				aio->a_expire_ok = false;
+				rv               = 0;
+			} else {
+				rv = NNG_ETIMEDOUT;
+			}
 
 			nni_aio_cancel_fn cancel_fn  = aio->a_cancel_fn;
 			void             *cancel_arg = aio->a_cancel_arg;
@@ -639,18 +647,22 @@ nni_aio_expire_loop(void *arg)
 			// If there is no cancellation function, then we cannot
 			// terminate the aio - we've tried, but it has to run
 			// to its natural conclusion.
-			if (cancel_fn != NULL) {
+			//
+			// For the special case of sleeping, we don't need to
+			// drop the lock and call the cancel function, we are
+			// already doing it right here!
+			if (aio->a_sleep) {
+				aio->a_result = rv;
+				aio->a_sleep  = false;
+				nni_task_dispatch(&aio->a_task);
+			} else if (cancel_fn != NULL) {
 				nni_mtx_unlock(mtx);
 				cancel_fn(aio, cancel_arg, rv);
 				nni_mtx_lock(mtx);
 			}
 			aio->a_expiring = false;
 		}
 		nni_cv_wake(cv);
-
-		if (now < q->eq_next) {
-			nni_cv_until(cv, q->eq_next);
-		}
 	}
 }
 
@@ -768,12 +780,30 @@ nni_sleep_aio(nng_duration ms, nng_aio *aio)
 	}
 }
 
+static bool
+nni_aio_expire_q_stop(nni_aio_expire_q *eq)
+{
+	bool result = false;
+	if (eq != NULL) {
+		nni_mtx_lock(&eq->eq_mtx);
+		eq->eq_stop = true;
+		nni_cv_wake(&eq->eq_cv);
+		while (!nni_list_empty(&eq->eq_list)) {
+			result = true;
+			nni_cv_wait(&eq->eq_cv);
+		}
+		nni_mtx_unlock(&eq->eq_mtx);
+	}
+	return (result);
+}
+
 static void
 nni_aio_expire_q_free(nni_aio_expire_q *eq)
 {
 	if (eq == NULL) {
 		return;
 	}
+	NNI_ASSERT(eq->eq_stop);
 	if (!eq->eq_exit) {
 		nni_mtx_lock(&eq->eq_mtx);
 		eq->eq_exit = true;
@@ -810,6 +840,18 @@ nni_aio_expire_q_alloc(void)
 	return (eq);
 }
 
+bool
+nni_aio_sys_drain(void)
+{
+	bool result = false;
+	for (int i = 0; i < nni_aio_expire_q_cnt; i++) {
+		if (nni_aio_expire_q_stop(nni_aio_expire_q_list[i])) {
+			result = true;
+		}
+	}
+	return (result);
+}
+
 void
 nni_aio_sys_fini(void)
 {

src/core/aio.h

@@ -193,6 +193,7 @@ extern void nni_aio_completions_add(
     nni_aio_completions *, nni_aio *, int, size_t);
 
 extern int  nni_aio_sys_init(nng_init_params *);
+extern bool nni_aio_sys_drain(void);
 extern void nni_aio_sys_fini(void);
 
 typedef struct nni_aio_expire_q nni_aio_expire_q;

src/core/aio_test.c

@@ -72,6 +72,46 @@ test_sleep_timeout(void)
 	nng_aio_free(aio);
 }
 
+static void
+sleep_reap(void *arg)
+{
+	nng_aio *aio = *(nng_aio **) arg;
+	if (nng_aio_result(aio) != NNG_ECANCELED) {
+		NUTS_FAIL(nng_aio_result(aio), NNG_ECANCELED);
+	}
+	nng_aio_reap(aio);
+}
+
+static void
+test_sleep_fini(void)
+{
+	static nng_aio *aio;
+	NUTS_TRUE(nng_aio_alloc(&aio, sleep_reap, &aio) == 0);
+	nng_sleep_aio(20000, aio);
+	nng_msleep(1);
+	// intentionally we do not free the aio here. reap should clean it.
+	nng_fini();
+	nng_init(NULL); // so that TEST_FINI will reap
+}
+
+static void
+test_sleep_fini_many(void)
+{
+#define NIOS 2000
+	static nng_aio *aios[NIOS];
+	for (int i = 0; i < NIOS; i++) {
+		int rv = nng_aio_alloc(&(aios[i]), sleep_reap, &(aios[i]));
+		if (rv != 0) {
+			NUTS_ASSERT(rv == 0);
+		}
+	}
+	for (int i = 0; i < NIOS; i++) {
+		nng_sleep_aio(20000, aios[i]);
+	}
+	nng_fini();
+	nng_init(NULL);
+}
+
 void
 test_insane_nio(void)
 {
@@ -400,6 +440,8 @@ test_aio_busy(void)
 NUTS_TESTS = {
 	{ "sleep", test_sleep },
 	{ "sleep timeout", test_sleep_timeout },
+	{ "sleep fini", test_sleep_fini },
+	{ "sleep fini many", test_sleep_fini_many },
 	{ "insane nio", test_insane_nio },
 	{ "provider cancel", test_provider_cancel },
 	{ "consumer cancel", test_consumer_cancel },

src/core/init.c

@@ -135,12 +135,19 @@ nng_fini(void)
 	}
 	nni_sock_closeall();
 	nni_sp_tran_sys_fini();
+
+	// Drain everything. This is important because some of
+	// these subsystems can dispatch things to other ones.
+	// So we need them *all* to be empty before proceeding.
+	while ((nni_aio_sys_drain() || nni_taskq_sys_drain() ||
+	    nni_reap_sys_drain())) {
+		continue;
+	}
 	nni_tls_sys_fini();
-	nni_reap_drain();
-	nni_aio_sys_fini();
 	nni_taskq_sys_fini();
-	nni_reap_sys_fini(); // must be before timer and aio (expire)
+	nni_aio_sys_fini();
 	nni_id_map_sys_fini();
+	nni_reap_sys_fini(); // must be near the end
 	nni_plat_fini();
 	nni_atomic_flag_reset(&init_busy);
 }

src/core/reap.c

@@ -18,9 +18,9 @@
 static nni_reap_list *reap_list = NULL;
 static nni_thr        reap_thr;
 static bool           reap_exit = false;
-static nni_mtx        reap_mtx = NNI_MTX_INITIALIZER;
+static nni_mtx        reap_mtx  = NNI_MTX_INITIALIZER;
 static bool           reap_empty;
-static nni_cv         reap_work_cv = NNI_CV_INITIALIZER(&reap_mtx);
+static nni_cv         reap_work_cv  = NNI_CV_INITIALIZER(&reap_mtx);
 static nni_cv         reap_empty_cv = NNI_CV_INITIALIZER(&reap_mtx);
 
 static void
@@ -90,14 +90,17 @@ nni_reap(nni_reap_list *rl, void *item)
 	nni_mtx_unlock(&reap_mtx);
 }
 
-void
-nni_reap_drain(void)
+bool
+nni_reap_sys_drain(void)
 {
+	bool result = false;
 	nni_mtx_lock(&reap_mtx);
 	while (!reap_empty) {
+		result = true;
 		nni_cv_wait(&reap_empty_cv);
 	}
 	nni_mtx_unlock(&reap_mtx);
+	return (result);
 }
 
 int

src/core/reap.h

@@ -1,5 +1,5 @@
 //
-// Copyright 2020 Staysail Systems, Inc. <[email protected]>
+// Copyright 2024 Staysail Systems, Inc. <[email protected]>
 // Copyright 2018 Capitar IT Group BV <[email protected]>
 //
 // This software is supplied under the terms of the MIT License, a
@@ -56,8 +56,9 @@ struct nni_reap_list {
 
 extern void nni_reap(nni_reap_list *, void *);
 
-// nni_reap_drain waits for the reap queue to be drained.
-extern void nni_reap_drain(void);
+// nni_reap_sys_drain waits for the reap queue to be drained.
+// It returns true if it found anything to wait for.
+extern bool nni_reap_sys_drain(void);
 
 extern int  nni_reap_sys_init(void);
 extern void nni_reap_sys_fini(void);

src/core/taskq.c

@@ -59,6 +59,7 @@ nni_taskq_thread(void *self)
 			continue;
 		}
 
+		nni_cv_wake(&tq->tq_wait_cv);
 		if (!tq->tq_run) {
 			break;
 		}
@@ -127,6 +128,19 @@ nni_taskq_fini(nni_taskq *tq)
 	NNI_FREE_STRUCT(tq);
 }
 
+bool
+nni_taskq_drain(nni_taskq *tq)
+{
+	bool result = false;
+	nni_mtx_lock(&tq->tq_mtx);
+	while (!nni_list_empty(&tq->tq_tasks)) {
+		result = true;
+		nni_cv_wait(&tq->tq_wait_cv);
+	}
+	nni_mtx_unlock(&tq->tq_mtx);
+	return (result);
+}
+
 void
 nni_task_exec(nni_task *task)
 {
@@ -263,6 +277,12 @@ nni_taskq_sys_init(nng_init_params *params)
 	return (nni_taskq_init(&nni_taskq_systq, (int) num_thr));
 }
 
+bool
+nni_taskq_sys_drain(void)
+{
+	return (nni_taskq_drain(nni_taskq_systq));
+}
+
 void
 nni_taskq_sys_fini(void)
 {

src/core/taskq.h

@@ -62,6 +62,7 @@ extern void nni_task_init(nni_task *, nni_taskq *, nni_cb, void *);
 extern void nni_task_fini(nni_task *);
 
 extern int  nni_taskq_sys_init(nng_init_params *);
+extern bool nni_taskq_sys_drain(void);
 extern void nni_taskq_sys_fini(void);
 
 // nni_task implementation details are not to be used except by the

src/platform/windows/win_ipcdial.c

@@ -295,7 +295,7 @@ nni_win_ipc_sysfini(void)
 {
 	ipc_dial_work *worker = &ipc_connector;
 
-	nni_reap_drain(); // so that listeners get cleaned up.
+	nni_reap_sys_drain(); // so that listeners get cleaned up.
 
 	nni_mtx_lock(&worker->mtx);
 	worker->exit = 1;