Report the in-memory counts of the kinds without a daily counter again

limit/service.go

@@ -200,11 +200,17 @@ func (s *Service) Tokens(key string, rule Rule) int64 {
 // sent. The bucket cannot exceed its burst, and the daily count cannot go below zero.
 func (s *Service) Refund(key string, kind Kind, n int64) {
 	s.refundTokens(key, kind, n)
+	start := periodStart(key, s.dbTime(s.now()))
 	i := dailyIndex(kind)
 	if i < 0 {
+		bs := s.buckets[shardIndex(key)]
+		bs.mu.Lock()
+		defer bs.mu.Unlock()
+		if e, ok := bs.entries[key]; ok && e.periodStart >= start {
+			e.counts[kind] = max(0, e.counts[kind]-n)
+		}
 		return
 	}
-	start := periodStart(key, s.dbTime(s.now()))
 	cs := s.counters[shardIndex(key)]
 	cs.mu.Lock()
 	defer cs.mu.Unlock()
@@ -213,19 +219,24 @@ func (s *Service) Refund(key string, kind Kind, n int64) {
 	}
 }
 
-// Usage returns the key's daily counts in its current period, from memory
+// Usage returns the key's counts in its current period, from memory: the daily counts (all
+// nodes' as of the last sync, plus this node's since), and this node's units of the other kinds
 func (s *Service) Usage(key string) Usage {
 	start := periodStart(key, s.dbTime(s.now()))
+	var u Usage
+	bs := s.buckets[shardIndex(key)]
+	bs.mu.Lock()
+	if e, ok := bs.entries[key]; ok && e.periodStart >= start {
+		u.counts = e.counts
+	}
+	bs.mu.Unlock()
 	cs := s.counters[shardIndex(key)]
 	cs.mu.Lock()
 	defer cs.mu.Unlock()
-	c, ok := cs.counters[key]
-	if !ok || c.periodStart < start {
-		return Usage{}
-	}
-	var u Usage
-	for i := 0; i < numDaily; i++ {
-		u.counts[i] = c.base[i] + c.unflushed[i]
+	if c, ok := cs.counters[key]; ok && c.periodStart >= start {
+		for i, kind := range dailyKinds {
+			u.counts[kind] = c.base[i] + c.unflushed[i]
+		}
 	}
 	return u
 }
@@ -482,6 +493,7 @@ func (s *Service) takeTokens(key string, rule *Rule, n int64, now time.Time) Dec
 	e := bs.entryNoLock(key, now)
 	e.seen = now.UnixNano()
 	if rule.Burst <= 0 {
+		e.countNoLock(key, rule.Kind, n, s.dbTime(now))
 		return Decision{OK: true}
 	}
 	b := &e.buckets[rule.Kind]
@@ -493,6 +505,7 @@ func (s *Service) takeTokens(key string, rule *Rule, n int64, now time.Time) Dec
 	}
 	b.tokens -= float64(n)
 	e.noteFull(b, rule)
+	e.countNoLock(key, rule.Kind, n, s.dbTime(now))
 	return Decision{OK: true}
 }
 
