prometheus

1
// Copyright 2015 The Prometheus Authors
2
// Licensed under the Apache License, Version 2.0 (the "License");
3
// you may not use this file except in compliance with the License.
4
// You may obtain a copy of the License at
5
//
6
// http://www.apache.org/licenses/LICENSE-2.0
7
//
8
// Unless required by applicable law or agreed to in writing, software
9
// distributed under the License is distributed on an "AS IS" BASIS,
10
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11
// See the License for the specific language governing permissions and
12
// limitations under the License.
13

14
package zookeeper
15

16
import (
17
	"context"
18
	"encoding/json"
19
	"errors"
20
	"fmt"
21
	"net"
22
	"strconv"
23
	"strings"
24
	"time"
25

26
	"github.com/go-kit/log"
27
	"github.com/go-zookeeper/zk"
28
	"github.com/prometheus/client_golang/prometheus"
29
	"github.com/prometheus/common/model"
30

31
	"github.com/prometheus/prometheus/discovery"
32
	"github.com/prometheus/prometheus/discovery/targetgroup"
33
	"github.com/prometheus/prometheus/util/strutil"
34
	"github.com/prometheus/prometheus/util/treecache"
35
)
36

37
var (
38
	// DefaultServersetSDConfig is the default Serverset SD configuration.
39
	DefaultServersetSDConfig = ServersetSDConfig{
40
		Timeout: model.Duration(10 * time.Second),
41
	}
42
	// DefaultNerveSDConfig is the default Nerve SD configuration.
43
	DefaultNerveSDConfig = NerveSDConfig{
44
		Timeout: model.Duration(10 * time.Second),
45
	}
46
)
47

48
func init() {
49
	discovery.RegisterConfig(&ServersetSDConfig{})
50
	discovery.RegisterConfig(&NerveSDConfig{})
51
}
52

53
// ServersetSDConfig is the configuration for Twitter serversets in Zookeeper based discovery.
54
type ServersetSDConfig struct {
55
	Servers []string       `yaml:"servers"`
56
	Paths   []string       `yaml:"paths"`
57
	Timeout model.Duration `yaml:"timeout,omitempty"`
58
}
59

60
// NewDiscovererMetrics implements discovery.Config.
61
func (*ServersetSDConfig) NewDiscovererMetrics(reg prometheus.Registerer, rmi discovery.RefreshMetricsInstantiator) discovery.DiscovererMetrics {
62
	return &discovery.NoopDiscovererMetrics{}
63
}
64

65
// Name returns the name of the Config.
66
func (*ServersetSDConfig) Name() string { return "serverset" }
67

68
// NewDiscoverer returns a Discoverer for the Config.
69
func (c *ServersetSDConfig) NewDiscoverer(opts discovery.DiscovererOptions) (discovery.Discoverer, error) {
70
	return NewServersetDiscovery(c, opts.Logger)
71
}
72

73
// UnmarshalYAML implements the yaml.Unmarshaler interface.
74
func (c *ServersetSDConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
75
	*c = DefaultServersetSDConfig
76
	type plain ServersetSDConfig
77
	err := unmarshal((*plain)(c))
78
	if err != nil {
79
		return err
80
	}
81
	if len(c.Servers) == 0 {
82
		return errors.New("serverset SD config must contain at least one Zookeeper server")
83
	}
84
	if len(c.Paths) == 0 {
85
		return errors.New("serverset SD config must contain at least one path")
86
	}
87
	for _, path := range c.Paths {
88
		if !strings.HasPrefix(path, "/") {
89
			return fmt.Errorf("serverset SD config paths must begin with '/': %s", path)
90
		}
91
	}
92
	return nil
93
}
94

95
// NerveSDConfig is the configuration for AirBnB's Nerve in Zookeeper based discovery.
96
type NerveSDConfig struct {
97
	Servers []string       `yaml:"servers"`
98
	Paths   []string       `yaml:"paths"`
99
	Timeout model.Duration `yaml:"timeout,omitempty"`
100
}
101

102
// NewDiscovererMetrics implements discovery.Config.
103
func (*NerveSDConfig) NewDiscovererMetrics(reg prometheus.Registerer, rmi discovery.RefreshMetricsInstantiator) discovery.DiscovererMetrics {
104
	return &discovery.NoopDiscovererMetrics{}
105
}
106

107
// Name returns the name of the Config.
108
func (*NerveSDConfig) Name() string { return "nerve" }
109

