Race condition fixes in pub/sub.

2 files changed, 17 insertions, 14 deletions

diff --git a/src/protocol/pubsub/pub.c b/src/protocol/pubsub/pub.c
index e1808c2b..ad0b9580 100644
--- a/src/protocol/pubsub/pub.c
+++ b/src/protocol/pubsub/pub.c
@@ -138,13 +138,11 @@ nni_pub_pipe_fini(void *arg)
 {
 	nni_pub_pipe *pp = arg;
 
-	if (pp != NULL) {
-		nni_msgq_fini(pp->sendq);
-		nni_aio_fini(&pp->aio_getq);
-		nni_aio_fini(&pp->aio_send);
-		nni_aio_fini(&pp->aio_recv);
-		NNI_FREE_STRUCT(pp);
-	}
+	nni_msgq_fini(pp->sendq);
+	nni_aio_fini(&pp->aio_getq);
+	nni_aio_fini(&pp->aio_send);
+	nni_aio_fini(&pp->aio_recv);
+	NNI_FREE_STRUCT(pp);
 }
 
 
@@ -160,7 +158,9 @@ nni_pub_pipe_add(void *arg)
 	nni_list_append(&pub->pipes, pp);
 
 	// Start the receiver and the queue reader.
+	nni_pipe_incref(pp->pipe);
 	nni_pipe_aio_recv(pp->pipe, &pp->aio_recv);
+	nni_pipe_incref(pp->pipe);
 	nni_msgq_aio_get(pp->sendq, &pp->aio_getq);
 
 	return (0);
@@ -174,7 +174,7 @@ nni_pub_pipe_rem(void *arg)
 	nni_pub_sock *pub = pp->pub;
 
 	nni_list_remove(&pub->pipes, pp);
-	nni_msgq_aio_cancel(pp->sendq, &pp->aio_getq);
+	nni_msgq_close(pp->sendq);
 }
 
 
@@ -190,7 +190,6 @@ nni_pub_sock_getq_cb(void *arg)
 	nni_pub_pipe *last;
 	int rv;
 
-
 	if (nni_aio_result(&pub->aio_getq) != 0) {
 		return;
 	}
@@ -230,6 +229,7 @@ nni_pub_pipe_recv_cb(void *arg)
 
 	if (nni_aio_result(&pp->aio_recv) != 0) {
 		nni_pipe_close(pp->pipe);
+		nni_pipe_decref(pp->pipe);
 		return;
 	}
 
@@ -246,6 +246,7 @@ nni_pub_pipe_getq_cb(void *arg)
 
 	if (nni_aio_result(&pp->aio_getq) != 0) {
 		nni_pipe_close(pp->pipe);
+		nni_pipe_decref(pp->pipe);
 		return;
 	}
 
@@ -265,6 +266,7 @@ nni_pub_pipe_send_cb(void *arg)
 		nni_msg_free(pp->aio_send.a_msg);
 		pp->aio_send.a_msg = NULL;
 		nni_pipe_close(pp->pipe);
+		nni_pipe_decref(pp->pipe);
 		return;
 	}
 
diff --git a/src/protocol/pubsub/sub.c b/src/protocol/pubsub/sub.c
index 864a80c7..6fae1962 100644
--- a/src/protocol/pubsub/sub.c
+++ b/src/protocol/pubsub/sub.c
@@ -108,11 +108,9 @@ nni_sub_pipe_fini(void *arg)
 {
 	nni_sub_pipe *sp = arg;
 
-	if (sp != NULL) {
-		nni_aio_fini(&sp->aio_putq);
-		nni_aio_fini(&sp->aio_recv);
-		NNI_FREE_STRUCT(sp);
-	}
+	nni_aio_fini(&sp->aio_putq);
+	nni_aio_fini(&sp->aio_recv);
+	NNI_FREE_STRUCT(sp);
 }
 
 
@@ -121,6 +119,7 @@ nni_sub_pipe_add(void *arg)
 {
 	nni_sub_pipe *sp = arg;
 
+	nni_pipe_incref(sp->pipe);
 	nni_pipe_aio_recv(sp->pipe, &sp->aio_recv);
 	return (0);
 }
@@ -145,6 +144,7 @@ nni_sub_recv_cb(void *arg)
 
 	if (nni_aio_result(&sp->aio_recv) != 0) {
 		nni_pipe_close(sp->pipe);
+		nni_pipe_decref(sp->pipe);
 		return;
 	}
 
@@ -163,6 +163,7 @@ nni_sub_putq_cb(void *arg)
 		nni_msg_free(sp->aio_putq.a_msg);
 		sp->aio_putq.a_msg = NULL;
 		nni_pipe_close(sp->pipe);
+		nni_pipe_decref(sp->pipe);
 		return;
 	}