1package backchannel
2
3import (
4	"fmt"
5	"sync"
6
7	"google.golang.org/grpc"
8)
9
10// ID is a monotonically increasing number that uniquely identifies a peer connection.
11type ID uint64
12
13// Registry is a thread safe registry for backchannels. It enables accessing the backchannels via a
14// unique ID.
15type Registry struct {
16	m            sync.RWMutex
17	currentID    ID
18	backchannels map[ID]*grpc.ClientConn
19}
20
21// NewRegistry returns a new Registry.
22func NewRegistry() *Registry { return &Registry{backchannels: map[ID]*grpc.ClientConn{}} }
23
24// Backchannel returns a backchannel for the ID. Returns an error if no backchannel is registered
25// for the ID.
26func (r *Registry) Backchannel(id ID) (*grpc.ClientConn, error) {
27	r.m.RLock()
28	defer r.m.RUnlock()
29	backchannel, ok := r.backchannels[id]
30	if !ok {
31		return nil, fmt.Errorf("no backchannel for peer %d", id)
32	}
33
34	return backchannel, nil
35}
36
37// RegisterBackchannel registers a new backchannel and returns its unique ID.
38func (r *Registry) RegisterBackchannel(conn *grpc.ClientConn) ID {
39	r.m.Lock()
40	defer r.m.Unlock()
41	r.currentID++
42	r.backchannels[r.currentID] = conn
43	return r.currentID
44}
45
46// RemoveBackchannel removes a backchannel from the registry.
47func (r *Registry) RemoveBackchannel(id ID) {
48	r.m.Lock()
49	defer r.m.Unlock()
50	delete(r.backchannels, id)
51}
52