4 files changed, 63 insertions, 10 deletions

diff --git a/src/supplemental/http/http_client.c b/src/supplemental/http/http_client.c
index a2427009..aba72715 100644
--- a/src/supplemental/http/http_client.c
+++ b/src/supplemental/http/http_client.c
@@ -241,11 +241,16 @@ http_connect_cancel(nni_aio *aio, int rv)
 void
 nni_http_client_connect(nni_http_client *c, nni_aio *aio)
 {
+	int rv;
 	if (nni_aio_begin(aio) != 0) {
 		return;
 	}
 	nni_mtx_lock(&c->mtx);
-	nni_aio_schedule(aio, http_connect_cancel, aio);
+	if ((rv = nni_aio_schedule(aio, http_connect_cancel, aio)) != 0) {
+		nni_mtx_unlock(&c->mtx);
+		nni_aio_finish_error(aio, rv);
+		return;
+	}
 	nni_list_append(&c->aios, aio);
 	if (nni_list_first(&c->aios) == aio) {
 		http_conn_start(c);
diff --git a/src/supplemental/http/http_conn.c b/src/supplemental/http/http_conn.c
index 0de40e10..f3b16370 100644
--- a/src/supplemental/http/http_conn.c
+++ b/src/supplemental/http/http_conn.c
@@ -365,6 +365,8 @@ http_rd_cancel(nni_aio *aio, int rv)
 static void
 http_rd_submit(nni_http_conn *conn, nni_aio *aio)
 {
+	int rv;
+
 	if (nni_aio_begin(aio) != 0) {
 		return;
 	}
@@ -372,7 +374,10 @@ http_rd_submit(nni_http_conn *conn, nni_aio *aio)
 		nni_aio_finish_error(aio, NNG_ECLOSED);
 		return;
 	}
-	nni_aio_schedule(aio, http_rd_cancel, conn);
+	if ((rv = nni_aio_schedule(aio, http_rd_cancel, conn)) != 0) {
+		nni_aio_finish_error(aio, rv);
+		return;
+	}
 	nni_list_append(&conn->rdq, aio);
 	if (conn->rd_uaio == NULL) {
 		http_rd_start(conn);
@@ -479,6 +484,8 @@ http_wr_cancel(nni_aio *aio, int rv)
 static void
 http_wr_submit(nni_http_conn *conn, nni_aio *aio)
 {
+	int rv;
+
 	if (nni_aio_begin(aio) != 0) {
 		return;
 	}
@@ -486,8 +493,12 @@ http_wr_submit(nni_http_conn *conn, nni_aio *aio)
 		nni_aio_finish_error(aio, NNG_ECLOSED);
 		return;
 	}
+	if ((rv = nni_aio_schedule(aio, http_wr_cancel, conn)) != 0) {
+		nni_aio_finish_error(aio, rv);
+		return;
+	}
 	nni_list_append(&conn->wrq, aio);
-	nni_aio_schedule(aio, http_wr_cancel, conn);
+
 	if (conn->wr_uaio == NULL) {
 		http_wr_start(conn);
 	}
diff --git a/src/supplemental/tls/mbedtls/tls.c b/src/supplemental/tls/mbedtls/tls.c
index 5246e06c..f721d4aa 100644
--- a/src/supplemental/tls/mbedtls/tls.c
+++ b/src/supplemental/tls/mbedtls/tls.c
@@ -566,6 +566,8 @@ nni_tls_net_recv(void *ctx, unsigned char *buf, size_t len)
 void
 nni_tls_send(nni_tls *tp, nni_aio *aio)
 {
+	int rv;
+
 	if (nni_aio_begin(aio) != 0) {
 		return;
 	}
@@ -575,7 +577,11 @@ nni_tls_send(nni_tls *tp, nni_aio *aio)
 		nni_aio_finish_error(aio, NNG_ECLOSED);
 		return;
 	}
-	nni_aio_schedule(aio, nni_tls_cancel, tp);
+	if ((rv = nni_aio_schedule(aio, nni_tls_cancel, tp)) != 0) {
+		nni_mtx_unlock(&tp->lk);
+		nni_aio_finish_error(aio, rv);
+		return;
+	}
 	nni_list_append(&tp->sends, aio);
 	nni_tls_do_send(tp);
 	nni_mtx_unlock(&tp->lk);
@@ -584,6 +590,8 @@ nni_tls_send(nni_tls *tp, nni_aio *aio)
 void
 nni_tls_recv(nni_tls *tp, nni_aio *aio)
 {
+	int rv;
+
 	if (nni_aio_begin(aio) != 0) {
 		return;
 	}
@@ -593,7 +601,12 @@ nni_tls_recv(nni_tls *tp, nni_aio *aio)
 		nni_aio_finish_error(aio, NNG_ECLOSED);
 		return;
 	}
-	nni_aio_schedule(aio, nni_tls_cancel, tp);
+	if ((rv = nni_aio_schedule(aio, nni_tls_cancel, tp)) != 0) {
+		nni_mtx_unlock(&tp->lk);
+		nni_aio_finish_error(aio, rv);
+		return;
+	}
+
 	nni_list_append(&tp->recvs, aio);
 	nni_tls_do_recv(tp);
 	nni_mtx_unlock(&tp->lk);
diff --git a/src/supplemental/websocket/websocket.c b/src/supplemental/websocket/websocket.c
index e562b209..9a3f5519 100644
--- a/src/supplemental/websocket/websocket.c
+++ b/src/supplemental/websocket/websocket.c
@@ -667,8 +667,12 @@ ws_send_close(nni_ws *ws, uint16_t code)
 		return;
 	}
 	// Close frames get priority!
+	if ((rv = nni_aio_schedule(aio, ws_cancel_close, ws)) != 0) {
+		ws->wclose = false;
+		nni_aio_finish_error(aio, rv);
+		return;
+	}
 	nni_list_prepend(&ws->txmsgs, wm);
-	nni_aio_schedule(aio, ws_cancel_close, ws);
 	ws_start_write(ws);
 }
 
@@ -737,7 +741,12 @@ nni_ws_send_msg(nni_ws *ws, nni_aio *aio)
 		ws_msg_fini(wm);
 		return;
 	}
