From 034b1a0377260fcb91656bf2f37749dffa56b04c Mon Sep 17 00:00:00 2001
From: Garrett D'Amore <garrett@damore.org>
Date: Mon, 9 Sep 2024 19:58:51 -0700
Subject: [PATCH] performance: Simplify locking for synchronous completion

If we're not going to issue a callback, then there is no
reason to drop and reacquire the lock.  This might help
some simple cases (such as immediate queuing of PUB messages)
reduce bus pressure and possible extra context switching.
---
 src/core/aio.c   |  5 ++++-
 src/core/aio.h   |  3 ++-
 src/core/taskq.c | 10 ++--------
 3 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/core/aio.c b/src/core/aio.c
index 8aee5f954..2f548d745 100644
--- a/src/core/aio.c
+++ b/src/core/aio.c
@@ -449,7 +449,10 @@ nni_aio_finish_impl(
 	aio->a_use_expire = false;
 	nni_mtx_unlock(&eq->eq_mtx);
 
-	if (sync) {
+	// If there is no callback, we can complete synchronously.
+	// If there's a callback, we have to be careful because
+	// the callback could try to acquire a lock held by the submitter.
+	if (sync || aio->a_task.task_cb == NULL) {
 		nni_task_exec(&aio->a_task);
 	} else {
 		nni_task_dispatch(&aio->a_task);
diff --git a/src/core/aio.h b/src/core/aio.h
index 958987018..a123e74d7 100644
--- a/src/core/aio.h
+++ b/src/core/aio.h
@@ -1,5 +1,5 @@
 //
-// Copyright 2023 Staysail Systems, Inc. <info@staysail.tech>
+// Copyright 2024 Staysail Systems, Inc. <info@staysail.tech>
 // Copyright 2018 Capitar IT Group BV <info@capitar.com>
 //
 // This software is supplied under the terms of the MIT License, a
@@ -13,6 +13,7 @@
 
 #include "defs.h"
 #include "list.h"
+#include "platform.h"
 #include "reap.h"
 #include "taskq.h"
 #include "thread.h"
diff --git a/src/core/taskq.c b/src/core/taskq.c
index ae1ddbdd5..75f95a761 100644
--- a/src/core/taskq.c
+++ b/src/core/taskq.c
@@ -135,13 +135,13 @@ nni_task_exec(nni_task *task)
 	} else {
 		task->task_busy++;
 	}
-	nni_mtx_unlock(&task->task_mtx);
 
 	if (task->task_cb != NULL) {
+		nni_mtx_unlock(&task->task_mtx);
 		task->task_cb(task->task_arg);
+		nni_mtx_lock(&task->task_mtx);
 	}
 
-	nni_mtx_lock(&task->task_mtx);
 	task->task_busy--;
 	if (task->task_busy == 0) {
 		nni_cv_wake(&task->task_cv);
@@ -154,12 +154,6 @@ nni_task_dispatch(nni_task *task)
 {
 	nni_taskq *tq = task->task_tq;
 
-	// If there is no callback to perform, then do nothing!
-	// The user will be none the wiser.
-	if (task->task_cb == NULL) {
-		nni_task_exec(task);
-		return;
-	}
 	nni_mtx_lock(&task->task_mtx);
 	if (task->task_prep) {
 		task->task_prep = false;