fix: improved cache

Author: turip

openmeter/entitlement/balanceworker/entitlementhandler.go

@@ -24,6 +24,10 @@ import (
 	"github.com/samber/lo"
 )
 
+const (
+	metricAttributeAction string = "action"
+)
+
 type handleEntitlementEventOptions struct {
 	// Source is the source of the event, e.g. the "subject" field from the upstream cloudevents event causing the change
 	source string
@@ -181,15 +185,17 @@ func (w *Worker) processEntitlementEntity(ctx context.Context, entitlementEntity
 type recalcAction string
 
 const (
-	recalcActionRecalculate recalcAction = "recalculate"
-	recalcActionUseCache    recalcAction = "use-cache"
+	recalcActionRecalculate      recalcAction = "recalculate"
+	recalcActionRecalculateOnHit recalcAction = "recalculate-on-hit"
+	recalcActionUseCache         recalcAction = "use-cache"
 )
 
 func (w *Worker) createSnapshotEventEstimator(ctx context.Context, entitlementEntity *entitlement.Entitlement, calculatedAt time.Time, opts handleEntitlementEventOptions) (marshaler.Event, error) {
 	if len(opts.rawIngestedEvents) == 0 {
 		return nil, fmt.Errorf("no raw ingested events provided")
 	}
 
+	// TODO: These should come from outside of the call
 	feature, err := w.entitlement.Feature.GetFeature(ctx, entitlementEntity.Namespace, entitlementEntity.FeatureID, feature.IncludeArchivedFeatureTrue)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get feature: %w", err)
@@ -213,6 +219,8 @@ func (w *Worker) createSnapshotEventEstimator(ctx context.Context, entitlementEn
 		Meter:       meterEntity,
 	}
 
+	w.metricEstimatorRequestsTotal.Add(ctx, 1)
+
 	action := recalcActionUseCache
 
 	ent, err := w.estimator.HandleEntitlementEvent(ctx, estimator.IngestEventInput{
@@ -245,7 +253,7 @@ func (w *Worker) createSnapshotEventEstimator(ctx context.Context, entitlementEn
 		}
 
 		if thHit {
-			action = recalcActionRecalculate
+			action = recalcActionRecalculateOnHit
 		}
 	}
 
@@ -254,7 +262,11 @@ func (w *Worker) createSnapshotEventEstimator(ctx context.Context, entitlementEn
 	// this should be a percentage and later we can say that for 1% of the events we should validate
 	// cache consistency
 
-	if action == recalcActionRecalculate {
+	w.metricEstimatorActionTotal.Add(ctx, 1, metric.WithAttributes(
+		attribute.String(metricAttributeAction, string(action)),
+	))
+
+	if action == recalcActionRecalculate || action == recalcActionRecalculateOnHit {
 		value, err := w.estimator.HandleRecalculation(ctx, target, func(ctx context.Context) (*snapshot.EntitlementValue, error) {
 			res, err := w.entitlement.Entitlement.GetEntitlementValue(ctx, entitlementEntity.Namespace, entitlementEntity.SubjectKey, entitlementEntity.ID, calculatedAt)
 			if err != nil {
@@ -302,6 +314,7 @@ func (w *Worker) createSnapshotEventEstimator(ctx context.Context, entitlementEn
 			balance := estimator.NewInfDecimal(lo.FromPtr(value.Balance))
 
 			if balance.GreaterThan(cacheEntry.ApproxUsage) {
+				w.metricEstimatorValidationErrorsTotal.Add(ctx, 1)
 				w.opts.Logger.Error("cache entry is inconsistent", "error", err)
 			}
 		}

openmeter/entitlement/balanceworker/estimator/backend.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"log/slog"
 	"time"
 
 	"github.com/go-redsync/redsync/v4"
@@ -23,11 +24,13 @@ type redisCacheBackend struct {
 
 	redsync     *redsync.Redsync
 	lockTimeout time.Duration
+	logger      *slog.Logger
 }
 
 type RedisCacheBackendOptions struct {
 	RedisURL    string
 	LockTimeout time.Duration
+	Logger      *slog.Logger
 }
 
 func (o *RedisCacheBackendOptions) Validate() error {
@@ -39,6 +42,10 @@ func (o *RedisCacheBackendOptions) Validate() error {
 		return errors.New("lockTimeout must be greater than 0")
 	}
 
+	if o.Logger == nil {
+		return errors.New("logger is required")
+	}
+
 	return nil
 }
 
@@ -58,6 +65,7 @@ func NewRedisCacheBackend(in RedisCacheBackendOptions) (CacheBackend, error) {
 		redis:       redis,
 		redsync:     redsync.New(goredis.NewPool(redis)),
 		lockTimeout: in.LockTimeout,
+		logger:      in.Logger,
 	}, nil
 }
 

openmeter/entitlement/balanceworker/estimator/estimator.go

@@ -62,7 +62,7 @@ func (t *TargetEntitlement) GetEntryHash() hasher.Hash {
 	val := strings.Join(
 		[]string{
 			t.Entitlement.ID, t.Entitlement.UpdatedAt.String(), lo.FromPtr(t.Entitlement.DeletedAt).String(),
-			// TODO: NR of voided and active grants
+			// TODO: NR of voided and active grants !!!!
 			// Entitlement usage period change
 			t.Entitlement.CurrentUsagePeriod.From.String(), t.Entitlement.CurrentUsagePeriod.To.String(),
 			// Feature changes
@@ -89,6 +89,7 @@ type Cache interface {
 type CacheOptions struct {
 	RedisURL    string
 	LockTimeout time.Duration
+	Logger      *slog.Logger
 }
 
 func (o *CacheOptions) Validate() error {
@@ -100,6 +101,10 @@ func (o *CacheOptions) Validate() error {
 		return errors.New("lockTimeout must be greater than 0")
 	}
 
+	if o.Logger == nil {
+		return errors.New("logger is required")
+	}
+
 	return nil
 }
 
@@ -111,20 +116,22 @@ func NewCache(in CacheOptions) (Cache, error) {
 	backend, err := NewRedisCacheBackend(RedisCacheBackendOptions{
 		RedisURL:    in.RedisURL,
 		LockTimeout: in.LockTimeout,
+		Logger:      in.Logger,
 	})
 	if err != nil {
 		return nil, err
 	}
 
 	return &cache{
 		backend: backend,
+		logger:  in.Logger,
 	}, nil
 }
 
 type cache struct {
 	backend CacheBackend
 
-	validationRate float64
+	logger *slog.Logger
 }
 
 var _ Cache = (*cache)(nil)
@@ -151,7 +158,7 @@ func (c *cache) HandleEntitlementEvent(ctx context.Context, event IngestEventInp
 			getChange, err := c.getChange(event.Target.Meter, event.DedupedEvents)
 			if err != nil {
 				// TODO: more details
-				slog.Warn("failed to get change for entitlement", "error", err)
+				c.logger.Warn("failed to get change for entitlement", "error", err)
 				// We should not fail the whole calculation, instead if we have any pending
 				// thresholds, this will force a recalculation of the entitlement balance.
 				approxIncrease = infinite

openmeter/entitlement/balanceworker/worker.go

@@ -158,6 +158,11 @@ type Worker struct {
 	metricRecalculationTime       metric.Int64Histogram
 	metricHighWatermarkCacheStats metric.Int64Counter
 
+	// Estimator metrics
+	metricEstimatorRequestsTotal         metric.Int64Counter
+	metricEstimatorValidationErrorsTotal metric.Int64Counter
+	metricEstimatorActionTotal           metric.Int64Counter
+
 	// Handlers
 	nonPublishingHandler *grouphandler.NoPublishingHandler
 }