github.com/polvi/nsproxy@v0.0.0-20140119202807-96ac732fa7f3/backends.go (about)

     1  package main
     2  
     3  import (
     4  	"encoding/json"
     5  	"log"
     6  	"net"
     7  	"strconv"
     8  	"sync"
     9  
    10  	"github.com/coreos/go-etcd/etcd"
    11  )
    12  
    13  type service struct {
    14  	Host string `json:"host"`
    15  	Port int `json:"port"`
    16  }
    17  
    18  type host struct {
    19  	key string
    20  	addr string
    21  }
    22  
    23  type backends struct {
    24  	path       string
    25  	hosts      []host
    26  	lastIndex  int
    27  	watchIndex uint64
    28  	lock       sync.RWMutex
    29  }
    30  
    31  func (b *backends) Dump(action string ) {
    32  	for _, v := range(b.hosts) {
    33  		log.Printf("Dump after %s %s -> %s", action, v.key, v.addr)
    34  	}
    35  }
    36  
    37  func (b *backends) Remove(key string) {
    38  	match := -1
    39  	for k, v := range(b.hosts) {
    40  		if v.key == key {
    41  			match = k
    42  		}
    43  	}
    44  
    45  	b.hosts = append(b.hosts[:match], b.hosts[match+1:]...)
    46  	b.Dump("remove")
    47  }
    48  
    49  func (b *backends) Update(node *etcd.Node, action string) {
    50  	b.lock.Lock()
    51  	defer b.lock.Unlock()
    52  
    53  	log.Printf("key: %s action: %s value: %s", node.Key, action, string(node.Value))
    54  
    55  	s := &service{}
    56  	if action == "delete" || action == "expire" {
    57  		b.Remove(node.Key)
    58  		return
    59  	}
    60  
    61  	err := json.Unmarshal([]byte(node.Value), s)
    62  	if err != nil {
    63  		panic(err)
    64  	}
    65  
    66  	addr := net.JoinHostPort(s.Host, strconv.Itoa(s.Port))
    67  
    68  
    69  	// TODO: create an actual set data structure
    70  	for _, v := range(b.hosts) {
    71  		if v.key == node.Key {
    72  			b.Dump(action)
    73  			return
    74  		}
    75  	}
    76  
    77  	b.hosts = append(b.hosts, host{addr: addr, key: node.Key})
    78  }
    79  
    80  func (b *backends) Watch(client *etcd.Client) {
    81  	receiver := make(chan *etcd.Response)
    82  	go client.Watch(b.path, uint64(b.watchIndex), true, receiver, nil)
    83  
    84  	for {
    85  		resp := <-receiver
    86  		b.Update(resp.Node, resp.Action)
    87  	}
    88  }
    89  
    90  func (b *backends) Sync(client *etcd.Client) error {
    91  	resp, err := client.Get(b.path, false, true)
    92  
    93  	if err != nil {
    94  		return err
    95  	}
    96  
    97  	for _, n := range(resp.Node.Nodes) {
    98  		b.Update(&n, resp.Action)
    99  	}
   100  
   101  	// Begin the watch after this sync from the next sync
   102  	b.watchIndex = resp.EtcdIndex + 1
   103  
   104  	return nil
   105  }
   106  
   107  func (b *backends) Next() string {
   108  	b.lock.RLock()
   109  	defer b.lock.RUnlock()
   110  
   111  	if len(b.hosts) == 0 {
   112  		return ""
   113  	}
   114  
   115  	index := (b.lastIndex + 1) % len(b.hosts)
   116  	b.lastIndex = index
   117  
   118  	return b.hosts[index].addr
   119  }