diff --git a/peer.go b/peer.go
index 63566a8e..a9833e4b 100644
--- a/peer.go
+++ b/peer.go
@@ -738,6 +738,7 @@ func (ms *msgStream) Stop() {
 func (ms *msgStream) msgConsumer() {
 	defer ms.wg.Done()
 	defer peerLog.Tracef(ms.stopMsg)
+	defer atomic.StoreInt32(&ms.streamShutdown, 1)
 
 	peerLog.Tracef(ms.startMsg)
 
@@ -754,7 +755,6 @@ func (ms *msgStream) msgConsumer() {
 			select {
 			case <-ms.quit:
 				ms.msgCond.L.Unlock()
-				atomic.StoreInt32(&ms.streamShutdown, 1)
 				return
 			default:
 			}
@@ -778,7 +778,6 @@ func (ms *msgStream) msgConsumer() {
 		select {
 		case ms.producerSema <- struct{}{}:
 		case <-ms.quit:
-			atomic.StoreInt32(&ms.streamShutdown, 1)
 			return
 		}
 	}