fix(v2): shard limit overflow (#4125)

Author: kolesnikovae

pkg/experiment/distributor/placement/adaptive_placement/config.go

@@ -34,8 +34,8 @@ func (o *PlacementLimits) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet
 	f.Uint64Var(&o.TenantShards, prefix+"tenant-shards", 0, "Number of shards per tenant. If 0, the limit is not applied.")
 	f.Uint64Var(&o.DefaultDatasetShards, prefix+"default-dataset-shards", 1, "Default number of shards per dataset.")
 	f.Uint64Var(&o.MinDatasetShards, prefix+"min-dataset-shards", 1, "Minimum number of shards per dataset.")
-	f.Uint64Var(&o.MaxDatasetShards, prefix+"max-dataset-shards", 32, "Maximum number of shards per dataset.")
-	f.Uint64Var(&o.UnitSizeBytes, prefix+"unit-size-bytes", 64<<10, "Shards are allocated based on the utilisation of units per second. The option specifies the unit size in bytes.")
+	f.Uint64Var(&o.MaxDatasetShards, prefix+"max-dataset-shards", 1<<10, "Maximum number of shards per dataset.")
+	f.Uint64Var(&o.UnitSizeBytes, prefix+"unit-size-bytes", 128<<10, "Shards are allocated based on the utilisation of units per second. The option specifies the unit size in bytes.")
 	f.DurationVar(&o.BurstWindow, prefix+"burst-window", 17*time.Minute, "Duration of the burst window. During this time, scale-outs are more aggressive.")
 	f.DurationVar(&o.DecayWindow, prefix+"decay-window", 19*time.Minute, "Duration of the decay window. During this time, scale-ins are delayed.")
 }

pkg/experiment/distributor/placement/adaptive_placement/shard_allocator.go

@@ -69,9 +69,11 @@ func (a *shardAllocator) observe(usage uint64, now int64) int {
 			// Reset multiplier if burst window has passed.
 			a.multiplier = 1
 		} else {
-			// Increase multiplier on consecutive
-			// scale-outs within burst window.
-			a.multiplier *= 2
+			// Increase multiplier on consecutive scale-outs within burst window.
+			// Limiting the multiplier here allow us to not worry about overflows.
+			if a.multiplier < 16 {
+				a.multiplier *= 2
+			}
 			scaled := target + int(math.Ceil(float64(delta)*a.multiplier))
 			target = min(2*target, scaled)
 		}

pkg/experiment/distributor/placement/adaptive_placement/shard_allocator_test.go

@@ -38,3 +38,22 @@ func Test_shard_allocator(t *testing.T) {
 		require.Equal(t, test.want, a.observe(test.usage, test.now), fmt.Sprint(i))
 	}
 }
+
+func Test_shard_limit(t *testing.T) {
+	a := &shardAllocator{
+		unitSize:    128 << 10,
+		min:         1,
+		max:         10,
+		burstWindow: 1e9 * 10,
+		decayWindow: 1e9 * 10 * 5,
+	}
+
+	var now int64
+	var hi int
+	for i := uint64(0); i < 100; i++ {
+		old := hi
+		hi = a.observe(2*a.unitSize*i, now)
+		require.GreaterOrEqual(t, hi, old)
+		now += 1e9 * 10
+	}
+}