Protocols keep their own reference counts.

1 files changed, 43 insertions, 28 deletions

diff --git a/src/protocol/pipeline/push.c b/src/protocol/pipeline/push.c
index 9554b2be..e69ebdbf 100644
--- a/src/protocol/pipeline/push.c
+++ b/src/protocol/pipeline/push.c
@@ -39,6 +39,8 @@ struct nni_push_pipe {
 	nni_aio		aio_recv;
 	nni_aio		aio_send;
 	nni_aio		aio_getq;
+	int		refcnt;
+	nni_mtx		mtx;
 };
 
 static int
@@ -70,6 +72,19 @@ nni_push_sock_fini(void *arg)
 }
 
 
+static void
+nni_push_pipe_fini(void *arg)
+{
+	nni_push_pipe *pp = arg;
+
+	nni_aio_fini(&pp->aio_recv);
+	nni_aio_fini(&pp->aio_send);
+	nni_aio_fini(&pp->aio_getq);
+	nni_mtx_fini(&pp->mtx);
+	NNI_FREE_STRUCT(pp);
+}
+
+
 static int
 nni_push_pipe_init(void **ppp, nni_pipe *pipe, void *psock)
 {
@@ -80,19 +95,17 @@ nni_push_pipe_init(void **ppp, nni_pipe *pipe, void *psock)
 		return (NNG_ENOMEM);
 	}
 	if ((rv = nni_aio_init(&pp->aio_recv, nni_push_recv_cb, pp)) != 0) {
-		NNI_FREE_STRUCT(pp);
-		return (rv);
+		goto fail;
 	}
 	if ((rv = nni_aio_init(&pp->aio_send, nni_push_send_cb, pp)) != 0) {
-		nni_aio_fini(&pp->aio_recv);
-		NNI_FREE_STRUCT(pp);
+		goto fail;
 		return (rv);
 	}
 	if ((rv = nni_aio_init(&pp->aio_getq, nni_push_getq_cb, pp)) != 0) {
-		nni_aio_fini(&pp->aio_send);
-		nni_aio_fini(&pp->aio_recv);
-		NNI_FREE_STRUCT(pp);
-		return (rv);
+		goto fail;
+	}
+	if ((rv = nni_mtx_init(&pp->mtx)) != 0) {
+		goto fail;
 	}
 
 	NNI_LIST_NODE_INIT(&pp->node);
@@ -100,18 +113,10 @@ nni_push_pipe_init(void **ppp, nni_pipe *pipe, void *psock)
 	pp->push = psock;
 	*ppp = pp;
 	return (0);
-}
 
-
-static void
-nni_push_pipe_fini(void *arg)
-{
-	nni_push_pipe *pp = arg;
-
-	nni_aio_fini(&pp->aio_recv);
-	nni_aio_fini(&pp->aio_send);
-	nni_aio_fini(&pp->aio_getq);
-	NNI_FREE_STRUCT(pp);
+fail:
+	nni_push_pipe_fini(pp);
+	return (rv);
 }
 
 
@@ -125,6 +130,10 @@ nni_push_pipe_start(void *arg)
 		return (NNG_EPROTO);
 	}
 
+	nni_mtx_lock(&pp->mtx);
+	pp->refcnt = 2;
+	nni_mtx_unlock(&pp->mtx);
+
 	// Schedule a receiver.  This is mostly so that we can detect
 	// a closed transport pipe.
 	nni_pipe_hold(pp->pipe);
@@ -139,12 +148,22 @@ nni_push_pipe_start(void *arg)
 
 
 static void
-nni_push_pipe_stop(void *arg)
+nni_push_pipe_stop(nni_push_pipe *pp)
 {
-	nni_push_pipe *pp = arg;
 	nni_push_sock *push = pp->push;
+	int refcnt;
 
 	nni_msgq_aio_cancel(push->uwq, &pp->aio_getq);
+
+	nni_mtx_lock(&pp->mtx);
+	NNI_ASSERT(pp->refcnt > 0);
+	pp->refcnt--;
+	refcnt = pp->refcnt;
+	nni_mtx_unlock(&pp->mtx);
+
+	if (refcnt == 0) {
+		nni_pipe_remove(pp->pipe);
+	}
 }
 
 
@@ -156,8 +175,7 @@ nni_push_recv_cb(void *arg)
 	// We normally expect to receive an error.  If a pipe actually
 	// sends us data, we just discard it.
 	if (nni_aio_result(&pp->aio_recv) != 0) {
-		nni_pipe_close(pp->pipe);
-		nni_pipe_rele(pp->pipe);
+		nni_push_pipe_stop(pp);
 		return;
 	}
 	nni_msg_free(pp->aio_recv.a_msg);
@@ -175,8 +193,7 @@ nni_push_send_cb(void *arg)
 	if (nni_aio_result(&pp->aio_send) != 0) {
 		nni_msg_free(pp->aio_send.a_msg);
 		pp->aio_send.a_msg = NULL;
-		nni_pipe_close(pp->pipe);
-		nni_pipe_rele(pp->pipe);
+		nni_push_pipe_stop(pp);
 		return;
 	}
 
@@ -192,8 +209,7 @@ nni_push_getq_cb(void *arg)
 
 	if (nni_aio_result(aio) != 0) {
 		// If the socket is closing, nothing else we can do.
-		nni_pipe_close(pp->pipe);
-		nni_pipe_rele(pp->pipe);
+		nni_push_pipe_stop(pp);
 		return;
 	}
 
@@ -244,7 +260,6 @@ static nni_proto_pipe_ops nni_push_pipe_ops = {
 	.pipe_init	= nni_push_pipe_init,
 	.pipe_fini	= nni_push_pipe_fini,
 	.pipe_start	= nni_push_pipe_start,
-	.pipe_stop	= nni_push_pipe_stop,
 };
 
 static nni_proto_sock_ops nni_push_sock_ops = {