github.com/stackdocker/rkt@v0.10.1-0.20151109095037-1aa827478248/Godeps/_workspace/src/google.golang.org/grpc/naming/etcd/etcd.go (about)

     1  package etcd
     2  
     3  import (
     4  	"log"
     5  	"sync"
     6  
     7  	etcdcl "github.com/coreos/etcd/client"
     8  	"github.com/coreos/rkt/Godeps/_workspace/src/golang.org/x/net/context"
     9  	"github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc/naming"
    10  )
    11  
    12  type kv struct {
    13  	key, value string
    14  }
    15  
    16  // recvBuffer is an unbounded channel of *kv to record all the pending changes from etcd server.
    17  type recvBuffer struct {
    18  	c        chan *kv
    19  	mu       sync.Mutex
    20  	stopping bool
    21  	backlog  []*kv
    22  }
    23  
    24  func newRecvBuffer() *recvBuffer {
    25  	b := &recvBuffer{
    26  		c: make(chan *kv, 1),
    27  	}
    28  	return b
    29  }
    30  
    31  func (b *recvBuffer) put(r *kv) {
    32  	b.mu.Lock()
    33  	defer b.mu.Unlock()
    34  	if b.stopping {
    35  		return
    36  	}
    37  	b.backlog = append(b.backlog, r)
    38  	select {
    39  	case b.c <- b.backlog[0]:
    40  		b.backlog = b.backlog[1:]
    41  	default:
    42  	}
    43  }
    44  
    45  func (b *recvBuffer) load() {
    46  	b.mu.Lock()
    47  	defer b.mu.Unlock()
    48  	if b.stopping || len(b.backlog) == 0 {
    49  		return
    50  	}
    51  	select {
    52  	case b.c <- b.backlog[0]:
    53  		b.backlog = b.backlog[1:]
    54  	default:
    55  	}
    56  }
    57  
    58  func (b *recvBuffer) get() <-chan *kv {
    59  	return b.c
    60  }
    61  
    62  // stop terminates the recvBuffer. After it is called, the recvBuffer is not usable any more.
    63  func (b *recvBuffer) stop() {
    64  	b.mu.Lock()
    65  	b.stopping = true
    66  	close(b.c)
    67  	b.mu.Unlock()
    68  }
    69  
    70  type etcdNR struct {
    71  	kAPI   etcdcl.KeysAPI
    72  	recv   *recvBuffer
    73  	ctx    context.Context
    74  	cancel context.CancelFunc
    75  }
    76  
    77  // NewETCDNR creates an etcd NameResolver.
    78  func NewETCDNR(cfg etcdcl.Config) (naming.Resolver, error) {
    79  	c, err := etcdcl.New(cfg)
    80  	if err != nil {
    81  		return nil, err
    82  	}
    83  	kAPI := etcdcl.NewKeysAPI(c)
    84  	ctx, cancel := context.WithCancel(context.Background())
    85  	return &etcdNR{
    86  		kAPI:   kAPI,
    87  		recv:   newRecvBuffer(),
    88  		ctx:    ctx,
    89  		cancel: cancel,
    90  	}, nil
    91  }
    92  
    93  // getNode builds the resulting key-value map starting from node recursively.
    94  func getNode(node *etcdcl.Node, res map[string]string) {
    95  	if !node.Dir {
    96  		res[node.Key] = node.Value
    97  		return
    98  	}
    99  	for _, val := range node.Nodes {
   100  		getNode(val, res)
   101  	}
   102  }
   103  
   104  func (nr *etcdNR) Get(target string) map[string]string {
   105  	resp, err := nr.kAPI.Get(nr.ctx, target, &etcdcl.GetOptions{Recursive: true, Sort: true})
   106  	if err != nil {
   107  		log.Printf("etcdNR.Get(_) stopped: %v", err)
   108  		return nil
   109  	}
   110  	res := make(map[string]string)
   111  	getNode(resp.Node, res)
   112  	return res
   113  }
   114  
   115  func (nr *etcdNR) Watch(target string) {
   116  	watcher := nr.kAPI.Watcher(target, &etcdcl.WatcherOptions{Recursive: true})
   117  	for {
   118  		resp, err := watcher.Next(nr.ctx)
   119  		if err != nil {
   120  			log.Printf("etcdNR.Watch(_) stopped: %v", err)
   121  			break
   122  		}
   123  		if resp.Node.Dir {
   124  			continue
   125  		}
   126  		entry := &kv{key: resp.Node.Key, value: resp.Node.Value}
   127  		nr.recv.put(entry)
   128  	}
   129  }
   130  
   131  func (nr *etcdNR) GetUpdate() (string, string) {
   132  	i := <-nr.recv.get()
   133  	nr.recv.load()
   134  	if i == nil {
   135  		return "", ""
   136  	}
   137  	// returns key and the corresponding value of the updated kv pair
   138  	return i.key, i.value
   139  
   140  }
   141  
   142  func (nr *etcdNR) Stop() {
   143  	nr.recv.stop()
   144  	nr.cancel()
   145  }