110
// NewDiscoverer returns a Discoverer for the Config.
111
func (c *NerveSDConfig) NewDiscoverer(opts discovery.DiscovererOptions) (discovery.Discoverer, error) {
112
	return NewNerveDiscovery(c, opts.Logger)
113
}
114

115
// UnmarshalYAML implements the yaml.Unmarshaler interface.
116
func (c *NerveSDConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
117
	*c = DefaultNerveSDConfig
118
	type plain NerveSDConfig
119
	err := unmarshal((*plain)(c))
120
	if err != nil {
121
		return err
122
	}
123
	if len(c.Servers) == 0 {
124
		return errors.New("nerve SD config must contain at least one Zookeeper server")
125
	}
126
	if len(c.Paths) == 0 {
127
		return errors.New("nerve SD config must contain at least one path")
128
	}
129
	for _, path := range c.Paths {
130
		if !strings.HasPrefix(path, "/") {
131
			return fmt.Errorf("nerve SD config paths must begin with '/': %s", path)
132
		}
133
	}
134
	return nil
135
}
136

137
// Discovery implements the Discoverer interface for discovering
138
// targets from Zookeeper.
139
type Discovery struct {
140
	conn *zk.Conn
141

142
	sources map[string]*targetgroup.Group
143

144
	updates     chan treecache.ZookeeperTreeCacheEvent
145
	pathUpdates []chan treecache.ZookeeperTreeCacheEvent
146
	treeCaches  []*treecache.ZookeeperTreeCache
147

148
	parse  func(data []byte, path string) (model.LabelSet, error)
149
	logger log.Logger
150
}
151

152
// NewNerveDiscovery returns a new Discovery for the given Nerve config.
153
func NewNerveDiscovery(conf *NerveSDConfig, logger log.Logger) (*Discovery, error) {
154
	return NewDiscovery(conf.Servers, time.Duration(conf.Timeout), conf.Paths, logger, parseNerveMember)
155
}
156

157
// NewServersetDiscovery returns a new Discovery for the given serverset config.
158
func NewServersetDiscovery(conf *ServersetSDConfig, logger log.Logger) (*Discovery, error) {
159
	return NewDiscovery(conf.Servers, time.Duration(conf.Timeout), conf.Paths, logger, parseServersetMember)
160
}
161

162
// NewDiscovery returns a new discovery along Zookeeper parses with
163
// the given parse function.
164
func NewDiscovery(
165
	srvs []string,
166
	timeout time.Duration,
167
	paths []string,
168
	logger log.Logger,
169
	pf func(data []byte, path string) (model.LabelSet, error),
170
) (*Discovery, error) {
171
	if logger == nil {
172
		logger = log.NewNopLogger()
173
	}
174

175
	conn, _, err := zk.Connect(
176
		srvs, timeout,
177
		func(c *zk.Conn) {
178
			c.SetLogger(treecache.NewZookeeperLogger(logger))
179
		})
180
	if err != nil {
181
		return nil, err
182
	}
183
	updates := make(chan treecache.ZookeeperTreeCacheEvent)
184
	sd := &Discovery{
185
		conn:    conn,
186
		updates: updates,
187
		sources: map[string]*targetgroup.Group{},
188
		parse:   pf,
189
		logger:  logger,
190
	}
191
	for _, path := range paths {
192
		pathUpdate := make(chan treecache.ZookeeperTreeCacheEvent)
193
		sd.pathUpdates = append(sd.pathUpdates, pathUpdate)
194
		sd.treeCaches = append(sd.treeCaches, treecache.NewZookeeperTreeCache(conn, path, pathUpdate, logger))
195
	}
196
	return sd, nil
197
}
198

199
// Run implements the Discoverer interface.
200
func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) {
201
	defer func() {
202
		for _, tc := range d.treeCaches {
203
			tc.Stop()
204
		}
205
		for _, pathUpdate := range d.pathUpdates {
206
			// Drain event channel in case the treecache leaks goroutines otherwise.
207
			for range pathUpdate {
208
			}
209
		}
210
		d.conn.Close()
211
	}()
212

213
	for _, pathUpdate := range d.pathUpdates {
214
		go func(update chan treecache.ZookeeperTreeCacheEvent) {
215
			for event := range update {
216
				select {
217
				case d.updates <- event:
218
				case <-ctx.Done():
219
					return
220
				}
221
			}
222
		}(pathUpdate)
223
	}
224

225
	for {
226
		select {
227
		case <-ctx.Done():
228
			return
229
		case event := <-d.updates:
230
			tg := &targetgroup.Group{
231
				Source: event.Path,
232
			}
233
			if event.Data != nil {
234
				labelSet, err := d.parse(*event.Data, event.Path)
235
				if err == nil {
236
					tg.Targets = []model.LabelSet{labelSet}
237
					d.sources[event.Path] = tg
238
				} else {
239
					delete(d.sources, event.Path)
240
				}
241
			} else {
242
				delete(d.sources, event.Path)
243
			}
244
			select {
245
			case <-ctx.Done():
246
				return
247
			case ch <- []*targetgroup.Group{tg}:
248
			}
249
		}
250
	}
251
}
252

