watchset: Settle time instead of timeout for Wait

Wait() should block until either context cancelled or channel(s) close.
This way we can block until there is actually interesting work rather
than just waiting for a timeout and returning empty set of closed channels.

Signed-off-by: Jussi Maki <[email protected]>

Author: joamaki

watchset.go

@@ -5,7 +5,6 @@ package statedb
 
 import (
 	"context"
-	"fmt"
 	"maps"
 	"slices"
 	"sync"
@@ -75,13 +74,12 @@ func (ws *WatchSet) Merge(other *WatchSet) {
 	}
 }
 
-// Wait for channels in the watch set to close until context is cancelled or timeout reached.
+// Wait for channels in the watch set to close or the context is cancelled.
+// After the first closed channel is seen Wait will wait [settleTime] for
+// more closed channels.
 // Returns the closed channels and removes them from the set.
-func (ws *WatchSet) Wait(ctx context.Context, timeout time.Duration) ([]<-chan struct{}, error) {
-	if timeout <= 0 {
-		return nil, fmt.Errorf("bad timeout %d, must be >0", timeout)
-	}
-	innerCtx, cancel := context.WithTimeout(ctx, timeout)
+func (ws *WatchSet) Wait(ctx context.Context, settleTime time.Duration) ([]<-chan struct{}, error) {
+	innerCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	ws.mu.Lock()
@@ -91,7 +89,7 @@ func (ws *WatchSet) Wait(ctx context.Context, timeout time.Duration) ([]<-chan s
 
 	// No channels to watch? Just watch the context.
 	if len(ws.chans) == 0 {
-		<-innerCtx.Done()
+		<-ctx.Done()
 		return nil, ctx.Err()
 	}
 
@@ -101,30 +99,44 @@ func (ws *WatchSet) Wait(ctx context.Context, timeout time.Duration) ([]<-chan s
 	chunkSize := 16
 	roundedSize := len(chans) + (chunkSize - len(chans)%chunkSize)
 	chans = slices.Grow(chans, roundedSize)[:roundedSize]
-
-	if len(ws.chans) <= chunkSize {
-		watch16(closedChannels, innerCtx.Done(), chans)
-		return closedChannels.chans, ctx.Err()
-	}
+	haveResult := make(chan struct{}, 1)
 
 	var wg sync.WaitGroup
-	for chunk := range slices.Chunk(chans, chunkSize) {
+	chunks := slices.Chunk(chans, chunkSize)
+	for chunk := range chunks {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			watch16(closedChannels, innerCtx.Done(), chunk)
+			watch16(haveResult, closedChannels, innerCtx.Done(), chunk)
 		}()
 	}
+
+	// Wait for the first closed channel to be seen. If [settleTime] is set,
+	// then wait a bit longer for more.
+	select {
+	case <-haveResult:
+		if settleTime > 0 {
+			select {
+			case <-time.After(settleTime):
+			case <-ctx.Done():
+			}
+		}
+	case <-ctx.Done():
+	}
+
+	// Stop waiting for more channels to close
+	cancel()
 	wg.Wait()
 
+	// Remove the closed channels from the watch set.
 	for _, ch := range closedChannels.chans {
 		delete(ws.chans, ch)
 	}
 
 	return closedChannels.chans, ctx.Err()
 }
 
-func watch16(closedChannels *closedChannelsSlice, stop <-chan struct{}, chans []<-chan struct{}) {
+func watch16(haveClosed chan struct{}, closedChannels *closedChannelsSlice, stop <-chan struct{}, chans []<-chan struct{}) {
 	for {
 		closedIndex := -1
 		select {
@@ -165,6 +177,13 @@ func watch16(closedChannels *closedChannelsSlice, stop <-chan struct{}, chans []
 		}
 		closedChannels.append(chans[closedIndex])
 		chans[closedIndex] = nil
+		if haveClosed != nil {
+			select {
+			case haveClosed <- struct{}{}:
+				haveClosed = nil
+			default:
+			}
+		}
 	}
 }
 

watchset_test.go

@@ -20,7 +20,7 @@ func TestWatchSet(t *testing.T) {
 
 	// Empty watch set, cancelled context.
 	ctx, cancel := context.WithCancel(context.Background())
-	go cancel()
+	cancel()
 	ch, err := ws.Wait(ctx, time.Second)
 	require.ErrorIs(t, err, context.Canceled)
 	require.Nil(t, ch)
@@ -31,7 +31,7 @@ func TestWatchSet(t *testing.T) {
 	ch3 := make(chan struct{})
 	ws.Add(ch1, ch2, ch3)
 	ctx, cancel = context.WithCancel(context.Background())
-	go cancel()
+	cancel()
 	ch, err = ws.Wait(ctx, time.Second)
 	require.ErrorIs(t, err, context.Canceled)
 	require.Nil(t, ch)
@@ -96,10 +96,6 @@ func TestWatchSetInQueries(t *testing.T) {
 	_, _, watch2, _ := table.GetWatch(txn, idIndex.Query(2))
 	_, _, watch3, _ := table.GetWatch(txn, idIndex.Query(3))
 
-	closed, err = ws.Wait(context.Background(), time.Millisecond)
-	require.NoError(t, err)
-	require.Empty(t, closed)
-
 	wtxn = db.WriteTxn(table)
 	table.Insert(wtxn, testObject{ID: 1, Tags: part.NewSet("foo")})
 	wtxn.Commit()