2 files changed, 28 insertions, 25 deletions

diff --git a/src/sp/protocol/bus0/bus.c b/src/sp/protocol/bus0/bus.c
index 5c7249dd..4822ffd6 100644
--- a/src/sp/protocol/bus0/bus.c
+++ b/src/sp/protocol/bus0/bus.c
@@ -269,14 +269,13 @@ bus0_sock_send(void *arg, nni_aio *aio)
 	uint32_t   sender = 0;
 	size_t     len;
 
-	if (nni_aio_begin(aio) != 0) {
-		return;
-	}
-
 	msg = nni_aio_get_msg(aio);
 	len = nni_msg_len(msg);
 	nni_aio_set_msg(aio, NULL);
 
+	// this test is so that we detect when the aio itself is terminated,
+	// otherwise we could loop forever.
+
 	if (s->raw) {
 		// In raw mode, we look for the message header, to see if it
 		// is being resent from another pipe (e.g. via a device).
@@ -290,6 +289,12 @@ bus0_sock_send(void *arg, nni_aio *aio)
 	}
 
 	nni_mtx_lock(&s->mtx);
+
+	if (!nni_aio_start(aio, NULL, NULL)) {
+		nni_mtx_unlock(&s->mtx);
+		return;
+	}
+
 	NNI_LIST_FOREACH (&s->pipes, pipe) {
 
 		if (s->raw && nni_pipe_id(pipe->pipe) == sender) {
@@ -331,17 +336,11 @@ bus0_sock_recv(void *arg, nni_aio *aio)
 	bus0_sock *s = arg;
 	nni_msg   *msg;
 
-	if (nni_aio_begin(aio) != 0) {
-		return;
-	}
-
 	nni_mtx_lock(&s->mtx);
 again:
 	if (nni_lmq_empty(&s->recv_msgs)) {
-		int rv;
-		if ((rv = nni_aio_schedule(aio, bus0_recv_cancel, s)) != 0) {
+		if (!nni_aio_start(aio, bus0_recv_cancel, s)) {
 			nni_mtx_unlock(&s->mtx);
-			nni_aio_finish_error(aio, rv);
 			return;
 		}
 		nni_list_append(&s->recv_wait, aio);
diff --git a/src/sp/protocol/bus0/bus_test.c b/src/sp/protocol/bus0/bus_test.c
index c9952d81..94ba8329 100644
--- a/src/sp/protocol/bus0/bus_test.c
+++ b/src/sp/protocol/bus0/bus_test.c
@@ -179,24 +179,28 @@ static void
 test_bus_aio_stopped(void)
 {
 	nng_socket s1;
-	nng_aio   *aio;
+	nng_aio   *aio1;
+	nng_aio   *aio2;
 	nng_msg   *msg;
 
 	NUTS_PASS(nng_bus0_open(&s1));
 	NUTS_PASS(nng_msg_alloc(&msg, 0));
-	NUTS_PASS(nng_aio_alloc(&aio, NULL, NULL));
-	nng_aio_stop(aio);
-
-	nng_recv_aio(s1, aio);
-	nng_aio_wait(aio);
-	NUTS_FAIL(nng_aio_result(aio), NNG_ESTOPPED);
-
-	nng_aio_set_msg(aio, msg);
-	nng_send_aio(s1, aio);
-	nng_aio_wait(aio);
-	NUTS_FAIL(nng_aio_result(aio), NNG_ESTOPPED);
-
-	nng_aio_free(aio);
+	NUTS_PASS(nng_aio_alloc(&aio1, NULL, NULL));
+	NUTS_PASS(nng_aio_alloc(&aio2, NULL, NULL));
+	nng_aio_stop(aio1);
+	nng_aio_stop(aio2);
+
+	nng_recv_aio(s1, aio1);
+	nng_aio_wait(aio1);
+	NUTS_FAIL(nng_aio_result(aio1), NNG_ESTOPPED);
+
+	nng_aio_set_msg(aio2, msg);
+	nng_send_aio(s1, aio2);
+	nng_aio_wait(aio2);
+	NUTS_FAIL(nng_aio_result(aio2), NNG_ESTOPPED);
+
+	nng_aio_free(aio1);
+	nng_aio_free(aio2);
 	nng_msg_free(msg);
 	NUTS_CLOSE(s1);
 }