watchset: Fix WatchSet.Wait() to not return empty closed channels

[ upstream commit da3849c ]

As the comment on Wait() implied the intended semantics for Wait() was to
wait for [settleTime] after the first channel has closed, but instead it
returned with empty closed channels when [settleTime] expired.

Fix this by first watching for 'ctx.Done()' and when a first closed channel
is encountered switch to a context with a [settleTime] timeout to gather
more channels.

Fixes: f0c8822 ("Use reflect.Select for WatchSet")
Signed-off-by: Jussi Maki <jussi.maki@isovalent.com>

Author: joamaki

watchset.go

@@ -77,11 +77,9 @@ func (ws *WatchSet) Merge(other *WatchSet) {
 // 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.
+// If [settleTime] is 0 waits until [ctx] cancelled or any channel closes.
 // Returns the closed channels and removes them from the set.
 func (ws *WatchSet) Wait(ctx context.Context, settleTime time.Duration) ([]<-chan struct{}, error) {
-	innerCtx, cancel := context.WithTimeout(ctx, settleTime)
-	defer cancel()
-
 	ws.mu.Lock()
 	defer ws.mu.Unlock()
 
@@ -94,10 +92,14 @@ func (ws *WatchSet) Wait(ctx context.Context, settleTime time.Duration) ([]<-cha
 	// Construct []SelectCase slice. Reuse the previous allocation.
 	ws.cases = slices.Grow(ws.cases, 1+len(ws.chans))
 	cases := ws.cases[:1+len(ws.chans)]
+
+	// Add [ctx.Done()] to stop when [ctx] is cancelled.
 	cases[0] = reflect.SelectCase{
 		Dir:  reflect.SelectRecv,
-		Chan: reflect.ValueOf(innerCtx.Done()),
+		Chan: reflect.ValueOf(ctx.Done()),
 	}
+
+	// Add the cases from the watch set.
 	casesIndex := 1
 	for ch := range ws.chans {
 		cases[casesIndex] = reflect.SelectCase{
@@ -109,21 +111,42 @@ func (ws *WatchSet) Wait(ctx context.Context, settleTime time.Duration) ([]<-cha
 
 	var closedChannels []<-chan struct{}
 
-	// Collect closed channels until [innerCtx] is cancelled.
-	for {
+	// At the end remove the closed channels from the watch set.
+	defer func() {
+		for _, ch := range closedChannels {
+			delete(ws.chans, ch)
+		}
+	}()
+
+	// Wait for the first channel to close and shift it out from [cases]
+	chosen, _, _ := reflect.Select(cases)
+	if chosen == 0 {
+		return nil, ctx.Err()
+	}
+	closedChannels = append(closedChannels, cases[chosen].Chan.Interface().(<-chan struct{}))
+	cases[chosen] = cases[len(cases)-1]
+	cases = cases[:len(cases)-1]
+
+	// If nothing else than context channel remains or we don't want to wait for further channels
+	// to close then we're done.
+	if len(cases) == 1 || settleTime == 0 {
+		return closedChannels, nil
+	}
+
+	// Swap out the 'ctx.Done()' to a context that times out when [settleTime] expires.
+	settleCtx, cancel := context.WithTimeout(ctx, settleTime)
+	defer cancel()
+	cases[0].Chan = reflect.ValueOf(settleCtx.Done())
+
+	for len(cases) > 1 {
 		chosen, _, _ := reflect.Select(cases)
-		if chosen == 0 /* == innerCtx.Done() */ {
+		if chosen == 0 /* settleCtx.Done() */ {
 			break
 		}
 		closedChannels = append(closedChannels, cases[chosen].Chan.Interface().(<-chan struct{}))
 		cases[chosen] = cases[len(cases)-1]
 		cases = cases[:len(cases)-1]
 	}
 
-	// Remove the closed channels from the watch set.
-	for _, ch := range closedChannels {
-		delete(ws.chans, ch)
-	}
-
 	return closedChannels, ctx.Err()
 }

watchset_test.go

@@ -22,9 +22,9 @@ func TestWatchSet(t *testing.T) {
 	// Empty watch set, cancelled context.
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
-	ch, err := ws.Wait(ctx, time.Second)
+	chs, err := ws.Wait(ctx, time.Second)
 	require.ErrorIs(t, err, context.Canceled)
-	require.Nil(t, ch)
+	require.Empty(t, chs)
 
 	// Few channels, cancelled context.
 	ch1 := make(chan struct{})
@@ -33,9 +33,35 @@ func TestWatchSet(t *testing.T) {
 	ws.Add(ch1, ch2, ch3)
 	ctx, cancel = context.WithCancel(context.Background())
 	cancel()
-	ch, err = ws.Wait(ctx, time.Second)
+	chs, err = ws.Wait(ctx, time.Second)
 	require.ErrorIs(t, err, context.Canceled)
-	require.Nil(t, ch)
+	require.Empty(t, chs)
+
+	// Few channels, timed out context. With tiny 'settleTime' we wait for the context to cancel.
+	duration := 10 * time.Millisecond
+	ctx, cancel = context.WithTimeout(context.Background(), duration)
+	t0 := time.Now()
+	chs, err = ws.Wait(ctx, time.Nanosecond)
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Empty(t, chs)
+	require.True(t, time.Since(t0) > duration, "expected to wait until context cancels")
+	cancel()
+
+	// One closed channel. Should wait until 'settleTime' expires.
+	close(ch1)
+	t0 = time.Now()
+	chs, err = ws.Wait(context.Background(), duration)
+	require.NoError(t, err)
+	require.ElementsMatch(t, chs, []<-chan struct{}{ch1})
+	require.True(t, time.Since(t0) > duration, "expected to wait until settle time expires")
+
+	// One closed channel, 0 wait time.
+	ws = NewWatchSet()
+	ws.Add(ch2)
+	close(ch2)
+	chs, err = ws.Wait(context.Background(), 0)
+	require.NoError(t, err)
+	require.ElementsMatch(t, chs, []<-chan struct{}{ch2})
 
 	// Many channels
 	for _, numChans := range []int{2, 16, 31, 1024} {