diff --git a/pkg/service/coordinator/allocation.go b/pkg/service/coordinator/allocation.go index abefe5d3..dcc42b07 100644 --- a/pkg/service/coordinator/allocation.go +++ b/pkg/service/coordinator/allocation.go @@ -7,19 +7,19 @@ import ( "github.com/google/uuid" - "go.atoms.co/splitter/lib/service/location" "go.atoms.co/lib/mapx" - "go.atoms.co/slicex" "go.atoms.co/lib/uuidx" + "go.atoms.co/slicex" + "go.atoms.co/splitter/lib/service/location" + splitterpb "go.atoms.co/splitter/pb" "go.atoms.co/splitter/pkg/allocation" "go.atoms.co/splitter/pkg/core" "go.atoms.co/splitter/pkg/model" splitteruuidx "go.atoms.co/splitter/pkg/util/uuidx" - splitterpb "go.atoms.co/splitter/pb" ) const ( - defaultShardLoad = allocation.Load(50) + defaultShardLoad = allocation.Load(defaultShardScore) unitShardLoad = defaultShardLoad * 5 regionAffinityLoad = defaultShardLoad * 2 namedShardsLoad = defaultShardLoad * 2 @@ -39,12 +39,12 @@ type ( // TODO(herohde) 11/12/2023: intra-domain anti-affinity to spread out domains evenly. Similar to general LB. -func newAllocation(id location.InstanceID, tenant model.TenantInfo, info model.ServiceInfoEx, placements []core.InternalPlacementInfo, activation time.Time) *Allocation { - return allocation.New(id, findPlacements(tenant, info), findColocations(info), findWork(info, placements), activation) +func newAllocation(id location.InstanceID, tenant model.TenantInfo, info model.ServiceInfoEx, placements []core.InternalPlacementInfo, trackers map[model.QualifiedDomainName]*domainLoadTracker, activation time.Time) *Allocation { + return allocation.New(id, findPlacements(tenant, info), findColocations(info), findWork(info, placements, trackers), activation) } -func updateAllocation(a *Allocation, tenant model.TenantInfo, info model.ServiceInfoEx, namedShards []model.Shard, placements []core.InternalPlacementInfo, activation time.Time) (*Allocation, []Grant) { - return allocation.Update(a, findPlacements(tenant, info, namedShards...), findColocations(info), findWork(info, placements), activation) +func updateAllocation(a *Allocation, tenant model.TenantInfo, info model.ServiceInfoEx, namedShards []model.Shard, placements []core.InternalPlacementInfo, trackers map[model.QualifiedDomainName]*domainLoadTracker, activation time.Time) (*Allocation, []Grant) { + return allocation.Update(a, findPlacements(tenant, info, namedShards...), findColocations(info), findWork(info, placements, trackers), activation) } // NamedShards handles named shard placement @@ -248,7 +248,7 @@ func findColocations(info model.ServiceInfoEx) []Colocation { return slicex.New[Colocation](a) } -func findWork(state model.ServiceInfoEx, placements []core.InternalPlacementInfo) []Work { +func findWork(state model.ServiceInfoEx, placements []core.InternalPlacementInfo, trackers map[model.QualifiedDomainName]*domainLoadTracker) []Work { var ret []Work m := mapx.New(placements, func(v core.InternalPlacementInfo) model.PlacementName { @@ -260,6 +260,8 @@ func findWork(state model.ServiceInfoEx, placements []core.InternalPlacementInfo for _, domain := range state.Domains() { switch domain.Type() { case model.Unit: + // TODO: (xuhui) 07/29/2026 Support load-aware allocation for unit domains. + // Unit domains currently use a fixed load and are excluded from load balancing. locations := toLocations(domain.Regions()...) if len(locations) == 0 { locations = defaultLocations @@ -293,15 +295,16 @@ func findWork(state model.ServiceInfoEx, placements []core.InternalPlacementInfo for _, shard := range shards { region := provider.Find(model.Key(shard.From())) + unit := model.Shard{ + Domain: domain.Name(), + Type: model.Global, + From: model.Key(shard.From()), + To: model.Key(shard.To()), + } w := Work{ - Unit: model.Shard{ - Domain: domain.Name(), - Type: model.Global, - From: model.Key(shard.From()), - To: model.Key(shard.To()), - }, + Unit: unit, Data: slicex.New(location.Location{Region: region}), - Load: defaultShardLoad, + Load: shardWorkLoad(state, trackers, unit), } ret = append(ret, w) } @@ -312,15 +315,16 @@ func findWork(state model.ServiceInfoEx, placements []core.InternalPlacementInfo locations = defaultLocations } for _, shard := range shards { + unit := model.Shard{ + Domain: domain.Name(), + Type: model.Global, + From: model.Key(shard.From()), + To: model.Key(shard.To()), + } w := Work{ - Unit: model.Shard{ - Domain: domain.Name(), - Type: model.Global, - From: model.Key(shard.From()), - To: model.Key(shard.To()), - }, + Unit: unit, Data: locations, - Load: defaultShardLoad, + Load: shardWorkLoad(state, trackers, unit), } ret = append(ret, w) } @@ -331,16 +335,17 @@ func findWork(state model.ServiceInfoEx, placements []core.InternalPlacementInfo shards := findShardsForRegion(domain, region) for _, shard := range shards { + unit := model.Shard{ + Region: region, + Type: model.Regional, + Domain: domain.Name(), + From: model.Key(shard.From()), + To: model.Key(shard.To()), + } ret = append(ret, Work{ - Unit: model.Shard{ - Region: region, - Type: model.Regional, - Domain: domain.Name(), - From: model.Key(shard.From()), - To: model.Key(shard.To()), - }, + Unit: unit, Data: slicex.New(location.Location{Region: region}), - Load: defaultShardLoad, + Load: shardWorkLoad(state, trackers, unit), }) } } @@ -353,6 +358,18 @@ func findWork(state model.ServiceInfoEx, placements []core.InternalPlacementInfo return ret } +func shardWorkLoad(state model.ServiceInfoEx, trackers map[model.QualifiedDomainName]*domainLoadTracker, shard model.Shard) allocation.Load { + if !state.Service().Config().TrackLoad() { + return defaultShardLoad + } + + shardScore := defaultShardScore + if tracker, ok := trackers[shard.Domain]; ok { + shardScore = tracker.shardScoreOrDefault(core.NewShard(shard.From, shard.To, shard.Region)) + } + return allocation.Load(max(1, shardScore)) +} + func findShards(domain model.Domain) []uuidx.Range { policy := domain.Config().ShardingPolicy() targetShardCount := policy.Shards() diff --git a/pkg/service/coordinator/allocation_test.go b/pkg/service/coordinator/allocation_test.go index 269a13b1..41bc925b 100644 --- a/pkg/service/coordinator/allocation_test.go +++ b/pkg/service/coordinator/allocation_test.go @@ -8,13 +8,14 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "go.atoms.co/splitter/lib/service/location" "go.atoms.co/lib/testing/assertx" - "go.atoms.co/slicex" "go.atoms.co/lib/uuidx" + "go.atoms.co/slicex" + "go.atoms.co/splitter/lib/service/location" + splitterpb "go.atoms.co/splitter/pb" "go.atoms.co/splitter/pkg/allocation" + "go.atoms.co/splitter/pkg/core" "go.atoms.co/splitter/pkg/model" - splitterpb "go.atoms.co/splitter/pb" "go.atoms.co/splitter/testing/prefab" ) @@ -143,6 +144,81 @@ func TestNamedShards(t *testing.T) { assert.Equal(t, namedShardsLoad, load) } +func TestFindWork_LoadScores(t *testing.T) { + serviceName := model.QualifiedServiceName{Tenant: "tenant", Service: "service"} + + newInfo := func(t *testing.T, domainType model.DomainType, cfg model.ServiceConfig, domainOpts ...model.DomainOption) model.ServiceInfoEx { + t.Helper() + + service, err := model.NewService(serviceName, time.Time{}, model.WithServiceConfig(cfg)) + require.NoError(t, err) + domain, err := model.NewDomain(model.QualifiedDomainName{Service: serviceName, Domain: "domain"}, domainType, time.Time{}, domainOpts...) + require.NoError(t, err) + return model.NewServiceInfoEx(model.NewServiceInfo(service, 1, time.Time{}), []model.Domain{domain}) + } + + globalDomainOpts := []model.DomainOption{model.WithDomainConfig(model.NewDomainConfig(model.WithDomainShardingPolicy(model.NewShardingPolicy(1))))} + + t.Run("tracking disabled", func(t *testing.T) { + info := newInfo(t, model.Global, model.NewServiceConfig(), globalDomainOpts...) + work := findWork(info, nil, nil) + require.Len(t, work, 1) + require.Equal(t, defaultShardLoad, work[0].Load) + }) + + t.Run("score unavailable", func(t *testing.T) { + info := newInfo(t, model.Global, model.NewServiceConfig(model.WithTrackLoad(true)), globalDomainOpts...) + work := findWork(info, nil, nil) + require.Len(t, work, 1) + require.Equal(t, allocation.Load(defaultShardScore), work[0].Load) + }) + + t.Run("published global score", func(t *testing.T) { + info := newInfo(t, model.Global, model.NewServiceConfig(model.WithTrackLoad(true)), globalDomainOpts...) + initial := findWork(info, nil, nil) + require.Len(t, initial, 1) + shard := initial[0].Unit + tracker := newDomainLoadTracker(time.Time{}, shard.Domain.Domain) + tracker.quantile = &domainQuantileInfo{ + domainQuantile: 20, + shardQuantiles: map[core.Shard]float64{ + core.NewShard(shard.From, shard.To, shard.Region): 60, + }, + } + + work := findWork(info, nil, map[model.QualifiedDomainName]*domainLoadTracker{shard.Domain: tracker}) + require.Len(t, work, 1) + require.Equal(t, allocation.Load(75), work[0].Load) + }) + + t.Run("published regional score includes region", func(t *testing.T) { + domainCfg := model.NewDomainConfig(model.WithDomainShardingPolicy(model.NewShardingPolicy(1)), model.WithDomainRegions("centralus")) + info := newInfo(t, model.Regional, model.NewServiceConfig(model.WithTrackLoad(true)), model.WithDomainConfig(domainCfg)) + initial := findWork(info, nil, nil) + require.Len(t, initial, 1) + shard := initial[0].Unit + tracker := newDomainLoadTracker(time.Time{}, shard.Domain.Domain) + tracker.quantile = &domainQuantileInfo{ + domainQuantile: 20, + shardQuantiles: map[core.Shard]float64{ + core.NewShard(shard.From, shard.To, ""): 1, + core.NewShard(shard.From, shard.To, shard.Region): 60, + }, + } + + work := findWork(info, nil, map[model.QualifiedDomainName]*domainLoadTracker{shard.Domain: tracker}) + require.Len(t, work, 1) + require.Equal(t, allocation.Load(75), work[0].Load) + }) + + t.Run("unit load unchanged", func(t *testing.T) { + info := newInfo(t, model.Unit, model.NewServiceConfig(model.WithTrackLoad(true))) + work := findWork(info, nil, nil) + require.Len(t, work, 1) + require.Equal(t, unitShardLoad, work[0].Load) + }) +} + func TestDomainState(t *testing.T) { t1, err := model.NewTenant("tenant1", time.Time{}) require.NoError(t, err) diff --git a/pkg/service/coordinator/coordinator.go b/pkg/service/coordinator/coordinator.go index 0ebc3488..e502ad69 100644 --- a/pkg/service/coordinator/coordinator.go +++ b/pkg/service/coordinator/coordinator.go @@ -393,6 +393,7 @@ func (c *coordinator) connect(ctx context.Context, sid session.ID, origin locati lease := now.Add(leaseDuration) s.TrySend(ctx, model.NewExtend(lease)) // grants will be covered under this lease + //TODO: (xuhui) 07/27/2026 revisit impact to services that set capacity. capacity := allocation.Load(int64(limit) * int64(defaultShardLoad)) // Set capacity shard limit * shard load (0 for no capacity) if capacity > 0 { log.Infof(ctx, "Consumer %v connected with non-zero capacity limit: %v", consumer, capacity) @@ -542,12 +543,12 @@ func (c *coordinator) init(ctx context.Context, state core.State, updates <-chan } now := time.Now() - c.alloc = newAllocation(c.self.ID(), tenant, info, c.cache.Placements(c.name.Tenant), now.Add(delay)) + c.restoreLoadTrackers(ctx) + + c.alloc = newAllocation(c.self.ID(), tenant, info, c.cache.Placements(c.name.Tenant), c.trackers, now.Add(delay)) c.noLb = c.findUnitDomains() c.cluster = model.NewClusterMap(model.NewClusterID(c.self, now), c.alloc.Units()) - c.restoreLoadTrackers(ctx) - log.Infof(ctx, "Coordinator %v/%v initialized, #shards=%v", c.name, c.self, c.alloc.Size()) c.recordAction(ctx, "init", "ok") c.recordActionLatency(ctx, "init", start) @@ -629,6 +630,7 @@ steady: // (1) Refresh allocation, (2) allocate, (3) broadcast cluster change now := time.Now() + trackLoad := c.info.Service().Config().TrackLoad() if err := c.cache.Update(upd, false); err != nil { log.Errorf(ctx, "Internal: invalid state update %v", err) @@ -649,6 +651,10 @@ steady: } c.info = info + if trackLoad && !info.Service().Config().TrackLoad() { + clear(c.trackers) + } + oldShards := c.alloc.Units() c.refresh(ctx, c.refreshDelay) @@ -707,9 +713,7 @@ steady: break } - for _, t := range c.trackers { - t.rotateIfNeeded(now) - } + c.rotateTrackerAndRefreshIfNeeded(ctx, now) // (1) emit domain load and shard load metrics c.emitLoadMetrics(ctx) @@ -755,6 +759,19 @@ steady: log.Infof(ctx, "Coordinator %v draining, #consumer=%v", c.self, len(c.consumers)) } +func (c *coordinator) rotateTrackerAndRefreshIfNeeded(ctx context.Context, now time.Time) { + var updated bool + for _, t := range c.trackers { + updated = t.rotateIfNeeded(now) || updated + } + + if !updated { + return + } + + c.refresh(ctx, c.refreshDelay) +} + func (c *coordinator) refresh(ctx context.Context, delay time.Duration) { now := time.Now() @@ -784,7 +801,7 @@ func (c *coordinator) refresh(ctx context.Context, delay time.Duration) { } } - upd, rejected := updateAllocation(c.alloc, c.tenant, c.info, namedShards, c.cache.Placements(c.name.Tenant), now.Add(delay)) + upd, rejected := updateAllocation(c.alloc, c.tenant, c.info, namedShards, c.cache.Placements(c.name.Tenant), c.trackers, now.Add(delay)) c.alloc = upd c.noLb = c.findUnitDomains() diff --git a/pkg/service/coordinator/coordinator_test.go b/pkg/service/coordinator/coordinator_test.go index a7fb1e96..65984f29 100644 --- a/pkg/service/coordinator/coordinator_test.go +++ b/pkg/service/coordinator/coordinator_test.go @@ -20,6 +20,7 @@ import ( "go.atoms.co/splitter/lib/service/session" splitterpb "go.atoms.co/splitter/pb" splitterprivatepb "go.atoms.co/splitter/pb/private" + "go.atoms.co/splitter/pkg/allocation" "go.atoms.co/splitter/pkg/core" "go.atoms.co/splitter/pkg/model" ) @@ -509,7 +510,7 @@ func TestCoordinator_NamedKeyDisconnectDropsNamedShardPenalty(t *testing.T) { readFn(t, out2, isClusterSnapshot) _, load = c.alloc.Load() - require.EqualValues(t, namedShardsLoad, load.Place, "matching shard should be marked as named") + require.Equal(t, namedShardsLoad, load.Place, "matching shard should be marked as named") close(in2) synctest.Wait() @@ -1229,6 +1230,145 @@ func TestCoordinator_RestoresDomainLoadTrackers(t *testing.T) { }) } +func TestCoordinator_ClearsLoadTrackersWhenTrackingDisabled(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + domainCfg := model.NewDomainConfig(model.WithDomainShardingPolicy(model.NewShardingPolicy(1))) + domain, err := model.NewDomain(domainName, model.Global, time.Now(), model.WithDomainConfig(domainCfg)) + require.NoError(t, err) + + shardRange := findShards(domain)[0] + shard := model.Shard{ + Domain: domainName, + Type: model.Global, + From: model.Key(shardRange.From()), + To: model.Key(shardRange.To()), + } + tracker := newDomainLoadTracker(testStart(), domain1) + tracker.quantile = &domainQuantileInfo{ + domainQuantile: 20, + shardQuantiles: map[core.Shard]float64{ + core.NewShard(shard.From, shard.To, shard.Region): 60, + }, + } + + status := core.NewServiceStatus(core.NewServiceLoadInfo(serviceName, []core.DomainLoadInfo{tracker.snapshot()})) + cfg := model.NewServiceConfig(model.WithTrackLoad(true)) + coord, _, updates := setupWithServiceConfigAndStatusesAndUpdates(ctx, t, []model.Domain{domain}, cfg, []core.ServiceStatus{status}) + defer coord.Close() + + c := coord.(*coordinator) + require.Contains(t, c.trackers, domainName) + work, ok := c.alloc.Unit(shard) + require.True(t, ok) + require.Equal(t, allocation.Load(75), work.Load) + + updateTrackLoad := func(enabled bool, version model.Version) { + serviceCfg := model.NewServiceConfig(model.WithTrackLoad(enabled)) + service, err := model.NewService(serviceName, time.Now(), model.WithServiceConfig(serviceCfg)) + require.NoError(t, err) + updates <- core.NewServiceUpdate(model.NewServiceInfo(service, version, time.Now())) + synctest.Wait() + } + + // disable TrackLoad + updateTrackLoad(false, 2) + require.Empty(t, c.trackers) + _, ok = c.cache.ServiceStatus(serviceName) + require.False(t, ok) + work, ok = c.alloc.Unit(shard) + require.True(t, ok) + require.Equal(t, defaultShardLoad, work.Load) + + // enable TrackLoad again, allocation should start with default value. + updateTrackLoad(true, 3) + require.Empty(t, c.trackers) + work, ok = c.alloc.Unit(shard) + require.True(t, ok) + require.Equal(t, defaultShardLoad, work.Load) + }) +} + +func TestCoordinator_RestoredLoadScoreUsedByInitialAllocation(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + domainCfg := model.NewDomainConfig(model.WithDomainShardingPolicy(model.NewShardingPolicy(2))) + domain, err := model.NewDomain(domainName, model.Global, time.Now(), model.WithDomainConfig(domainCfg)) + require.NoError(t, err) + + shardRange := findShards(domain)[0] + shard := model.Shard{ + Domain: domainName, + Type: model.Global, + From: model.Key(shardRange.From()), + To: model.Key(shardRange.To()), + } + expected := newDomainLoadTracker(testStart(), domain1) + for range 5 { + expected.add(shard, model.Load(20)) + } + expected.quantile = &domainQuantileInfo{ + domainQuantile: 20, + shardQuantiles: map[core.Shard]float64{ + core.NewShard(shard.From, shard.To, shard.Region): 60, + }, + } + + status := core.NewServiceStatus(core.NewServiceLoadInfo(serviceName, []core.DomainLoadInfo{expected.snapshot()})) + cfg := model.NewServiceConfig(model.WithTrackLoad(true)) + coord, _ := setupWithServiceConfigAndStatuses(ctx, t, []model.Domain{domain}, cfg, []core.ServiceStatus{status}) + defer coord.Close() + + c := coord.(*coordinator) + work, ok := c.alloc.Unit(shard) + require.True(t, ok) + require.Equal(t, allocation.Load(75), work.Load) + }) +} + +func TestCoordinator_AppliesLoadScoresAfterTrackerRotation(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + + domainCfg := model.NewDomainConfig(model.WithDomainShardingPolicy(model.NewShardingPolicy(2))) + domain, err := model.NewDomain(domainName, model.Global, time.Now(), model.WithDomainConfig(domainCfg)) + require.NoError(t, err) + + cfg := model.NewServiceConfig(model.WithTrackLoad(true)) + coord, _ := setupWithServiceConfig(ctx, t, []model.Domain{domain}, cfg) + defer coord.Close() + c := coord.(*coordinator) + + require.NoError(t, c.txn(ctx, func() error { + work := c.alloc.Work() + require.Len(t, work, 2) + require.Equal(t, allocation.Load(defaultShardScore), work[0].Load) + require.Equal(t, allocation.Load(defaultShardScore), work[1].Load) + + now := time.Now() + tracker := newDomainLoadTracker(now.Add(-defaultRotationInterval-time.Second), domain1) + for range 10 { + tracker.add(work[0].Unit, model.Load(10)) + tracker.add(work[1].Unit, model.Load(30)) + } + c.trackers[domainName] = tracker + + c.rotateTrackerAndRefreshIfNeeded(ctx, now) + updated := c.alloc.Work() + require.Len(t, updated, 2) + + loads := []allocation.Load{updated[0].Load, updated[1].Load} + for _, oldWork := range work { + score := tracker.shardScoreOrDefault(core.NewShard(oldWork.Unit.From, oldWork.Unit.To, oldWork.Unit.Region)) + require.Contains(t, loads, allocation.Load(score)) + } + return nil + })) + }) +} + // updateCreatedAt updates the createdAt of domainLoadTrackers to simulate time advancing and avoid a long sleep (24 hours). // With synctest, time.Sleep triggers all tickers to fire within the synctest bubble; // long sleeps slow the test. @@ -1265,6 +1405,13 @@ func setupWithServiceConfig(ctx context.Context, t *testing.T, domains []model.D func setupWithServiceConfigAndStatuses(ctx context.Context, t *testing.T, domains []model.Domain, cfg model.ServiceConfig, statuses []core.ServiceStatus, opts ...Option) (Coordinator, <-chan core.ServiceStatusMessage) { t.Helper() + c, out, _ := setupWithServiceConfigAndStatusesAndUpdates(ctx, t, domains, cfg, statuses, opts...) + return c, out +} + +func setupWithServiceConfigAndStatusesAndUpdates(ctx context.Context, t *testing.T, domains []model.Domain, cfg model.ServiceConfig, statuses []core.ServiceStatus, opts ...Option) (Coordinator, <-chan core.ServiceStatusMessage, chan<- core.Update) { + t.Helper() + loc := location.New("centralus", "splitter-0") tenant, err := model.NewTenant(tenant1, time.Now()) @@ -1285,7 +1432,7 @@ func setupWithServiceConfigAndStatuses(ctx context.Context, t *testing.T, domain c, out := New(ctx, loc, serviceName, state, updates, opts...) <-c.Initialized().Closed() - return c, out + return c, out, updates } func isAssign(msg model.ConsumerMessage) (model.AssignMessage, bool) { diff --git a/pkg/service/coordinator/tracking.go b/pkg/service/coordinator/tracking.go index 4ba52498..87d788f9 100644 --- a/pkg/service/coordinator/tracking.go +++ b/pkg/service/coordinator/tracking.go @@ -9,18 +9,21 @@ import ( "go.atoms.co/splitter/pkg/util/p2quantile" ) +// score represents the load score of a shard. It is in range (0, 100) +type score float64 + const ( // P50 quantile to track median value median = 0.5 // Score is ranged in [0, scoreRange) scoreRange = 100.0 + // defaultShardScore is used until a shard has a published score. + // It is also used as defaultShardLoad by services that doesn't enable load tracking. + defaultShardScore score = scoreRange * median // defaultRotationInterval defines the interval a domainLoadTracker lives before rotation. defaultRotationInterval = 24 * time.Hour ) -// score represents the load score of a shard. It is in range (0, 100) -type score float64 - // domainQuantileInfo holds published quantile values for a domain and its shards. // Mutable version of core.DomainQuantileInfo. type domainQuantileInfo struct { @@ -198,13 +201,17 @@ func newDomainLoadTracker(now time.Time, domain model.DomainName) *domainLoadTra } // rotateIfNeeded seals current load tracker and creates a new one if needed. -func (t *domainLoadTracker) rotateIfNeeded(now time.Time) { +// It returns whether new quantiles were published. +func (t *domainLoadTracker) rotateIfNeeded(now time.Time) bool { + updated := false if t.tracker.needsRotation(now) { if q, ok := t.tracker.quantileInfo(); ok { t.quantile = q + updated = true } t.tracker = newDomainTracker(now) } + return updated } // add adds an observation of a shard load. @@ -225,14 +232,13 @@ func (t *domainLoadTracker) shardLoad() map[core.Shard]model.Load { // shardScoreOrDefault returns score of a shard if it has been tracked. // Or, default score that equals to (scoreRange / 2). func (t *domainLoadTracker) shardScoreOrDefault(shard core.Shard) score { - defaultScore := score(scoreRange / 2) if t.quantile == nil { - return defaultScore + return defaultShardScore } s, ok := t.quantile.score(shard) if !ok { - return defaultScore + return defaultShardScore } return s } diff --git a/pkg/service/coordinator/tracking_test.go b/pkg/service/coordinator/tracking_test.go index 69bd6ae8..c5128444 100644 --- a/pkg/service/coordinator/tracking_test.go +++ b/pkg/service/coordinator/tracking_test.go @@ -82,7 +82,8 @@ func TestLoadTracker_TryRotatePublishesMetrics(t *testing.T) { tr.add(shard, model.Load(10)) } require.Nil(t, tr.quantile) - tr.rotateIfNeeded(start.Add(defaultRotationInterval + time.Second)) + require.True(t, tr.rotateIfNeeded(start.Add(defaultRotationInterval+time.Second))) + require.False(t, tr.rotateIfNeeded(start.Add(defaultRotationInterval+2*time.Second))) require.NotNil(t, tr.quantile) require.True(t, tr.snapshot().HasQuantileInfo())