blob: 407d770f918a0ffbc131efedc24d54fb791b5b3b (
plain)
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
|
package backchannel
import (
"fmt"
"sync"
"google.golang.org/grpc"
)
// ID is a monotonically increasing number that uniquely identifies a peer connection.
type ID uint64
// Registry is a thread safe registry for backchannels. It enables accessing the backchannels via a
// unique ID.
type Registry struct {
m sync.RWMutex
currentID ID
backchannels map[ID]*grpc.ClientConn
}
// NewRegistry returns a new Registry.
func NewRegistry() *Registry { return &Registry{backchannels: map[ID]*grpc.ClientConn{}} }
// Backchannel returns a backchannel for the ID. Returns an error if no backchannel is registered
// for the ID.
func (r *Registry) Backchannel(id ID) (*grpc.ClientConn, error) {
r.m.RLock()
defer r.m.RUnlock()
backchannel, ok := r.backchannels[id]
if !ok {
return nil, fmt.Errorf("no backchannel for peer %d", id)
}
return backchannel, nil
}
// RegisterBackchannel registers a new backchannel and returns its unique ID.
func (r *Registry) RegisterBackchannel(conn *grpc.ClientConn) ID {
r.m.Lock()
defer r.m.Unlock()
r.currentID++
r.backchannels[r.currentID] = conn
return r.currentID
}
// RemoveBackchannel removes a backchannel from the registry.
func (r *Registry) RemoveBackchannel(id ID) {
r.m.Lock()
defer r.m.Unlock()
delete(r.backchannels, id)
}
|