/
forward.go
60 lines (52 loc) · 1.15 KB
/
forward.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
package cerebrum
import (
"net"
log "github.com/mgutz/logxi/v1"
"github.com/blacklabeldata/grim"
"github.com/blacklabeldata/namedtuple"
"github.com/blacklabeldata/yamuxer"
"golang.org/x/net/context"
)
const (
connForward yamuxer.StreamType = 0x01
connRaft = 0x02
)
type ForwardingHandler struct {
applier Applier
logger log.Logger
}
func (f *ForwardingHandler) Handle(c context.Context, conn net.Conn) {
f.logger.Info("Accepted raft connection", "addr", conn.RemoteAddr().String())
g := grim.ReaperWithContext(c)
defer g.Wait()
messages := make(chan namedtuple.Tuple, 1)
g.SpawnFunc(func(ctx context.Context) {
defer conn.Close()
for {
select {
case <-ctx.Done():
case msg, ok := <-messages:
if !ok {
return
}
if err := f.applier.Apply(msg); err != nil {
f.logger.Warn("error applying message")
}
}
}
})
g.SpawnFunc(func(ctx context.Context) {
defer close(messages)
decoder := namedtuple.NewDecoder(namedtuple.DefaultRegistry, conn)
for {
tuple, err := decoder.Decode()
if err != nil {
return
}
if tuple.Is(nodeStatus) {
messages <- tuple
}
}
})
return
}