253
const (
254
	serversetLabelPrefix         = model.MetaLabelPrefix + "serverset_"
255
	serversetStatusLabel         = serversetLabelPrefix + "status"
256
	serversetPathLabel           = serversetLabelPrefix + "path"
257
	serversetEndpointLabelPrefix = serversetLabelPrefix + "endpoint"
258
	serversetShardLabel          = serversetLabelPrefix + "shard"
259
)
260

261
type serversetMember struct {
262
	ServiceEndpoint     serversetEndpoint
263
	AdditionalEndpoints map[string]serversetEndpoint
264
	Status              string `json:"status"`
265
	Shard               int    `json:"shard"`
266
}
267

268
type serversetEndpoint struct {
269
	Host string
270
	Port int
271
}
272

273
func parseServersetMember(data []byte, path string) (model.LabelSet, error) {
274
	member := serversetMember{}
275

276
	if err := json.Unmarshal(data, &member); err != nil {
277
		return nil, fmt.Errorf("error unmarshaling serverset member %q: %w", path, err)
278
	}
279

280
	labels := model.LabelSet{}
281
	labels[serversetPathLabel] = model.LabelValue(path)
282
	labels[model.AddressLabel] = model.LabelValue(
283
		net.JoinHostPort(member.ServiceEndpoint.Host, strconv.Itoa(member.ServiceEndpoint.Port)))
284

285
	labels[serversetEndpointLabelPrefix+"_host"] = model.LabelValue(member.ServiceEndpoint.Host)
286
	labels[serversetEndpointLabelPrefix+"_port"] = model.LabelValue(strconv.Itoa(member.ServiceEndpoint.Port))
287

288
	for name, endpoint := range member.AdditionalEndpoints {
289
		cleanName := model.LabelName(strutil.SanitizeLabelName(name))
290
		labels[serversetEndpointLabelPrefix+"_host_"+cleanName] = model.LabelValue(
291
			endpoint.Host)
292
		labels[serversetEndpointLabelPrefix+"_port_"+cleanName] = model.LabelValue(
293
			strconv.Itoa(endpoint.Port))
294
	}
295

296
	labels[serversetStatusLabel] = model.LabelValue(member.Status)
297
	labels[serversetShardLabel] = model.LabelValue(strconv.Itoa(member.Shard))
298

299
	return labels, nil
300
}
301

302
const (
303
	nerveLabelPrefix         = model.MetaLabelPrefix + "nerve_"
304
	nervePathLabel           = nerveLabelPrefix + "path"
305
	nerveEndpointLabelPrefix = nerveLabelPrefix + "endpoint"
306
)
307

308
type nerveMember struct {
309
	Host string `json:"host"`
310
	Port int    `json:"port"`
311
	Name string `json:"name"`
312
}
313

314
func parseNerveMember(data []byte, path string) (model.LabelSet, error) {
315
	member := nerveMember{}
316
	err := json.Unmarshal(data, &member)
317
	if err != nil {
318
		return nil, fmt.Errorf("error unmarshaling nerve member %q: %w", path, err)
319
	}
320

321
	labels := model.LabelSet{}
322
	labels[nervePathLabel] = model.LabelValue(path)
323
	labels[model.AddressLabel] = model.LabelValue(
324
		net.JoinHostPort(member.Host, strconv.Itoa(member.Port)))
325

326
	labels[nerveEndpointLabelPrefix+"_host"] = model.LabelValue(member.Host)
327
	labels[nerveEndpointLabelPrefix+"_port"] = model.LabelValue(strconv.Itoa(member.Port))
328
	labels[nerveEndpointLabelPrefix+"_name"] = model.LabelValue(member.Name)
329

330
	return labels, nil
331
}
332