open-consul/agent/proxycfg/internal/watch/watchmap.go

124 lines
3.0 KiB
Go
Raw Normal View History

package watch
import "context"
// Map safely stores and retrieves values by validating that
// there is a live watch for a key. InitWatch must be called
// to associate a key with its cancel function before any
// Set's are called.
type Map[K comparable, V any] struct {
M map[K]watchedVal[V]
}
type watchedVal[V any] struct {
Val *V
// keeping cancel private has a beneficial side effect:
// copying Map with copystructure.Copy will zero out
// cancel, preventing it from being called by the
// receiver of a proxy config snapshot.
cancel context.CancelFunc
}
func NewMap[K comparable, V any]() Map[K, V] {
return Map[K, V]{M: make(map[K]watchedVal[V])}
}
// InitWatch associates a cancel function with a key,
// allowing Set to be called for the key. The cancel
// function is allowed to be nil.
//
// Any existing data for a key will be cancelled and
// overwritten.
func (m Map[K, V]) InitWatch(key K, cancel func()) {
if _, present := m.M[key]; present {
m.CancelWatch(key)
}
m.M[key] = watchedVal[V]{
cancel: cancel,
}
}
// CancelWatch first calls the cancel function
// associated with the key then deletes the key
// from the map. No-op if key is not present.
func (m Map[K, V]) CancelWatch(key K) {
if entry, ok := m.M[key]; ok {
if entry.cancel != nil {
entry.cancel()
}
delete(m.M, key)
}
}
// IsWatched returns true if InitWatch has been
// called for key and has not been cancelled by
// CancelWatch.
func (m Map[K, V]) IsWatched(key K) bool {
if _, present := m.M[key]; present {
return true
}
return false
}
// Set stores V if K exists in the map.
// No-op if the key never was initialized with InitWatch
// or if the entry got cancelled by CancelWatch.
func (m Map[K, V]) Set(key K, val V) bool {
if entry, ok := m.M[key]; ok {
entry.Val = &val
m.M[key] = entry
return true
}
return false
}
// Get returns the underlying value for a key.
// If an entry has been set, returns (V, true).
// Otherwise, returns the zero value (V, false).
//
// Note that even if InitWatch has been called
// for a key, unless Set has been called this
// function will return false.
func (m Map[K, V]) Get(key K) (V, bool) {
if entry, ok := m.M[key]; ok {
if entry.Val != nil {
return *entry.Val, true
}
}
var empty V
return empty, false
}
func (m Map[K, V]) Len() int {
return len(m.M)
}
// ForEachKey iterates through the map, calling f
// for each iteration. It is up to the caller to
// Get the value and nil-check if required.
// Stops iterating if f returns false.
// Order of iteration is non-deterministic.
func (m Map[K, V]) ForEachKey(f func(K) bool) {
for k := range m.M {
if ok := f(k); !ok {
return
}
}
}
// ForEachKeyE iterates through the map, calling f
// for each iteration. It is up to the caller to
// Get the value and nil-check if required.
// If a non-nil error is returned by f, iterating
// stops and the error is returned.
// Order of iteration is non-deterministic.
func (m Map[K, V]) ForEachKeyE(f func(K) error) error {
for k := range m.M {
if err := f(k); err != nil {
return err
}
}
return nil
}