|
|
- // 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()
- }
|