// Wait blocks until i) the new transport is up or ii) ctx is done or iii) cc is closed.
func (cc *Conn) Wait(ctx context.Context) (transport.ClientTransport, error) {
for {
cc.mu.Lock()
switch {
case cc.state == Shutdown:
cc.mu.Unlock()
return nil, ErrClientConnClosing
case cc.state == Ready:
cc.mu.Unlock()
return cc.transport, nil
default:
ready := cc.ready
if ready == nil {
ready = make(chan struct{})
cc.ready = ready
}
cc.mu.Unlock()
select {
case <-ctx.Done():
return nil, transport.ContextErr(ctx.Err())
// Wait until the new transport is ready or failed.
case <-ready:
}
}
}
}
// When wait returns, either the new transport is up or ClientConn is
// closing. Used to avoid working on a dying transport. It updates and
// returns the transport and its version when there is no error.
func (cc *ClientConn) wait(ctx context.Context, ts int) (transport.ClientTransport, int, error) {
for {
cc.mu.Lock()
switch {
case cc.state == Shutdown:
cc.mu.Unlock()
return nil, 0, ErrClientConnClosing
case ts < cc.transportSeq:
// Worked on a dying transport. Try the new one immediately.
defer cc.mu.Unlock()
return cc.transport, cc.transportSeq, nil
default:
ready := cc.ready
if ready == nil {
ready = make(chan struct{})
cc.ready = ready
}
cc.mu.Unlock()
select {
case <-ctx.Done():
return nil, 0, transport.ContextErr(ctx.Err())
// Wait until the new transport is ready or failed.
case <-ready:
}
}
}
}
// NewClientStream creates a new Stream for the client side. This is called
// by generated code.
func NewClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (ClientStream, error) {
var (
t transport.ClientTransport
err error
)
t, err = cc.dopts.picker.Pick(ctx)
if err != nil {
return nil, toRPCErr(err)
}
// TODO(zhaoq): CallOption is omitted. Add support when it is needed.
callHdr := &transport.CallHdr{
Host: cc.authority,
Method: method,
}
cs := &clientStream{
desc: desc,
codec: cc.dopts.codec,
tracing: EnableTracing,
}
if cs.tracing {
cs.trInfo.tr = trace.New("grpc.Sent."+methodFamily(method), method)
cs.trInfo.firstLine.client = true
if deadline, ok := ctx.Deadline(); ok {
cs.trInfo.firstLine.deadline = deadline.Sub(time.Now())
}
cs.trInfo.tr.LazyLog(&cs.trInfo.firstLine, false)
ctx = trace.NewContext(ctx, cs.trInfo.tr)
}
s, err := t.NewStream(ctx, callHdr)
if err != nil {
cs.finish(err)
return nil, toRPCErr(err)
}
cs.t = t
cs.s = s
cs.p = &parser{s: s}
// Listen on ctx.Done() to detect cancellation when there is no pending
// I/O operations on this stream.
go func() {
select {
case <-t.Error():
// Incur transport error, simply exit.
case <-s.Context().Done():
err := s.Context().Err()
cs.finish(err)
cs.closeTransportStream(transport.ContextErr(err))
}
}()
return cs, nil
}
func TestContextErr(t *testing.T) {
for _, test := range []struct {
// input
errIn error
// outputs
errOut transport.StreamError
}{
{context.DeadlineExceeded, transport.StreamErrorf(codes.DeadlineExceeded, "%v", context.DeadlineExceeded)},
{context.Canceled, transport.StreamErrorf(codes.Canceled, "%v", context.Canceled)},
} {
err := transport.ContextErr(test.errIn)
if err != test.errOut {
t.Fatalf("ContextErr{%v} = %v \nwant %v", test.errIn, err, test.errOut)
}
}
}