-
Notifications
You must be signed in to change notification settings - Fork 10
/
register.go
78 lines (65 loc) · 2.06 KB
/
register.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
package service
import (
"encoding/json"
"fmt"
"time"
etcd "github.com/coreos/etcd/client"
"golang.org/x/net/context"
)
const (
HEARTBEAT_DURATION = 5
)
type Infos struct {
Name string `json:"name"`
Critical bool `json:"critical"` // Is the service critical to the infrastructure health?
}
func Register(service string, host *Host, infos *Infos, stop chan struct{}) (chan struct{}, error) {
if len(host.Name) == 0 {
host.Name = hostname
}
if infos != nil {
infos.Name = service
}
registered := make(chan struct{}, 1)
hostKey := fmt.Sprintf("/services/%s/%s", service, host.Name)
hostJson, _ := json.Marshal(&host)
hostValue := string(hostJson)
serviceKey := fmt.Sprintf("/services_infos/%s", service)
serviceJson, _ := json.Marshal(infos)
serviceValue := string(serviceJson)
go func() {
ticker := time.NewTicker((HEARTBEAT_DURATION - 1) * time.Second)
KAPI().Set(context.Background(), hostKey, hostValue, &etcd.SetOptions{TTL: HEARTBEAT_DURATION * time.Second})
if infos != nil {
KAPI().Set(context.Background(), serviceKey, serviceValue, nil)
}
registered <- struct{}{}
for {
select {
case <-stop:
_, err := KAPI().Delete(context.Background(), hostKey, &etcd.DeleteOptions{Recursive: false})
if err != nil {
logger.Println("fail to remove key", hostKey)
}
ticker.Stop()
return
case <-ticker.C:
_, err := KAPI().Set(context.Background(), hostKey, hostValue, &etcd.SetOptions{TTL: HEARTBEAT_DURATION * time.Second})
// If for any random reason, there is an error,
// we retry every second until it's ok.
for err != nil {
logger.Printf("lost registration of '%v': %v (%v)", service, err, Client().Endpoints())
time.Sleep(1 * time.Second)
if infos != nil {
KAPI().Set(context.Background(), serviceKey, serviceValue, nil)
}
_, err = KAPI().Set(context.Background(), hostKey, hostValue, &etcd.SetOptions{TTL: HEARTBEAT_DURATION * time.Second})
if err == nil {
logger.Printf("recover registration of '%v'", service)
}
}
}
}
}()
return registered, nil
}