zk/zk.go

// Package zk is an in-memory ZooKeeper-style coordination store: a tree of
// znodes with versioned data, ephemeral and sequential nodes, sessions, and
// one-shot watches.
package zk

import (
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
)

var (
	ErrNoNode                  = errors.New("zk: node does not exist")
	ErrNodeExists              = errors.New("zk: node already exists")
	ErrBadVersion              = errors.New("zk: version mismatch")
	ErrNotEmpty                = errors.New("zk: node has children")
	ErrNoChildrenForEphemerals = errors.New("zk: ephemeral nodes cannot have children")
	ErrBadPath                 = errors.New("zk: invalid path")
	ErrSessionClosed           = errors.New("zk: session closed")
)

// AnyVersion skips the version check in Set and Delete.
const AnyVersion int32 = -1

type CreateMode int

const (
	Persistent CreateMode = iota
	Ephemeral
	PersistentSequential
	EphemeralSequential
)

func (m CreateMode) ephemeral() bool  { return m == Ephemeral || m == EphemeralSequential }
func (m CreateMode) sequential() bool { return m == PersistentSequential || m == EphemeralSequential }

type EventType int

const (
	NodeCreated EventType = iota + 1
	NodeDeleted
	NodeDataChanged
	NodeChildrenChanged
)

func (t EventType) String() string {
	return [...]string{"", "NodeCreated", "NodeDeleted", "NodeDataChanged", "NodeChildrenChanged"}[t]
}

type Event struct {
	Type EventType
	Path string
}

type Stat struct {
	Czxid          int64 // zxid of the create
	Mzxid          int64 // zxid of the last data change
	Pzxid          int64 // zxid of the last child create/delete
	Version        int32 // data version
	Cversion       int32 // child list version
	EphemeralOwner int64 // session id, 0 if persistent
	NumChildren    int
	DataLength     int
}

type node struct {
	data     []byte
	stat     Stat
	children map[string]*node
	seq      int32 // next sequence suffix for sequential children
}

// Store is the shared tree. Clients talk to it through a Session.
type Store struct {
	mu          sync.Mutex
	root        *node
	zxid        int64
	nextSession int64
	ephemerals  map[int64]map[string]struct{} // session -> paths it owns
	dataWatches map[string][]chan Event       // Get / Exists watches
	childWatch  map[string][]chan Event       // Children watches
}

func New() *Store {
	return &Store{
		root:        &node{children: map[string]*node{}},
		ephemerals:  map[int64]map[string]struct{}{},
		dataWatches: map[string][]chan Event{},
		childWatch:  map[string][]chan Event{},
	}
}

// Session owns ephemeral nodes; closing it deletes them.
type Session struct {
	s  *Store
	id int64
}

func (s *Store) NewSession() *Session {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.nextSession++
	s.ephemerals[s.nextSession] = map[string]struct{}{}
	return &Session{s: s, id: s.nextSession}
}

func (c *Session) ID() int64 { return c.id }

// Close deletes the session's ephemeral nodes, deepest first.
func (c *Session) Close() {
	s := c.s
	s.mu.Lock()
	defer s.mu.Unlock()
	owned, ok := s.ephemerals[c.id]
	if !ok {
		return
	}
	paths := make([]string, 0, len(owned))
	for p := range owned {
		paths = append(paths, p)
	}
	sort.Slice(paths, func(i, j int) bool { return len(paths[i]) > len(paths[j]) })
	for _, p := range paths {
		_ = s.deleteLocked(p, AnyVersion)
	}
	delete(s.ephemerals, c.id)
}

func (c *Session) alive() error {
	if _, ok := c.s.ephemerals[c.id]; !ok {
		return ErrSessionClosed
	}
	return nil
}

func split(path string) ([]string, error) {
	if path == "/" {
		return nil, nil
	}
	if !strings.HasPrefix(path, "/") || strings.HasSuffix(path, "/") {
		return nil, ErrBadPath
	}
	parts := strings.Split(path[1:], "/")
	for _, p := range parts {
		if p == "" || p == "." || p == ".." {
			return nil, ErrBadPath
		}
	}
	return parts, nil
}

func parentOf(path string) (parent, name string) {
	i := strings.LastIndex(path, "/")
	if i == 0 {
		return "/", path[1:]
	}
	return path[:i], path[i+1:]
}

func (s *Store) lookup(path string) (*node, error) {
	parts, err := split(path)
	if err != nil {
		return nil, err
	}
	n := s.root
	for _, p := range parts {
		if n = n.children[p]; n == nil {
			return nil, ErrNoNode
		}
	}
	return n, nil
}

func (s *Store) fire(m map[string][]chan Event, path string, t EventType) {
	for _, ch := range m[path] {
		ch <- Event{Type: t, Path: path} // buffered(1), fires once
	}
	delete(m, path)
}

func (s *Store) watch(m map[string][]chan Event, path string) <-chan Event {
	ch := make(chan Event, 1)
	m[path] = append(m[path], ch)
	return ch
}

func snapshot(n *node) Stat {
	st := n.stat
	st.NumChildren = len(n.children)
	st.DataLength = len(n.data)
	return st
}

