forked from nats-io/graft
/
nats_rpc.go
162 lines (140 loc) · 4.3 KB
/
nats_rpc.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
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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
// Copyright 2013-2015 Apcera Inc. All rights reserved.
package graft
import (
"errors"
"fmt"
"github.com/nats-io/nats"
)
// The subject space for the nats rpc driver is based on the
// cluster name, which is filled in below on the heartbeats
// and vote requests. The vote responses are directed by
// using the node.Id().
const (
HEARTBEAT_SUB = "graft.%s.heartbeat"
VOTE_REQ_SUB = "graft.%s.vote_request"
VOTE_RESP_SUB = "graft.%s.vote_response"
)
var (
NotInitializedErr = errors.New("graft(nats_rpc): Driver is not properly initialized")
)
type NatsRpcDriver struct {
// NATS encoded connection.
ec *nats.EncodedConn
// Heartbeat subscription.
hbSub *nats.Subscription
// Vote request subscription.
vreqSub *nats.Subscription
// Vote response subscription.
vrespSub *nats.Subscription
// Graft node.
node *Node
}
// Create a new instance of the driver. The NATS encoded connection
// will use json and the options passed in.
func NewNatsRpc(opts *nats.Options) (*NatsRpcDriver, error) {
nc, err := opts.Connect()
ec, _ := nats.NewEncodedConn(nc, "json")
if err != nil {
return nil, err
}
return &NatsRpcDriver{ec: ec}, nil
}
// Initialize via the Graft node.
func (rpc *NatsRpcDriver) Init(n *Node) (err error) {
rpc.node = n
// Create the heartbeat subscription.
hbSub := fmt.Sprintf(HEARTBEAT_SUB, n.ClusterInfo().Name)
rpc.hbSub, err = rpc.ec.Subscribe(hbSub, rpc.HeartbeatCallback)
if err != nil {
return err
}
// Create the voteRequest subscription.
rpc.vreqSub, err = rpc.ec.Subscribe(rpc.vreqSubject(), rpc.VoteRequestCallback)
if err != nil {
return err
}
return nil
}
// Close down the subscriptions and the NATS encoded connection.
// Will nil everything out.
func (rpc *NatsRpcDriver) Close() {
if rpc.hbSub != nil {
rpc.hbSub.Unsubscribe()
rpc.hbSub = nil
}
if rpc.vreqSub != nil {
rpc.vreqSub.Unsubscribe()
rpc.vreqSub = nil
}
if rpc.vrespSub != nil {
rpc.vrespSub.Unsubscribe()
rpc.vrespSub = nil
}
if rpc.ec != nil {
rpc.ec.Close()
}
}
// Convenience function for generating the directed response
// subject for vote requests. We will use the candidate's id
// to form a directed response
func (rpc *NatsRpcDriver) vrespSubject(candidate string) string {
return fmt.Sprintf(VOTE_RESP_SUB, candidate)
}
// Convenience funstion for generating the vote request subject.
func (rpc *NatsRpcDriver) vreqSubject() string {
return fmt.Sprintf(VOTE_REQ_SUB, rpc.node.ClusterInfo().Name)
}
// Heartbeat callback which will place the heartbeat on the Graft
// node's appropriate channel.
func (rpc *NatsRpcDriver) HeartbeatCallback(hb *Heartbeat) {
rpc.node.HeartBeats <- hb
}
// VoteRequest callback which will place the request on the Graft
// node's appropriate channel.
func (rpc *NatsRpcDriver) VoteRequestCallback(vreq *VoteRequest) {
// Don't respond to our own request.
if vreq.Candidate != rpc.node.Id() {
rpc.node.VoteRequests <- vreq
}
}
// VoteResponse callback which will place the response on the Graft
// node's appropriate channel.
func (rpc *NatsRpcDriver) VoteResponseCallback(vresp *VoteResponse) {
rpc.node.VoteResponses <- vresp
}
// RequestVote is sent from the Graft node when it has become a
// candidate.
func (rpc *NatsRpcDriver) RequestVote(vr *VoteRequest) error {
// Create a new response subscription for each outstanding
// RequestVote and cancel the previous.
if rpc.vrespSub != nil {
rpc.vrespSub.Unsubscribe()
rpc.vrespSub = nil
}
inbox := rpc.vrespSubject(rpc.node.Id())
sub, err := rpc.ec.Subscribe(inbox, rpc.VoteResponseCallback)
if err != nil {
return err
}
// If we can auto-unsubscribe to max number of expected responses
// which will be the cluster size.
if size := rpc.node.ClusterInfo().Size; size > 0 {
sub.AutoUnsubscribe(size)
}
// hold to cancel later.
rpc.vrespSub = sub
// Fire off the request.
return rpc.ec.PublishRequest(rpc.vreqSubject(), inbox, vr)
}
// Heartbeat is called from the Graft node to send out a heartbeat
// while it is a LEADER.
func (rpc *NatsRpcDriver) HeartBeat(hb *Heartbeat) error {
if rpc.hbSub == nil {
return NotInitializedErr
}
return rpc.ec.Publish(rpc.hbSub.Subject, hb)
}
// SendVoteResponse is called from the Graft node to respond to a vote request.
func (rpc *NatsRpcDriver) SendVoteResponse(id string, vresp *VoteResponse) error {
return rpc.ec.Publish(rpc.vrespSubject(id), vresp)
}