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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
|
package sidechannel
import (
"fmt"
"net"
"sync"
)
// sidechannelID is the type of ID used to differeniate sidechannel connections
// in the same registry
type sidechannelID int64
// Registry manages sidechannel connections. It allows the RPC
// handlers to wait for the secondary incoming connection made by the client.
type Registry struct {
nextID sidechannelID
waiters map[sidechannelID]*Waiter
mu sync.Mutex
}
// Waiter lets the caller waits until a connection with matched id is pushed
// into the registry, then execute the callback
type Waiter struct {
id sidechannelID
registry *Registry
errC chan error
accept chan net.Conn
callback func(net.Conn) error
}
// NewRegistry returns a new Registry instance
func NewRegistry() *Registry {
return &Registry{
waiters: make(map[sidechannelID]*Waiter),
}
}
// Register registers the caller into the waiting list. The caller must provide
// a callback function. The caller receives a waiter instance. After the
// connection arrives, the callback function is executed with arrived
// connection in a new goroutine. The caller receives execution result via
// waiter.Wait().
func (s *Registry) Register(callback func(net.Conn) error) *Waiter {
s.mu.Lock()
defer s.mu.Unlock()
waiter := &Waiter{
id: s.nextID,
registry: s,
errC: make(chan error),
accept: make(chan net.Conn),
callback: callback,
}
s.nextID++
go waiter.run()
s.waiters[waiter.id] = waiter
return waiter
}
// receive looks into the registry for a waiter with the given ID. If
// there is an associated ID, the waiter is removed from the registry, and the
// connection is pushed into the waiter's accept channel. After the callback is done, the
// connection is closed. When the ID is not found, an error is returned and the
// connection is closed immediately.
func (s *Registry) receive(id sidechannelID, conn net.Conn) (err error) {
s.mu.Lock()
defer func() {
s.mu.Unlock()
if err != nil {
conn.Close()
}
}()
waiter, exist := s.waiters[id]
if !exist {
return fmt.Errorf("sidechannel registry: ID not registered")
}
delete(s.waiters, waiter.id)
waiter.accept <- conn
return nil
}
func (s *Registry) removeWaiter(waiter *Waiter) {
s.mu.Lock()
defer s.mu.Unlock()
if _, exist := s.waiters[waiter.id]; exist {
delete(s.waiters, waiter.id)
close(waiter.accept)
}
}
func (s *Registry) waiting() int {
s.mu.Lock()
defer s.mu.Unlock()
return len(s.waiters)
}
func (w *Waiter) run() {
defer close(w.errC)
if conn := <-w.accept; conn != nil {
defer conn.Close()
w.errC <- w.callback(conn)
}
}
// Close cleans the waiter, removes it from the registry. If the callback is
// executing, this method is blocked until the callback is done.
func (w *Waiter) Close() error {
w.registry.removeWaiter(w)
return <-w.errC
}
// Wait waits until either the callback is executed, or the waiter is closed
func (w *Waiter) Wait() error {
return <-w.errC
}
|
