github.com/insionng/yougam@v0.0.0-20170714101924-2bc18d833463/libraries/ngaut/go-zookeeper/zk/lock.go (about)

     1  package zk
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"strconv"
     7  	"strings"
     8  )
     9  
    10  var (
    11  	ErrDeadlock  = errors.New("zk: trying to acquire a lock twice")
    12  	ErrNotLocked = errors.New("zk: not locked")
    13  )
    14  
    15  type Lock struct {
    16  	c        *Conn
    17  	path     string
    18  	acl      []ACL
    19  	lockPath string
    20  	seq      int
    21  }
    22  
    23  func NewLock(c *Conn, path string, acl []ACL) *Lock {
    24  	return &Lock{
    25  		c:    c,
    26  		path: path,
    27  		acl:  acl,
    28  	}
    29  }
    30  
    31  func parseSeq(path string) (int, error) {
    32  	parts := strings.Split(path, "-")
    33  	return strconv.Atoi(parts[len(parts)-1])
    34  }
    35  
    36  func (l *Lock) Lock() error {
    37  	if l.lockPath != "" {
    38  		return ErrDeadlock
    39  	}
    40  
    41  	prefix := fmt.Sprintf("%s/lock-", l.path)
    42  
    43  	path := ""
    44  	var err error
    45  	for i := 0; i < 3; i++ {
    46  		path, err = l.c.CreateProtectedEphemeralSequential(prefix, []byte{}, l.acl)
    47  		if err == ErrNoNode {
    48  			// Create parent node.
    49  			parts := strings.Split(l.path, "/")
    50  			pth := ""
    51  			for _, p := range parts[1:] {
    52  				pth += "/" + p
    53  				_, err := l.c.Create(pth, []byte{}, 0, l.acl)
    54  				if err != nil && err != ErrNodeExists {
    55  					return err
    56  				}
    57  			}
    58  		} else if err == nil {
    59  			break
    60  		} else {
    61  			return err
    62  		}
    63  	}
    64  	if err != nil {
    65  		return err
    66  	}
    67  
    68  	seq, err := parseSeq(path)
    69  	if err != nil {
    70  		return err
    71  	}
    72  
    73  	for {
    74  		children, _, err := l.c.Children(l.path)
    75  		if err != nil {
    76  			return err
    77  		}
    78  
    79  		lowestSeq := seq
    80  		prevSeq := 0
    81  		prevSeqPath := ""
    82  		for _, p := range children {
    83  			s, err := parseSeq(p)
    84  			if err != nil {
    85  				return err
    86  			}
    87  			if s < lowestSeq {
    88  				lowestSeq = s
    89  			}
    90  			if s < seq && s > prevSeq {
    91  				prevSeq = s
    92  				prevSeqPath = p
    93  			}
    94  		}
    95  
    96  		if seq == lowestSeq {
    97  			// Acquired the lock
    98  			break
    99  		}
   100  
   101  		// Wait on the node next in line for the lock
   102  		_, _, ch, err := l.c.GetW(l.path + "/" + prevSeqPath)
   103  		if err != nil && err != ErrNoNode {
   104  			return err
   105  		} else if err != nil && err == ErrNoNode {
   106  			// try again
   107  			continue
   108  		}
   109  
   110  		ev := <-ch
   111  		if ev.Err != nil {
   112  			return ev.Err
   113  		}
   114  	}
   115  
   116  	l.seq = seq
   117  	l.lockPath = path
   118  	return nil
   119  }
   120  
   121  func (l *Lock) Unlock() error {
   122  	if l.lockPath == "" {
   123  		return ErrNotLocked
   124  	}
   125  	if err := l.c.Delete(l.lockPath, -1); err != nil {
   126  		return err
   127  	}
   128  	l.lockPath = ""
   129  	l.seq = 0
   130  	return nil
   131  }