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

147 lines
3.7 KiB
Go

package watch
import "context"
// DeepCopyable describes a type that implements the DeepCopy
// method to get a copy of itself that is safe to pass around
// without worrying about receivers modifying the original.
type DeepCopyable[T any] interface {
DeepCopy() T
}
// 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 DeepCopyable[V]] 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 DeepCopyable[V]]() Map[K, V] {
return Map[K, V]{M: make(map[K]watchedVal[V])}
}
// DeepCopy returns a copy of the Map that is safe to be passed
// around without worrying about receivers modifying the original
// or canceling its watches.
func (m Map[K, V]) DeepCopy() Map[K, V] {
dup := make(map[K]watchedVal[V], len(m.M))
for k, v := range m.M {
var val *V
if v.Val != nil {
dc := (*v.Val).DeepCopy()
val = &dc
}
dup[k] = watchedVal[V]{Val: val}
}
return Map[K, V]{M: dup}
}
// 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
}