ExpiringErrorCounter maintains a registry of counter errors that expires after a specified amount of time. This package also supports registering even handlers that will be called if the counter of a particular error crosses a threshold in the specified amount of time.

expiringerrorcounter.go

package threadsafe

import (
	"sync"
	"time"
)

// Based on https://gist.github.com/clarkmcc/2c40db4bb7b3aea412d78c8a490f030e
// ExpiringErrorCounter maintains a registry of counter errors that expires
// after a specified amount of time. For example if the expiration were set
// at ten seconds, and ten of the same error were counted, on second then,
// the counter would show a count of ten for that error, while a second later
// the counter would be reset to 0. In addition to maintaining a expiring
// counter of errors, this package also supports registering even handlers
// that will be called if the counter of a particular error crosses a threshold
// in the specified amount of time.
type ExpiringErrorCounter struct {
	lock sync.Mutex
	globalExpiration time.Duration
	errorSpecificExpiration map[string]time.Duration
	errorSpecificHandlers map[string]*ErrorThresholdHandler
	errors map[string]*ExpiringCounter
}

// ErrorThresholdHandler holds the threshold of errors as well as the func
// that should be executed if that threshold is met.
type ErrorThresholdHandler struct {
	handler   ExpiringErrorCounterThresholdFunc
	threshold int
}

func NewThresholdHandler(threshold int, handler ExpiringErrorCounterThresholdFunc) *ErrorThresholdHandler {
	return &ErrorThresholdHandler{threshold: threshold, handler: handler}
}

type ExpiringErrorCounterThresholdFunc func(err error)

// NewExpiringErrorCounter returns a new expiring error counter
func NewExpiringErrorCounter() *ExpiringErrorCounter {
	return &ExpiringErrorCounter{
		globalExpiration: 0,
		errors: map[string]*ExpiringCounter{},
		errorSpecificHandlers: map[string]*ErrorThresholdHandler{},
		errorSpecificExpiration: map[string]time.Duration{},
	}
}

// WithExpiration sets the global expiration for all counted errors,
// this value can be overridden by error specific expirations.
func (e *ExpiringErrorCounter) WithExpiration(expiration time.Duration) *ExpiringErrorCounter {
	e.lock.Lock()
	defer e.lock.Unlock()
	e.globalExpiration = expiration
	return e
}

// Count counts the error and returns the current value of the counter
func (e *ExpiringErrorCounter) Count(err error) int {
	e.lock.Lock()
	defer e.lock.Unlock()
	if counter, ok := e.errors[err.Error()]; ok {
		// Assuming the counter already exists for this error
		counter.Inc()

		// Check for a handler for this error
		if handler, ok := e.errorSpecificHandlers[err.Error()]; ok {
			// If there is a handler, check if the counter value is higher
			// than the handler threshold.
			if counter.Value() >= handler.threshold {
				// If it is higher, execute the handler
				handler.handler(err)
				// And reset the counter
				counter.Reset()
			}
		}

		// Call counter.Value again in case it changed
		return counter.Value()
	}

	// Check to see if theres a specific expiration for this error
	// and give that expiration priority over the global expiration
	var expiration time.Duration
	var ok bool
	if expiration, ok = e.errorSpecificExpiration[err.Error()]; !ok {
		expiration = e.globalExpiration
	}

	// Assuming we don't have a counter for this error yet
	counter := NewCounter().
		WithExpiration(expiration)
	counter.Inc()
	// Save the counter to the error counter
	e.errors[err.Error()] = counter
	// Return the error counter value
	return counter.Value()
}

// RegisterThresholdHandler registers a threshold handler, this handler will get called anytime
// there counter for the provided error crosses the provided threshold.
func (e *ExpiringErrorCounter) RegisterThresholdHandler(err error, threshold int, handler ExpiringErrorCounterThresholdFunc) {
	e.lock.Lock()
	defer e.lock.Unlock()
	e.errorSpecificHandlers[err.Error()] = NewThresholdHandler(threshold, handler)
}

// SetErrorBasedExpiration sets the expiration to use for any counted errors
// if the error hasn't been counted yet, it will be created with an expiring
// counter that expires with the provided expiration time.Duration here.
func (e *ExpiringErrorCounter) SetErrorBasedExpiration(err error, expiration time.Duration) {
	e.lock.Lock()
	defer e.lock.Unlock()
	// Set the expiration for any new counters if this error type
	e.errorSpecificExpiration[err.Error()] = expiration
	// Check for any existing counters of this error type and change the
	// expiration for them as well
	if counter, ok := e.errors[err.Error()]; ok {
		counter.WithExpiration(expiration)
	}
}