[Prometheus] Counter with reset interval

interval.go

package main

import (
	"context"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

type IntervalFunc = func(now time.Time) time.Duration

// StaticInterval allows to specify a static duration.
func StaticInterval(duration time.Duration) IntervalFunc {
	return func (_ time.Time) time.Duration {
		return duration
	}
}

// Daily returns duration for a day (time until 00:00)
func Daily(now time.Time) time.Duration {
	year, month, day := now.Date()
	rounded := time.Date(year, month, day + 1, 0, 0, 0, 0, now.Location())
	return rounded.Sub(now)
}

// IntervalCounter is a prometheus.Metric that acts like a counter or gauge which
// will be reset after a specified interval.
//
// Metric can be used to keep value only for a specific period of time.
//
// Note: you should call IntervalCounter.Countdown in order to start the reset timer.
type IntervalCounter struct {
	prometheus.Gauge
	intervalFunc IntervalFunc
}

// NewIntervalCounter creates a new Gauge-based counter which will be reset to zero after
// specified period of time.
//
// Reset interval is returned from a function or can be static using StaticInterval.
func NewIntervalCounter(opts prometheus.GaugeOpts, intervalFunc IntervalFunc) IntervalCounter {
	return IntervalCounter{
		Gauge:        prometheus.NewGauge(opts),
		intervalFunc: intervalFunc,
	}
}

// Countdown starts countdown and will reset counter after specified interval.
func (r IntervalCounter) Countdown(ctx context.Context) {
	interval := r.intervalFunc(time.Now())
	t := time.NewTicker(interval)
	for {
		select {
		case <-ctx.Done():
			t.Stop()
			return
		case <-t.C:
			r.Set(0)
			if newInterval := r.intervalFunc(time.Now()); newInterval != interval {
				// Update interval
				t.Reset(newInterval)
			}
		}
	}
}

// Example
func main() {
	ctx := context.TODO()
	requestsPerMinute := NewIntervalCounter(prometheus.GaugeOpts{
		Name: "requests-per-minute",
	}, StaticInterval(time.Minute))
	go requestsPerMinute.Countdown(ctx)
	
	requestsPerDay := NewIntervalCounter(prometheus.GaugeOpts{Name: "per-day"}, Daily)
	go requestsPerDay.Countdown(ctx)
	
	prometheus.DefaultRegisterer.MustRegister(requestsPerMinute, requestsPerDay)
}