forked from hashicorp/nomad
-
Notifications
You must be signed in to change notification settings - Fork 0
/
consul.go
341 lines (296 loc) · 9.96 KB
/
consul.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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
package client
import (
"crypto/tls"
"fmt"
"log"
"net/http"
"net/url"
"strings"
"sync"
"time"
consul "github.com/hashicorp/consul/api"
"github.com/hashicorp/go-multierror"
"github.com/hashicorp/nomad/nomad/structs"
)
const (
syncInterval = 5 * time.Second
)
// consulApi is the interface which wraps the actual consul api client
type consulApi interface {
CheckRegister(check *consul.AgentCheckRegistration) error
CheckDeregister(checkID string) error
ServiceRegister(service *consul.AgentServiceRegistration) error
ServiceDeregister(ServiceID string) error
Services() (map[string]*consul.AgentService, error)
Checks() (map[string]*consul.AgentCheck, error)
}
// consulApiClient is the actual implementation of the consulApi which
// talks to the consul agent
type consulApiClient struct {
client *consul.Client
}
func (a *consulApiClient) CheckRegister(check *consul.AgentCheckRegistration) error {
return a.client.Agent().CheckRegister(check)
}
func (a *consulApiClient) CheckDeregister(checkID string) error {
return a.client.Agent().CheckDeregister(checkID)
}
func (a *consulApiClient) ServiceRegister(service *consul.AgentServiceRegistration) error {
return a.client.Agent().ServiceRegister(service)
}
func (a *consulApiClient) ServiceDeregister(serviceId string) error {
return a.client.Agent().ServiceDeregister(serviceId)
}
func (a *consulApiClient) Services() (map[string]*consul.AgentService, error) {
return a.client.Agent().Services()
}
func (a *consulApiClient) Checks() (map[string]*consul.AgentCheck, error) {
return a.client.Agent().Checks()
}
// trackedTask is a Task that we are tracking for changes in service and check
// definitions and keep them sycned with Consul Agent
type trackedTask struct {
allocID string
task *structs.Task
}
// ConsulService is the service which tracks tasks and syncs the services and
// checks defined in them with Consul Agent
type ConsulService struct {
client consulApi
logger *log.Logger
shutdownCh chan struct{}
trackedTasks map[string]*trackedTask
serviceStates map[string]string
trackedTskLock sync.Mutex
}
// A factory method to create new consul service
func NewConsulService(logger *log.Logger, consulAddr string, token string,
auth string, enableSSL bool, verifySSL bool) (*ConsulService, error) {
var err error
var c *consul.Client
cfg := consul.DefaultConfig()
cfg.Address = consulAddr
if token != "" {
cfg.Token = token
}
if auth != "" {
var username, password string
if strings.Contains(auth, ":") {
split := strings.SplitN(auth, ":", 2)
username = split[0]
password = split[1]
} else {
username = auth
}
cfg.HttpAuth = &consul.HttpBasicAuth{
Username: username,
Password: password,
}
}
if enableSSL {
cfg.Scheme = "https"
}
if enableSSL && !verifySSL {
cfg.HttpClient.Transport = &http.Transport{
TLSClientConfig: &tls.Config{
InsecureSkipVerify: true,
},
}
}
if c, err = consul.NewClient(cfg); err != nil {
return nil, err
}
consulService := ConsulService{
client: &consulApiClient{client: c},
logger: logger,
trackedTasks: make(map[string]*trackedTask),
serviceStates: make(map[string]string),
shutdownCh: make(chan struct{}),
}
return &consulService, nil
}
// Register starts tracking a task for changes to it's services and tasks and
// adds/removes services and checks associated with it.
func (c *ConsulService) Register(task *structs.Task, allocID string) error {
var mErr multierror.Error
c.trackedTskLock.Lock()
tt := &trackedTask{allocID: allocID, task: task}
c.trackedTasks[fmt.Sprintf("%s-%s", allocID, task.Name)] = tt
c.trackedTskLock.Unlock()
for _, service := range task.Services {
c.logger.Printf("[INFO] consul: Registering service %s with Consul.", service.Name)
if err := c.registerService(service, task, allocID); err != nil {
mErr.Errors = append(mErr.Errors, err)
}
}
return mErr.ErrorOrNil()
}
// Deregister stops tracking a task for changes to it's services and checks and
// removes all the services and checks associated with the Task
func (c *ConsulService) Deregister(task *structs.Task, allocID string) error {
var mErr multierror.Error
c.trackedTskLock.Lock()
delete(c.trackedTasks, fmt.Sprintf("%s-%s", allocID, task.Name))
c.trackedTskLock.Unlock()
for _, service := range task.Services {
if service.Id == "" {
continue
}
c.logger.Printf("[INFO] consul: De-Registering service %v with Consul", service.Name)
if err := c.deregisterService(service.Id); err != nil {
c.logger.Printf("[DEBUG] consul: Error in de-registering service %v from Consul", service.Name)
mErr.Errors = append(mErr.Errors, err)
}
}
return mErr.ErrorOrNil()
}
func (c *ConsulService) ShutDown() {
close(c.shutdownCh)
}
// SyncWithConsul is a long lived function that performs calls to sync
// checks and services periodically with Consul Agent
func (c *ConsulService) SyncWithConsul() {
sync := time.After(syncInterval)
for {
select {
case <-sync:
c.performSync()
sync = time.After(syncInterval)
case <-c.shutdownCh:
c.logger.Printf("[INFO] Shutting down Consul Client")
return
}
}
}
// performSync syncs checks and services with Consul and removed tracked
// services which are no longer present in tasks
func (c *ConsulService) performSync() {
// Get the list of the services and that Consul knows about
consulServices, _ := c.client.Services()
consulChecks, _ := c.client.Checks()
delete(consulServices, "consul")
knownChecks := make(map[string]struct{})
knownServices := make(map[string]struct{})
// Add services and checks which Consul doesn't know about
for _, trackedTask := range c.trackedTasks {
for _, service := range trackedTask.task.Services {
// Add new services which Consul agent isn't aware of
knownServices[service.Id] = struct{}{}
if _, ok := consulServices[service.Id]; !ok {
c.registerService(service, trackedTask.task, trackedTask.allocID)
continue
}
// If a service has changed, re-register it with Consul agent
if service.Hash() != c.serviceStates[service.Id] {
c.registerService(service, trackedTask.task, trackedTask.allocID)
continue
}
// Add new checks that Consul isn't aware of
for _, check := range service.Checks {
knownChecks[check.Id] = struct{}{}
if _, ok := consulChecks[check.Id]; !ok {
host, port := trackedTask.task.FindHostAndPortFor(service.PortLabel)
cr := c.makeCheck(service, check, host, port)
c.registerCheck(cr)
}
}
}
}
// Remove services from the service tracker which no longer exists
for serviceId := range c.serviceStates {
if _, ok := knownServices[serviceId]; !ok {
delete(c.serviceStates, serviceId)
}
}
// Remove services that are not present anymore
for _, consulService := range consulServices {
if _, ok := knownServices[consulService.ID]; !ok {
delete(c.serviceStates, consulService.ID)
c.deregisterService(consulService.ID)
}
}
// Remove checks that are not present anymore
for _, consulCheck := range consulChecks {
if _, ok := knownChecks[consulCheck.CheckID]; !ok {
c.deregisterCheck(consulCheck.CheckID)
}
}
}
// registerService registers a Service with Consul
func (c *ConsulService) registerService(service *structs.Service, task *structs.Task, allocID string) error {
var mErr multierror.Error
service.Id = fmt.Sprintf("%s-%s", allocID, service.Name)
host, port := task.FindHostAndPortFor(service.PortLabel)
if host == "" || port == 0 {
return fmt.Errorf("consul: The port:%s marked for registration of service: %s couldn't be found", service.PortLabel, service.Name)
}
c.serviceStates[service.Id] = service.Hash()
asr := &consul.AgentServiceRegistration{
ID: service.Id,
Name: service.Name,
Tags: service.Tags,
Port: port,
Address: host,
}
if err := c.client.ServiceRegister(asr); err != nil {
c.logger.Printf("[DEBUG] consul: Error while registering service %v with Consul: %v", service.Name, err)
mErr.Errors = append(mErr.Errors, err)
}
for _, check := range service.Checks {
cr := c.makeCheck(service, check, host, port)
if err := c.registerCheck(cr); err != nil {
c.logger.Printf("[ERROR] consul: Error while registerting check %v with Consul: %v", check.Name, err)
mErr.Errors = append(mErr.Errors, err)
}
}
return mErr.ErrorOrNil()
}
// registerCheck registers a check with Consul
func (c *ConsulService) registerCheck(check *consul.AgentCheckRegistration) error {
c.logger.Printf("[DEBUG] Registering Check with ID: %v for Service: %v", check.ID, check.ServiceID)
return c.client.CheckRegister(check)
}
// deregisterCheck de-registers a check with a specific ID from Consul
func (c *ConsulService) deregisterCheck(checkID string) error {
c.logger.Printf("[DEBUG] Removing check with ID: %v", checkID)
return c.client.CheckDeregister(checkID)
}
// deregisterService de-registers a Service with a specific id from Consul
func (c *ConsulService) deregisterService(serviceId string) error {
delete(c.serviceStates, serviceId)
if err := c.client.ServiceDeregister(serviceId); err != nil {
return err
}
return nil
}
// makeCheck creates a Consul Check Registration struct
func (c *ConsulService) makeCheck(service *structs.Service, check *structs.ServiceCheck, ip string, port int) *consul.AgentCheckRegistration {
if check.Name == "" {
check.Name = fmt.Sprintf("service: %q%s%q check", service.Name)
}
check.Id = check.Hash(service.Id)
cr := &consul.AgentCheckRegistration{
ID: check.Id,
Name: check.Name,
ServiceID: service.Id,
}
cr.Interval = check.Interval.String()
cr.Timeout = check.Timeout.String()
switch check.Type {
case structs.ServiceCheckHTTP:
if check.Protocol == "" {
check.Protocol = "http"
}
url := url.URL{
Scheme: check.Protocol,
Host: fmt.Sprintf("%s:%d", ip, port),
Path: check.Path,
}
cr.HTTP = url.String()
case structs.ServiceCheckTCP:
cr.TCP = fmt.Sprintf("%s:%d", ip, port)
case structs.ServiceCheckScript:
cr.Script = check.Script // TODO This needs to include the path of the alloc dir and based on driver types
}
return cr
}