// Create makes a node and returns its actual path (which differs from path
// for sequential nodes).
func (c *Session) Create(path string, data []byte, mode CreateMode) (string, error) {
	s := c.s
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := c.alive(); err != nil {
		return "", err
	}
	if path == "/" {
		return "", ErrNodeExists
	}
	if _, err := split(path); err != nil {
		return "", err
	}
	parentPath, name := parentOf(path)
	parent, err := s.lookup(parentPath)
	if err != nil {
		return "", err
	}
	if parent.stat.EphemeralOwner != 0 {
		return "", ErrNoChildrenForEphemerals
	}
	if mode.sequential() {
		name = fmt.Sprintf("%s%010d", name, parent.seq)
		path = strings.TrimSuffix(parentPath, "/") + "/" + name
	}
	if _, ok := parent.children[name]; ok {
		return "", ErrNodeExists
	}

	s.zxid++
	n := &node{
		data:     append([]byte(nil), data...),
		children: map[string]*node{},
		stat:     Stat{Czxid: s.zxid, Mzxid: s.zxid, Pzxid: s.zxid},
	}
	if mode.ephemeral() {
		n.stat.EphemeralOwner = c.id
		s.ephemerals[c.id][path] = struct{}{}
	}
	parent.children[name] = n
	parent.seq++
	parent.stat.Cversion++
	parent.stat.Pzxid = s.zxid

	s.fire(s.dataWatches, path, NodeCreated)
	s.fire(s.childWatch, parentPath, NodeChildrenChanged)
	return path, nil
}

// Get returns a copy of the node's data. If watch is true, the returned
// channel receives one event on the next data change or delete.
func (c *Session) Get(path string, watch bool) ([]byte, Stat, <-chan Event, error) {
	s := c.s
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := c.alive(); err != nil {
		return nil, Stat{}, nil, err
	}
	n, err := s.lookup(path)
	if err != nil {
		return nil, Stat{}, nil, err
	}
	var ch <-chan Event
	if watch {
		ch = s.watch(s.dataWatches, path)
	}
	return append([]byte(nil), n.data...), snapshot(n), ch, nil
}

// Exists reports the node's stat, or nil if absent. A watch set here also
// fires on creation of a missing node.
func (c *Session) Exists(path string, watch bool) (*Stat, <-chan Event, error) {
	s := c.s
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := c.alive(); err != nil {
		return nil, nil, err
	}
	var ch <-chan Event
	n, err := s.lookup(path)
	if err != nil && err != ErrNoNode {
		return nil, nil, err
	}
	if watch {
		ch = s.watch(s.dataWatches, path)
	}
	if n == nil {
		return nil, ch, nil
	}
	st := snapshot(n)
	return &st, ch, nil
}

// Set replaces the data if version matches (or is AnyVersion).
func (c *Session) Set(path string, data []byte, version int32) (Stat, error) {
	s := c.s
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := c.alive(); err != nil {
		return Stat{}, err
	}
	n, err := s.lookup(path)
	if err != nil {
		return Stat{}, err
	}
	if version != AnyVersion && version != n.stat.Version {
		return Stat{}, ErrBadVersion
	}
	s.zxid++
	n.data = append([]byte(nil), data...)
	n.stat.Version++
	n.stat.Mzxid = s.zxid
	s.fire(s.dataWatches, path, NodeDataChanged)
	return snapshot(n), nil
}

// Delete removes a childless node if version matches (or is AnyVersion).
func (c *Session) Delete(path string, version int32) error {
	s := c.s
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := c.alive(); err != nil {
		return err
	}
	return s.deleteLocked(path, version)
}

func (s *Store) deleteLocked(path string, version int32) error {
	if path == "/" {
		return ErrBadPath
	}
	n, err := s.lookup(path)
	if err != nil {
		return err
	}
	if version != AnyVersion && version != n.stat.Version {
		return ErrBadVersion
	}
	if len(n.children) > 0 {
		return ErrNotEmpty
	}
	parentPath, name := parentOf(path)
	parent, _ := s.lookup(parentPath)

	s.zxid++
	delete(parent.children, name)
	parent.stat.Cversion++
	parent.stat.Pzxid = s.zxid
	if owner := n.stat.EphemeralOwner; owner != 0 {
		delete(s.ephemerals[owner], path)
	}

	s.fire(s.dataWatches, path, NodeDeleted)
	s.fire(s.childWatch, path, NodeDeleted)
	s.fire(s.childWatch, parentPath, NodeChildrenChanged)
	return nil
}

// Children returns the sorted child names. If watch is true, the channel
// fires when a child is added/removed or the node itself is deleted.
func (c *Session) Children(path string, watch bool) ([]string, <-chan Event, error) {
	s := c.s
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := c.alive(); err != nil {
		return nil, nil, err
	}
	n, err := s.lookup(path)
	if err != nil {
		return nil, nil, err
	}
	names := make([]string, 0, len(n.children))
	for k := range n.children {
		names = append(names, k)
	}
	sort.Strings(names)
	var ch <-chan Event
	if watch {
		ch = s.watch(s.childWatch, path)
	}
	return names, ch, nil
}