watcher.go

// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package hcat

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/hashicorp/hcat/dep"
	"github.com/hashicorp/hcat/events"
	idep "github.com/hashicorp/hcat/internal/dependency"
	"github.com/pkg/errors"
)

// dataBufferSize is the default number of views to process in a batch.
const DefaultDataBufferSize = 2048

// standard error returned when you try to register the same notifier twice
var ErrRegistry = fmt.Errorf("duplicate watcher registry entry")
var ErrStop = fmt.Errorf("Stop")

// RetryFunc defines the function type used to determine how many and how often
// to retry calls to the external services.
type RetryFunc func(int) (bool, time.Duration)

// Cacher defines the interface required by the watcher for caching data
// retreived from external services. It is implemented by Store.
type Cacher interface {
	Save(key string, value interface{})
	Recall(key string) (value interface{}, found bool)
	Delete(key string)
	Reset()
}

// Make sure we implement Collector
var _ Collector = (*Watcher)(nil)

// Watcher is a manager for views that poll external sources for data.
type Watcher struct {
	// clients is the collection of API clients to talk to upstreams.
	clients Looker
	// cache stores the data fetched from remote sources
	cache Cacher
	// event holds the callback for event processing
	event events.EventHandler

	// dataCh is the chan where Views will be published.
	dataCh chan *view
	// errCh is the chan where any errors will be published.
	errCh chan error

	// stopCh is the chan used internally to notify of Stop calls
	stopCh drainableChan
	// waitingCh is used internally to test when Wait is waiting
	waitingCh chan struct{}

	// tracker tracks template<->dependencies (see bottom of this file)
	tracker *tracker

	// bufferTimers manages the buffer period per notifier to accumulate
	// dependency changes.
	bufferTimers *timers
	// bufferTrigger is the notification channel for template IDs that have
	// completed their active buffer period.
	bufferTrigger chan string

	// Consul related
	retryFuncConsul RetryFunc
	// blockWaitTime is how long to block on consul's blocking queries
	blockWaitTime time.Duration
	// maxStale passed to consul to control staleness
	maxStale time.Duration

	// Vault related
	retryFuncVault RetryFunc
	// defaultLease is used for non-renewable leases when secret has no lease
	defaultLease time.Duration
}

type WatcherInput struct {
	// Clients is the client set to communicate with upstreams.
	Clients Looker
	// Cache is the Cacher for caching watched values
	Cache Cacher

	// EventHandler takes the callback for event processing
	EventHandler events.EventHandler

	// Optional Vault specific parameters
	// Default non-renewable secret duration
	VaultDefaultLease time.Duration
	// RetryFun for Vault
	VaultRetryFunc RetryFunc

	// Optional Consul specific parameters
	// MaxStale is the max time Consul will return a stale value.
	ConsulMaxStale time.Duration
	// BlockWait is amount of time Consul will block on a query.
	ConsulBlockWait time.Duration
	// RetryFun for Consul
	ConsulRetryFunc RetryFunc

	// Override the default data buffer size (for testing)
	DataBufferSize *int
}

type drainableChan chan struct{}

func (s drainableChan) drain() {
	for {
		select {
		case <-s:
		default:
			return
		}
	}
}

// NewWatcher creates a new watcher using the given API client.
func NewWatcher(i WatcherInput) *Watcher {
	cache := i.Cache
	if cache == nil {
		cache = NewStore()
	}
	clients := i.Clients
	if clients == nil {
		clients = DefaultClientSet()
	}
	eventHandler := i.EventHandler
	if eventHandler == nil {
		eventHandler = func(events.Event) {}
	}

	var dataBufferSize int
	switch i.DataBufferSize {
	case nil:
		dataBufferSize = DefaultDataBufferSize
	default:
		dataBufferSize = *i.DataBufferSize
	}

	bufferTriggerCh := make(chan string, dataBufferSize/2)
	w := &Watcher{
		clients:         clients,
		cache:           cache,
		event:           eventHandler,
		dataCh:          make(chan *view, dataBufferSize),
		errCh:           make(chan error),
		waitingCh:       make(chan struct{}, 1),
		stopCh:          make(chan struct{}, 1),
		tracker:         newTracker(),
		bufferTrigger:   bufferTriggerCh,
		bufferTimers:    newTimers(),
		retryFuncConsul: i.ConsulRetryFunc,
		maxStale:        i.ConsulMaxStale,
		blockWaitTime:   i.ConsulBlockWait,
		retryFuncVault:  i.VaultRetryFunc,
		defaultLease:    i.VaultDefaultLease,
	}

	go w.bufferTimers.Run(bufferTriggerCh)

	return w
}

