chore: remove redis based highwatermark cache (#3605)

Author: turip

app/common/openmeter_balanceworker.go

@@ -32,7 +32,6 @@ var BalanceWorker = wire.NewSet(
 	NewBalanceWorkerOptions,
 	NewBalanceWorker,
 	BalanceWorkerGroup,
-	NewBalanceWorkerFilterStateStorage,
 )
 
 var BalanceWorkerAdapter = wire.NewSet(
@@ -94,7 +93,7 @@ func NewBalanceWorkerOptions(
 	subjectService subject.Service,
 	customerService customer.Service,
 	logger *slog.Logger,
-	filterStateStorage balanceworker.FilterStateStorage,
+	balanceWorkerConfiguration config.BalanceWorkerConfiguration,
 ) balanceworker.WorkerOptions {
 	return balanceworker.WorkerOptions{
 		SystemEventsTopic:        eventConfig.SystemEvents.Topic,
@@ -109,7 +108,7 @@ func NewBalanceWorkerOptions(
 		Logger:                   logger,
 		MetricMeter:              routerOptions.MetricMeter,
 		NotificationService:      notificationService,
-		FilterStateStorage:       filterStateStorage,
+		HighWatermarkCacheSize:   balanceWorkerConfiguration.StateStorage.HighWatermarkCache.LRUCacheSize,
 	}
 }
 
@@ -122,30 +121,6 @@ func NewBalanceWorker(workerOptions balanceworker.WorkerOptions) (*balanceworker
 	return worker, nil
 }
 
-func NewBalanceWorkerFilterStateStorage(conf config.BalanceWorkerConfiguration) (balanceworker.FilterStateStorage, error) {
-	switch conf.StateStorage.Driver {
-	case config.BalanceWorkerStateStorageDriverRedis:
-		redis, err := conf.StateStorage.GetRedisBackendConfiguration()
-		if err != nil {
-			return balanceworker.FilterStateStorage{}, fmt.Errorf("failed to get redis backend configuration: %w", err)
-		}
-
-		client, err := redis.NewClient()
-		if err != nil {
-			return balanceworker.FilterStateStorage{}, fmt.Errorf("failed to create redis client: %w", err)
-		}
-
-		return balanceworker.NewFilterStateStorage(balanceworker.FilterStateStorageRedis{
-			Client:     client,
-			Expiration: redis.Expiration,
-		})
-	case config.BalanceWorkerStateStorageDriverInMemory:
-		return balanceworker.NewFilterStateStorage(balanceworker.FilterStateStorageInMemory{})
-	default:
-		return balanceworker.FilterStateStorage{}, fmt.Errorf("unsupported state storage driver: %s", conf.StateStorage.Driver)
-	}
-}
-
 func BalanceWorkerGroup(
 	ctx context.Context,
 	worker *balanceworker.Worker,

app/config/balanceworker.go

@@ -2,16 +2,10 @@ package config
 
 import (
 	"errors"
-	"fmt"
-	"slices"
-	"time"
 
-	"github.com/mitchellh/mapstructure"
 	"github.com/spf13/viper"
 
 	"github.com/openmeterio/openmeter/pkg/errorsx"
-	"github.com/openmeterio/openmeter/pkg/models"
-	"github.com/openmeterio/openmeter/pkg/redis"
 )
 
 type BalanceWorkerConfiguration struct {
@@ -46,115 +40,35 @@ type rawBalanceWorkerStateStorageConfiguration struct {
 }
 
 type BalanceWorkerStateStorageConfiguration struct {
-	Driver BalanceWorkerStateStorageDriver
-
-	BalanceWorkerStateStorageBackendConfiguration
-}
-
-func (c *BalanceWorkerStateStorageConfiguration) DecodeMap(v map[string]any) error {
-	var raw rawBalanceWorkerStateStorageConfiguration
-
-	if err := mapstructure.Decode(v, &raw); err != nil {
-		return err
-	}
-
-	c.Driver = raw.Driver
-
-	switch c.Driver {
-	case BalanceWorkerStateStorageDriverRedis:
-		var redisConfig BalanceWorkerStateStorageRedisBackendConfiguration
-
-		decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
-			Metadata:         nil,
-			Result:           &redisConfig,
-			WeaklyTypedInput: true,
-			DecodeHook: mapstructure.ComposeDecodeHookFunc(
-				mapstructure.StringToTimeDurationHookFunc(),
-			),
-		})
-		if err != nil {
-			return err
-		}
-
-		if err := decoder.Decode(raw.Config); err != nil {
-			return err
-		}
-
-		c.BalanceWorkerStateStorageBackendConfiguration = redisConfig
-	case BalanceWorkerStateStorageDriverInMemory:
-		// no config
-	}
-
-	return nil
+	HighWatermarkCache BalanceWorkerHighWatermarkCacheConfiguration
 }
 
 func (c BalanceWorkerStateStorageConfiguration) Validate() error {
-	errs := []error{}
-	if !slices.Contains([]BalanceWorkerStateStorageDriver{
-		BalanceWorkerStateStorageDriverRedis,
-		BalanceWorkerStateStorageDriverInMemory,
-	}, c.Driver) {
-		errs = append(errs, fmt.Errorf("invalid driver: %s", c.Driver))
-	}
+	var errs []error
 
-	if c.Driver == BalanceWorkerStateStorageDriverRedis {
-		if c.BalanceWorkerStateStorageBackendConfiguration == nil {
-			errs = append(errs, errors.New("state storage backend configuration is required"))
-		} else {
-			if err := c.BalanceWorkerStateStorageBackendConfiguration.Validate(); err != nil {
-				errs = append(errs, fmt.Errorf("state storage backend: %w", err))
-			}
-		}
+	if err := c.HighWatermarkCache.Validate(); err != nil {
+		errs = append(errs, errorsx.WithPrefix(err, "high watermark cache"))
 	}
 
 	return errors.Join(errs...)
 }
 
-func (c BalanceWorkerStateStorageConfiguration) GetRedisBackendConfiguration() (BalanceWorkerStateStorageRedisBackendConfiguration, error) {
-	if c.Driver != BalanceWorkerStateStorageDriverRedis {
-		return BalanceWorkerStateStorageRedisBackendConfiguration{}, fmt.Errorf("driver is not redis")
-	}
-
-	if c.BalanceWorkerStateStorageBackendConfiguration == nil {
-		return BalanceWorkerStateStorageRedisBackendConfiguration{}, errors.New("state storage backend configuration is required")
-	}
-
-	redisConfig, ok := c.BalanceWorkerStateStorageBackendConfiguration.(BalanceWorkerStateStorageRedisBackendConfiguration)
-	if !ok {
-		return BalanceWorkerStateStorageRedisBackendConfiguration{}, fmt.Errorf("state storage backend configuration is not a redis configuration")
-	}
-
-	return redisConfig, nil
+type BalanceWorkerHighWatermarkCacheConfiguration struct {
+	LRUCacheSize int
 }
 
-type BalanceWorkerStateStorageBackendConfiguration interface {
-	models.Validator
-}
-
-type BalanceWorkerStateStorageRedisBackendConfiguration struct {
-	redis.Config `mapstructure:",squash"`
-	Expiration   time.Duration
-}
-
-func (c BalanceWorkerStateStorageRedisBackendConfiguration) Validate() error {
-	errs := []error{}
-
-	if c.Expiration <= 0 {
-		errs = append(errs, errors.New("expiration should be greater than 0"))
-	}
-
-	if err := c.Config.Validate(); err != nil {
-		errs = append(errs, fmt.Errorf("redis: %w", err))
+func (c BalanceWorkerHighWatermarkCacheConfiguration) Validate() error {
+	if c.LRUCacheSize <= 0 {
+		return errors.New("LRU cache size must be positive")
 	}
 
-	return errors.Join(errs...)
+	return nil
 }
 
 func ConfigureBalanceWorker(v *viper.Viper) {
 	ConfigureConsumer(v, "balanceWorker")
 	v.SetDefault("balanceWorker.dlq.topic", "om_sys.balance_worker_dlq")
 	v.SetDefault("balanceWorker.consumerGroupName", "om_balance_worker")
 
-	v.SetDefault("balanceWorker.stateStorage.driver", BalanceWorkerStateStorageDriverInMemory)
-	v.SetDefault("balanceWorker.stateStorage.config.expiration", 24*time.Hour)
+	v.SetDefault("balanceWorker.stateStorage.highWatermarkCache.lruCacheSize", 100_000)
 }

app/config/config_test.go

@@ -341,12 +341,8 @@ func TestComplete(t *testing.T) {
 				ConsumerGroupName: "om_balance_worker",
 			},
 			StateStorage: BalanceWorkerStateStorageConfiguration{
-				Driver: BalanceWorkerStateStorageDriverRedis,
-				BalanceWorkerStateStorageBackendConfiguration: BalanceWorkerStateStorageRedisBackendConfiguration{
-					Expiration: 23 * time.Hour,
-					Config: redis.Config{
-						Address: "127.0.0.1:6379",
-					},
+				HighWatermarkCache: BalanceWorkerHighWatermarkCacheConfiguration{
+					LRUCacheSize: 100_000,
 				},
 			},
 		},

cmd/balance-worker/wire_gen.go

@@ -14,20 +14,15 @@ func NewRecalculateBalanceSnapshotsCommand() *cobra.Command {
 		Use:   "recalculate-balance-snapshots",
 		Short: "Recalculate balance snapshots and send the resulting events into the eventbus",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			filterStateStorage, err := balanceworker.NewFilterStateStorage(balanceworker.FilterStateStorageInMemory{})
-			if err != nil {
-				return err
-			}
-
 			recalculator, err := balanceworker.NewRecalculator(balanceworker.RecalculatorOptions{
-				Entitlement:         internal.App.EntitlementRegistry,
-				EventBus:            internal.App.EventPublisher,
-				MetricMeter:         internal.App.Meter,
-				NotificationService: internal.App.NotificationService,
-				FilterStateStorage:  filterStateStorage,
-				Logger:              internal.App.Logger,
-				Customer:            internal.App.Customer,
-				Subject:             internal.App.Subject,
+				Entitlement:            internal.App.EntitlementRegistry,
+				EventBus:               internal.App.EventPublisher,
+				MetricMeter:            internal.App.Meter,
+				NotificationService:    internal.App.NotificationService,
+				HighWatermarkCacheSize: 100_000,
+				Logger:                 internal.App.Logger,
+				Customer:               internal.App.Customer,
+				Subject:                internal.App.Subject,
 			})
 			if err != nil {
 				return err

cmd/jobs/entitlement/recalculatesnapshots.go

@@ -38,8 +38,9 @@ var _ models.Validator = (*EntitlementFiltersConfig)(nil)
 type EntitlementFiltersConfig struct {
 	NotificationService notification.Service
 	MetricMeter         metric.Meter
-	StateStorage        FilterStateStorage
 	Logger              *slog.Logger
+
+	HighWatermarkCacheSize int
 }
 
 func (c EntitlementFiltersConfig) Validate() error {
@@ -53,8 +54,8 @@ func (c EntitlementFiltersConfig) Validate() error {
 		errs = append(errs, fmt.Errorf("metric meter is required"))
 	}
 
-	if err := c.StateStorage.Validate(); err != nil {
-		errs = append(errs, fmt.Errorf("state storage: %w", err))
+	if c.HighWatermarkCacheSize <= 0 {
+		errs = append(errs, fmt.Errorf("high watermark cache size must be positive"))
 	}
 
 	if c.Logger == nil {
@@ -76,6 +77,10 @@ type EntitlementFilters struct {
 }
 
 func NewEntitlementFilters(cfg EntitlementFiltersConfig) (*EntitlementFilters, error) {
+	if err := cfg.Validate(); err != nil {
+		return nil, err
+	}
+
 	notificationFilter, err := filters.NewNotificationsFilter(filters.NotificationsFilterConfig{
 		NotificationService: cfg.NotificationService,
 		CacheTTL:            defaultCacheTTL,
@@ -85,33 +90,7 @@ func NewEntitlementFilters(cfg EntitlementFiltersConfig) (*EntitlementFilters, e
 		return nil, err
 	}
 
-	var highWatermarkCacheBackend filters.HighWatermarkBackend
-
-	switch cfg.StateStorage.Driver() {
-	case FilterStateStorageDriverRedis:
-		redis, err := cfg.StateStorage.Redis()
-		if err != nil {
-			return nil, err
-		}
-
-		highWatermarkCacheBackend, err = filters.NewHighWatermarkRedisBackend(filters.HighWatermarkRedisBackendConfig{
-			Redis:      redis.Client,
-			Logger:     cfg.Logger,
-			Expiration: redis.Expiration,
-		})
-		if err != nil {
-			return nil, err
-		}
-	case FilterStateStorageDriverInMemory:
-		highWatermarkCacheBackend, err = filters.NewHighWatermarkInMemoryBackend(defaultLRUCacheSize)
-		if err != nil {
-			return nil, err
-		}
-	default:
-		return nil, fmt.Errorf("unsupported state storage driver: %s", cfg.StateStorage.Driver())
-	}
-
-	highWatermarkCache, err := filters.NewHighWatermarkCache(highWatermarkCacheBackend)
+	highWatermarkCache, err := filters.NewHighWatermarkCache(cfg.HighWatermarkCacheSize)
 	if err != nil {
 		return nil, err
 	}