Merge pull request #20614 from apullo777/progress-notify

cache: make the cache progress-aware

Author: serathius

cache/cache.go

@@ -271,13 +271,7 @@ func (c *Cache) applyStorage(storeW *watcher) error {
 			if !ok {
 				return nil
 			}
-			if resp.Canceled {
-				return nil
-			}
-			if len(resp.Events) == 0 {
-				continue
-			}
-			if err := c.store.Apply(resp.Events); err != nil {
+			if err := c.store.Apply(resp); err != nil {
 				return err
 			}
 		}

cache/demux.go

@@ -16,9 +16,11 @@ package cache
 
 import (
 	"context"
+	"errors"
 	"sync"
 	"time"
 
+	"go.etcd.io/etcd/api/v3/etcdserverpb"
 	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
@@ -108,6 +110,9 @@ func (d *demux) Unregister(w *watcher) {
 func (d *demux) Init(minRev int64) {
 	d.mu.Lock()
 	defer d.mu.Unlock()
+	if minRev == 0 {
+		return
+	}
 	if d.minRev == 0 {
 		// Watch started for empty demux
 		d.minRev = minRev
@@ -129,23 +134,35 @@ func (d *demux) Init(minRev int64) {
 }
 
 func (d *demux) Broadcast(resp clientv3.WatchResponse) error {
-	events := resp.Events
-	if len(events) == 0 {
-		return nil
-	}
-
 	d.mu.Lock()
 	defer d.mu.Unlock()
-	err := validateRevisions(events, d.maxRev)
+	if d.minRev == 0 {
+		return errors.New("demux: not initialized")
+	}
+	err := validateRevisions(resp, d.maxRev)
 	if err != nil {
 		return err
 	}
-	d.updateStoreLocked(events)
-	d.broadcastLocked(events)
+	d.updateStoreLocked(resp)
+	d.broadcastLocked(resp)
 	return nil
 }
 
-func (d *demux) updateStoreLocked(events []*clientv3.Event) {
+func (d *demux) LatestRev() int64 {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+	return d.maxRev
+}
+
+func (d *demux) updateStoreLocked(resp clientv3.WatchResponse) {
+	if resp.IsProgressNotify() {
+		d.maxRev = resp.Header.Revision
+		return
+	}
+	if len(resp.Events) == 0 {
+		return
+	}
+	events := resp.Events
 	batchStart := 0
 	for end := 1; end < len(events); end++ {
 		if events[end].Kv.ModRevision != events[batchStart].Kv.ModRevision {
@@ -167,7 +184,33 @@ func (d *demux) updateStoreLocked(events []*clientv3.Event) {
 	d.maxRev = events[len(events)-1].Kv.ModRevision
 }
 
-func (d *demux) broadcastLocked(events []*clientv3.Event) {
+func (d *demux) broadcastLocked(resp clientv3.WatchResponse) {
+	switch {
+	case resp.IsProgressNotify():
+		d.broadcastProgressLocked(resp.Header.Revision)
+	case len(resp.Events) != 0:
+		d.broadcastEventsLocked(resp.Events)
+	default:
+	}
+}
+
+func (d *demux) broadcastProgressLocked(progressRev int64) {
+	for w, nextRev := range d.activeWatchers {
+		if nextRev >= progressRev {
+			continue
+		}
+		resp := clientv3.WatchResponse{
+			Header: etcdserverpb.ResponseHeader{
+				Revision: progressRev,
+			},
+		}
+		if w.enqueueResponse(resp) {
+			d.activeWatchers[w] = progressRev + 1
+		}
+	}
+}
+
+func (d *demux) broadcastEventsLocked(events []*clientv3.Event) {
 	firstRev := events[0].Kv.ModRevision
 	lastRev := events[len(events)-1].Kv.ModRevision
 
@@ -177,18 +220,22 @@ func (d *demux) broadcastLocked(events []*clientv3.Event) {
 			delete(d.activeWatchers, w)
 			continue
 		}
+
 		sendStart := len(events)
 		for i, ev := range events {
 			if ev.Kv.ModRevision >= nextRev {
 				sendStart = i
 				break
 			}
 		}
+
 		if sendStart == len(events) {
 			continue
 		}
 
-		if !w.enqueueEvent(events[sendStart:]) { // overflow → lagging
+		if !w.enqueueResponse(clientv3.WatchResponse{
+			Events: events[sendStart:],
+		}) { // overflow → lagging
 			d.laggingWatchers[w] = nextRev
 			delete(d.activeWatchers, w)
 		} else {
@@ -241,17 +288,26 @@ func (d *demux) resyncLaggingWatchers() {
 			continue
 		}
 		// TODO: re-enable key‐predicate in Filter when non‐zero startRev or performance tuning is needed
-		enqueueFailed := false
+		resyncSuccess := true
 		d.history.AscendGreaterOrEqual(nextRev, func(rev int64, eventBatch []*clientv3.Event) bool {
-			if !w.enqueueEvent(eventBatch) { // buffer overflow: watcher still lagging
-				enqueueFailed = true
+			resp := clientv3.WatchResponse{
+				Events: eventBatch,
+			}
+			if !w.enqueueResponse(resp) { // buffer overflow: watcher still lagging
+				resyncSuccess = false
 				return false
 			}
 			nextRev = rev + 1
 			return true
 		})
-
-		if !enqueueFailed {
+		// Send progress to just resync.
+		if resyncSuccess {
+			resp := clientv3.WatchResponse{
+				Header: etcdserverpb.ResponseHeader{Revision: d.maxRev},
+			}
+			if d.maxRev > nextRev && w.enqueueResponse(resp) {
+				nextRev = d.maxRev + 1
+			}
 			delete(d.laggingWatchers, w)
 			d.activeWatchers[w] = nextRev
 		} else {

cache/demux_test.go

@@ -257,6 +257,7 @@ func TestBroadcastBatching(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			d := newDemux(16, 10*time.Millisecond)
 			w := newWatcher(len(tt.input)+1, nil)
+			d.Init(1)
 			d.Register(w, 0)
 
 			d.Broadcast(respWithEventRevs(tt.input...))
@@ -305,6 +306,7 @@ func TestSlowWatcherResync(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			d := newDemux(16, 10*time.Millisecond)
 			w := newWatcher(1, nil)
+			d.Init(1)
 			d.Register(w, 0)
 
 			d.Broadcast(respWithEventRevs(tt.input...))

cache/store.go

@@ -15,6 +15,7 @@
 package cache
 
 import (
+	"errors"
 	"fmt"
 	"sync"
 
@@ -92,6 +93,10 @@ func (s *store) getSnapshot(rev int64) (*snapshot, int64, error) {
 		targetSnapshot = snap
 		return false
 	})
+	// If s.history < rev < s.latest.rev serve latest.
+	if targetSnapshot == nil {
+		targetSnapshot = &s.latest
+	}
 
 	return targetSnapshot, s.latest.rev, nil
 }
@@ -110,14 +115,36 @@ func (s *store) Restore(kvs []*mvccpb.KeyValue, rev int64) {
 	s.history.Append(newClonedSnapshot(rev, s.latest.tree))
 }
 
-func (s *store) Apply(events []*clientv3.Event) error {
+func (s *store) Apply(resp clientv3.WatchResponse) error {
+	if resp.Canceled {
+		return errors.New("canceled")
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
-
-	if err := validateRevisions(events, s.latest.rev); err != nil {
+	if err := validateRevisions(resp, s.latest.rev); err != nil {
 		return err
 	}
 
+	switch {
+	case resp.IsProgressNotify():
+		s.applyProgressNotifyLocked(resp.Header.Revision)
+		return nil
+	case len(resp.Events) != 0:
+		return s.applyEventsLocked(resp.Events)
+	default:
+		return nil
+	}
+}
+
+func (s *store) applyProgressNotifyLocked(revision int64) {
+	if s.latest.rev == 0 {
+		return
+	}
+	s.latest.rev = revision
+}
+
+func (s *store) applyEventsLocked(events []*clientv3.Event) error {
 	for i := 0; i < len(events); {
 		rev := events[i].Kv.ModRevision
 
@@ -145,7 +172,14 @@ func (s *store) LatestRev() int64 {
 	return s.latest.rev
 }
 
-func validateRevisions(events []*clientv3.Event, latestRev int64) error {
+func validateRevisions(resp clientv3.WatchResponse, latestRev int64) error {
+	if resp.IsProgressNotify() {
+		if resp.Header.Revision < latestRev {
+			return fmt.Errorf("cache: progress notification out of order (progress %d < latest %d)", resp.Header.Revision, latestRev)
+		}
+		return nil
+	}
+	events := resp.Events
 	if len(events) == 0 {
 		return nil
 	}

cache/store_test.go

@@ -281,7 +281,8 @@ func TestStoreApply(t *testing.T) {
 
 			var gotErr error
 			for batchIndex, batch := range test.eventBatches {
-				if err := s.Apply(batch); err != nil {
+				resp := clientv3.WatchResponse{Events: batch}
+				if err := s.Apply(resp); err != nil {
 					gotErr = err
 					if !test.expectErr {
 						t.Fatalf("Apply(batch %d) unexpected error: %v", batchIndex, err)
@@ -397,7 +398,8 @@ func TestRestoreAppendCloneImmutability(t *testing.T) {
 				s.Restore(test.initialKVs, test.initialRev)
 			}
 			if len(test.events) > 0 {
-				if err := s.Apply(test.events); err != nil {
+				resp := clientv3.WatchResponse{Events: test.events}
+				if err := s.Apply(resp); err != nil {
 					t.Fatalf("Apply failed: %v", err)
 				}
 			}

cache/watcher.go

@@ -38,21 +38,21 @@ func newWatcher(bufSize int, pred KeyPredicate) *watcher {
 
 // true  -> events delivered (or filtered/duplicate)
 // false -> buffer full (caller should mark watcher “lagging”)
-func (w *watcher) enqueueEvent(eventBatch []*clientv3.Event) bool {
-	if w.keyPred != nil {
-		filtered := make([]*clientv3.Event, 0, len(eventBatch))
-		for _, event := range eventBatch {
+func (w *watcher) enqueueResponse(resp clientv3.WatchResponse) bool {
+	if !resp.IsProgressNotify() && w.keyPred != nil {
+		filtered := make([]*clientv3.Event, 0, len(resp.Events))
+		for _, event := range resp.Events {
 			if w.keyPred(event.Kv.Key) {
 				filtered = append(filtered, event)
 			}
 		}
 		if len(filtered) == 0 {
 			return true
 		}
-		eventBatch = filtered
+		resp.Events = filtered
 	}
 	select {
-	case w.respCh <- clientv3.WatchResponse{Events: eventBatch}:
+	case w.respCh <- resp:
 		return true
 	default:
 		return false