fix(store/watch): fix the slow consumer bug

store/watcher.go

@@ -21,6 +21,7 @@ type Watcher struct {
 	stream     bool
 	recursive  bool
 	sinceIndex uint64
+	hub        *watcherHub
 	removed    bool
 	remove     func()
 }
@@ -51,8 +52,9 @@ func (w *Watcher) notify(e *Event, originalPath bool, deleted bool) bool {
 		select {
 		case w.EventChan <- e:
 		default:
-			// We have missed a notification. Close the channel to indicate this situation.
-			close(w.EventChan)
+			// We have missed a notification. Remove the watcher.
+			// Removing the watcher also closes the EventChan.
+			w.remove()
 		}
 		return true
 	}
@@ -62,11 +64,9 @@ func (w *Watcher) notify(e *Event, originalPath bool, deleted bool) bool {
 // Remove removes the watcher from watcherHub
 // The actual remove function is guaranteed to only be executed once
 func (w *Watcher) Remove() {
-	if w.remove != nil {
-		w.remove()
-	} else {
-		// We attached a remove function to watcher
-		// Other pkg cannot change it, so this should not happen
-		panic("missing Watcher remove function")
-	}
+	w.hub.mutex.Lock()
+	defer w.hub.mutex.Unlock()
+
+	close(w.EventChan)
+	w.remove()
 }

store/watcher_hub.go

@@ -50,6 +50,7 @@ func (wh *watcherHub) watch(key string, recursive, stream bool, index uint64) (*
 		recursive:  recursive,
 		stream:     stream,
 		sinceIndex: index,
+		hub:        wh,
 	}
 
 	if event != nil {
@@ -77,10 +78,6 @@ func (wh *watcherHub) watch(key string, recursive, stream bool, index uint64) (*
 		if w.removed { // avoid remove it twice
 			return
 		}
-
-		wh.mutex.Lock()
-		defer wh.mutex.Unlock()
-
 		w.removed = true
 		l.Remove(elem)
 		atomic.AddInt64(&wh.count, -1)