feat: add TryRecv and RecvWithContext functions

Both of this should simplify places where we are using `select` currently.
Also fix flaky tests while we are at it.

Signed-off-by: Dmitriy Matrenichev <[email protected]>

Author: Dmitriy Matrenichev

channel/recv.go

@@ -0,0 +1,55 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package channel
+
+import "context"
+
+// RecvWithContext tries to receive a value from a channel which is aborted if the context is canceled.
+//
+// Function returns true if the value was received, false if the context was canceled or the channel was closed.
+func RecvWithContext[T any](ctx context.Context, ch <-chan T) (T, bool) {
+	select {
+	case <-ctx.Done():
+		var zero T
+
+		return zero, false
+	case val, ok := <-ch:
+		if !ok {
+			return val, false
+		}
+
+		return val, true
+	}
+}
+
+// RecvState is the state of a channel after receiving a value.
+type RecvState int
+
+const (
+	// StateRecv means that a value was received from the channel.
+	StateRecv RecvState = iota
+	// StateEmpty means that the channel was empty.
+	StateEmpty
+	// StateClosed means that the channel was closed.
+	StateClosed
+)
+
+// TryRecv tries to receive a value from a channel.
+//
+// Function returns the value and the state of the channel.
+func TryRecv[T any](ch <-chan T) (T, RecvState) {
+	var zero T
+
+	select {
+	case val, ok := <-ch:
+		if !ok {
+			return zero, StateClosed
+		}
+
+		return val, StateRecv
+	default:
+		return zero, StateEmpty
+	}
+}

channel/recv_test.go

@@ -0,0 +1,72 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package channel_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/siderolabs/gen/channel"
+)
+
+func TestRecvWithContext(t *testing.T) {
+	t.Parallel()
+
+	ch := make(chan int, 1)
+	ch <- 42
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	val, ok := channel.RecvWithContext(ctx, ch)
+	assert.Equal(t, 42, val)
+	assert.True(t, ok)
+
+	cancel()
+
+	val, ok = channel.RecvWithContext(ctx, ch)
+	assert.Zero(t, val)
+	assert.False(t, ok)
+}
+
+func TestRecvWithContextCloseCh(t *testing.T) {
+	t.Parallel()
+
+	ch := make(chan int, 1)
+	ch <- 42
+
+	ctx := context.Background()
+
+	val, ok := channel.RecvWithContext(ctx, ch)
+	assert.Equal(t, 42, val)
+	assert.True(t, ok)
+
+	close(ch)
+	val, ok = channel.RecvWithContext(ctx, ch)
+	assert.Zero(t, val)
+	assert.False(t, ok)
+}
+
+func TestTryRecv(t *testing.T) {
+	t.Parallel()
+
+	ch := make(chan int, 1)
+	ch <- 42
+
+	val, state := channel.TryRecv(ch)
+	assert.Equal(t, 42, val)
+	assert.Equal(t, channel.StateRecv, state)
+
+	val, state = channel.TryRecv(ch)
+	assert.Zero(t, val)
+	assert.Equal(t, channel.StateEmpty, state)
+
+	close(ch)
+	val, state = channel.TryRecv(ch)
+	assert.Zero(t, val)
+	assert.Equal(t, channel.StateClosed, state)
+}

channel/send_test.go

@@ -14,6 +14,8 @@ import (
 )
 
 func TestSendWithContext(t *testing.T) {
+	t.Parallel()
+
 	ch := make(chan int, 1)
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)

containers/lazymap_test.go

@@ -7,6 +7,7 @@ package containers_test
 
 import (
 	"fmt"
+	"sort"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -132,6 +133,9 @@ func TestLazyBiMap(t *testing.T) {
 			values = append(values, v)
 		})
 
+		sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
+		sort.Slice(values, func(i, j int) bool { return values[i] < values[j] })
+
 		assert.Equal(t, []int{1, 2}, keys)
 		assert.Equal(t, []int{100, 200}, values)
 	})
@@ -227,6 +231,9 @@ func TestLazyMap(t *testing.T) {
 			values = append(values, v)
 		})
 
+		sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
+		sort.Slice(values, func(i, j int) bool { return values[i] < values[j] })
+
 		assert.Equal(t, []int{4, 5}, keys)
 		assert.Equal(t, []int{400, 500}, values)
 	})