nng_sleep_aio should honor aio timeout.

The first problem was that using nng_sleep_aio
was found to reset the timeout, and this caused subsequent
operations to start failing with timeouts when reusing
the AIO for other operations.

The second thing is that we think it would be nicer if the
presence of real aio timeouts were still honored, so that
if the timeout is shorter than the sleep time, then we get
back an NNG_ETIMEDOUT like every other operation, and we
get back a 0 if the logical sleep operation completes
normally.

3 files changed, 59 insertions, 23 deletions

diff --git a/docs/man/nng_sleep_aio.adoc b/docs/man/nng_sleep_aio.adoc
index dff00219..8181046d 100644
--- a/docs/man/nng_sleep_aio.adoc
+++ b/docs/man/nng_sleep_aio.adoc
@@ -26,12 +26,11 @@ void nng_sleep_aio(nng_duration msec, nng_aio *aio);
 
 The `nng_sleep_aio()` function performs an asynchronous "`sleep``",
 causing the callback for _aio_ to be executed after _msec_ milliseconds.
-This is logically the equivalent of starting an asynchronous operation
-that does nothing at all, but expires after _msec_ duration, _except_ that
-the completion result will be zero rather `NNG_ETIMEDOUT`.
+If the sleep finishes completely, the result will always be zero.
 
-NOTE: This overrides and replaces any timeout on the _aio_ set with
-<<nng_aio_set_timeout#,nng_aio_set_timeout(3)>>.
+NOTE: If a timeout is set on _aio_ using
+<<nng_aio_timeout#,nng_aio_set_timeout(3)>>, and it is shorter than _msec_,
+then the sleep will wake up early, with a result code of `NNG_ETIMEDOUT`.
 
 == RETURN VALUES
 
diff --git a/src/core/aio.c b/src/core/aio.c
index e0206830..1aaffe12 100644
--- a/src/core/aio.c
+++ b/src/core/aio.c
@@ -321,19 +321,24 @@ nni_aio_start(nni_aio *aio, nni_aio_cancelfn cancelfn, void *data)
 	}
 
 	// Convert the relative timeout to an absolute timeout.
-	switch (aio->a_timeout) {
-	case NNG_DURATION_ZERO:
-		aio->a_expire = NNI_TIME_ZERO;
-		nni_aio_expire_add(aio);
-		break;
-	case NNG_DURATION_INFINITE:
-	case NNG_DURATION_DEFAULT:
-		aio->a_expire = NNI_TIME_NEVER;
-		break;
-	default:
-		aio->a_expire = nni_clock() + aio->a_timeout;
+	if (aio->a_sleep) {
+		// expire node already set.
 		nni_aio_expire_add(aio);
-		break;
+	} else {
+		switch (aio->a_timeout) {
+		case NNG_DURATION_ZERO:
+			aio->a_expire = NNI_TIME_ZERO;
+			nni_aio_expire_add(aio);
+			break;
+		case NNG_DURATION_INFINITE:
+		case NNG_DURATION_DEFAULT:
+			aio->a_expire = NNI_TIME_NEVER;
+			break;
+		default:
+			aio->a_expire = nni_clock() + aio->a_timeout;
+			nni_aio_expire_add(aio);
+			break;
+		}
 	}
 	nni_mtx_unlock(&nni_aio_lk);
 	return (0);
@@ -517,10 +522,11 @@ nni_aio_expire_loop(void *arg)
 		// already being canceled or finished.
 		if (cancelfn != NULL) {
 			nni_mtx_unlock(&nni_aio_lk);
-			cancelfn(aio, aio->a_sleep ? 0 : NNG_ETIMEDOUT);
+			cancelfn(aio, rv);
 			nni_mtx_lock(&nni_aio_lk);
 		} else {
-			aio->a_pend = 1;
+			aio->a_pend   = 1;
+			aio->a_result = rv;
 		}
 
 		NNI_ASSERT(aio->a_pend); // nni_aio_finish was run
@@ -621,11 +627,23 @@ nni_aio_iov_advance(nni_aio *aio, size_t n)
 void
 nni_sleep_aio(nng_duration ms, nng_aio *aio)
 {
-	nni_aio_set_timeout(aio, ms);
-	aio->a_sleep = 1;
-	if (nni_aio_start(aio, NULL, NULL)) {
-		return;
+	switch (aio->a_timeout) {
+	case NNG_DURATION_DEFAULT:
+	case NNG_DURATION_INFINITE:
+		// No premature timeout, honor our expected values.
+		break;
+	default:
+		// If the timeout on the aio is shorter than our sleep time,
+		// then let it still wake up early, but with NNG_ETIMEDOUT.
+		if (ms > aio->a_timeout) {
+			aio->a_sleep = 0;
+			(void) nni_aio_start(aio, NULL, NULL);
+			return;
+		}
 	}
+	aio->a_sleep  = 1;
+	aio->a_expire = nni_clock() + ms;
+	(void) nni_aio_start(aio, NULL, NULL);
 }
 
 void
diff --git a/tests/aio.c b/tests/aio.c
index b57c5a69..244ae3b9 100644
--- a/tests/aio.c
+++ b/tests/aio.c
@@ -56,6 +56,25 @@ Main({
 			So((nng_clock() - start) <= 1000);
 			nng_aio_free(saio);
 		});
+
+		Convey("Sleep timeout works", {
+			nng_time start = 0;
+			nng_time end   = 0;
+			nng_aio *saio;
+			So(nng_aio_alloc(&saio, sleepdone, &end) == 0);
+			nng_aio_set_timeout(saio, 100);
+			start = nng_clock();
+			nng_sleep_aio(2000, saio);
+			nng_aio_wait(saio);
+			So(nng_aio_result(saio) == NNG_ETIMEDOUT);
+			So(end != 0);
+			So((end - start) >= 100);
+			So((end - start) <= 1000);
+			So((nng_clock() - start) >= 100);
+			So((nng_clock() - start) <= 1000);
+			nng_aio_free(saio);
+		});
+
 		Convey("Given a connected pair of sockets", {
 			nng_socket s1;
 			nng_socket s2;