-	nni_aio_schedule(aio, ws_write_cancel, ws);
+	if ((rv = nni_aio_schedule(aio, ws_write_cancel, ws)) != 0) {
+		nni_mtx_unlock(&ws->mtx);
+		nni_aio_finish_error(aio, rv);
+		ws_msg_fini(wm);
+		return;
+	}
 	nni_aio_set_prov_extra(aio, 0, wm);
 	nni_list_append(&ws->sendq, aio);
 	nni_list_append(&ws->txmsgs, wm);
@@ -1060,7 +1069,12 @@ nni_ws_recv_msg(nni_ws *ws, nni_aio *aio)
 		return;
 	}
 	nni_mtx_lock(&ws->mtx);
-	nni_aio_schedule(aio, ws_read_cancel, ws);
+	if ((rv = nni_aio_schedule(aio, ws_read_cancel, ws)) != 0) {
+		nni_mtx_unlock(&ws->mtx);
+		ws_msg_fini(wm);
+		nni_aio_finish_error(aio, rv);
+		return;
+	}
 	nni_aio_set_prov_extra(aio, 0, wm);
 	nni_list_append(&ws->recvq, aio);
 	nni_list_append(&ws->rxmsgs, wm);
@@ -1647,6 +1661,7 @@ void
 nni_ws_listener_accept(nni_ws_listener *l, nni_aio *aio)
 {
 	nni_ws *ws;
+	int     rv;
 
 	if (nni_aio_begin(aio) != 0) {
 		return;
@@ -1669,7 +1684,11 @@ nni_ws_listener_accept(nni_ws_listener *l, nni_aio *aio)
 		nni_aio_finish(aio, 0, 0);
 		return;
 	}
-	nni_aio_schedule(aio, ws_accept_cancel, l);
+	if ((rv = nni_aio_schedule(aio, ws_accept_cancel, l)) != 0) {
+		nni_aio_finish_error(aio, rv);
+		nni_mtx_unlock(&l->mtx);
+		return;
+	}
 	nni_list_append(&l->aios, aio);
 	nni_mtx_unlock(&l->mtx);
 }
@@ -1995,11 +2014,16 @@ nni_ws_dialer_dial(nni_ws_dialer *d, nni_aio *aio)
 		ws_fini(ws);
 		return;
 	}
+	if ((rv = nni_aio_schedule(aio, ws_dial_cancel, ws)) != 0) {
+		nni_mtx_unlock(&d->mtx);
+		nni_aio_finish_error(aio, rv);
+		ws_fini(ws);
+		return;
+	}
 	ws->dialer  = d;
 	ws->useraio = aio;
 	ws->mode    = NNI_EP_MODE_DIAL;
 	nni_list_append(&d->wspend, ws);
-	nni_aio_schedule(aio, ws_dial_cancel, ws);
 	nni_http_client_connect(d->client, ws->connaio);
 	nni_mtx_unlock(&d->mtx);
 }