jackhammer

channel.go (3521B)

---

 package channel
 
 import (
 	"context"
 	"reflect"
 	"slices"
 	"time"
 )
 
 // Drain the provided channel.
 func Drain[T any](ch <-chan T) {
 	for range ch {
 	}
 }
 
 // TryReceive selects on a channel returning true if there was a value.
 func TryReceive[T any](ch <-chan T) (t T, ok bool) {
 	select {
 	case t = <-ch:
 		ok = true
 	default:
 	}
 	return t, ok
 }
 
 // TrySend attempts a sends on a channel returning true if success.
 func TrySend[T any](ch chan<- T, t T) bool {
 	select {
 	case ch <- t:
 		return true
 	default:
 		return false
 	}
 }
 
 // ReceiveIn selects on the provided channel up to duration `d` and returns the value.
 // If the duration elapses, the function returns the zero value of `T` and false.
 func ReceiveIn[T any](ch <-chan T, d time.Duration) (T, bool) {
 	timer := time.NewTimer(d)
 	select {
 	case t := <-ch:
 		return t, true
 	case <-timer.C:
 		return *new(T), false
 	}
 }
 
 // Collect all values from a channel into a slice.
 func Collect[T any](in <-chan T) []T {
 	out := []T{}
 	for t := range in {
 		out = append(out, t)
 	}
 	return out
 }
 
 // CollectAllCtx buffers every value received value into a slice until all input channels are closed
 // or the context is cancelled.
 func CollectAllCtx[T any](ctx context.Context, chs ...chan T) []T {
 	out := []T{}
 	for v := range FanInCtx(ctx, chs...) {
 		out = append(out, v)
 	}
 	return out
 }
 
 // CollectAll buffers every value received value into a slice until all input channels are closed.
 func CollectAll[T any](chs ...chan T) (out []T) {
 	return CollectAllCtx(context.Background(), chs...)
 }
 
 // BufferedSend uses slice allocation to buffer items providing a dynamically
 // buffered channel send. This send will never block.
 //
 // If the channel is open, we send on it.
 // If the channel is full, we select off the slice, append to it and send it
 // back onto the channel.
 //
 // The loop will drain the channel until it becomes available to send. After
 // the function returns, the value v is guaranteed to be sent on the channel,
 // however it is not guaranteed that this is the only value sent.
 //
 // The input channel MUST have capacity of 1.
 func BufferedSend[S ~[]E, E any](in chan S, v E) {
 	if cap(in) != 1 {
 		panic("BufferedSend is only valid for single buffered channels")
 	}
 
 	var items S = []E{v}
 
 	for {
 		select {
 		case buf := <-in:
 			items = append(buf, items...)
 		case in <- items:
 			return
 		}
 	}
 }
 
 // FanIn returns a channel that emits every value received from the input channel set.
 // The output channel is closed when all input channels are closed.
 func FanIn[T any](chs ...chan T) chan T {
 	return FanInCtx(context.Background(), chs...)
 }
 
 // FanInCtx returns a channel that emits every value received from the input channel set.
 // The output channel is closed when all input channels are closed, or when the context
 // is cancelled.
 func FanInCtx[T any](ctx context.Context, chs ...chan T) chan T {
 	out := make(chan T)
 
 	go func() {
 		defer close(out)
 
 		cases := make([]reflect.SelectCase, 0, len(chs)+1)
 
 		cases = append(cases, reflect.SelectCase{
 			Dir:  reflect.SelectRecv,
 			Chan: reflect.ValueOf(ctx.Done()),
 		})
 
 		for _, ch := range chs {
 			cases = append(cases, reflect.SelectCase{
 				Dir:  reflect.SelectRecv,
 				Chan: reflect.ValueOf(ch),
 			})
 		}
 
 		for len(cases) > 1 {
 			ii, recv, ok := reflect.Select(cases)
 			if ok {
 				out <- recv.Interface().(T)
 			} else {
 				cases = slices.Delete(cases, ii, ii+1)
 			}
 			if ii == 0 {
 				// This is always the ctx.Done channel.
 				return
 			}
 		}
 	}()
 
 	return out
 }