Add support for synchronous AIO completions.

We add a flag (auto-clearing) that can be set on an AIO to indicate
that the AIO should not processed asynchronously on a taskq.  This
can be used to enhance performance in some cases, but it can also
be used to permit an AIO be destroyed from a completion callback.
(For the latter, the callback must execute the new nni_aio_fini_cb()
routine, which destroys the AIO without waiting for it to finish.)

2 files changed, 48 insertions, 3 deletions

diff --git a/src/core/aio.c b/src/core/aio.c
index 9b308cd1..568cd633 100644
--- a/src/core/aio.c
+++ b/src/core/aio.c
@@ -68,6 +68,9 @@ nni_aio_init(nni_aio **aiop, nni_cb cb, void *arg)
 	nni_cv_init(&aio->a_cv, &nni_aio_lk);
 	aio->a_expire = NNI_TIME_NEVER;
 	aio->a_init   = 1;
+	if (arg == NULL) {
+		arg = aio;
+	}
 	nni_task_init(NULL, &aio->a_task, cb, arg);
 	*aiop = aio;
 	return (0);
@@ -86,6 +89,13 @@ nni_aio_fini(nni_aio *aio)
 	}
 }
 
+void
+nni_aio_fini_cb(nni_aio *aio)
+{
+	nni_cv_fini(&aio->a_cv);
+	NNI_FREE_STRUCT(aio);
+}
+
 // nni_aio_stop cancels any oustanding operation, and waits for the
 // callback to complete, if still running.  It also marks the AIO as
 // stopped, preventing further calls to nni_aio_start from succeeding.
@@ -175,10 +185,18 @@ nni_aio_count(nni_aio *aio)
 }
 
 void
+nni_aio_set_synch(nni_aio *aio)
+{
+	aio->a_synch = 1;
+}
+
+void
 nni_aio_wait(nni_aio *aio)
 {
 	nni_mtx_lock(&nni_aio_lk);
-	while ((aio->a_active) && (!aio->a_done)) {
+	// Wait until we're done, and the synchronous completion flag
+	// is cleared (meaning any synch completion is finished).
+	while ((aio->a_active) && ((!aio->a_done) || (aio->a_synch))) {
 		aio->a_waiting = 1;
 		nni_cv_wait(&aio->a_cv);
 	}
@@ -396,11 +414,18 @@ nni_aio_expire_loop(void *arg)
 		NNI_ASSERT(aio->a_prov_cancel == NULL);
 		aio->a_expiring = 0;
 		aio->a_done     = 1;
+		if (!aio->a_synch) {
+			nni_task_dispatch(&aio->a_task);
+		} else {
+			nni_mtx_unlock(&nni_aio_lk);
+			aio->a_task.task_cb(aio->a_task.task_arg);
+			nni_mtx_lock(&nni_aio_lk);
+			aio->a_synch = 0;
+		}
 		if (aio->a_waiting) {
 			aio->a_waiting = 0;
 			nni_cv_wake(&aio->a_cv);
 		}
-		nni_task_dispatch(&aio->a_task);
 		nni_mtx_unlock(&nni_aio_lk);
 	}
 }
diff --git a/src/core/aio.h b/src/core/aio.h
index eae17446..717e7995 100644
--- a/src/core/aio.h
+++ b/src/core/aio.h
@@ -35,7 +35,8 @@ struct nni_aio {
 	unsigned a_active : 1;   // aio was started
 	unsigned a_expiring : 1; // expiration callback in progress
 	unsigned a_waiting : 1;  // a thread is waiting for this to finish
-	unsigned a_pad : 26;     // ensure 32-bit alignment
+	unsigned a_synch : 1;    // run completion synchronously
+	unsigned a_pad : 25;     // ensure 32-bit alignment
 	nni_task a_task;
 
 	// Read/write operations.
@@ -76,6 +77,12 @@ extern int nni_aio_init(nni_aio **, nni_cb, void *);
 // on zero'd memory.
 extern void nni_aio_fini(nni_aio *);
 
+// nni_aio_fini_cb finalizes the aio WITHOUT waiting for it to complete.
+// This is intended exclusively for finalizing an AIO from a completion
+// callack for that AIO. It is important that the caller ensure that nothing
+// else might be waiting for that AIO or using it.
+extern void nni_aio_fini_cb(nni_aio *);
+
 // nni_aio_stop cancels any unfinished I/O, running completion callbacks,
 // but also prevents any new operations from starting (nni_aio_start will
 // return NNG_ESTATE).  This should be called before nni_aio_fini().  The
@@ -102,6 +109,19 @@ extern void *   nni_aio_get_pipe(nni_aio *);
 extern void     nni_aio_set_ep(nni_aio *, void *);
 extern void *   nni_aio_get_ep(nni_aio *);
 
+// nni_aio_set_synch sets a synchronous completion flag on the AIO.
+// When this is set, the next time the AIO is completed, the callback
+// be run synchronously, from the thread calling the finish routine.
+// It is important that this only be set when the provider knows that
+// it is not holding any locks or resources when completing the operation,
+// or when the consumer knows that the callback routine does not acquire
+// any locks.  Use with caution to avoid deadlocks.  The flag is cleared
+// automatically when the completion callback is executed.  Some care has
+// been taken so that other aio operations like aio_wait will work,
+// although it is still an error to try waiting for an aio from that aio's
+// completion callback.
+void nni_aio_set_synch(nni_aio *);
+
 // nni_aio_set_timeout sets the timeout (absolute) when the AIO will
 // be canceled.  The cancelation does not happen until after nni_aio_start
 // is called.