diff --git a/automod/countstore/countstore.go b/automod/countstore/countstore.go index 21004739..bde0d72c 100644 --- a/automod/countstore/countstore.go +++ b/automod/countstore/countstore.go @@ -28,6 +28,7 @@ const ( // a count in each supported period bucket size. // In other words, one call to CountStore.Increment causes three increments internally: // one to the count for the hour, one to the count for the day, and one to thte all-time count. +// The "IncrementPeriod" method allows only incrementing a single period bucket. Care must be taken to match the "GetCount" period with the incremented period when using this variant. // // The exact implementation and precision of the "*Distinct" methods may vary: // in the MemCountStore implementation, it is precise (it's based on large maps); @@ -40,10 +41,10 @@ const ( // only the all-time counts go without expiration. // The MemCountStore grows without bound (it's intended to be used in testing // and other non-production operations). -// type CountStore interface { GetCount(ctx context.Context, name, val, period string) (int, error) Increment(ctx context.Context, name, val string) error + IncrementPeriod(ctx context.Context, name, val, period string) error // TODO: batch increment method GetCountDistinct(ctx context.Context, name, bucket, period string) (int, error) IncrementDistinct(ctx context.Context, name, bucket, val string) error diff --git a/automod/countstore/countstore_mem.go b/automod/countstore/countstore_mem.go index 3c265a68..f33c076e 100644 --- a/automod/countstore/countstore_mem.go +++ b/automod/countstore/countstore_mem.go @@ -33,14 +33,21 @@ func (s MemCountStore) GetCount(ctx context.Context, name, val, period string) ( func (s MemCountStore) Increment(ctx context.Context, name, val string) error { for _, p := range []string{PeriodTotal, PeriodDay, PeriodHour} { - k := periodBucket(name, val, p) - s.Counts.Compute(k, func(oldVal int, _ bool) (int, bool) { - return oldVal+1, false - }) + if err := s.IncrementPeriod(ctx, name, val, p); err != nil { + return err + } } return nil } +func (s MemCountStore) IncrementPeriod(ctx context.Context, name, val, period string) error { + k := periodBucket(name, val, period) + s.Counts.Compute(k, func(oldVal int, _ bool) (int, bool) { + return oldVal + 1, false + }) + return nil +} + func (s MemCountStore) GetCountDistinct(ctx context.Context, name, bucket, period string) (int, error) { v, ok := s.DistinctCounts.Load(periodBucket(name, bucket, period)) if !ok { @@ -52,7 +59,7 @@ func (s MemCountStore) GetCountDistinct(ctx context.Context, name, bucket, perio func (s MemCountStore) IncrementDistinct(ctx context.Context, name, bucket, val string) error { for _, p := range []string{PeriodTotal, PeriodDay, PeriodHour} { k := periodBucket(name, bucket, p) - s.DistinctCounts.Compute(k,func(nested *xsync.MapOf[string, bool], _ bool) (*xsync.MapOf[string, bool], bool) { + s.DistinctCounts.Compute(k, func(nested *xsync.MapOf[string, bool], _ bool) (*xsync.MapOf[string, bool], bool) { if nested == nil { nested = xsync.NewMapOf[string, bool]() } diff --git a/automod/countstore/countstore_redis.go b/automod/countstore/countstore_redis.go index d4ef1c76..bd7d5f21 100644 --- a/automod/countstore/countstore_redis.go +++ b/automod/countstore/countstore_redis.go @@ -71,7 +71,7 @@ func (s *RedisCountStore) IncrementPeriod(ctx context.Context, name, val, period // multiple ops in a single redis round-trip multi := s.Client.Pipeline() - key := redisCountPrefix + PeriodBucket(name, val, period) + key := redisCountPrefix + periodBucket(name, val, period) multi.Incr(ctx, key) switch period {