forked from vitessio/vitess
-
Notifications
You must be signed in to change notification settings - Fork 0
/
fake_healthcheck.go
194 lines (168 loc) · 5.04 KB
/
fake_healthcheck.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
package discovery
import (
"sync"
"github.com/youtube/vitess/go/sync2"
"github.com/youtube/vitess/go/vt/tabletserver/sandboxconn"
"github.com/youtube/vitess/go/vt/tabletserver/tabletconn"
"github.com/youtube/vitess/go/vt/topo"
querypb "github.com/youtube/vitess/go/vt/proto/query"
topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)
// This file contains the definitions for a FakeHealthCheck class to
// simulate a HealthCheck module. Note it is not in a sub-package because
// otherwise it couldn't be used in this package's tests because of
// circular dependencies.
// NewFakeHealthCheck returns the fake healthcheck object.
func NewFakeHealthCheck() *FakeHealthCheck {
return &FakeHealthCheck{
items: make(map[string]*fhcItem),
}
}
// FakeHealthCheck implements discovery.HealthCheck.
type FakeHealthCheck struct {
// mu protects the items map
mu sync.RWMutex
items map[string]*fhcItem
// GetStatsFromTargetCounter counts GetTabletStatsFromTarget() being called.
// (it can be accessed concurrently by 'multiGo', so using atomic)
GetStatsFromTargetCounter sync2.AtomicInt32
// GetStatsFromKeyspaceShardCounter counts GetTabletStatsFromKeyspaceShard() being called.
GetStatsFromKeyspaceShardCounter int
}
type fhcItem struct {
ts *TabletStats
conn tabletconn.TabletConn
}
//
// discovery.HealthCheck interface methods
//
// RegisterStats is not implemented.
func (fhc *FakeHealthCheck) RegisterStats() {
}
// SetListener is not implemented.
func (fhc *FakeHealthCheck) SetListener(listener HealthCheckStatsListener) {
}
// AddTablet adds the tablet.
func (fhc *FakeHealthCheck) AddTablet(cell, name string, tablet *topodatapb.Tablet) {
key := TabletToMapKey(tablet)
item := &fhcItem{
ts: &TabletStats{
Tablet: tablet,
Name: name,
},
}
fhc.mu.Lock()
defer fhc.mu.Unlock()
fhc.items[key] = item
}
// RemoveTablet removes the tablet.
func (fhc *FakeHealthCheck) RemoveTablet(tablet *topodatapb.Tablet) {
fhc.mu.Lock()
defer fhc.mu.Unlock()
key := TabletToMapKey(tablet)
delete(fhc.items, key)
}
// GetTabletStatsFromKeyspaceShard returns all TabletStats for the given keyspace/shard.
func (fhc *FakeHealthCheck) GetTabletStatsFromKeyspaceShard(keyspace, shard string) []*TabletStats {
fhc.mu.RLock()
defer fhc.mu.RUnlock()
fhc.GetStatsFromKeyspaceShardCounter++
var res []*TabletStats
for _, item := range fhc.items {
if item.ts.Target == nil {
continue
}
if item.ts.Target.Keyspace == keyspace && item.ts.Target.Shard == shard {
res = append(res, item.ts)
}
}
return res
}
// GetTabletStatsFromTarget returns all TabletStats for the given target.
func (fhc *FakeHealthCheck) GetTabletStatsFromTarget(keyspace, shard string, tabletType topodatapb.TabletType) []*TabletStats {
fhc.GetStatsFromTargetCounter.Add(1)
fhc.mu.RLock()
defer fhc.mu.RUnlock()
var res []*TabletStats
for _, item := range fhc.items {
if item.ts.Target == nil {
continue
}
if item.ts.Target.Keyspace == keyspace && item.ts.Target.Shard == shard && item.ts.Target.TabletType == tabletType {
res = append(res, item.ts)
}
}
return res
}
// GetConnection returns the TabletConn of the given tablet.
func (fhc *FakeHealthCheck) GetConnection(tablet *topodatapb.Tablet) tabletconn.TabletConn {
fhc.mu.RLock()
defer fhc.mu.RUnlock()
key := TabletToMapKey(tablet)
if item := fhc.items[key]; item != nil {
return item.conn
}
return nil
}
// CacheStatus is not implemented.
func (fhc *FakeHealthCheck) CacheStatus() TabletsCacheStatusList {
return nil
}
// Close is not implemented.
func (fhc *FakeHealthCheck) Close() error {
return nil
}
//
// Management methods
//
// Reset cleans up the internal state.
func (fhc *FakeHealthCheck) Reset() {
fhc.mu.Lock()
defer fhc.mu.Unlock()
fhc.GetStatsFromTargetCounter.Set(0)
fhc.GetStatsFromKeyspaceShardCounter = 0
fhc.items = make(map[string]*fhcItem)
}
// AddTestTablet inserts a fake entry into FakeHealthCheck.
// The Tablet can be talked to using the provided connection.
func (fhc *FakeHealthCheck) AddTestTablet(cell, host string, port int32, keyspace, shard string, tabletType topodatapb.TabletType, serving bool, reparentTS int64, err error) *sandboxconn.SandboxConn {
t := topo.NewTablet(0, cell, host)
t.Keyspace = keyspace
t.Shard = shard
t.Type = tabletType
t.PortMap["vt"] = port
key := TabletToMapKey(t)
fhc.mu.Lock()
defer fhc.mu.Unlock()
item := fhc.items[key]
if item == nil {
item = &fhcItem{
ts: &TabletStats{
Tablet: t,
},
}
fhc.items[key] = item
}
item.ts.Target = &querypb.Target{
Keyspace: keyspace,
Shard: shard,
TabletType: tabletType,
}
item.ts.Serving = serving
item.ts.TabletExternallyReparentedTimestamp = reparentTS
item.ts.Stats = &querypb.RealtimeStats{}
item.ts.LastError = err
conn := sandboxconn.NewSandboxConn(t)
item.conn = conn
return conn
}
// GetAllTablets returns all the tablets we have.
func (fhc *FakeHealthCheck) GetAllTablets() map[string]*topodatapb.Tablet {
res := make(map[string]*topodatapb.Tablet)
fhc.mu.RLock()
defer fhc.mu.RUnlock()
for key, t := range fhc.items {
res[key] = t.ts.Tablet
}
return res
}