157 lines
3.6 KiB
Go
157 lines
3.6 KiB
Go
package state
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/armon/go-radix"
|
|
)
|
|
|
|
// Watch is the external interface that's common to all the different flavors.
|
|
type Watch interface {
|
|
// Wait registers the given channel and calls it back when the watch
|
|
// fires.
|
|
Wait(notifyCh chan struct{})
|
|
|
|
// Clear deregisters the given channel.
|
|
Clear(notifyCh chan struct{})
|
|
}
|
|
|
|
// FullTableWatch implements a single notify group for a table.
|
|
type FullTableWatch struct {
|
|
group NotifyGroup
|
|
}
|
|
|
|
// NewFullTableWatch returns a new full table watch.
|
|
func NewFullTableWatch() *FullTableWatch {
|
|
return &FullTableWatch{}
|
|
}
|
|
|
|
// See Watch.
|
|
func (w *FullTableWatch) Wait(notifyCh chan struct{}) {
|
|
w.group.Wait(notifyCh)
|
|
}
|
|
|
|
// See Watch.
|
|
func (w *FullTableWatch) Clear(notifyCh chan struct{}) {
|
|
w.group.Clear(notifyCh)
|
|
}
|
|
|
|
// Notify wakes up all the watchers registered for this table.
|
|
func (w *FullTableWatch) Notify() {
|
|
w.group.Notify()
|
|
}
|
|
|
|
// DumbWatchManager is a wrapper that allows nested code to arm full table
|
|
// watches multiple times but fire them only once.
|
|
type DumbWatchManager struct {
|
|
tableWatches map[string]*FullTableWatch
|
|
armed map[string]bool
|
|
}
|
|
|
|
// NewDumbWatchManager returns a new dumb watch manager.
|
|
func NewDumbWatchManager(tableWatches map[string]*FullTableWatch) *DumbWatchManager {
|
|
return &DumbWatchManager{
|
|
tableWatches: tableWatches,
|
|
armed: make(map[string]bool),
|
|
}
|
|
}
|
|
|
|
// Arm arms the given table's watch.
|
|
func (d *DumbWatchManager) Arm(table string) {
|
|
if _, ok := d.armed[table]; !ok {
|
|
d.armed[table] = true
|
|
}
|
|
}
|
|
|
|
// Notify fires watches for all the armed tables.
|
|
func (d *DumbWatchManager) Notify() {
|
|
for table, _ := range d.armed {
|
|
d.tableWatches[table].Notify()
|
|
}
|
|
}
|
|
|
|
// PrefixWatch maintains a notify group for each prefix, allowing for much more
|
|
// fine-grained watches.
|
|
type PrefixWatch struct {
|
|
// watches has the set of notify groups, organized by prefix.
|
|
watches *radix.Tree
|
|
|
|
// lock protects the watches tree.
|
|
lock sync.Mutex
|
|
}
|
|
|
|
// NewPrefixWatch returns a new prefix watch.
|
|
func NewPrefixWatch() *PrefixWatch {
|
|
return &PrefixWatch{watches: radix.New()}
|
|
}
|
|
|
|
// GetSubwatch returns the notify group for the given prefix.
|
|
func (w *PrefixWatch) GetSubwatch(prefix string) *NotifyGroup {
|
|
w.lock.Lock()
|
|
defer w.lock.Unlock()
|
|
|
|
if raw, ok := w.watches.Get(prefix); ok {
|
|
return raw.(*NotifyGroup)
|
|
}
|
|
|
|
group := &NotifyGroup{}
|
|
w.watches.Insert(prefix, group)
|
|
return group
|
|
}
|
|
|
|
// Notify wakes up all the watchers associated with the given prefix. If subtree
|
|
// is true then we will also notify all the tree under the prefix, such as when
|
|
// a key is being deleted.
|
|
func (w *PrefixWatch) Notify(prefix string, subtree bool) {
|
|
w.lock.Lock()
|
|
defer w.lock.Unlock()
|
|
|
|
var cleanup []string
|
|
fn := func(k string, v interface{}) bool {
|
|
group := v.(*NotifyGroup)
|
|
group.Notify()
|
|
if k != "" {
|
|
cleanup = append(cleanup, k)
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Invoke any watcher on the path downward to the key.
|
|
w.watches.WalkPath(prefix, fn)
|
|
|
|
// If the entire prefix may be affected (e.g. delete tree),
|
|
// invoke the entire prefix.
|
|
if subtree {
|
|
w.watches.WalkPrefix(prefix, fn)
|
|
}
|
|
|
|
// Delete the old notify groups.
|
|
for i := len(cleanup) - 1; i >= 0; i-- {
|
|
w.watches.Delete(cleanup[i])
|
|
}
|
|
}
|
|
|
|
// MultiWatch wraps several watches and allows any of them to trigger the
|
|
// caller.
|
|
type MultiWatch struct {
|
|
watches []Watch
|
|
}
|
|
|
|
func NewMultiWatch(watches ...Watch) *MultiWatch {
|
|
return &MultiWatch{watches: watches}
|
|
}
|
|
|
|
// See Watch.
|
|
func (w *MultiWatch) Wait(notifyCh chan struct{}) {
|
|
for _, watch := range w.watches {
|
|
watch.Wait(notifyCh)
|
|
}
|
|
}
|
|
|
|
// See Watch.
|
|
func (w *MultiWatch) Clear(notifyCh chan struct{}) {
|
|
for _, watch := range w.watches {
|
|
watch.Clear(notifyCh)
|
|
}
|
|
}
|