sync: Add Select, Combine

sync/gate.go

@@ -0,0 +1,94 @@
+package usync
+
+import "sync"
+
+// Error defines errors that this package can produce
+type Error string; const (
+	ErrAlreadyClosed Error = "AlreadyClosed"
+)
+
+// Error fullfills the error interface.
+func (err Error) Error () string {
+	return string(err)
+}
+
+// Gate wraps a channel and allows the receiver to abruptly stop receiving
+// messages without causing the sender to lock up.
+type Gate[T any] struct {
+	channel    chan T
+	lock       sync.RWMutex
+	open       bool
+	bufferSize int
+}
+
+// NewGate creates a new gate with no buffer.
+func NewGate[T any] () Gate[T] {
+	return Gate[T] {
+		channel: make(chan T),
+		open:    true,
+	}
+}
+
+// NewBufferedGate creates a new gate with a buffer.
+func NewBufferedGate[T any] (buffer int) Gate[T] {
+	return Gate[T] {
+		channel:    make(chan T, buffer),
+		open:       true,
+		bufferSize: buffer,
+	}
+}
+
+// Send sends and item to the channel, returning whether the item was sent.
+func (this *Gate[T]) Send (item T) bool {
+	this.lock.RLock()
+	defer this.lock.RUnlock()
+	if !this.open { return false }
+	this.channel <- item
+	return true
+}
+
+// Receive returns a receive-only channel that can be used to receive items.
+func (this *Gate[T]) Receive () <- chan T {
+	this.lock.RLock()
+	defer this.lock.RUnlock()
+	return this.channel
+}
+
+// Close closes the gate, drains all remaining messages, and closes the channel.
+func (this *Gate[T]) Close () error {
+	this.lock.Lock()
+	if !this.open { return ErrAlreadyClosed }
+	this.open = false
+	this.lock.Unlock()
+	for len(this.channel) > 0 { <- this.channel }
+	close(this.channel)
+	return nil
+}
+
+// Reset re-opens the gate if it is closed, and creates a new channel.
+func (this *Gate[T]) Reset () error {
+	this.lock.Lock()
+	defer this.lock.Unlock()
+	this.open = true
+	if this.channel != nil { close(this.channel) }
+	this.channel = make(chan T, this.bufferSize)
+	return nil
+}
+
+// Open returns whether the gate is open.
+func (this *Gate[T]) Open () bool {
+	this.lock.RLock()
+	defer this.lock.RUnlock()
+	return this.open
+}
+
+// Len returns the amount of items in the channel.
+func (this *Gate[T]) Len () int {
+	return len(this.channel)
+}
+
+// Cap returns the amount of items the channel can hold.
+func (this *Gate[T]) Cap () int {
+	return cap(this.channel)
+}
+

sync/select.go

@@ -0,0 +1,78 @@
+package usync
+
+import "slices"
+import "context"
+import "reflect"
+
+// A type-safe wrapper around reflect.Select. Taken from:
+// https://stackoverflow.com/questions/19992334
+func Select[T any] (ctx context.Context, channels ...chan T) (int, T, bool) {
+	var zero T
+
+	// add all channels as select cases
+	cases := make([]reflect.SelectCase, len(channels) + 1)
+	for i, ch := range channels {
+		cases[i] = reflect.SelectCase {
+			Dir:  reflect.SelectRecv,
+			Chan: reflect.ValueOf(ch),
+		}
+	}
+
+	// add ctx.Done() as another select case to stop listening when the
+	// context is closed
+	cases[len(channels)] = reflect.SelectCase {
+		Dir:  reflect.SelectRecv,
+		Chan: reflect.ValueOf(ctx.Done()),
+	}
+	
+	// read from the channel
+	chosen, value, ok := reflect.Select(cases)
+	if !ok {
+		if ctx.Err() != nil {
+			return -1, zero, false
+		}
+		return chosen, zero, false
+	}
+
+	// cast return value
+	if ret, ok := value.Interface().(T); ok {
+		return chosen, ret, true
+	}
+	return chosen, zero, false
+}
+
+// Combine returns a channel that continuously returns the result of the select
+// until all input sources are exhauste, or the context is canceled.
+func Combine[T any] (ctx context.Context, channels ...chan T) <- chan T {
+	channel := make(chan T)
+	// our silly slection routine
+	go func () {
+		for {
+			if len(channels) < 2 {
+				// only the context is left, stop everything
+				close(channel)
+				return
+			}
+			// read new value
+			chosen, value, ok := Select(ctx, channels...)
+			if ok {
+				// we have a value
+				channel <- value
+			} else {
+				// a channel has been closed and we need to do
+				// something about it
+				if chosen == len(channels) - 1 {
+					// the context has expired, stop
+					// everything
+					close(channel)
+					return
+				} else {
+					// a normal channel has closed, remove
+					// it from the list
+					channels = slices.Delete(channels, chosen, chosen + 1)
+				}
+			}
+		}
+	} ()
+	return channel
+}

sync/select_test.go

@@ -0,0 +1,24 @@
+package usync
+
+import "time"
+import "testing"
+import "context"
+
+func TestSelect (test *testing.T) {
+	// https://stackoverflow.com/questions/19992334
+	c1 := make(chan int)
+	c2 := make(chan int)
+	c3 := make(chan int)
+	chs := []chan int { c1, c2, c3 }
+	go func () {
+		time.Sleep(time.Second)
+		c2 <- 42
+	} ()
+	ctx, done := context.WithTimeout(context.Background(), 5 * time.Second)
+	defer done()
+
+	chosen, val, ok := Select(ctx, chs...)
+	if !ok          { test.Fatal("not ok")                  }
+	if 1  != chosen { test.Fatal("expected 1, got", chosen) }
+	if 42 != val    { test.Fatal("expected 42, got", val)   }
+}