const vaultTokenDummyTemplateID = "dummy.watcher.vault-token.id"

// WatchVaultToken takes a vault token and watches it to keep it updated.
// This is a specialized method as this token can be required without being in
// a template. I hope to generalize this idea so you can watch arbitrary
// dependencies in the future.
func (w *Watcher) WatchVaultToken(token string) error {
	// Start a watcher for the Vault renew if that config was specified
	if token != "" {
		vt, err := idep.NewVaultTokenQuery(token)
		if err != nil {
			return errors.Wrap(err, "watcher")
		}
		// fakeNotifier is defined near end of file
		n := fakeNotifier(vaultTokenDummyTemplateID)
		if err := w.Register(n); err != nil {
			return err
		}
		w.Track(n, vt)
		w.Poll(vt)
	}
	return nil
}

// Clients returns the Looker/ClientSet to give easy access to the clients
// after initial setup.
func (w *Watcher) Clients() Looker {
	return w.clients
}

// WaitCh returns an error channel and runs Wait sending the result down
// the channel. Useful for when you need to use Wait in a select block.
func (w *Watcher) WaitCh(ctx context.Context) <-chan error {
	errCh := make(chan error)
	go func() {
		errCh <- w.Wait(ctx)
	}()
	return errCh
}

// Wait blocks until new a watched value changes or until context is closed
// or exceeds its deadline.
func (w *Watcher) Wait(ctx context.Context) error {
	w.stopCh.drain() // in case Stop was already called

	for {
		notifiers, err := w.wait(ctx)
		switch err {
		case nil: // continue
		case context.DeadlineExceeded, ErrStop:
			return nil
		default:
			return err
		}

		if len(notifiers) > 0 {
			return nil
		}
	}
}

// Watch will continuously watch for data updates and send template IDs to
// the provided channel which templates have changes to be rendered. Useful
// for when caller wants to process templates asynchronously. Only one Watch
// should be called at given time and should not be called with Wait.
func (w *Watcher) Watch(ctx context.Context, notifierCh chan string) error {
	w.stopCh.drain()

	// send waiting notification, only used for testing
	select {
	case w.waitingCh <- struct{}{}:
	default:
	}

	for {
		notifiers, err := w.wait(ctx)
		switch err {
		case nil: // continue
		case context.DeadlineExceeded, ErrStop:
			return nil
		default:
			return err
		}

		for nID := range notifiers {
			notifierCh <- nID
		}
	}
}

