grpcproxy: reliably track rid in watchergroups

Couldn't find watcher group from rid on server stream close, leading to
the watcher group sending on a closed channel.

Also got rid of send closing the watcher stream if the buffer is full,
this could lead to a send after close while broadcasting to all receivers.
Instead, if a send times out, have the caller issue a cancel().

Fixes #6739

proxy/grpcproxy/watch.go

@@ -192,18 +192,15 @@ func (sws *serverWatchStream) addCoalescedWatcher(w watcher) {
 }
 
 func (sws *serverWatchStream) addDedicatedWatcher(w watcher, rev int64) {
-	sws.mu.Lock()
-	defer sws.mu.Unlock()
-
 	ctx, cancel := context.WithCancel(sws.proxyCtx)
-
 	wch := sws.cw.Watch(ctx,
 		w.wr.key, clientv3.WithRange(w.wr.end),
 		clientv3.WithRev(rev),
 		clientv3.WithProgressNotify(),
 		clientv3.WithCreatedNotify(),
 	)
-
+	sws.mu.Lock()
+	defer sws.mu.Unlock()
 	ws := newWatcherSingle(wch, cancel, w, sws)
 	sws.singles[w.id] = ws
 	go ws.run()
@@ -213,12 +210,11 @@ func (sws *serverWatchStream) maybeCoalesceWatcher(ws watcherSingle) bool {
 	sws.mu.Lock()
 	defer sws.mu.Unlock()
 
-	rid := receiverID{streamID: sws.id, watcherID: ws.w.id}
 	// do not add new watchers when stream is closing
 	if sws.inGroups == nil {
 		return false
 	}
-	if sws.groups.maybeJoinWatcherSingle(rid, ws) {
+	if sws.groups.maybeJoinWatcherSingle(ws) {
 		delete(sws.singles, ws.w.id)
 		sws.inGroups[ws.w.id] = struct{}{}
 		return true

proxy/grpcproxy/watcher.go

@@ -37,9 +37,9 @@ type watcher struct {
 	ch       chan<- *pb.WatchResponse
 }
 
-func (w *watcher) send(wr clientv3.WatchResponse) {
+func (w *watcher) send(wr clientv3.WatchResponse) bool {
 	if wr.IsProgressNotify() && !w.progress {
-		return
+		return true
 	}
 
 	events := make([]*mvccpb.Event, 0, len(wr.Events))
@@ -77,7 +77,7 @@ func (w *watcher) send(wr clientv3.WatchResponse) {
 
 	// all events are filtered out?
 	if !wr.IsProgressNotify() && !wr.Created && len(events) == 0 {
-		return
+		return true
 	}
 
 	pbwr := &pb.WatchResponse{
@@ -89,8 +89,7 @@ func (w *watcher) send(wr clientv3.WatchResponse) {
 	select {
 	case w.ch <- pbwr:
 	case <-time.After(50 * time.Millisecond):
-		// close the watch chan will notify the stream sender.
-		// the stream will gc all its watchers.
-		close(w.ch)
+		return false
 	}
+	return true
 }

proxy/grpcproxy/watcher_group.go

@@ -63,7 +63,10 @@ func (wg *watcherGroup) broadcast(wr clientv3.WatchResponse) {
 
 	wg.rev = wr.Header.Revision
 	for _, r := range wg.receivers {
-		r.send(wr)
+		if !r.send(wr) {
+			wg.cancel()
+			break
+		}
 	}
 }
 

proxy/grpcproxy/watcher_groups.go

@@ -96,24 +96,24 @@ func (wgs *watchergroups) removeWatcher(rid receiverID) (int64, bool) {
 	return -1, false
 }
 
-func (wgs *watchergroups) maybeJoinWatcherSingle(rid receiverID, ws watcherSingle) bool {
+func (wgs *watchergroups) maybeJoinWatcherSingle(ws watcherSingle) bool {
 	wgs.mu.Lock()
 	defer wgs.mu.Unlock()
 
+	rid := receiverID{streamID: ws.sws.id, watcherID: ws.w.id}
 	group, ok := wgs.groups[ws.w.wr]
 	if ok {
-		return group.add(receiverID{streamID: ws.sws.id, watcherID: ws.w.id}, ws.w) != -1
+		return group.add(rid, ws.w) != -1
 	}
-
-	if ws.canPromote() {
-		wg := newWatchergroup(ws.ch, ws.cancel)
-		wgs.groups[ws.w.wr] = wg
-		wg.add(receiverID{streamID: ws.sws.id, watcherID: ws.w.id}, ws.w)
-		go wg.run()
-		return true
+	if !ws.canPromote() {
+		return false
 	}
-
-	return false
+	wg := newWatchergroup(ws.ch, ws.cancel)
+	wgs.groups[ws.w.wr] = wg
+	wgs.idToGroup[rid] = wg
+	wg.add(rid, ws.w)
+	go wg.run()
+	return true
 }
 
 func (wgs *watchergroups) stop() {

proxy/grpcproxy/watcher_single.go

@@ -53,8 +53,10 @@ func (ws watcherSingle) run() {
 
 	for wr := range ws.ch {
 		ws.lastStoreRev = wr.Header.Revision
-		ws.w.send(wr)
-
+		if !ws.w.send(wr) {
+			ws.cancel()
+			return
+		}
 		if ws.sws.maybeCoalesceWatcher(ws) {
 			return
 		}