Pipeline is now race free (and simpler to boot!)

1 files changed, 6 insertions, 0 deletions

diff --git a/src/protocol/pipeline/pull.c b/src/protocol/pipeline/pull.c
index ec66fab6..622bd958 100644
--- a/src/protocol/pipeline/pull.c
+++ b/src/protocol/pipeline/pull.c
@@ -107,6 +107,7 @@ nni_pull_pipe_start(void *arg)
 	nni_pull_pipe *pp = arg;
 
 	// Start the pending pull...
+	nni_pipe_incref(pp->pipe);
 	nni_pull_recv(pp);
 
 	return (0);
@@ -133,6 +134,7 @@ nni_pull_recv_cb(void *arg)
 	if (nni_aio_result(aio) != 0) {
 		// Failed to get a message, probably the pipe is closed.
 		nni_pipe_close(pp->pipe);
+		nni_pipe_decref(pp->pipe);
 		return;
 	}
 
@@ -153,7 +155,10 @@ nni_pull_putq_cb(void *arg)
 	if (nni_aio_result(aio) != 0) {
 		// If we failed to put, probably NNG_ECLOSED, nothing else
 		// we can do.  Just close the pipe.
+		nni_msg_free(aio->a_msg);
+		aio->a_msg = NULL;
 		nni_pipe_close(pp->pipe);
+		nni_pipe_decref(pp->pipe);
 		return;
 	}
 
@@ -168,6 +173,7 @@ nni_pull_recv(nni_pull_pipe *pp)
 	// Schedule the aio with callback.
 	if (nni_pipe_aio_recv(pp->pipe, &pp->recv_aio) != 0) {
 		nni_pipe_close(pp->pipe);
+		nni_pipe_decref(pp->pipe);
 	}
 }