2 files changed, 67 insertions, 1 deletions

diff --git a/src/core/msgqueue.c b/src/core/msgqueue.c
index 931fda68..9ef48f2d 100644
--- a/src/core/msgqueue.c
+++ b/src/core/msgqueue.c
@@ -25,6 +25,8 @@ struct nni_msgq {
 	int		mq_get;
 	int		mq_put;
 	int		mq_closed;
+	int		mq_puterr;
+	int		mq_geterr;
 	int		mq_rwait;       // readers waiting (unbuffered)
 	int		mq_wwait;
 	nni_msg **	mq_msgs;
@@ -120,6 +122,44 @@ nni_msgq_fini(nni_msgq *mq)
 }
 
 
+void
+nni_msgq_set_put_error(nni_msgq *mq, int error)
+{
+	nni_mtx_lock(&mq->mq_lock);
+	mq->mq_puterr = error;
+	if (error) {
+		nni_cv_wake(&mq->mq_writeable);
+	}
+	nni_mtx_unlock(&mq->mq_lock);
+}
+
+
+void
+nni_msgq_set_get_error(nni_msgq *mq, int error)
+{
+	nni_mtx_lock(&mq->mq_lock);
+	mq->mq_geterr = error;
+	if (error) {
+		nni_cv_wake(&mq->mq_readable);
+	}
+	nni_mtx_unlock(&mq->mq_lock);
+}
+
+
+void
+nni_msgq_set_error(nni_msgq *mq, int error)
+{
+	nni_mtx_lock(&mq->mq_lock);
+	mq->mq_geterr = error;
+	mq->mq_puterr = error;
+	if (error) {
+		nni_cv_wake(&mq->mq_readable);
+		nni_cv_wake(&mq->mq_writeable);
+	}
+	nni_mtx_unlock(&mq->mq_lock);
+}
+
+
 // nni_msgq_signal raises a signal on the signal object. This allows a
 // waiter to be signaled, so that it can be woken e.g. due to a pipe closing.
 // Note that the signal object must be *zero* if no signal is raised.
@@ -150,6 +190,11 @@ nni_msgq_put_(nni_msgq *mq, nni_msg *msg, nni_time expire, nni_signal *sig)
 			return (NNG_ECLOSED);
 		}
 
+		if ((rv = mq->mq_puterr) != 0) {
+			nni_mtx_unlock(&mq->mq_lock);
+			return (rv);
+		}
+
 		// room in the queue?
 		if (mq->mq_len < mq->mq_cap) {
 			break;
@@ -185,7 +230,6 @@ nni_msgq_put_(nni_msgq *mq, nni_msg *msg, nni_time expire, nni_signal *sig)
 	}
 
 	// Writeable!  Yay!!
-
 	mq->mq_msgs[mq->mq_put] = msg;
 	mq->mq_put++;
 	if (mq->mq_put == mq->mq_alloc) {
@@ -251,6 +295,10 @@ nni_msgq_get_(nni_msgq *mq, nni_msg **msgp, nni_time expire, nni_signal *sig)
 			nni_mtx_unlock(&mq->mq_lock);
 			return (NNG_ECLOSED);
 		}
+		if ((rv = mq->mq_geterr) != 0) {
+			nni_mtx_unlock(&mq->mq_lock);
+			return (rv);
+		}
 		if (expire == NNI_TIME_ZERO) {
 			nni_mtx_unlock(&mq->mq_lock);
 			return (NNG_EAGAIN);
diff --git a/src/core/msgqueue.h b/src/core/msgqueue.h
index 09d7fa0c..450c55a4 100644
--- a/src/core/msgqueue.h
+++ b/src/core/msgqueue.h
@@ -87,6 +87,24 @@ extern int nni_msgq_get_sig(nni_msgq *, nni_msg **, nni_signal *);
 // It modifies the turnstile's value under the lock to a non-zero value.
 extern void nni_msgq_signal(nni_msgq *, nni_signal *);
 
+// nni_msgq_set_error sets an error condition on the message queue,
+// which causes all current and future readers/writes to return the
+// given error condition (if non-zero).  Threads waiting to put or get
+// are woken as well, if non-zero.  If zero, then any present error
+// condition is cleared, and waiters are not woken (there shouldn't be
+// any waiters unless it was already zero.)
+extern void nni_msgq_set_error(nni_msgq *, int);
+
+// nni_msgq_set_put_error sets an error condition on the put side of the
+// message queue, and for that side behaves like nni_msgq_set_error.
+// Readers (nni_msgq_get*) are unaffected.
+extern void nni_msgq_set_put_error(nni_msgq *, int);
+
+// nni_msgq_set_get_error sets an error condition on the get side of the
+// message queue, and for that side behaves like nni_msgq_set_error.
+// Readers (nni_msgq_put*) are unaffected.
+extern void nni_msgq_set_get_error(nni_msgq *, int);
+
 // nni_msgq_close closes the queue.  After this all operates on the
 // message queue will return NNG_ECLOSED.  Messages inside the queue
 // are freed.  Unlike closing a go channel, this operation is idempotent.