@@ -515,6 +528,7 @@ func (s *Service) reset(key string, kind Kind) {
 	bs.mu.Lock()
 	if e, ok := bs.entries[key]; ok {
 		e.buckets[kind].tokens = float64(e.buckets[kind].burst)
+		e.counts[kind] = 0
 	}
 	bs.mu.Unlock()
 	i := dailyIndex(kind)
@@ -558,6 +572,19 @@ func (bs *bucketShard) entryNoLock(key string, now time.Time) *bucketEntry {
 	return e
 }
 
+// countNoLock adds n units of a kind without a counter to the entry's counts of the current
+// period, starting them over if the period has ended. The caller must hold the shard's lock.
+func (e *bucketEntry) countNoLock(key string, kind Kind, n int64, dbNow time.Time) {
+	if dailyIndex(kind) >= 0 {
+		return // Counted in the key's counter
+	}
+	if start := periodStart(key, dbNow); e.periodStart < start {
+		e.counts = [numKinds]int64{}
+		e.periodStart = start
+	}
+	e.counts[kind] += n
+}
+
 // noteFull moves the time from which every bucket is full again to when this bucket is full, if
 // that is later. A bucket without refill never fills, so it does not hold the key.
 func (e *bucketEntry) noteFull(b *bucket, rule *Rule) {

limit/service_test.go

@@ -127,7 +127,7 @@ func TestService_Allow_N(t *testing.T) {
 	require.Equal(t, 200*time.Millisecond, d.RetryAfter)
 	clock.advance(200 * time.Millisecond)
 	require.True(t, s.Allow("k", rule, 600).OK)
-	require.Equal(t, int64(0), s.Usage("k").Count(KindBandwidth)) // Buckets only, no daily count
+	require.Equal(t, int64(1200), s.Usage("k").Count(KindBandwidth))
 
 	// More than the burst can never be allowed, so there is no retry time
 	d = s.Allow("k", rule, 1001)
@@ -198,7 +198,7 @@ func TestService_Allow_Unlimited(t *testing.T) {
 	for i := 0; i < 1000; i++ {
 		require.True(t, s.Allow("k", Rule{Kind: KindTopics}, 1).OK)
 	}
-	require.Equal(t, 1, s.Size())
+	require.Equal(t, int64(1000), s.Usage("k").Count(KindTopics))
 }
 
 func TestService_Allow_NeverRefills(t *testing.T) {
@@ -443,6 +443,7 @@ func TestService_Allow_ConcurrentBucketOnly(t *testing.T) {
 	close(start)
 	wg.Wait()
 	require.Equal(t, rule.Burst, admitted.Load())
+	require.Equal(t, admitted.Load(), s.Usage("hot").Count(rule.Kind))
 	require.Zero(t, s.Tokens("hot", rule))
 }
 
@@ -485,6 +486,7 @@ func TestService_Prune_ConcurrentAllow(t *testing.T) {
 	close(start)
 	wg.Wait()
 	require.Zero(t, admitted.Load())
+	require.Equal(t, int64(1), s.Usage("hot").Count(rule.Kind))
 	require.Equal(t, 1, s.Size())
 }
 

limit/types.go

@@ -49,7 +49,7 @@ type Decision struct {
 
 // Usage is a key's daily counts in its current period
 type Usage struct {
-	counts [numDaily]int64
+	counts [numKinds]int64
 }
 
 // Config configures the service
@@ -67,9 +67,11 @@ type counter struct {
 
 // bucketEntry is one key's token buckets
 type bucketEntry struct {
-	buckets [numKinds]bucket
-	seen    int64 // Last use, unix nanoseconds
-	full    int64 // When every bucket is full again, unix nanoseconds
+	buckets     [numKinds]bucket
+	counts      [numKinds]int64 // Units used in the current period of the kinds without a counter (bandwidth, ...), memory only
+	periodStart int64           // The period of counts (see periodStart)
+	seen        int64           // Last use, unix nanoseconds
+	full        int64           // When every bucket is full again, unix nanoseconds
 }
 
 // bucket is a token bucket. It is created full: last is zero until the first use, and tokens
@@ -107,8 +109,8 @@ func ParseKind(name string) (Kind, bool) {
 
 // Count returns the usage of the given kind; 0 for a kind without a daily count
 func (u Usage) Count(kind Kind) int64 {
-	if i := dailyIndex(kind); i >= 0 {
-		return u.counts[i]
+	if kind < numKinds {
+		return u.counts[kind]
 	}
 	return 0
 }

limit/types_test.go

@@ -32,10 +32,23 @@ func TestUsage_Count(t *testing.T) {
 	}, 42).OK)
 	u := s.Usage("k")
 	require.Equal(t, int64(3), u.Count(KindMessages))
-	require.Equal(t, int64(0), u.Count(KindBandwidth)) // Buckets only, no daily count
+	require.Equal(t, int64(42), u.Count(KindBandwidth)) // In memory only, but counted
 	require.Equal(t, int64(0), u.Count(KindEmails))
 }
 
+func TestLimiter_Value_Bandwidth(t *testing.T) {
+	s := newTestService(t)
+	l := s.Limiter("k", Rule{
+		Kind:   KindBandwidth,
+		Burst:  1000,
+		Refill: time.Second,
+	})
+	require.True(t, l.AllowN(300))
+	require.Equal(t, int64(300), l.Value())
+	require.True(t, l.AllowN(-100))
+	require.Equal(t, int64(200), l.Value())
+}
+
 func TestLimiter_AllowNRefundsNegative(t *testing.T) {
 	s := newTestService(t)
 	l := s.Limiter("k", Rule{

limit/util.go

@@ -23,6 +23,9 @@ func shardIndex(key string) int {
 	return int(h % shardCount)
 }
 
+// dailyKinds are the kinds with a counter, in the order of the counter's arrays (see dailyIndex)
+var dailyKinds = [numDaily]Kind{KindMessages, KindEmails, KindCalls}
+
 // dailyIndex returns the index of the kind in a counter's arrays, or -1 if the kind has no daily count
 func dailyIndex(kind Kind) int {
 	switch kind {