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

// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: BUSL-1.1

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
}