cache: remove early compaction check from Cache.Watch

Signed-off-by: Peter Chang <[email protected]>

Author: apullo777

cache/cache.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	pb "go.etcd.io/etcd/api/v3/etcdserverpb"
+	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
 	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
@@ -95,18 +96,6 @@ func (c *Cache) Watch(ctx context.Context, key string, opts ...clientv3.OpOption
 	op := clientv3.OpWatch(key, opts...)
 	startRev := op.Rev()
 
-	if startRev != 0 {
-		if oldest := c.demux.PeekOldest(); oldest != 0 && startRev < oldest {
-			ch := make(chan clientv3.WatchResponse, 1)
-			ch <- clientv3.WatchResponse{
-				Canceled:        true,
-				CompactRevision: startRev,
-			}
-			close(ch)
-			return ch
-		}
-	}
-
 	pred, err := c.validateWatch(key, op)
 	if err != nil {
 		ch := make(chan clientv3.WatchResponse, 1)
@@ -132,6 +121,13 @@ func (c *Cache) Watch(ctx context.Context, key string, opts ...clientv3.OpOption
 				return
 			case events, ok := <-w.eventQueue:
 				if !ok {
+					if w.cancelResp != nil {
+						select {
+						case <-ctx.Done():
+						case <-c.internalCtx.Done():
+						case responseChan <- *w.cancelResp:
+						}
+					}
 					return
 				}
 				select {
@@ -299,7 +295,11 @@ func (c *Cache) watchEvents(watchCh clientv3.WatchChan, applyErr <-chan error, r
 			readyOnce.Do(func() { c.ready.Set() })
 			if err := resp.Err(); err != nil {
 				c.ready.Reset()
-				c.demux.Purge()
+				if errors.Is(err, rpctypes.ErrCompacted) || resp.CompactRevision > 0 {
+					c.demux.Compact(resp.CompactRevision)
+				} else {
+					c.demux.Purge()
+				}
 				return err
 			}
 			c.demux.Broadcast(resp.Events)

cache/demux.go

@@ -151,7 +151,7 @@ func (d *demux) Broadcast(events []*clientv3.Event) {
 	}
 }
 
-// Purge is called when etcd compaction invalidates our cached history, so clients should resubscribe.
+// Purge stops all watchers and rebase history on watch errors
 func (d *demux) Purge() {
 	d.mu.Lock()
 	defer d.mu.Unlock()
@@ -166,6 +166,21 @@ func (d *demux) Purge() {
 	d.laggingWatchers = make(map[*watcher]int64)
 }
 
+// Compact is called when etcd reports a compaction at compactRev to rebase history;
+// it keeps provably-too-old watchers for later cancellation, stops others, and clients should resubscribe.
+func (d *demux) Compact(compactRev int64) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.history.RebaseHistory()
+
+	for w, next := range d.activeWatchers {
+		if next != 0 && next <= compactRev {
+			delete(d.activeWatchers, w)
+			d.laggingWatchers[w] = next
+		}
+	}
+}
+
 func (d *demux) resyncLaggingWatchers() {
 	d.mu.Lock()
 	defer d.mu.Unlock()
@@ -177,7 +192,7 @@ func (d *demux) resyncLaggingWatchers() {
 
 	for w, nextRev := range d.laggingWatchers {
 		if nextRev < oldestRev {
-			w.Stop()
+			w.Compact(nextRev)
 			delete(d.laggingWatchers, w)
 			continue
 		}
@@ -200,9 +215,3 @@ func (d *demux) resyncLaggingWatchers() {
 		}
 	}
 }
-
-func (d *demux) PeekOldest() int64 {
-	d.mu.RLock()
-	defer d.mu.RUnlock()
-	return d.history.PeekOldest()
-}

cache/watcher.go

@@ -15,24 +15,24 @@
 package cache
 
 import (
-	"sync/atomic"
+	"sync"
 
+	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
 	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 // watcher holds one client’s buffered stream of events.
 type watcher struct {
 	eventQueue chan []*clientv3.Event
+	cancelResp *clientv3.WatchResponse
 	keyPred    KeyPredicate
-	stopped    int32
-	done       chan struct{} // closed together with Stop()
+	stopOnce   sync.Once
 }
 
 func newWatcher(bufSize int, pred KeyPredicate) *watcher {
 	return &watcher{
 		eventQueue: make(chan []*clientv3.Event, bufSize),
 		keyPred:    pred,
-		done:       make(chan struct{}),
 	}
 }
 
@@ -59,12 +59,21 @@ func (w *watcher) enqueueEvent(eventBatch []*clientv3.Event) bool {
 	}
 }
 
+func (w *watcher) Compact(compactRev int64) {
+	resp := &clientv3.WatchResponse{
+		Canceled:        true,
+		CompactRevision: compactRev,
+		CancelReason:    rpctypes.ErrCompacted.Error(),
+	}
+	w.stopOnce.Do(func() {
+		w.cancelResp = resp
+		close(w.eventQueue)
+	})
+}
+
 // Stop closes the event channel atomically.
 func (w *watcher) Stop() {
-	if atomic.CompareAndSwapInt32(&w.stopped, 0, 1) {
+	w.stopOnce.Do(func() {
 		close(w.eventQueue)
-		close(w.done)
-	}
+	})
 }
-
-func (w *watcher) Done() <-chan struct{} { return w.done }