udp: more nni_aio_start

2 files changed, 31 insertions, 27 deletions

diff --git a/src/sp/transport/udp/udp.c b/src/sp/transport/udp/udp.c
index f2fc6889..ac3d3f1a 100644
--- a/src/sp/transport/udp/udp.c
+++ b/src/sp/transport/udp/udp.c
@@ -956,14 +956,6 @@ udp_pipe_send(void *arg, nni_aio *aio)
 	nng_msg    *msg;
 	size_t      count = 0;
 
-	if (nni_aio_begin(aio) != 0) {
-		// No way to give the message back to the  protocol,
-		// so we just discard it silently to prevent it from leaking.
-		nni_msg_free(nni_aio_get_msg(aio));
-		nni_aio_set_msg(aio, NULL);
-		return;
-	}
-
 	msg = nni_aio_get_msg(aio);
 	ep  = p->ep;
 
@@ -971,6 +963,7 @@ udp_pipe_send(void *arg, nni_aio *aio)
 		count = nni_msg_len(msg) + nni_msg_header_len(msg);
 	}
 
+	nni_aio_reset(aio);
 	nni_mtx_lock(&ep->mtx);
 	if ((nni_msg_len(msg) + nni_msg_header_len(msg)) > p->sndmax) {
 		nni_mtx_unlock(&ep->mtx);
@@ -1020,20 +1013,16 @@ udp_pipe_recv(void *arg, nni_aio *aio)
 {
 	udp_pipe *p  = arg;
 	udp_ep   *ep = p->ep;
-	int       rv;
 
-	if (nni_aio_begin(aio) != 0) {
-		return;
-	}
+	nni_aio_reset(aio);
 	nni_mtx_lock(&ep->mtx);
 	if (p->closed) {
 		nni_mtx_unlock(&ep->mtx);
 		nni_aio_finish_error(aio, NNG_ECLOSED);
 		return;
 	}
-	if ((rv = nni_aio_schedule(aio, udp_pipe_recv_cancel, p)) != 0) {
+	if (!nni_aio_start(aio, udp_pipe_recv_cancel, p)) {
 		nni_mtx_unlock(&ep->mtx);
-		nni_aio_finish_error(aio, rv);
 		return;
 	}
 
@@ -1740,20 +1729,16 @@ static void
 udp_ep_accept(void *arg, nni_aio *aio)
 {
 	udp_ep *ep = arg;
-	int     rv;
 
-	if (nni_aio_begin(aio) != 0) {
-		return;
-	}
+	nni_aio_reset(aio);
 	nni_mtx_lock(&ep->mtx);
 	if (ep->closed) {
 		nni_mtx_unlock(&ep->mtx);
 		nni_aio_finish_error(aio, NNG_ECLOSED);
 		return;
 	}
-	if ((rv = nni_aio_schedule(aio, udp_ep_cancel, ep)) != 0) {
+	if (!nni_aio_start(aio, udp_ep_cancel, ep)) {
 		nni_mtx_unlock(&ep->mtx);
-		nni_aio_finish_error(aio, rv);
 		return;
 	}
 	nni_aio_list_append(&ep->connaios, aio);
diff --git a/src/sp/transport/udp/udp_tran_test.c b/src/sp/transport/udp/udp_tran_test.c
index 941b28e0..014c36dc 100644
--- a/src/sp/transport/udp/udp_tran_test.c
+++ b/src/sp/transport/udp/udp_tran_test.c
@@ -223,14 +223,32 @@ test_udp_multi_send_recv(void)
 	NUTS_PASS(nng_dialer_start(d, 0));
 	nng_msleep(100);
 
-	for (int i = 0; i < 1000; i++) {
-		NUTS_PASS(nng_send(s1, msg, 95, 0));
-		NUTS_PASS(nng_recv(s0, buf, &sz, 0));
-		NUTS_TRUE(sz == 95);
-		NUTS_PASS(nng_send(s0, msg, 95, 0));
-		NUTS_PASS(nng_recv(s1, buf, &sz, 0));
-		NUTS_TRUE(sz == 95);
+	int rv;
+	int i;
+	for (i = 0; i < 1000; i++) {
+		if ((rv = nng_send(s1, msg, 95, 0)) != 0) {
+			;
+			break;
+		}
+		if ((rv = nng_recv(s0, buf, &sz, 0)) != 0) {
+			break;
+		}
+		if (sz != 95) {
+			break;
+		}
+		if ((rv = nng_send(s0, msg, 95, 0)) != 0) {
+			break;
+		}
+		if ((rv = nng_recv(s1, buf, &sz, 0)) != 0) {
+			break;
+		}
+		if (sz != 95) {
+			break;
+		}
 	}
+	NUTS_PASS(rv);
+	NUTS_TRUE(i == 1000);
+	NUTS_TRUE(sz == 95);
 	NUTS_CLOSE(s0);
 	NUTS_CLOSE(s1);
 }
@@ -255,6 +273,7 @@ udp_recv_count_cb(void *arg)
 	switch (rv) {
 	case NNG_ECLOSED:
 	case NNG_ECANCELED:
+	case NNG_ESTOPPED:
 		c->fail++;
 		return;
 	case NNG_ETIMEDOUT: