github.com/safing/portbase@v0.19.5/metrics/persistence.go (about)

     1  package metrics
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/tevino/abool"
    10  
    11  	"github.com/safing/portbase/database"
    12  	"github.com/safing/portbase/database/record"
    13  	"github.com/safing/portbase/log"
    14  )
    15  
    16  var (
    17  	storage       *metricsStorage
    18  	storageKey    string
    19  	storageInit   = abool.New()
    20  	storageLoaded = abool.New()
    21  
    22  	db = database.NewInterface(&database.Options{
    23  		Local:    true,
    24  		Internal: true,
    25  	})
    26  
    27  	// ErrAlreadyInitialized is returned when trying to initialize an option
    28  	// more than once or if the time window for initializing is over.
    29  	ErrAlreadyInitialized = errors.New("already initialized")
    30  )
    31  
    32  type metricsStorage struct {
    33  	sync.Mutex
    34  	record.Base
    35  
    36  	Start    time.Time
    37  	Counters map[string]uint64
    38  }
    39  
    40  // EnableMetricPersistence enables metric persistence for metrics that opted
    41  // for it. They given key is the database key where the metric data will be
    42  // persisted.
    43  // This call also directly loads the stored data from the database.
    44  // The returned error is only about loading the metrics, not about enabling
    45  // persistence.
    46  // May only be called once.
    47  func EnableMetricPersistence(key string) error {
    48  	// Check if already initialized.
    49  	if !storageInit.SetToIf(false, true) {
    50  		return ErrAlreadyInitialized
    51  	}
    52  
    53  	// Set storage key.
    54  	storageKey = key
    55  
    56  	// Load metrics from storage.
    57  	var err error
    58  	storage, err = getMetricsStorage(storageKey)
    59  	switch {
    60  	case err == nil:
    61  		// Continue.
    62  	case errors.Is(err, database.ErrNotFound):
    63  		return nil
    64  	default:
    65  		return err
    66  	}
    67  	storageLoaded.Set()
    68  
    69  	// Load saved state for all counter metrics.
    70  	registryLock.RLock()
    71  	defer registryLock.RUnlock()
    72  
    73  	for _, m := range registry {
    74  		counter, ok := m.(*Counter)
    75  		if ok {
    76  			counter.loadState()
    77  		}
    78  	}
    79  
    80  	return nil
    81  }
    82  
    83  func (c *Counter) loadState() {
    84  	// Check if we can and should load the state.
    85  	if !storageLoaded.IsSet() || !c.Opts().Persist {
    86  		return
    87  	}
    88  
    89  	c.Set(storage.Counters[c.LabeledID()])
    90  }
    91  
    92  func storePersistentMetrics() {
    93  	// Check if persistence is enabled.
    94  	if !storageInit.IsSet() || storageKey == "" {
    95  		return
    96  	}
    97  
    98  	// Create new storage.
    99  	newStorage := &metricsStorage{
   100  		// TODO: This timestamp should be taken from previous save, if possible.
   101  		Start:    time.Now(),
   102  		Counters: make(map[string]uint64),
   103  	}
   104  	newStorage.SetKey(storageKey)
   105  	// Copy values from previous version.
   106  	if storageLoaded.IsSet() {
   107  		newStorage.Start = storage.Start
   108  	}
   109  
   110  	registryLock.RLock()
   111  	defer registryLock.RUnlock()
   112  
   113  	// Export all counter metrics.
   114  	for _, m := range registry {
   115  		if m.Opts().Persist {
   116  			counter, ok := m.(*Counter)
   117  			if ok {
   118  				newStorage.Counters[m.LabeledID()] = counter.Get()
   119  			}
   120  		}
   121  	}
   122  
   123  	// Save to database.
   124  	err := db.Put(newStorage)
   125  	if err != nil {
   126  		log.Warningf("metrics: failed to save metrics storage to db: %s", err)
   127  	}
   128  }
   129  
   130  func getMetricsStorage(key string) (*metricsStorage, error) {
   131  	r, err := db.Get(key)
   132  	if err != nil {
   133  		return nil, err
   134  	}
   135  
   136  	// unwrap
   137  	if r.IsWrapped() {
   138  		// only allocate a new struct, if we need it
   139  		newStorage := &metricsStorage{}
   140  		err = record.Unwrap(r, newStorage)
   141  		if err != nil {
   142  			return nil, err
   143  		}
   144  		return newStorage, nil
   145  	}
   146  
   147  	// or adjust type
   148  	newStorage, ok := r.(*metricsStorage)
   149  	if !ok {
   150  		return nil, fmt.Errorf("record not of type *metricsStorage, but %T", r)
   151  	}
   152  	return newStorage, nil
   153  }