Merge pull request #144 from Iceber/stream_recv_channel

First process the pending messages in recv channel
This commit is contained in:
Fu Wei 2023-05-09 13:11:40 +08:00 committed by GitHub
commit ac26f8cbea
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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
case msg := <-s.recv:
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)
// If recv has a pending message, process that first
select {
case msg = <-s.recv:
default:
return s.recvErr
}
// return the payload buffer for reuse
c.channel.putmbuf(msg.payload)
return err
case msg = <-s.recv:
}
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
}