registry.go 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180
  1. // Copyright 2015 CoreOS, Inc.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package subnet
  15. import (
  16. "fmt"
  17. golog "log"
  18. "os"
  19. "path"
  20. "sync"
  21. "time"
  22. "github.com/coreos/flannel/Godeps/_workspace/src/github.com/coreos/go-etcd/etcd"
  23. log "github.com/coreos/flannel/Godeps/_workspace/src/github.com/golang/glog"
  24. "github.com/coreos/flannel/Godeps/_workspace/src/golang.org/x/net/context"
  25. )
  26. type Registry interface {
  27. getConfig(ctx context.Context, network string) (*etcd.Response, error)
  28. getSubnets(ctx context.Context, network string) (*etcd.Response, error)
  29. createSubnet(ctx context.Context, network, sn, data string, ttl uint64) (*etcd.Response, error)
  30. updateSubnet(ctx context.Context, network, sn, data string, ttl uint64) (*etcd.Response, error)
  31. watchSubnets(ctx context.Context, network string, since uint64) (*etcd.Response, error)
  32. }
  33. type EtcdConfig struct {
  34. Endpoints []string
  35. Keyfile string
  36. Certfile string
  37. CAFile string
  38. Prefix string
  39. }
  40. type etcdSubnetRegistry struct {
  41. mux sync.Mutex
  42. cli *etcd.Client
  43. etcdCfg *EtcdConfig
  44. }
  45. func init() {
  46. etcd.SetLogger(golog.New(os.Stderr, "go-etcd", golog.LstdFlags))
  47. }
  48. func newEtcdClient(c *EtcdConfig) (*etcd.Client, error) {
  49. if c.Keyfile != "" || c.Certfile != "" || c.CAFile != "" {
  50. return etcd.NewTLSClient(c.Endpoints, c.Certfile, c.Keyfile, c.CAFile)
  51. } else {
  52. return etcd.NewClient(c.Endpoints), nil
  53. }
  54. }
  55. func newEtcdSubnetRegistry(config *EtcdConfig) (Registry, error) {
  56. r := &etcdSubnetRegistry{
  57. etcdCfg: config,
  58. }
  59. var err error
  60. r.cli, err = newEtcdClient(config)
  61. if err != nil {
  62. return nil, err
  63. }
  64. return r, nil
  65. }
  66. func (esr *etcdSubnetRegistry) getConfig(ctx context.Context, network string) (*etcd.Response, error) {
  67. key := path.Join(esr.etcdCfg.Prefix, network, "config")
  68. resp, err := esr.client().Get(key, false, false)
  69. if err != nil {
  70. return nil, err
  71. }
  72. return resp, nil
  73. }
  74. func (esr *etcdSubnetRegistry) getSubnets(ctx context.Context, network string) (*etcd.Response, error) {
  75. key := path.Join(esr.etcdCfg.Prefix, network, "subnets")
  76. return esr.client().Get(key, false, true)
  77. }
  78. func (esr *etcdSubnetRegistry) createSubnet(ctx context.Context, network, sn, data string, ttl uint64) (*etcd.Response, error) {
  79. key := path.Join(esr.etcdCfg.Prefix, network, "subnets", sn)
  80. resp, err := esr.client().Create(key, data, ttl)
  81. if err != nil {
  82. return nil, err
  83. }
  84. ensureExpiration(resp, ttl)
  85. return resp, nil
  86. }
  87. func (esr *etcdSubnetRegistry) updateSubnet(ctx context.Context, network, sn, data string, ttl uint64) (*etcd.Response, error) {
  88. key := path.Join(esr.etcdCfg.Prefix, network, "subnets", sn)
  89. resp, err := esr.client().Set(key, data, ttl)
  90. if err != nil {
  91. return nil, err
  92. }
  93. ensureExpiration(resp, ttl)
  94. return resp, nil
  95. }
  96. type watchResp struct {
  97. resp *etcd.Response
  98. err error
  99. }
  100. func (esr *etcdSubnetRegistry) watchSubnets(ctx context.Context, network string, since uint64) (*etcd.Response, error) {
  101. stop := make(chan bool)
  102. respCh := make(chan watchResp)
  103. go func() {
  104. for {
  105. key := path.Join(esr.etcdCfg.Prefix, network, "subnets")
  106. rresp, err := esr.client().RawWatch(key, since, true, nil, stop)
  107. if err != nil {
  108. respCh <- watchResp{nil, err}
  109. return
  110. }
  111. if len(rresp.Body) == 0 {
  112. // etcd timed out, go back but recreate the client as the underlying
  113. // http transport gets hosed (http://code.google.com/p/go/issues/detail?id=8648)
  114. esr.resetClient()
  115. continue
  116. }
  117. resp, err := rresp.Unmarshal()
  118. respCh <- watchResp{resp, err}
  119. }
  120. }()
  121. select {
  122. case <-ctx.Done():
  123. close(stop)
  124. <-respCh // Wait for f to return.
  125. return nil, ctx.Err()
  126. case wr := <-respCh:
  127. return wr.resp, wr.err
  128. }
  129. }
  130. func (esr *etcdSubnetRegistry) client() *etcd.Client {
  131. esr.mux.Lock()
  132. defer esr.mux.Unlock()
  133. return esr.cli
  134. }
  135. func (esr *etcdSubnetRegistry) resetClient() {
  136. esr.mux.Lock()
  137. defer esr.mux.Unlock()
  138. var err error
  139. esr.cli.Close()
  140. esr.cli, err = newEtcdClient(esr.etcdCfg)
  141. if err != nil {
  142. panic(fmt.Errorf("resetClient: error recreating etcd client: %v", err))
  143. }
  144. }
  145. func ensureExpiration(resp *etcd.Response, ttl uint64) {
  146. if resp.Node.Expiration == nil {
  147. // should not be but calc it ourselves in this case
  148. log.Info("Expiration field missing on etcd response, calculating locally")
  149. exp := time.Now().Add(time.Duration(ttl) * time.Second)
  150. resp.Node.Expiration = &exp
  151. }
  152. }