// wait implements the watchers select/data update blocking call along with the
// cache update and buffering check. It returns a list (as a map to dedup) of
// the notifiers that need to be notified (handled ball caller).
func (w *Watcher) wait(ctx context.Context) (map[string]struct{}, error) {
	type notifierMap map[string]struct{}
	empty := struct{}{}
	// send waiting notification, only used for testing
	select {
	case w.waitingCh <- empty:
	default:
	}

	dataUpdate := func(v *view, notifiers notifierMap) notifierMap {
		id := v.ID()
		w.cache.Save(id, v.Data())
		for _, n := range w.tracker.notifiersFor(v) {
			if n.Notify(v.Data()) && !w.Buffering(n) {
				notifiers[n.ID()] = empty
			}
		}
		return notifiers
	}

	// use map to dedup results
	notifiers := make(notifierMap)

	select {
	case v := <-w.dataCh:
		notifiers = dataUpdate(v, notifiers)
		// this case (<-w.dataCh) only happens if there is new/udpated data
		for drain := true; drain; {
			select {
			case v := <-w.dataCh:
				notifiers = dataUpdate(v, notifiers)
			default:
				drain = false
			}
		}
		return notifiers, nil
	case nID := <-w.bufferTrigger:
		notifiers[nID] = empty
		// A template is now ready to be rendered, though there might be a
		// few ready around the same time if they have the same dependencies.
		// Drain the channel similar for the dataCh above.
		for drain := true; drain; {
			select {
			case nID := <-w.bufferTrigger:
				notifiers[nID] = empty
			default:
				drain = false
			}
		}
		return notifiers, nil
	case <-w.stopCh:
		return nil, ErrStop

	case err := <-w.errCh:
		// Push the error back up the stack
		return nil, err

	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// Buffering sets the template to activate buffer and accumulate changes for a
// period. If the template has not been initalized or a buffer period is not
// configured for the template, it will skip the buffering.
// period.
func (w *Watcher) Buffering(n Notifier) bool {
	// to enable buffering only after the notifier/template has been completely
	// evaluated/rendered, you need to add '&& w.Complete(n)' to the line below
	if w.bufferTimers.isBuffering(n.ID()) {
		w.bufferTimers.tick(n.ID())
		return true
	}
	return false
}

// BufferReset resets an active buffer period to inactive.
func (w *Watcher) BufferReset(n Notifier) {
	w.bufferTimers.Reset(n.ID())
}

// Register registers one or more Notifiers with the Watcher for future use.
// Trying to register the same Notifier twice will result in an error and none
// of the Notifiers will be registered (all or nothing).
// Trying to use a Notifier without Registering it will result in a *panic*.
func (w *Watcher) Register(ns ...Notifier) error {
	return w.tracker.registerNotifiers(ns...)
}

// Deregister de-registers one or more Notifiers from the Watcher.
func (w *Watcher) Deregister(ns ...Notifier) {
	w.tracker.deregisterNotifiers(ns...)

	for _, n := range ns {
		w.tracker.markForSweep(n)
		w.tracker.sweep(n, w.cache)
	}
}

// Track is used to add dependencies to be monitored by the watcher. It sets
// everything up but stops short of running the polling, waiting for an
// explicit start (see Poll below).
// It calls Register as a convenience, but ignores the returned error so it can
// be used with already Registered Notifiers.
func (w *Watcher) Track(n Notifier, d dep.Dependency) {
	w.Register(n)
	w.track(n, d)
}

// track is the private form of Track that returns the new view.
// Returned view is useful internally and for testing.
// Private as we don't want `view` public at this point.
func (w *Watcher) track(n Notifier, d dep.Dependency) *view {
	w.tracker.markInUse(n, d)
	if v, ok := w.tracker.lookup(n, d); ok {
		return v
	}
	// Choose the correct retry function based off of the dependency's type.
	// NOTE: I would like to abstract this part out to not have type specific
	//       things embedded in general code.
	var retryFunc RetryFunc
	switch d.(type) {
	case idep.ConsulType:
		retryFunc = w.retryFuncConsul
	case idep.VaultType:
		retryFunc = w.retryFuncVault
	}

	v := newView(&newViewInput{
		Dependency:        d,
		Clients:           w.clients,
		EventHandler:      w.event,
		MaxStale:          w.maxStale,
		BlockWaitTime:     w.blockWaitTime,
		RetryFunc:         retryFunc,
		VaultDefaultLease: w.defaultLease,
	})
	w.event(events.TrackStart{ID: v.ID()})
	w.tracker.add(v, n)
	return v
}

// Poll starts any/all polling as needed.
// It is idepotent.
// If nothing is passed it checks all views (dependencies).
func (w *Watcher) Poll(deps ...dep.Dependency) {
	if len(deps) == 0 {
		for _, v := range w.tracker.views {
			deps = append(deps, v.Dependency())
		}
	}
	for _, d := range deps {
		if v := w.tracker.view(d.ID()); v != nil {
			// view.poll checks if it is already polling
			go v.poll(w.dataCh, w.errCh)
		}
	}
}

// Recaller returns a Recaller (function) that wraps the Store (cache)
// to enable tracking dependencies on the Watcher.
func (w *Watcher) Recaller(n Notifier) Recaller {
	return func(dep dep.Dependency) (interface{}, bool) {
		w.track(n, dep)
		data, ok := w.cache.Recall(dep.ID())
		switch {
		case ok:
			w.tracker.cacheAccessed(n, dep)
		default:
			w.Poll(dep)
		}
		return data, ok
	}
}

// Complete checks if all values in use have been fetched.
func (w *Watcher) Complete(n Notifier) bool {
	return w.tracker.complete(n)
}

// Mark-n-Sweep garbage-collector-like cleaning of views that are no in use.
// Stops the (garbage) views and removes all references.
// Should be used before/after the code that uses the dependencies (eg. template).
//
// Mark's all tracked dependencies as being *not* in use.
func (w *Watcher) MarkForSweep(notifier IDer) {
	w.tracker.markForSweep(notifier)
}

// Sweeps (stop and dereference) all views for dependencies marked as *not* in use.
func (w *Watcher) Sweep(notifier IDer) {
	w.tracker.sweep(notifier, w.cache)
}

// SetBufferPeriod sets a buffer period to accumulate dependency changes for
// a template.
func (w *Watcher) SetBufferPeriod(min, max time.Duration, notifierIDs ...string) {
	for _, id := range notifierIDs {
		w.bufferTimers.Add(min, max, id)
	}
}

// ID here is to meet the IDer interface and be used with events/logging
func (w *Watcher) ID() string {
	return fmt.Sprintf("watcher (%p)", w)
}

// Stop halts this watcher and any currently polling views immediately. If a
// view was in the middle of a poll, no data will be returned.
func (w *Watcher) Stop() {
	w.event(events.Trace{ID: w.ID(), Message: "stopping watcher"})
	w.bufferTimers.Stop()

	w.tracker.stopViews()

	w.stopCh.drain() // So calling Stop twice doesn't block
	w.stopCh <- struct{}{}

	// Empty cache
	if w.cache != nil {
		w.cache.Reset()
	}

	// Close any idle TCP connections
	if w.clients != nil {
		w.clients.Stop()
	}
}

// Size returns the number of views this watcher is watching.
func (w *Watcher) Size() int {
	return w.tracker.viewCount()
}

// Watching determines if the given dependency (id) is being watched.
func (w *Watcher) Watching(id string) bool {
	v := w.tracker.view(id)
	return (v != nil)
}

// view is a convenience function for accessing stored views by id
// note that dependency IDs and their corresponding view IDs are identical
func (w *Watcher) view(id string) *view {
	return w.tracker.view(id)
}

///////////////////////////////////////////////////////////////////////////
// internal structure used to track template <-> dependencies relationships

func newTracker() *tracker {
	return &tracker{
		tracked:   make([]trackedPair, 0, 8),
		views:     make(map[string]*view),
		notifiers: make(map[string]Notifier),
	}
}

// 1 view/notifier pair. Think many-2-many RDBMS table with annotations.
type trackedPair struct {
	// view: id of view watched, notify: id of notifier (eg. template)
	view, notify string
	// mark(-n-sweep) flag gets set off pre-render and back on at use
	mark bool
	// cacheAccessed is set when recalled from cache the first time
	cacheAccessed bool
}

// markUsed sets as mark (=true) and returns new pair to keep as value
func (tp trackedPair) markInUse() trackedPair {
	tp.mark = true
	return tp
}

// clearUse clears mark (=false) and returns new pair to keep as value
func (tp trackedPair) clearInUse() trackedPair {
	tp.mark = false
	return tp
}

// markDataUsed sets fetched data as being used
// this is only important for new dependencies, so only needs to be set once
func (tp trackedPair) markCacheAccessed() trackedPair {
	tp.cacheAccessed = true
	return tp
}

// IDer an interface that supports and ID
type IDer interface {
	ID() string
}

// Notifier indicates support for notifications
type Notifier interface {
	IDer
	Notify(interface{}) bool
}

// If performance of looping through tracked gets to be to much build 2 indexes
// of views/notifiers to their trackedPair entries and use that to accel lookups.
// It will require updating though, and complicates things. So wait.

type tracker struct {
	sync.Mutex
	// think in terms of a many-2-many DB relationship
	tracked []trackedPair
	// viewID -> view
	views map[string]*view
	// stringID -> Notifier (stringID is usually template-id)
	notifiers map[string]Notifier
}

// cacheAccessed records that the fetched data was used at least once
func (t *tracker) cacheAccessed(notifier IDer, d dep.Dependency) {
	notifierID, depID := notifier.ID(), d.ID()
	t.Lock()
	defer t.Unlock()
	for i, tp := range t.tracked {
		if !tp.cacheAccessed && tp.view == depID && tp.notify == notifierID {
			t.tracked[i] = tp.markCacheAccessed()
		}
	}
}

// viewCount returns the number of views watched
func (t *tracker) viewCount() int {
	t.Lock()
	defer t.Unlock()
	return len(t.views)
}

// registerNotifiers adds the notifiers to those tracked, it returns an error
// if a notifier (indexed by n.ID()) has already been registered. If an error
// occurs none of the notifiers will be added (all or nothing).
func (t *tracker) registerNotifiers(ns ...Notifier) error {
	t.Lock()
	defer t.Unlock()
	for _, n := range ns {
		if _, ok := t.notifiers[n.ID()]; ok {
			return ErrRegistry
		}
	}
	for _, n := range ns {
		t.notifiers[n.ID()] = n
	}
	return nil
}

// deregisterNotifiers removes the notifiers from those tracked
func (t *tracker) deregisterNotifiers(ns ...Notifier) {
	t.Lock()
	defer t.Unlock()
	for _, n := range ns {
		// Delete from notifier map
		delete(t.notifiers, n.ID())
	}
}

// notifierTracked tests if a registered notifier has been paired with a
// dependency (a tracked_pair added) and thus used at least once
func (t *tracker) notifierTracked(n Notifier) bool {
	t.Lock()
	defer t.Unlock()
	for _, tp := range t.tracked {
		if tp.notify == n.ID() {
			return true
		}
	}
	return false
}

// lookup returns the view and true, or nil and false
// true is returned if the notifier and depencency match a tracked pair
// returns the view as it is the 1 thing that you don't have yet
// note that a view's and dependency's IDs are interchangeable (identical)
func (t *tracker) lookup(notifier IDer, d dep.Dependency) (*view, bool) {
	notifierID, depID := notifier.ID(), d.ID()
	t.Lock()
	defer t.Unlock()
	for _, tp := range t.tracked {
		if tp.view == depID && tp.notify == notifierID {
			return t.views[tp.view], true
		}
	}
	return nil, false
}

// view returns the view (or nil)
// note that a view's and dependency's IDs are interchangeable (identical)
func (t *tracker) view(viewID string) *view {
	t.Lock()
	defer t.Unlock()
	return t.views[viewID]
}

// adds new tracked entry
func (t *tracker) add(v *view, n Notifier) {
	t.Lock()
	defer t.Unlock()
	if _, ok := t.views[v.ID()]; !ok {
		t.views[v.ID()] = v
	}
	if _, ok := t.notifiers[n.ID()]; !ok {
		panic("attempt to use an unregistered notifier")
	}
	t.tracked = append(t.tracked,
		trackedPair{view: v.ID(), notify: n.ID(), mark: true})
}

// Marks all trackedPairs w/ a view as having been used
func (t *tracker) markInUse(notifier IDer, d dep.Dependency) {
	notifierID, depID := notifier.ID(), d.ID()
	t.Lock()
	defer t.Unlock()
	for i, tp := range t.tracked {
		if tp.view == depID && tp.notify == notifierID {
			t.tracked[i] = tp.markInUse()
		}
	}
}

// stop all view from polling/watching
func (t *tracker) stopViews() {
	t.Lock()
	defer t.Unlock()
	for id, view := range t.views {
		delete(t.views, id)
		if view == nil {
			continue
		}
		view.stop()
	}
}

// Return all Notifiers for a view
func (t *tracker) notifiersFor(view IDer) []Notifier {
	viewID := view.ID()
	results := make([]Notifier, 0, 8)
	t.Lock()
	defer t.Unlock()
	for _, tp := range t.tracked {
		if tp.view == viewID {
			results = append(results, t.notifiers[tp.notify])
		}
	}
	return results
}

// complete returns true if every dependency used has been initialized
// ie. it returns true if all values have been fetched
func (t *tracker) complete(notifier IDer) bool {
	t.Lock()
	defer t.Unlock()
	for _, tp := range t.tracked {
		thisNotifier := tp.notify == notifier.ID()
		cacheNotAccessed := !tp.cacheAccessed
		if thisNotifier && cacheNotAccessed {
			return false
		}
	}
	return true
}

// Clean out un-used trackedPair entries and their views (if the last use).
// Checks based on passed in notifier, ignores others.
//
// mark all pairs used by this notifier not used (used with sweep)
func (t *tracker) markForSweep(notifier IDer) {
	t.Lock()
	defer t.Unlock()
	for i, tp := range t.tracked {
		if tp.notify == notifier.ID() && tp.mark {
			t.tracked[i] = tp.clearInUse()
		}
	}
}

// sweep (delete) unused pairs and views. It stops views before deleting their
// reference.
// Notifiers are not handled as they aren't internal objects.
func (t *tracker) sweep(notifier IDer, cache Cacher) {
	t.Lock()
	defer t.Unlock()
	used := make(map[string]struct{})
	tmp := t.tracked[:0]
	for _, tp := range t.tracked {
		otherNotifier := tp.notify != notifier.ID()
		if tp.mark || otherNotifier {
			tmp = append(tmp, tp)
			used[tp.view] = struct{}{}
			used[tp.notify] = struct{}{}
		}
	}
	t.tracked = tmp
	// remove views/notifiers no longer referenced
	for viewId, view := range t.views {
		if _, ok := used[viewId]; !ok {
			delete(t.views, viewId)
			view.stop()
			cache.Delete(viewId)
		}
	}
}

// dummy Notifier for use by vault token above and in tests
type dummyNotifier struct {
	name   string
	notify bool
	deps   chan interface{}
}

func fakeNotifier(name string) *dummyNotifier {
	return &dummyNotifier{
		name: name, notify: true,
		deps: make(chan interface{}, 100),
	}
}

func (n *dummyNotifier) Notify(d interface{}) bool {
	n.deps <- d
	return n.notify
}

func (n *dummyNotifier) ID() string {
	return string(n.name)
}

func (n *dummyNotifier) count() int {
	return len(n.deps)
}