3 files changed, 23 insertions, 4 deletions

diff --git a/src/core/dialer.c b/src/core/dialer.c
index 443571bc..7dd018f6 100644
--- a/src/core/dialer.c
+++ b/src/core/dialer.c
@@ -322,6 +322,7 @@ nni_dialer_rele(nni_dialer *d)
 	bool reap;
 
 	nni_mtx_lock(&dialers_lk);
+	NNI_ASSERT(d->d_ref > 0);
 	d->d_ref--;
 	reap = ((d->d_ref == 0) && (d->d_closed));
 	nni_mtx_unlock(&dialers_lk);
@@ -346,6 +347,7 @@ nni_dialer_close(nni_dialer *d)
 
 	nni_dialer_shutdown(d);
 
+	nni_sock_remove_dialer(d);
 	nni_dialer_rele(d);
 }
 
diff --git a/src/core/socket.c b/src/core/socket.c
index 316f3603..1f44ebfa 100644
--- a/src/core/socket.c
+++ b/src/core/socket.c
@@ -699,7 +699,6 @@ nni_sock_shutdown(nni_sock *sock)
 
 	while ((d = nni_list_first(&sock->s_dialers)) != NULL) {
 		nni_dialer_hold(d);
-		nni_list_node_remove(&d->d_node);
 		nni_mtx_unlock(&sock->s_mx);
 		nni_dialer_close(d);
 		nni_mtx_lock(&sock->s_mx);
@@ -941,10 +940,17 @@ int
 nni_sock_add_dialer(nni_sock *s, nni_dialer *d)
 {
 	nni_sockopt *sopt;
+	int rv;
+
+	// grab a hold on the dialer for the socket
+	if ((rv = nni_dialer_hold(d)) != 0) {
+		return (rv);
+	}
 
 	nni_mtx_lock(&s->s_mx);
 	if (s->s_closing) {
 		nni_mtx_unlock(&s->s_mx);
+		nni_dialer_rele(d);
 		return (NNG_ECLOSED);
 	}
 
@@ -968,6 +974,19 @@ nni_sock_add_dialer(nni_sock *s, nni_dialer *d)
 	return (0);
 }
 
+void
+nni_sock_remove_dialer(nni_dialer *d)
+{
+	nni_sock *s = d->d_sock;
+	nni_mtx_lock(&s->s_mx);
+	NNI_ASSERT(nni_list_node_active(&d->d_node));
+	nni_list_node_remove(&d->d_node);
+	nni_mtx_unlock(&s->s_mx);
+
+	// also drop the hold from the socket
+	nni_dialer_rele(d);
+}
+
 int
 nni_sock_setopt(
     nni_sock *s, const char *name, const void *v, size_t sz, nni_type t)
@@ -1541,7 +1560,6 @@ nni_dialer_shutdown(nni_dialer *d)
 	NNI_LIST_FOREACH (&d->d_pipes, p) {
 		nni_pipe_close(p);
 	}
-	nni_list_node_remove(&d->d_node);
 	nni_mtx_unlock(&s->s_mx);
 }
 
@@ -1575,8 +1593,6 @@ dialer_reap(void *arg)
 		return;
 	}
 
-	nni_list_node_remove(&d->d_node);
-
 	nni_mtx_unlock(&s->s_mx);
 
 	nni_sock_rele(s);
diff --git a/src/core/sockimpl.h b/src/core/sockimpl.h
index 9734d94a..62c6b27c 100644
--- a/src/core/sockimpl.h
+++ b/src/core/sockimpl.h
@@ -126,6 +126,7 @@ struct nni_pipe {
 extern int  nni_sock_add_dialer(nni_sock *, nni_dialer *);
 extern int  nni_sock_add_listener(nni_sock *, nni_listener *);
 extern void nni_sock_remove_listener(nni_listener *);
+extern void nni_sock_remove_dialer(nni_dialer *);
 
 extern void nni_dialer_add_pipe(nni_dialer *, void *);
 extern void nni_dialer_shutdown(nni_dialer *);