-
Notifications
You must be signed in to change notification settings - Fork 0
/
initialization.go
74 lines (60 loc) · 1.56 KB
/
initialization.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
package etcd
import (
"context"
"github.com/coreos/etcd/clientv3"
"github.com/coreos/etcd/clientv3/concurrency"
"github.com/sensu/sensu-go/backend/store"
)
const (
initializationLockKey = ".initialized.lock"
initializationKey = ".initialized"
)
// StoreInitializer ...
type StoreInitializer struct {
mutex *concurrency.Mutex
ctx context.Context
session *concurrency.Session
client *clientv3.Client
}
// NewInitializer returns a new store initializer
func (store *Store) NewInitializer() (store.Initializer, error) {
client := store.client
session, err := concurrency.NewSession(client) // TODO: move session into etcdStore?
if err != nil {
return nil, err
}
return &StoreInitializer{
mutex: concurrency.NewMutex(session, initializationLockKey),
session: session,
client: client,
ctx: context.TODO(),
}, nil
}
// Lock mutex to avoid competing writes
func (s *StoreInitializer) Lock() error {
return s.mutex.Lock(s.ctx)
}
// IsInitialized checks the state of the .initialized key
func (s *StoreInitializer) IsInitialized() (bool, error) {
r, err := s.client.Get(s.ctx, initializationKey)
if err != nil {
return false, err
}
return r.Count > 0, nil
}
// FlagAsInitialized - set .initialized key
func (s *StoreInitializer) FlagAsInitialized() error {
_, err := s.client.Put(s.ctx, initializationKey, "1")
return err
}
// Close session & unlock
func (s *StoreInitializer) Close() error {
if err := s.mutex.Unlock(s.ctx); err != nil {
return err
}
if err := s.session.Close(); err != nil {
return err
}
<-s.session.Done()
return nil
}