github.com/mmatczuk/gohan@v0.0.0-20170206152520-30e45d9bdb69/sync/etcdv3/etcd.go (about)

     1  // Copyright (C) 2015 NTT Innovation Institute, 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
    12  // implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package etcdv3
    17  
    18  import (
    19  	"encoding/json"
    20  	"errors"
    21  	"fmt"
    22  	"os"
    23  	"strings"
    24  	syn "sync"
    25  	"time"
    26  
    27  	"github.com/cloudwan/gohan/sync"
    28  	etcd "github.com/coreos/etcd/clientv3"
    29  	pb "github.com/coreos/etcd/mvcc/mvccpb"
    30  	cmap "github.com/streamrail/concurrent-map"
    31  	"github.com/twinj/uuid"
    32  	"golang.org/x/net/context"
    33  )
    34  
    35  const masterTTL = 10
    36  
    37  //Sync is struct for etcd based sync
    38  type Sync struct {
    39  	locks      cmap.ConcurrentMap
    40  	timeout    time.Duration
    41  	etcdClient *etcd.Client
    42  	processID  string
    43  }
    44  
    45  func (s *Sync) withTimeout() context.Context {
    46  	ctx, _ := context.WithTimeout(context.Background(), s.timeout)
    47  	return ctx
    48  }
    49  
    50  //NewSync initialize new etcd sync
    51  func NewSync(etcdServers []string, timeout time.Duration) (*Sync, error) {
    52  	sync := &Sync{
    53  		locks:   cmap.New(),
    54  		timeout: timeout,
    55  	}
    56  	client, err := etcd.New(
    57  		etcd.Config{
    58  			Endpoints:   etcdServers,
    59  			DialTimeout: timeout,
    60  		},
    61  	)
    62  	if err != nil {
    63  		return nil, err
    64  	}
    65  	sync.etcdClient = client
    66  	hostname, _ := os.Hostname()
    67  	sync.processID = hostname + uuid.NewV4().String()
    68  	return sync, nil
    69  }
    70  
    71  //Update sync update sync
    72  //When jsonString is empty, this method do nothing because
    73  //etcd v3 doesn't support directories.
    74  func (s *Sync) Update(key, jsonString string) error {
    75  	var err error
    76  	if jsonString == "" {
    77  		// do nothing, because clientv3 doesn't have directories
    78  		return nil
    79  	} else {
    80  		_, err = s.etcdClient.Put(s.withTimeout(), key, jsonString)
    81  	}
    82  	if err != nil {
    83  		log.Error(fmt.Sprintf("failed to sync with backend %s", err))
    84  		return err
    85  	}
    86  	return nil
    87  }
    88  
    89  //Delete sync update sync
    90  func (s *Sync) Delete(key string) error {
    91  	_, err := s.etcdClient.Delete(s.withTimeout(), key)
    92  	return err
    93  }
    94  
    95  //Fetch data from sync
    96  func (s *Sync) Fetch(key string) (*sync.Node, error) {
    97  	node, err := s.etcdClient.Get(s.withTimeout(), key, etcd.WithSort(etcd.SortByKey, etcd.SortAscend))
    98  	if err != nil {
    99  		return nil, err
   100  	}
   101  	dir, err := s.etcdClient.Get(s.withTimeout(), key+"/", etcd.WithPrefix(), etcd.WithSort(etcd.SortByKey, etcd.SortAscend))
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  
   106  	return s.recursiveFetch(key, node.Kvs, dir.Kvs)
   107  }
   108  
   109  func (s *Sync) recursiveFetch(rootKey string, node []*pb.KeyValue, children []*pb.KeyValue) (*sync.Node, error) {
   110  	if len(node) == 0 && len(children) == 0 {
   111  		return nil, errors.New("Not found")
   112  	}
   113  
   114  	subMap := make(map[string]*sync.Node, len(children))
   115  
   116  	rootNode := &sync.Node{}
   117  	subMap[""] = rootNode
   118  	if len(node) != 0 {
   119  		rootNode.Key = string(node[0].Key)
   120  		rootNode.Value = string(node[0].Value)
   121  		rootNode.Revision = node[0].ModRevision
   122  	} else {
   123  		rootNode.Key = rootKey
   124  	}
   125  
   126  	for _, kv := range children {
   127  		key := string(kv.Key)
   128  		n := &sync.Node{
   129  			Key: key,
   130  			Value: string(kv.Value),
   131  			Revision: kv.ModRevision,
   132  		}
   133  		path := strings.TrimPrefix(key, rootKey)
   134  		steps := strings.Split(path, "/")
   135  
   136  		for i := 1; i < len(steps); i++ {
   137  			parent, ok := subMap[strings.Join(steps[:len(steps)-i], "/")]
   138  			if ok {
   139  				for j := 1; j < i; j++ {
   140  					bridge := &sync.Node{
   141  						Key: rootKey + strings.Join(steps[:len(steps)-i+j], "/"),
   142  					}
   143  					parent.Children = []*sync.Node{bridge}
   144  					parent = bridge
   145  				}
   146  				if parent.Children == nil {
   147  					parent.Children = make([]*sync.Node, 0, 1)
   148  				}
   149  				parent.Children = append(parent.Children, n)
   150  				break
   151  			}
   152  		}
   153  	}
   154  
   155  	return rootNode, nil
   156  }
   157  
   158  //HasLock checks current process owns lock or not
   159  func (s *Sync) HasLock(path string) bool {
   160  	return s.locks.Has(path)
   161  }
   162  
   163  // Lock locks resources on sync
   164  // This call blocks until you can get lock
   165  func (s *Sync) Lock(path string, block bool) error {
   166  	for {
   167  		if s.HasLock(path) {
   168  			return nil
   169  		}
   170  		var err error
   171  		lease, err := s.etcdClient.Grant(s.withTimeout(), masterTTL)
   172  		var resp *etcd.TxnResponse
   173  		if err == nil {
   174  			cmp := etcd.Compare(etcd.CreateRevision(path), "=", 0)
   175  			put := etcd.OpPut(path, s.processID, etcd.WithLease(lease.ID))
   176  			resp, err = s.etcdClient.Txn(s.withTimeout()).If(cmp).Then(put).Commit()
   177  		}
   178  		if err != nil || !resp.Succeeded {
   179  			msg := fmt.Sprintf("failed to lock path %s", path)
   180  			if err != nil {
   181  				msg = fmt.Sprintf("failed to lock path %s: %s", path, err)
   182  			}
   183  			log.Notice(msg)
   184  
   185  			s.locks.Remove(path)
   186  			if !block {
   187  				return errors.New(msg)
   188  			}
   189  			time.Sleep(masterTTL * time.Second)
   190  			continue
   191  		}
   192  		log.Info("Locked %s", path)
   193  		s.locks.Set(path, lease.ID)
   194  		//Refresh master token
   195  		go func() {
   196  			defer func() {
   197  				log.Notice("releasing keepalive lock for %s", path)
   198  				s.locks.Remove(path)
   199  			}()
   200  			for s.HasLock(path) {
   201  				ch, err := s.etcdClient.KeepAlive(s.withTimeout(), lease.ID)
   202  				if err != nil {
   203  					log.Notice("failed to keepalive lock for %s %s", path, err)
   204  					return
   205  				}
   206  				for range ch {
   207  				}
   208  			}
   209  		}()
   210  
   211  		return nil
   212  	}
   213  }
   214  
   215  //Unlock path
   216  func (s *Sync) Unlock(path string) error {
   217  	leaseID, ok := s.locks.Get(path)
   218  	if !ok {
   219  		return nil
   220  	}
   221  	s.locks.Remove(path)
   222  	s.etcdClient.Revoke(s.withTimeout(), leaseID.(etcd.LeaseID))
   223  	log.Info("Unlocked path %s", path)
   224  	return nil
   225  }
   226  
   227  func eventsFromNode(action string, kvs []*pb.KeyValue, responseChan chan *sync.Event) {
   228  	for _, kv := range kvs {
   229  		event := &sync.Event{
   230  			Action: action,
   231  			Key:    string(kv.Key),
   232  			Revision: kv.ModRevision,
   233  		}
   234  		if kv.Value != nil {
   235  			err := json.Unmarshal(kv.Value, &event.Data)
   236  			if err != nil {
   237  				log.Warning("failed to unmarshal watch response: %s", err)
   238  				continue
   239  			}
   240  		}
   241  		responseChan <- event
   242  	}
   243  }
   244  
   245  //Watch keep watch update under the path
   246  func (s *Sync) Watch(path string, responseChan chan *sync.Event, stopChan chan bool, revision int64) error {
   247  	if revision == sync.RevisionCurrent {
   248  		node, err := s.etcdClient.Get(s.withTimeout(), path, etcd.WithSort(etcd.SortByKey, etcd.SortAscend))
   249  		if err != nil {
   250  			return err
   251  		}
   252  		eventsFromNode("get", node.Kvs, responseChan)
   253  
   254  		revision = node.Header.Revision + 1
   255  	}
   256  
   257  	dir, err := s.etcdClient.Get(s.withTimeout(), path+"/", etcd.WithPrefix(), etcd.WithSort(etcd.SortByKey, etcd.SortAscend))
   258  	if err != nil {
   259  		return err
   260  	}
   261  	eventsFromNode("get", dir.Kvs, responseChan)
   262  
   263  	ctx, cancel := context.WithCancel(context.Background())
   264  	errors := make(chan error, 2)
   265  	var wg syn.WaitGroup
   266  	wg.Add(1)
   267  	go func() {
   268  		defer wg.Done()
   269  		err := func() error {
   270  			rch := s.etcdClient.Watch(ctx, path, etcd.WithRev(revision))
   271  
   272  			for wresp := range rch {
   273  				err := wresp.Err()
   274  				if err != nil {
   275  					return err
   276  				}
   277  				for _, ev := range wresp.Events {
   278  					action := "unknown"
   279  					switch ev.Type {
   280  					case etcd.EventTypePut:
   281  						action = "set"
   282  					case etcd.EventTypeDelete:
   283  						action = "delete"
   284  					}
   285  					eventsFromNode(action, []*pb.KeyValue{ev.Kv}, responseChan)
   286  				}
   287  			}
   288  
   289  			return nil
   290  		}()
   291  		errors <- err
   292  	}()
   293  
   294  	wg.Add(1)
   295  	go func() {
   296  		defer wg.Done()
   297  		err := func() error {
   298  			rch := s.etcdClient.Watch(ctx, path+"/", etcd.WithPrefix(), etcd.WithRev(revision))
   299  
   300  			for wresp := range rch {
   301  				err := wresp.Err()
   302  				if err != nil {
   303  					return err
   304  				}
   305  				for _, ev := range wresp.Events {
   306  					action := "unknown"
   307  					switch ev.Type {
   308  					case etcd.EventTypePut:
   309  						action = "set"
   310  					case etcd.EventTypeDelete:
   311  						action = "delete"
   312  					}
   313  					eventsFromNode(action, []*pb.KeyValue{ev.Kv}, responseChan)
   314  				}
   315  			}
   316  
   317  			return nil
   318  		}()
   319  		errors <- err
   320  	}()
   321  
   322  	select {
   323  	case <-stopChan:
   324  		cancel()
   325  		wg.Wait()
   326  		return nil
   327  	case err := <-errors:
   328  		cancel()
   329  		wg.Wait()
   330  		return err
   331  	}
   332  }
   333  
   334  func (s *Sync) Close()  {
   335  	s.etcdClient.Close()
   336  }