Remove CAS/WatchKey layers from notification delay tests.

Author: seizethedave

kv/memberlist/memberlist_client.go

@@ -497,7 +497,7 @@ func (m *KV) running(ctx context.Context) error {
 		// Start delayed key notifications.
 		notifTicker := time.NewTicker(m.cfg.NotifyInterval)
 		defer notifTicker.Stop()
-		go m.sendKeyNotifications(ctx, notifTicker.C)
+		go m.monitorKeyNotifications(ctx, notifTicker.C)
 	}
 
 	var tickerChan <-chan time.Time
@@ -932,14 +932,26 @@ func (m *KV) notifyWatchers(key string) {
 	m.keyNotifications[key] = struct{}{}
 }
 
-// sendKeyNotifications sends accumulated notifications to all watchers of
+// monitorKeyNotifications sends accumulated notifications to all watchers of
 // respective keys when the given channel ticks.
-func (m *KV) sendKeyNotifications(ctx context.Context, tickChan <-chan time.Time) {
+func (m *KV) monitorKeyNotifications(ctx context.Context, tickChan <-chan time.Time) {
 	if m.cfg.NotifyInterval <= 0 {
 		panic("sendNotifications called with NotifyInterval <= 0")
 	}
 
+	for {
+		select {
+		case <-tickChan:
+			m.sendKeyNotifications()
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func (m *KV) sendKeyNotifications() {
 	newNotifs := func() map[string]struct{} {
+		// Grab and clear accumulated notifications.
 		m.notifMu.Lock()
 		defer m.notifMu.Unlock()
 
@@ -952,15 +964,8 @@ func (m *KV) sendKeyNotifications(ctx context.Context, tickChan <-chan time.Time
 		return notifs
 	}
 
-	for {
-		select {
-		case <-tickChan:
-			for key := range newNotifs() {
-				m.notifyWatchersSync(key)
-			}
-		case <-ctx.Done():
-			return
-		}
+	for key := range newNotifs() {
+		m.notifyWatchersSync(key)
 	}
 }
 

kv/memberlist/memberlist_client_test.go

@@ -1801,116 +1801,91 @@ func TestNotificationDelay(t *testing.T) {
 		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), kv))
 	})
 
-	cli, err := NewClient(kv, codec)
-	require.NoError(t, err)
-
-	casInterval := 1 * time.Second
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	// Arrange to do our own ticking.
-	tick := make(chan time.Time)
-	go kv.sendKeyNotifications(ctx, tick)
+	watchChan := make(chan string, 16)
 
-	// Mirror any WatchKey updates to our own map and verify that they
-	// eventually arrive.
+	// Add ourselves as a watcher.
+	kv.watchersMu.Lock()
+	kv.watchers["foo_123"] = append(kv.watchers["foo_123"], watchChan)
+	kv.watchers["foo_124"] = append(kv.watchers["foo_124"], watchChan)
+	kv.watchersMu.Unlock()
 
-	dbMu := sync.Mutex{}
-	db := make(map[string]*data)
-	keyCalls := make(map[string]int)
+	defer func() {
+		kv.watchersMu.Lock()
+		removeWatcherChannel("foo_123", watchChan, kv.watchers)
+		removeWatcherChannel("foo_124", watchChan, kv.watchers)
+		kv.watchersMu.Unlock()
+	}()
 
-	setData := func(key string, d *data) {
-		dbMu.Lock()
-		defer dbMu.Unlock()
-		db[key] = d
-		keyCalls[key]++
-	}
-	getData := func(key string) *data {
-		dbMu.Lock()
-		defer dbMu.Unlock()
-		return db[key]
-	}
-	callsForKey := func(key string) int {
-		dbMu.Lock()
-		defer dbMu.Unlock()
-		return keyCalls[key]
-	}
-	verifyVal := func(k, v string) bool {
-		d := getData(k)
-		if d == nil {
-			return false
+	verifyNotifs := func(expected map[string]int, comment string) {
+		observed := make(map[string]int, len(expected))
+		for kk := range expected {
+			observed[kk] = 0
+		}
+	loop:
+		for {
+			select {
+			case k := <-watchChan:
+				observed[k]++
+			default:
+				break loop
+			}
 		}
-		_, ok := d.Members[v]
-		return ok
+		require.Equal(t, expected, observed, comment)
 	}
 
-	watchData := func(key string) {
-		go cli.WatchKey(ctx, key, func(in any) bool {
-			setData(key, in.(*data))
-			return true
-		})
+	drainChan := func() {
+		for {
+			select {
+			case <-watchChan:
+			default:
+				return
+			}
+		}
 	}
 
-	watchData("foo_123")
-	watchData("foo_124")
-
-	assert.Equal(t, 0, callsForKey("foo_123"))
-	assert.Equal(t, 0, callsForKey("foo_124"))
-
-	err = cas(cli, "foo_123", updateFn("val1"))
-	require.NoError(t, err)
-
-	tick <- time.Now()
-
-	require.Eventually(t, func() bool {
-		return verifyVal("foo_123", "val1")
-	}, 3*casInterval, 25*time.Millisecond)
-
-	assert.Equal(t, 1, callsForKey("foo_123"))
-	assert.Equal(t, 0, callsForKey("foo_124"))
+	kv.notifyWatchers("foo_123")
+	kv.sendKeyNotifications()
+	verifyNotifs(map[string]int{"foo_123": 1}, "1 change 1 notification")
 
 	// Test coalescing of updates.
-	err = cas(cli, "foo_123", updateFn("val2"))
-	require.NoError(t, err)
-	err = cas(cli, "foo_123", updateFn("val3"))
-	require.NoError(t, err)
-	err = cas(cli, "foo_123", updateFn("val4"))
-	require.NoError(t, err)
-
-	assert.Equal(t, 1, callsForKey("foo_123"), "no flush -> no callback")
-	assert.Equal(t, 0, callsForKey("foo_124"), "no flush -> no callback")
-
-	tick <- time.Now()
-
-	require.Eventually(t, func() bool {
-		return verifyVal("foo_123", "val4")
-	}, 3*casInterval, 25*time.Millisecond, "multiple updates should be coalesced into the last one")
-
-	assert.Equal(t, 2, callsForKey("foo_123"))
-	assert.Equal(t, 0, callsForKey("foo_124"))
-
-	err = cas(cli, "foo_123", updateFn("val100"))
-	require.NoError(t, err)
-	err = cas(cli, "foo_124", updateFn("val101"))
-	require.NoError(t, err)
-
-	tick <- time.Now()
-
-	require.Eventually(t, func() bool {
-		return verifyVal("foo_123", "val100") && verifyVal("foo_124", "val101")
-	}, 3*casInterval, 25*time.Millisecond)
-
-	assert.Equal(t, 3, callsForKey("foo_123"))
-	assert.Equal(t, 1, callsForKey("foo_124"))
-
-	require.NotPanics(t, func() {
-		tick <- time.Now()
-		tick <- time.Now()
-		tick <- time.Now()
-	}, "shouldn't panic or anything like that when ticked without updates")
-
-	time.Sleep(100 * time.Millisecond)
-	assert.Equal(t, 3, callsForKey("foo_123"))
-	assert.Equal(t, 1, callsForKey("foo_124"))
+	drainChan()
+	verifyNotifs(map[string]int{"foo_123": 0}, "chan drained")
+	kv.notifyWatchers("foo_123")
+	verifyNotifs(map[string]int{"foo_123": 0}, "no flush -> no watcher notification")
+	kv.notifyWatchers("foo_123")
+	verifyNotifs(map[string]int{"foo_123": 0}, "no flush -> no watcher notification")
+	kv.notifyWatchers("foo_123")
+	verifyNotifs(map[string]int{"foo_123": 0}, "no flush -> no watcher notification")
+	kv.notifyWatchers("foo_123")
+	verifyNotifs(map[string]int{"foo_123": 0}, "no flush -> no watcher notification")
+	kv.notifyWatchers("foo_123")
+	verifyNotifs(map[string]int{"foo_123": 0}, "no flush -> no watcher notification")
+	kv.notifyWatchers("foo_123")
+	verifyNotifs(map[string]int{"foo_123": 0}, "no flush -> no watcher notification")
+	kv.sendKeyNotifications()
+	verifyNotifs(map[string]int{"foo_123": 1}, "flush should coalesce updates")
+
+	// multiple buffered updates
+	drainChan()
+	verifyNotifs(map[string]int{"foo_123": 0}, "chan drained")
+	kv.notifyWatchers("foo_123")
+	kv.sendKeyNotifications()
+	kv.notifyWatchers("foo_123")
+	kv.sendKeyNotifications()
+	verifyNotifs(map[string]int{"foo_123": 2}, "two buffered updates")
+
+	// multiple keys
+	drainChan()
+	kv.notifyWatchers("foo_123")
+	kv.notifyWatchers("foo_124")
+	kv.sendKeyNotifications()
+	verifyNotifs(map[string]int{"foo_123": 1, "foo_124": 1}, "2 changes 2 notifications")
+	kv.sendKeyNotifications()
+	verifyNotifs(map[string]int{"foo_123": 0, "foo_124": 0}, "no new notifications")
+
+	// and finally, sendKeyNotifications can be called repeatedly without new updates.
+	kv.sendKeyNotifications()
+	kv.sendKeyNotifications()
+	kv.sendKeyNotifications()
+	kv.sendKeyNotifications()
 }