fixes #1536 nng_sendmsg and nng_recvmsg could be faster

fixes #1535 Desire nng_ctx_sendmsg and nng_ctx_recvmsg

2 files changed, 157 insertions, 46 deletions

diff --git a/src/nng.c b/src/nng.c
index b1ebbd11..57f24a70 100644
--- a/src/nng.c
+++ b/src/nng.c
@@ -120,29 +120,33 @@ nng_recv(nng_socket s, void *buf, size_t *szp, int flags)
 int
 nng_recvmsg(nng_socket s, nng_msg **msgp, int flags)
 {
-	int      rv;
-	nng_aio *ap;
+	int       rv;
+	nni_sock *sock;
+	nni_aio   aio;
 
-	if ((rv = nng_aio_alloc(&ap, NULL, NULL)) != 0) {
+	if ((rv = nni_sock_find(&sock, s.id)) != 0) {
 		return (rv);
 	}
+
+	nni_aio_init(&aio, NULL, NULL);
 	if (flags & NNG_FLAG_NONBLOCK) {
-		nng_aio_set_timeout(ap, NNG_DURATION_ZERO);
+		nng_aio_set_timeout(&aio, NNG_DURATION_ZERO);
 	} else {
-		nng_aio_set_timeout(ap, NNG_DURATION_DEFAULT);
+		nng_aio_set_timeout(&aio, NNG_DURATION_DEFAULT);
 	}
+	nni_sock_recv(sock, &aio);
+	nni_sock_rele(sock);
 
-	nng_recv_aio(s, ap);
-	nng_aio_wait(ap);
+	nni_aio_wait(&aio);
 
-	if ((rv = nng_aio_result(ap)) == 0) {
-		*msgp = nng_aio_get_msg(ap);
+	if ((rv = nni_aio_result(&aio)) == 0) {
+		*msgp = nng_aio_get_msg(&aio);
 
 	} else if ((rv == NNG_ETIMEDOUT) &&
 	    ((flags & NNG_FLAG_NONBLOCK) == NNG_FLAG_NONBLOCK)) {
 		rv = NNG_EAGAIN;
 	}
-	nng_aio_free(ap);
+	nni_aio_fini(&aio);
 
 	return (rv);
 }
@@ -171,24 +175,31 @@ nng_send(nng_socket s, void *buf, size_t len, int flags)
 int
 nng_sendmsg(nng_socket s, nng_msg *msg, int flags)
 {
-	int      rv;
-	nng_aio *ap;
+	int       rv;
+	nni_aio   aio;
+	nni_sock *sock;
 
-	if ((rv = nng_aio_alloc(&ap, NULL, NULL)) != 0) {
+	if (msg == NULL) {
+		return (NNG_EINVAL);
+	}
+	if ((rv = nni_sock_find(&sock, s.id)) != 0) {
 		return (rv);
 	}
+
+	nni_aio_init(&aio, NULL, NULL);
 	if ((flags & NNG_FLAG_NONBLOCK) == NNG_FLAG_NONBLOCK) {
-		nng_aio_set_timeout(ap, NNG_DURATION_ZERO);
+		nni_aio_set_timeout(&aio, NNG_DURATION_ZERO);
 	} else {
-		nng_aio_set_timeout(ap, NNG_DURATION_DEFAULT);
+		nni_aio_set_timeout(&aio, NNG_DURATION_DEFAULT);
 	}
 
-	nng_aio_set_msg(ap, msg);
-	nng_send_aio(s, ap);
-	nng_aio_wait(ap);
+	nng_aio_set_msg(&aio, msg);
+	nni_sock_send(sock, &aio);
+	nni_sock_rele(sock);
 
-	rv = nng_aio_result(ap);
-	nng_aio_free(ap);
+	nni_aio_wait(&aio);
+	rv = nni_aio_result(&aio);
+	nni_aio_fini(&aio);
 
 	// Possibly massage nonblocking attempt.  Note that nonblocking is
 	// still done asynchronously, and the calling thread loses context.
@@ -242,7 +253,7 @@ int
 nng_ctx_open(nng_ctx *cp, nng_socket s)
 {
 	nni_sock *sock;
-	nni_ctx * ctx;
+	nni_ctx  *ctx;
 	int       rv;
 	nng_ctx   c;
 
@@ -280,6 +291,40 @@ nng_ctx_id(nng_ctx c)
 	return (((int) c.id > 0) ? (int) c.id : -1);
 }
 
+int
+nng_ctx_recvmsg(nng_ctx cid, nng_msg **msgp, int flags)
+{
+	int       rv;
+	nni_aio   aio;
+	nni_ctx  *ctx;
+
+	if ((rv = nni_ctx_find(&ctx, cid.id, false)) != 0) {
+		return (rv);
+	}
+
+	nni_aio_init(&aio, NULL, NULL);
+	if (flags & NNG_FLAG_NONBLOCK) {
+		nng_aio_set_timeout(&aio, NNG_DURATION_ZERO);
+	} else {
+		nng_aio_set_timeout(&aio, NNG_DURATION_DEFAULT);
+	}
+	nni_ctx_recv(ctx, &aio);
+	nni_ctx_rele(ctx);
+
+	nni_aio_wait(&aio);
+
+	if ((rv = nni_aio_result(&aio)) == 0) {
+		*msgp = nng_aio_get_msg(&aio);
+
+	} else if ((rv == NNG_ETIMEDOUT) &&
+	    ((flags & NNG_FLAG_NONBLOCK) == NNG_FLAG_NONBLOCK)) {
+		rv = NNG_EAGAIN;
+	}
+	nni_aio_fini(&aio);
+
+	return (rv);
+}
+
 void
 nng_ctx_recv(nng_ctx cid, nng_aio *aio)
 {
@@ -318,6 +363,45 @@ nng_ctx_send(nng_ctx cid, nng_aio *aio)
 	nni_ctx_rele(ctx);
 }
 
+int
+nng_ctx_sendmsg(nng_ctx cid, nng_msg *msg, int flags)
+{
+	int       rv;
+	nni_aio   aio;
+	nni_ctx *ctx;
+
+	if (msg == NULL) {
+		return (NNG_EINVAL);
+	}
+	if ((rv = nni_ctx_find(&ctx, cid.id, false)) != 0) {
+		return (rv);
+	}
+
+	nni_aio_init(&aio, NULL, NULL);
+	if ((flags & NNG_FLAG_NONBLOCK) == NNG_FLAG_NONBLOCK) {
+		nni_aio_set_timeout(&aio, NNG_DURATION_ZERO);
+	} else {
+		nni_aio_set_timeout(&aio, NNG_DURATION_DEFAULT);
+	}
+
+	nng_aio_set_msg(&aio, msg);
+	nni_ctx_send(ctx, &aio);
+	nni_ctx_rele(ctx);
+
+	nni_aio_wait(&aio);
+	rv = nni_aio_result(&aio);
+	nni_aio_fini(&aio);
+
+	// Possibly massage nonblocking attempt.  Note that nonblocking is
+	// still done asynchronously, and the calling thread loses context.
+	if ((rv == NNG_ETIMEDOUT) &&
+	    ((flags & NNG_FLAG_NONBLOCK) == NNG_FLAG_NONBLOCK)) {
+		rv = NNG_EAGAIN;
+	}
+
+	return (rv);
+}
+
 static int
 ctx_get(nng_ctx id, const char *n, void *v, size_t *szp, nni_type t)
 {
@@ -466,7 +550,7 @@ nng_dial(nng_socket sid, const char *addr, nng_dialer *dp, int flags)
 {
 	nni_dialer *d;
 	int         rv;
-	nni_sock *  s;
+	nni_sock   *s;
 
 	if ((rv = nni_sock_find(&s, sid.id)) != 0) {
 		return (rv);
@@ -492,7 +576,7 @@ int
 nng_listen(nng_socket sid, const char *addr, nng_listener *lp, int flags)
 {
 	int           rv;
-	nni_sock *    s;
+	nni_sock     *s;
 	nni_listener *l;
 
 	if ((rv = nni_sock_find(&s, sid.id)) != 0) {
@@ -519,7 +603,7 @@ nng_listen(nng_socket sid, const char *addr, nng_listener *lp, int flags)
 int
 nng_listener_create(nng_listener *lp, nng_socket sid, const char *addr)
 {
-	nni_sock *    s;
+	nni_sock     *s;
 	int           rv;
 	nni_listener *l;
 	nng_listener  lid;
@@ -560,7 +644,7 @@ nng_listener_id(nng_listener l)
 int
 nng_dialer_create(nng_dialer *dp, nng_socket sid, const char *addr)
 {
-	nni_sock *  s;
+	nni_sock   *s;
 	nni_dialer *d;
 	int         rv;
 	nng_dialer  did;
@@ -1247,7 +1331,7 @@ nng_socket
 nng_pipe_socket(nng_pipe p)
 {
 	nng_socket s = NNG_SOCKET_INITIALIZER;
-	nni_pipe * pipe;
+	nni_pipe  *pipe;
 
 	if ((nni_init() == 0) && (nni_pipe_find(&pipe, p.id) == 0)) {
 		s.id = nni_pipe_sock_id(pipe);
@@ -1260,7 +1344,7 @@ nng_dialer
 nng_pipe_dialer(nng_pipe p)
 {
 	nng_dialer d = NNG_DIALER_INITIALIZER;
-	nni_pipe * pipe;
+	nni_pipe  *pipe;
 	if ((nni_init() == 0) && (nni_pipe_find(&pipe, p.id) == 0)) {
 		d.id = nni_pipe_dialer_id(pipe);
 		nni_pipe_rele(pipe);
@@ -1272,7 +1356,7 @@ nng_listener
 nng_pipe_listener(nng_pipe p)
 {
 	nng_listener l = NNG_LISTENER_INITIALIZER;
-	nni_pipe *   pipe;
+	nni_pipe    *pipe;
 	if ((nni_init() == 0) && (nni_pipe_find(&pipe, p.id) == 0)) {
 		l.id = nni_pipe_listener_id(pipe);
 		nni_pipe_rele(pipe);
diff --git a/src/sp/protocol/survey0/survey_test.c b/src/sp/protocol/survey0/survey_test.c
index 95d27adf..1e1b8635 100644
--- a/src/sp/protocol/survey0/survey_test.c
+++ b/src/sp/protocol/survey0/survey_test.c
@@ -1,5 +1,5 @@
 //
-// Copyright 2020 Staysail Systems, Inc. <info@staysail.tech>
+// Copyright 2021 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
@@ -15,7 +15,7 @@ test_surv_identity(void)
 {
 	nng_socket s;
 	int        p;
-	char *     n;
+	char      *n;
 
 	NUTS_PASS(nng_surveyor0_open(&s));
 	NUTS_PASS(nng_socket_get_int(s, NNG_OPT_PROTO, &p));
@@ -74,7 +74,7 @@ test_surv_survey_time_option(void)
 	nng_duration d;
 	bool         b;
 	size_t       sz  = sizeof(b);
-	const char * opt = NNG_OPT_SURVEYOR_SURVEYTIME;
+	const char  *opt = NNG_OPT_SURVEYOR_SURVEYTIME;
 
 	NUTS_PASS(nng_surveyor0_open(&surv));
 
@@ -93,7 +93,7 @@ void
 test_surv_recv_bad_state(void)
 {
 	nng_socket surv;
-	nng_msg *  msg = NULL;
+	nng_msg   *msg = NULL;
 
 	NUTS_PASS(nng_surveyor0_open(&surv));
 	NUTS_FAIL(nng_recvmsg(surv, &msg, 0), NNG_ESTATE);
@@ -106,7 +106,7 @@ test_surv_recv_garbage(void)
 {
 	nng_socket resp;
 	nng_socket surv;
-	nng_msg *  m;
+	nng_msg   *m;
 	uint32_t   surv_id;
 
 	NUTS_PASS(nng_respondent0_open_raw(&resp));
@@ -216,7 +216,7 @@ test_surv_cancel(void)
 void
 test_surv_cancel_abort_recv(void)
 {
-	nng_aio *    aio;
+	nng_aio     *aio;
 	nng_duration time = SECOND * 10; // 10s (kind of never)
 	nng_socket   surv;
 	nng_socket   resp;
@@ -334,7 +334,7 @@ test_surv_poll_readable(void)
 	int        fd;
 	nng_socket surv;
 	nng_socket resp;
-	nng_msg *  msg;
+	nng_msg   *msg;
 
 	NUTS_PASS(nng_surveyor0_open(&surv));
 	NUTS_PASS(nng_respondent0_open(&resp));
@@ -392,8 +392,8 @@ test_surv_ctx_recv_nonblock(void)
 	nng_socket surv;
 	nng_socket resp;
 	nng_ctx    ctx;
-	nng_aio *  aio;
-	nng_msg *  msg;
+	nng_aio   *aio;
+	nng_msg   *msg;
 
 	NUTS_PASS(nng_surveyor0_open(&surv));
 	NUTS_PASS(nng_respondent0_open(&resp));
@@ -418,12 +418,38 @@ test_surv_ctx_recv_nonblock(void)
 }
 
 static void
+test_surv_ctx_send_recv_msg(void)
+{
+	nng_socket surv;
+	nng_socket resp;
+	nng_ctx    ctx1;
+	nng_ctx    ctx2;
+	nng_msg   *msg;
+
+	NUTS_PASS(nng_surveyor0_open(&surv));
+	NUTS_PASS(nng_respondent0_open(&resp));
+	NUTS_PASS(nng_ctx_open(&ctx1, surv));
+	NUTS_PASS(nng_ctx_open(&ctx2, resp));
+	NUTS_PASS(nng_msg_alloc(&msg, 0));
+
+	NUTS_MARRY(surv, resp);
+
+	NUTS_PASS(nng_ctx_sendmsg(ctx1, msg, 0));
+	NUTS_PASS(nng_ctx_recvmsg(ctx2, &msg, 0));
+	nng_msg_free(msg);
+	NUTS_PASS(nng_ctx_close(ctx1));
+	NUTS_PASS(nng_ctx_close(ctx2));
+	NUTS_CLOSE(surv);
+	NUTS_CLOSE(resp);
+}
+
+static void
 test_surv_ctx_send_nonblock(void)
 {
 	nng_socket surv;
 	nng_ctx    ctx;
-	nng_aio *  aio;
-	nng_msg *  msg;
+	nng_aio   *aio;
+	nng_msg   *msg;
 
 	NUTS_PASS(nng_surveyor0_open(&surv));
 	NUTS_PASS(nng_ctx_open(&ctx, surv));
@@ -490,8 +516,8 @@ test_surv_ctx_recv_close_socket(void)
 	nng_socket surv;
 	nng_socket resp;
 	nng_ctx    ctx;
-	nng_aio *  aio;
-	nng_msg *  m;
+	nng_aio   *aio;
+	nng_msg   *m;
 
 	NUTS_PASS(nng_surveyor0_open(&surv));
 	NUTS_PASS(nng_respondent0_open(&resp));
@@ -518,8 +544,8 @@ test_surv_context_multi(void)
 	nng_socket surv;
 	nng_socket resp;
 	nng_ctx    c[5];
-	nng_aio *  aio;
-	nng_msg *  m;
+	nng_aio   *aio;
+	nng_msg   *m;
 	int        cnt = sizeof(c) / sizeof(c[0]);
 
 	NUTS_PASS(nng_surveyor0_open(&surv));
@@ -575,9 +601,9 @@ static void
 test_surv_validate_peer(void)
 {
 	nng_socket s1, s2;
-	nng_stat * stats;
-	nng_stat * reject;
-	char *     addr;
+	nng_stat  *stats;
+	nng_stat  *reject;
+	char      *addr;
 
 	NUTS_ADDR(addr, "inproc");
 	NUTS_PASS(nng_surveyor0_open(&s1));
@@ -618,6 +644,7 @@ TEST_LIST = {
 	    test_surv_ctx_recv_close_socket },
 	{ "survey context recv nonblock", test_surv_ctx_recv_nonblock },
 	{ "survey context send nonblock", test_surv_ctx_send_nonblock },
+	{ "survey context send recv msg", test_surv_ctx_send_recv_msg },
 	{ "survey timeout", test_surv_survey_timeout },
 	{ "survey send best effort", test_surv_send_best_effort },
 	{ "survey context multi", test_surv_context_multi },