mirror of
1
Fork 0
forgejo/vendor/github.com/ngaut/zkhelper/fakezk.go

520 lines
12 KiB
Go

// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package fakezk is a pretty complete mock implementation of a
// Zookeper connection (see go/zk/zk.Conn). All operations
// work as expected with the exceptions of zk.Conn.ACL and
// zk.Conn.SetACL. zk.Conn.SetACL will succeed, but it is a noop (and
// the ACLs won't be respected). zk.Conn.ACL will panic. It is OK to
// access the connection from multiple goroutines, but the locking is
// very naive (every operation locks the whole connection).
package zkhelper
import (
"bytes"
"encoding/json"
"fmt"
"io/ioutil"
"strings"
"sync"
"time"
"github.com/ngaut/go-zookeeper/zk"
)
type zconn struct {
mu sync.Mutex
root *stat
zxid int64
existWatches map[string][]chan zk.Event
}
func (conn *zconn) getZxid() int64 {
conn.zxid++
return conn.zxid
}
func (conn *zconn) Seq2Str(seq int64) string {
return fmt.Sprintf("%0.10d", seq)
}
// NewConn returns a fake zk.Conn implementation. Data is stored in
// memory, and there's a global connection lock for concurrent access.
func NewConn() Conn {
return &zconn{
root: &stat{
name: "/",
children: make(map[string]*stat),
},
existWatches: make(map[string][]chan zk.Event)}
}
// NewConnFromFile returns a fake zk.Conn implementation, that is seeded
// with the json data extracted from the input file.
func NewConnFromFile(filename string) Conn {
result := &zconn{
root: &stat{
name: "/",
children: make(map[string]*stat),
},
existWatches: make(map[string][]chan zk.Event)}
data, err := ioutil.ReadFile(filename)
if err != nil {
panic(fmt.Errorf("NewConnFromFile failed to read file %v: %v", filename, err))
}
values := make(map[string]interface{})
if err := json.Unmarshal(data, &values); err != nil {
panic(fmt.Errorf("NewConnFromFile failed to json.Unmarshal file %v: %v", filename, err))
}
for k, v := range values {
jv, err := json.Marshal(v)
if err != nil {
panic(fmt.Errorf("NewConnFromFile failed to json.Marshal value %v: %v", k, err))
}
// CreateRecursive will work for a leaf node where the parent
// doesn't exist, but not for a node in the middle of a tree
// that already exists. So have to use 'Set' as a backup.
if _, err := CreateRecursive(result, k, string(jv), 0, nil); err != nil {
if ZkErrorEqual(err, zk.ErrNodeExists) {
_, err = result.Set(k, jv, -1)
}
if err != nil {
panic(fmt.Errorf("NewConnFromFile failed to zk.CreateRecursive value %v: %v", k, err))
}
}
}
return result
}
func (conn *zconn) GetACL(path string) ([]zk.ACL, zk.Stat, error) {
return nil, nil, nil
}
func (conn *zconn) Get(zkPath string) (data []byte, stat zk.Stat, err error) {
conn.mu.Lock()
defer conn.mu.Unlock()
node, _, rest, err := conn.getNode(zkPath, "get")
if err != nil {
return nil, nil, err
}
if len(rest) != 0 {
return nil, nil, zkError(zk.ErrNoNode, "get", zkPath)
}
return []byte(node.content), node, nil
}
func (conn *zconn) GetW(zkPath string) (data []byte, stat zk.Stat, watch <-chan zk.Event, err error) {
conn.mu.Lock()
defer conn.mu.Unlock()
node, _, rest, err := conn.getNode(zkPath, "getw")
if err != nil {
return nil, nil, nil, err
}
if len(rest) != 0 {
return nil, nil, nil, zkError(zk.ErrNoNode, "getw", zkPath)
}
c := make(chan zk.Event, 1)
node.changeWatches = append(node.changeWatches, c)
return []byte(node.content), node, c, nil
}
func (conn *zconn) Children(zkPath string) (children []string, stat zk.Stat, err error) {
conn.mu.Lock()
defer conn.mu.Unlock()
//println("Children:", conn.String())
node, _, rest, err := conn.getNode(zkPath, "children")
if err != nil {
return nil, nil, err
}
if len(rest) != 0 {
return nil, nil, zkError(zk.ErrNoNode, "children", zkPath)
}
for name := range node.children {
children = append(children, name)
}
return children, node, nil
}
func (conn *zconn) ChildrenW(zkPath string) (children []string, stat zk.Stat, watch <-chan zk.Event, err error) {
conn.mu.Lock()
defer conn.mu.Unlock()
//println("ChildrenW:", conn.String())
node, _, rest, err := conn.getNode(zkPath, "childrenw")
if err != nil {
return nil, nil, nil, err
}
if len(rest) != 0 {
return nil, nil, nil, zkError(zk.ErrNoNode, "childrenw", zkPath)
}
c := make(chan zk.Event, 1)
node.childrenWatches = append(node.childrenWatches, c)
for name := range node.children {
children = append(children, name)
}
return children, node, c, nil
}
func (conn *zconn) Exists(zkPath string) (exist bool, stat zk.Stat, err error) {
// FIXME(szopa): if the path is bad, Op will be "get."
exist = false
_, stat, err = conn.Get(zkPath)
if err != nil {
if ZkErrorEqual(err, zk.ErrNoNode) {
err = nil
}
} else {
exist = true
}
return exist, stat, err
}
func (conn *zconn) ExistsW(zkPath string) (exist bool, stat zk.Stat, watch <-chan zk.Event, err error) {
conn.mu.Lock()
defer conn.mu.Unlock()
exist = false
c := make(chan zk.Event, 1)
node, _, rest, err := conn.getNode(zkPath, "existsw")
if err != nil {
return exist, nil, nil, err
}
if len(rest) != 0 {
watches, ok := conn.existWatches[zkPath]
if !ok {
watches = make([]chan zk.Event, 0)
conn.existWatches[zkPath] = watches
}
conn.existWatches[zkPath] = append(watches, c)
return exist, nil, c, nil
}
exist = true
node.existWatches = append(node.existWatches, c)
return exist, node, c, nil
}
func (conn *zconn) Create(zkPath string, value []byte, flags int32, aclv []zk.ACL) (zkPathCreated string, err error) {
conn.mu.Lock()
defer conn.mu.Unlock()
node, _, rest, err := conn.getNode(zkPath, "create")
if err != nil {
return "", err
}
if len(rest) == 0 {
return "", zkError(zk.ErrNodeExists, "create", zkPath)
}
if len(rest) > 1 {
return "", zkError(zk.ErrNoNode, "create", zkPath)
}
zxid := conn.getZxid()
name := rest[0]
if (flags & zk.FlagSequence) != 0 {
sequence := node.nextSequence()
name += sequence
zkPath = zkPath + sequence
}
stat := &stat{
name: name,
content: string(value),
children: make(map[string]*stat),
acl: aclv,
mtime: time.Now(),
ctime: time.Now(),
czxid: zxid,
mzxid: zxid,
existWatches: make([]chan zk.Event, 0),
}
node.children[name] = stat
event := zk.Event{
Type: zk.EventNodeCreated,
Path: zkPath,
State: zk.StateConnected,
}
if watches, ok := conn.existWatches[zkPath]; ok {
delete(conn.existWatches, zkPath)
for _, watch := range watches {
watch <- event
}
}
childrenEvent := zk.Event{
Type: zk.EventNodeChildrenChanged,
Path: zkPath,
State: zk.StateConnected,
}
for _, watch := range node.childrenWatches {
watch <- childrenEvent
close(watch)
}
node.childrenWatches = nil
node.cversion++
return zkPath, nil
}
func (conn *zconn) Set(zkPath string, value []byte, version int32) (stat zk.Stat, err error) {
conn.mu.Lock()
defer conn.mu.Unlock()
node, _, rest, err := conn.getNode(zkPath, "set")
if err != nil {
return nil, err
}
if len(rest) != 0 {
return nil, zkError(zk.ErrNoNode, "set", zkPath)
}
if version != -1 && node.version != int(version) {
return nil, zkError(zk.ErrBadVersion, "set", zkPath)
}
node.content = string(value)
node.version++
for _, watch := range node.changeWatches {
watch <- zk.Event{
Type: zk.EventNodeDataChanged,
Path: zkPath,
State: zk.StateConnected,
}
}
node.changeWatches = nil
return node, nil
}
func (conn *zconn) Delete(zkPath string, version int32) (err error) {
conn.mu.Lock()
defer conn.mu.Unlock()
node, parent, rest, err := conn.getNode(zkPath, "delete")
if err != nil {
return err
}
if len(rest) > 0 {
return zkError(zk.ErrNoNode, "delete", zkPath)
}
if len(node.children) > 0 {
return zkError(zk.ErrNotEmpty, "delete", zkPath)
}
delete(parent.children, node.name)
event := zk.Event{
Type: zk.EventNodeDeleted,
Path: zkPath,
State: zk.StateConnected,
}
for _, watch := range node.existWatches {
watch <- event
}
for _, watch := range node.changeWatches {
watch <- event
}
node.existWatches = nil
node.changeWatches = nil
childrenEvent := zk.Event{
Type: zk.EventNodeChildrenChanged,
Path: zkPath,
State: zk.StateConnected}
for _, watch := range parent.childrenWatches {
watch <- childrenEvent
}
return nil
}
func (conn *zconn) Close() {
conn.mu.Lock()
defer conn.mu.Unlock()
for _, watches := range conn.existWatches {
for _, c := range watches {
close(c)
}
}
conn.root.closeAllWatches()
}
/*
func (conn *zconn) RetryChange(path string, flags int, acl []zk.ACL, changeFunc zk.ChangeFunc) error {
for {
oldValue, oldStat, err := conn.Get(path)
if err != nil && !ZkErrorEqual(err, zk.ErrNoNode) {
return err
}
newValue, err := changeFunc(oldValue, oldStat)
if err != nil {
return err
}
if oldStat == nil {
_, err := conn.Create(path, newValue, flags, acl)
if err == nil || !ZkErrorEqual(err, zk.ZNODEEXISTS) {
return err
}
continue
}
if newValue == oldValue {
return nil // Nothing to do.
}
_, err = conn.Set(path, newValue, oldStat.Version())
if err == nil || !ZkErrorEqual(err, zk.ZBADVERSION) && !ZkErrorEqual(err, zk.ErrNoNode) {
return err
}
}
}
*/
func (conn *zconn) SetACL(zkPath string, aclv []zk.ACL, version int32) (zk.Stat, error) {
return nil, nil
}
func (conn *zconn) getNode(zkPath string, op string) (node *stat, parent *stat, rest []string, err error) {
// FIXME(szopa): Make sure the path starts with /.
parts := strings.Split(zkPath, "/")
if parts[0] != "" {
//todo: fix this, error bad arguments
return nil, nil, nil, zkError(zk.ErrUnknown, op, zkPath)
}
elements := parts[1:]
parent = nil
current := conn.root
for i, el := range elements {
candidateParent := current
candidate, ok := current.children[el]
if !ok {
return current, parent, elements[i:], nil
}
current, parent = candidate, candidateParent
}
return current, parent, []string{}, nil
}
type ZkError struct {
Code error
Op string
Path string
}
func (ze *ZkError) Error() string {
return ze.Code.Error()
}
// zkError creates an appropriate error return from
// a ZooKeeper status
func zkError(code error, op, path string) error {
return &ZkError{
Op: op,
Code: code,
Path: path,
}
}
type stat struct {
name string
content string
children map[string]*stat
acl []zk.ACL
mtime time.Time
ctime time.Time
czxid int64
mzxid int64
pzxid int64
version int
cversion int
aversion int
sequence int
existWatches []chan zk.Event
changeWatches []chan zk.Event
childrenWatches []chan zk.Event
}
func (st stat) closeAllWatches() {
for _, c := range st.existWatches {
close(c)
}
for _, c := range st.changeWatches {
close(c)
}
for _, c := range st.childrenWatches {
close(c)
}
for _, child := range st.children {
child.closeAllWatches()
}
}
func (st stat) Czxid() int64 {
return st.czxid
}
func (st stat) Mzxid() int64 {
return st.mzxid
}
func (st stat) CTime() time.Time {
return st.ctime
}
func (st stat) MTime() time.Time {
return st.mtime
}
func (st stat) Version() int {
return st.version
}
func (st stat) CVersion() int {
return st.cversion
}
func (st stat) AVersion() int {
return st.aversion
}
func (st stat) EphemeralOwner() int64 {
return 0
}
func (st stat) DataLength() int {
return len(st.content)
}
func (st stat) NumChildren() int {
return len(st.children)
}
func (st stat) Pzxid() int64 {
return st.pzxid
}
func (st *stat) nextSequence() string {
st.sequence++
return fmt.Sprintf("%010d", st.sequence)
}
func (st stat) fprintRecursive(level int, buf *bytes.Buffer) {
start := strings.Repeat(" ", level)
fmt.Fprintf(buf, "%v-%v:\n", start, st.name)
if st.content != "" {
fmt.Fprintf(buf, "%v content: %q\n\n", start, st.content)
}
if len(st.children) > 0 {
for _, child := range st.children {
child.fprintRecursive(level+1, buf)
}
}
}
func (conn *zconn) String() string {
b := new(bytes.Buffer)
conn.root.fprintRecursive(0, b)
return b.String()
}