利用zookeeper实现分布式锁

package main

import (
	"time"
	"github.com/samuel/go-zookeeper/zk"
	"errors"
	"strings"
	"strconv"
	"fmt"
	"sync"
)

const (
	Pre_Child_Node = "child_"
	Default_Version = 0
	Root_Node = "/lock"
	)

func main()  {
	var hosts = []string{"localhost:2181"}
	conn,_,err := zk.Connect(hosts,time.Second*60)
	if err != nil{
		return
	}
	zkServer := ZKServer{Conn:conn,Nodes:make(map[int]Node)}
	_,err = zkServer.Conn.Create(Root_Node,[]byte(""),zk.FlagSequence,zk.WorldACL(zk.PermAll))

	children,_,_:= zkServer.Conn.Children(Root_Node)
	for _,childNode := range children{
		zkServer.Conn.Delete(fmt.Sprintf("%s/%s",Root_Node,childNode),Default_Version)
	}

	lock := DistributeLock{ZKServer:&zkServer}
	var waitGroutp = sync.WaitGroup{}
	//循环创建子节点
	for i:=0;i<10;i++{
		waitGroutp.Add(1)
		go func(i int) {
			err = lock.Acquire()
			children,_,_= zkServer.Conn.Children(Root_Node)
			fmt.Println(err)
			fmt.Println(fmt.Sprintf("goroute_%d",i))
			fmt.Println(fmt.Sprintf("done:%d",i))
			lock.Release()
			waitGroutp.Done()
		}(i)
	}
	waitGroutp.Wait()
	fmt.Println("end...")
}

type DistributeLock struct{
	ZKServer *ZKServer
	NodeNum int
}

//获取锁，没获取到一直阻塞
func (dLock *DistributeLock)Acquire()(error){
	zkServer := dLock.ZKServer
	node,err := dLock.ZKServer.CreateNode(fmt.Sprintf("%s/%s",Root_Node,Pre_Child_Node),"",zk.FlagSequence)
	ChildNodeNum := zkServer.GetChildNodeNumByName(node.Path)
	minChildNodeNum,err := zkServer.GetMinChildNum(Root_Node)
	if err != nil{
		return err
	}
	haveGetLock := false
	for{
		if ChildNodeNum == minChildNodeNum{
			haveGetLock = true
		}
		if haveGetLock {
			fmt.Println("777")
			break
		}

		lastNodeNum ,_:= zkServer.GetLastSortChildNode(ChildNodeNum)
		isExist,_,event_ch,_:= zkServer.Conn.ExistsW(zkServer.GetNodePathByNum(lastNodeNum))
		if !isExist{
			break
		}
		event := <- event_ch
		if event.Type == zk.EventNodeDeleted{
			//上一个节点删除则触发获得锁
			haveGetLock = true
		}
	}
	dLock.NodeNum = ChildNodeNum
	return nil
}

//释放锁
func (dLock *DistributeLock)Release()error{
	nodePath:= dLock.ZKServer.GetNodePathByNum(dLock.NodeNum)
	isExist,_,err := dLock.ZKServer.Conn.Exists(nodePath)
	if err != nil{
		return err
	}
	if isExist{
		err = dLock.ZKServer.Conn.Delete(nodePath,Default_Version)
	}
	return err
}

type ZKServer struct {
	Conn *zk.Conn
	Nodes map[int]Node
	Mtx sync.Mutex
}
//获取最小子节点
func (zkServer *ZKServer)GetMinChildNum(path string)(nodeNum int,err error){
	childNodeNums ,err := zkServer.GetAllChildrenNum(path)
	if err != nil{
		return 0,err
	}

	if len(childNodeNums)<=0{
		return 0,nil
	}

	minChildNum := childNodeNums[0]
	for _,childNodeNum :=range childNodeNums{
		if childNodeNum<minChildNum {
			minChildNum = childNodeNum
		}
	}

	return minChildNum,nil
}

//获取所有子节点编号
func (zkServer * ZKServer) GetAllChildrenNum(path string)([]int,error){
	if zkServer==nil && zkServer.Conn==nil{
		return nil,errors.New("fail to get zk conn")
	}
	childPaths,_,err := zkServer.Conn.Children(path)
	if err != nil{
		return nil,err
	}

	childNodeNums := make([]int,0)
	for _,childPath:=range childPaths{
		segmentPaths := strings.Split(childPath,Pre_Child_Node)
		if len(segmentPaths)!=2{
			return nil,errors.New("path is not correct")
		}
		childNodeNum ,_:= strconv.Atoi(segmentPaths[1])
		childNodeNums = append(childNodeNums,childNodeNum)
	}
	return childNodeNums,nil
}

type Node struct {
	Num int
	Path string
}

//创建节点
func (zkServer *ZKServer)CreateNode(path string,data string,nodeType int)(*Node,error){
	if zkServer==nil && zkServer.Conn==nil{
		return  nil,errors.New("fail to get zk conn")
	}
	name,err := zkServer.Conn.Create(path,[]byte(data),int32(nodeType),zk.WorldACL(zk.PermAll))
	num := zkServer.GetChildNodeNumByName(name)
	node := Node{
		Num:zkServer.GetChildNodeNumByName(name),
		Path:name,
	}
	zkServer.Mtx.Lock()
	zkServer.Nodes[num] = node
	zkServer.Mtx.Unlock()
   return &node,err
}

//删除节点
func(zkServer *ZKServer)DeleteNode(node *Node)error{
	if zkServer==nil && zkServer.Conn==nil{
		return errors.New("fail to get zk conn")
	}
	err := zkServer.Conn.Delete(node.Path,Default_Version)
	fmt.Println(err)
	return err
}

//根据根据名字获取节点编号
func (zkServer *ZKServer)GetChildNodeNumByName(nodeName string)int{
	segmentPaths := strings.Split(nodeName,Pre_Child_Node)
	if len(segmentPaths)!=2{
		return 0
	}
	childNodeNum ,_:= strconv.Atoi(segmentPaths[1])
	return childNodeNum
}

//获取上一个顺序节点
func (zkServer *ZKServer)GetLastSortChildNode(nodeNum int)(int,error){
	childNodeNums ,err:= zkServer.GetAllChildrenNum(Root_Node)
	if err != nil{
		return 0,err
	}

	allLessNodes := make([]int,0) //所有序号大于nodeNum节点序号
	for _,childNodeNum := range childNodeNums{
		if childNodeNum<nodeNum{
			allLessNodes = append(allLessNodes,childNodeNum)
		}
	}

	if len(allLessNodes)<=0{
		return 0,errors.New("get fail")
	}
	lastNodeNum := allLessNodes[0]
	for _,nodeNum:= range allLessNodes{
		if nodeNum>lastNodeNum{
			lastNodeNum = nodeNum
		}
	}
	return lastNodeNum,nil
}

func (zkServer *ZKServer) GetNodePathByNum( nodeNum int)string{
	node := zkServer.Nodes[nodeNum]
	return node.Path
}