forked from spf13/viper
/
remote.go
103 lines (92 loc) · 2.39 KB
/
remote.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
// Copyright © 2015 Steve Francia <spf@spf13.com>.
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
// Package remote integrates the remote features of Viper.
package remote
import (
"bytes"
"context"
"io"
"strings"
"time"
"github.com/Be-MobileNV/viper"
"github.com/coreos/etcd/clientv3"
)
type remoteConfigProvider struct {
client *clientv3.Client
}
func (rc remoteConfigProvider) Get(rp viper.RemoteProvider) (io.Reader, error) {
if rc.client == nil {
var err error
rc.client, err = getConfigManager(rp)
if err != nil {
return nil, err
}
}
r, err := rc.client.Get(context.Background(), rp.Path())
if err != nil {
return nil, err
}
return bytes.NewReader(r.Kvs[0].Value), nil
}
func (rc remoteConfigProvider) Watch(rp viper.RemoteProvider) (io.Reader, error) {
if rc.client == nil {
var err error
rc.client, err = getConfigManager(rp)
if err != nil {
return nil, err
}
}
w := rc.client.Watch(context.Background(), rp.Path())
resp := <-w
if resp.Err() != nil {
return nil, resp.Err()
}
val := resp.Events[0].Kv.Value
return bytes.NewReader(val), nil
}
func (rc remoteConfigProvider) WatchChannel(rp viper.RemoteProvider) (responseChannel <-chan *viper.RemoteResponse, quitwc chan bool, err error) {
quitwc = make(chan bool)
viperResponsCh := make(chan *viper.RemoteResponse)
if rc.client == nil {
var err error
rc.client, err = getConfigManager(rp)
if err != nil {
return nil, nil, err
}
}
ctx, cancel := context.WithCancel(context.Background())
w := rc.client.Watch(ctx, rp.Path())
go func(etcdResponseChannel clientv3.WatchChan, vr chan<- *viper.RemoteResponse, quitwc <-chan bool, cancel context.CancelFunc) {
for {
select {
case <-quitwc:
cancel()
return
case resp := <-etcdResponseChannel:
vr <- &viper.RemoteResponse{
Error: resp.Err(),
Value: resp.Events[0].Kv.Value,
}
}
}
}(w, viperResponsCh, quitwc, cancel)
return viperResponsCh, quitwc, nil
}
func getConfigManager(rp viper.RemoteProvider) (*clientv3.Client, error) {
etcdConfig := clientv3.Config{
Endpoints: strings.Split(rp.Endpoint(), ","),
DialTimeout: 5 * time.Second,
Username: rp.Username(),
Password: rp.Password(),
}
client, err := clientv3.New(etcdConfig)
if err != nil {
return nil, err
}
return client, nil
}
func init() {
viper.RemoteConfig = &remoteConfigProvider{}
}