diff --git a/clientv3/watch.go b/clientv3/watch.go index 3612d35f1..d3981bfd9 100644 --- a/clientv3/watch.go +++ b/clientv3/watch.go @@ -311,55 +311,63 @@ func (w *watcher) Watch(ctx context.Context, key string, opts ...OpOption) Watch ok := false ctxKey := streamKeyFromCtx(ctx) - // find or allocate appropriate grpc watch stream - w.mu.Lock() - if w.streams == nil { - // closed - w.mu.Unlock() - ch := make(chan WatchResponse) - close(ch) - return ch - } - wgs := w.streams[ctxKey] - if wgs == nil { - wgs = w.newWatcherGrpcStream(ctx) - w.streams[ctxKey] = wgs - } - donec := wgs.donec - reqc := wgs.reqc - w.mu.Unlock() - - // couldn't create channel; return closed channel - closeCh := make(chan WatchResponse, 1) - - // submit request - select { - case reqc <- wr: - ok = true - case <-wr.ctx.Done(): - case <-donec: - if wgs.closeErr != nil { - closeCh <- WatchResponse{Canceled: true, closeErr: wgs.closeErr} - break + var closeCh chan WatchResponse + for { + // find or allocate appropriate grpc watch stream + w.mu.Lock() + if w.streams == nil { + // closed + w.mu.Unlock() + ch := make(chan WatchResponse) + close(ch) + return ch } - // retry; may have dropped stream from no ctxs - return w.Watch(ctx, key, opts...) - } + wgs := w.streams[ctxKey] + if wgs == nil { + wgs = w.newWatcherGrpcStream(ctx) + w.streams[ctxKey] = wgs + } + donec := wgs.donec + reqc := wgs.reqc + w.mu.Unlock() - // receive channel - if ok { + // couldn't create channel; return closed channel + if closeCh == nil { + closeCh = make(chan WatchResponse, 1) + } + + // submit request select { - case ret := <-wr.retc: - return ret - case <-ctx.Done(): + case reqc <- wr: + ok = true + case <-wr.ctx.Done(): + ok = false case <-donec: + ok = false if wgs.closeErr != nil { closeCh <- WatchResponse{Canceled: true, closeErr: wgs.closeErr} break } // retry; may have dropped stream from no ctxs - return w.Watch(ctx, key, opts...) + continue } + + // receive channel + if ok { + select { + case ret := <-wr.retc: + return ret + case <-ctx.Done(): + case <-donec: + if wgs.closeErr != nil { + closeCh <- WatchResponse{Canceled: true, closeErr: wgs.closeErr} + break + } + // retry; may have dropped stream from no ctxs + continue + } + } + break } close(closeCh)