Merge pull request #144 from Iceber/stream_recv_channel

First process the pending messages in recv channel

client.go

@@ -483,23 +483,30 @@ func (c *Client) dispatch(ctx context.Context, req *Request, resp *Response) err
 	}
 	defer c.deleteStream(s)
 
+	var msg *streamMessage
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
 	case <-c.ctx.Done():
 		return ErrClosed
 	case <-s.recvClose:
-		return s.recvErr
+		// If recv has a pending message, process that first
-	case msg := <-s.recv:
+		select {
-		if msg.header.Type == messageTypeResponse {
+		case msg = <-s.recv:
-			err = proto.Unmarshal(msg.payload[:msg.header.Length], resp)
+		default:
-		} else {
+			return s.recvErr
-			err = fmt.Errorf("unexpected %q message received: %w", msg.header.Type, ErrProtocol)
 		}
-
+	case msg = <-s.recv:
-		// return the payload buffer for reuse
-		c.channel.putmbuf(msg.payload)
-
-		return err
 	}
+
+	if msg.header.Type == messageTypeResponse {
+		err = proto.Unmarshal(msg.payload[:msg.header.Length], resp)
+	} else {
+		err = fmt.Errorf("unexpected %q message received: %w", msg.header.Type, ErrProtocol)
+	}
+
+	// return the payload buffer for reuse
+	c.channel.putmbuf(msg.payload)
+
+	return err
 }