-
-
Notifications
You must be signed in to change notification settings - Fork 108
/
mpsc.go
68 lines (57 loc) · 1.57 KB
/
mpsc.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
// Package mpsc implements a multiple-producer, single-consumer queue.
package mpsc
// N.B. this is a trivial wrapper around the spsc package, which just adds
// a lock to the Tx end to make it multiple-producer.
import (
"context"
"sync"
"capnproto.org/go/capnp/v3/exp/spsc"
)
// A multiple-producer, single-consumer queue. Create one with New(),
// and send from many gorotuines with Tx.Send(). Only one gorotuine may
// call Rx.Recv().
type Queue[T any] struct {
Tx[T]
Rx[T]
}
// The receive end of a Queue.
type Rx[T any] struct {
rx spsc.Rx[T]
}
// The send/transmit end of a Queue.
type Tx[T any] struct {
mu sync.Mutex
tx spsc.Tx[T]
}
// Create a new, initially empty Queue.
func New[T any]() *Queue[T] {
q := spsc.New[T]()
return &Queue[T]{
Tx: Tx[T]{tx: q.Tx},
Rx: Rx[T]{rx: q.Rx},
}
}
// Send a message on the queue.
func (tx *Tx[T]) Send(v T) {
tx.mu.Lock()
defer tx.mu.Unlock()
tx.tx.Send(v)
}
// Close the queue. Calls to Recv on the other end will return io.EOF.
func (tx *Tx[T]) Close() error {
tx.mu.Lock()
defer tx.mu.Unlock()
return tx.tx.Close()
}
// Receive a message from the queue. Blocks if the queue is empty.
// If the context ends before the receive happens, this returns
// ctx.Err(). If Close is called on the corresponding Tx, this
// returns io.EOF
func (rx *Rx[T]) Recv(ctx context.Context) (T, error) {
return rx.rx.Recv(ctx)
}
// Try to receive a message from the queue. If successful, ok will be true.
// If the queue is empty, this will return immediately with ok = false.
func (rx *Rx[T]) TryRecv() (v T, ok bool) {
return rx.rx.TryRecv()
}