blob: c5369ab83cc98d7f0c50e2824777bf77836e2599 [file] [log] [blame]
// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package chans provides utility functions for working with channels.
package chans
import (
"context"
"runtime"
)
// ReadAll reads from c until the channel is closed or the context is
// canceled, returning all the values read.
func ReadAll[Elem any](ctx context.Context, c <-chan Elem) []Elem {
var r []Elem
for {
select {
case <-ctx.Done():
return r
case v, ok := <-c:
if !ok {
return r
}
r = append(r, v)
}
}
}
// Merge merges two channels into a single channel.
// This will leave a goroutine running until either both channels are closed
// or the context is canceled, at which point the returned channel is closed.
func Merge[Elem any](ctx context.Context, c1, c2 <-chan Elem) <-chan Elem {
r := make(chan Elem)
go func(ctx context.Context, c1, c2 <-chan Elem, r chan<- Elem) {
defer close(r)
for c1 != nil || c2 != nil {
select {
case <-ctx.Done():
return
case v1, ok := <-c1:
if ok {
r <- v1
} else {
c1 = nil
}
case v2, ok := <-c2:
if ok {
r <- v2
} else {
c2 = nil
}
}
}
}(ctx, c1, c2, r)
return r
}
// Filter calls f on each value read from c. If f returns true the value
// is sent on the returned channel. This will leave a goroutine running
// until c is closed or the context is canceled, at which point the
// returned channel is closed.
func Filter[Elem any](ctx context.Context, c <-chan Elem, f func(Elem) bool) <-chan Elem {
r := make(chan Elem)
go func(ctx context.Context, c <-chan Elem, f func(Elem) bool, r chan<- Elem) {
defer close(r)
for {
select {
case <-ctx.Done():
return
case v, ok := <-c:
if !ok {
return
}
if f(v) {
r <- v
}
}
}
}(ctx, c, f, r)
return r
}
// Sink returns a channel that discards all values sent to it.
// This will leave a goroutine running until the context is canceled
// or the returned channel is closed.
func Sink[Elem any](ctx context.Context) chan<- Elem {
r := make(chan Elem)
go func(ctx context.Context, r <-chan Elem) {
for {
select {
case <-ctx.Done():
return
case _, ok := <-r:
if !ok {
return
}
}
}
}(ctx, r)
return r
}
// An Exclusive is a value that may only be used by a single goroutine
// at a time. This is implemented using channels rather than a mutex.
type Exclusive[Val any] struct {
c chan Val
}
// MakeExclusive makes an initialized exclusive value.
func MakeExclusive[Val any](initial Val) *Exclusive[Val] {
r := &Exclusive[Val]{
c: make(chan Val, 1),
}
r.c <- initial
return r
}
// Acquire acquires the exclusive value for private use.
// It must be released using the Release method.
func (e *Exclusive[Val]) Acquire() Val {
return <-e.c
}
// TryAcquire attempts to acquire the value. The ok result reports whether
// the value was acquired. If the value is acquired, it must be released
// using the Release method.
func (e *Exclusive[Val]) TryAcquire() (v Val, ok bool) {
select {
case r := <-e.c:
return r, true
default:
return v, false
}
}
// Release updates and releases the value.
// This method panics if the value has not been acquired.
func (e *Exclusive[Val]) Release(v Val) {
select {
case e.c <- v:
default:
panic("Exclusive Release without Acquire")
}
}
// Ranger returns a Sender and a Receiver. The Receiver provides a
// Next method to retrieve values. The Sender provides a Send method
// to send values and a Close method to stop sending values. The Next
// method indicates when the Sender has been closed, and the Send
// method indicates when the Receiver has been freed.
//
// This is a convenient way to exit a goroutine sending values when
// the receiver stops reading them.
func Ranger[Elem any]() (*Sender[Elem], *Receiver[Elem]) {
c := make(chan Elem)
d := make(chan struct{})
s := &Sender[Elem]{
values: c,
done: d,
}
r := &Receiver[Elem] {
values: c,
done: d,
}
runtime.SetFinalizer(r, (*Receiver[Elem]).finalize)
return s, r
}
// A Sender is used to send values to a Receiver.
type Sender[Elem any] struct {
values chan<- Elem
done <-chan struct{}
}
// Send sends a value to the receiver. It reports whether the value was sent.
// The value will not be sent if the context is closed or the receiver
// is freed.
func (s *Sender[Elem]) Send(ctx context.Context, v Elem) bool {
select {
case <-ctx.Done():
return false
case s.values <- v:
return true
case <-s.done:
return false
}
}
// Close tells the receiver that no more values will arrive.
// After Close is called, the Sender may no longer be used.
func (s *Sender[Elem]) Close() {
close(s.values)
}
// A Receiver receives values from a Sender.
type Receiver[Elem any] struct {
values <-chan Elem
done chan<- struct{}
}
// Next returns the next value from the channel. The bool result indicates
// whether the value is valid.
func (r *Receiver[Elem]) Next(ctx context.Context) (v Elem, ok bool) {
select {
case <-ctx.Done():
case v, ok = <-r.values:
}
return v, ok
}
// finalize is a finalizer for the receiver.
func (r *Receiver[Elem]) finalize() {
close(r.done)
}