diff --git a/blobstore/scheduler/client/clustermgr.go b/blobstore/scheduler/client/clustermgr.go index 37a95fb54..7c0d74ab0 100644 --- a/blobstore/scheduler/client/clustermgr.go +++ b/blobstore/scheduler/client/clustermgr.go @@ -79,7 +79,7 @@ type ClusterMgrDiskAPI interface { type ClusterMgrServiceAPI interface { Register(ctx context.Context, info RegisterInfo) error - GetService(ctx context.Context, name string, clusterID proto.ClusterID) (hosts []string, err error) + GetService(ctx context.Context, name string, clusterID proto.ClusterID) (nodes []cmapi.ServiceNode, err error) } type ClusterMgrTaskAPI interface { @@ -871,18 +871,19 @@ func (c *clustermgrClient) Register(ctx context.Context, info RegisterInfo) erro return c.client.RegisterService(ctx, node, info.HeartbeatIntervalS, info.HeartbeatTicks, info.ExpiresTicks) } -// GetService returns services -func (c *clustermgrClient) GetService(ctx context.Context, name string, clusterID proto.ClusterID) (hosts []string, err error) { +// GetService returns service nodes (with host and idc) for the given name and cluster. +func (c *clustermgrClient) GetService(ctx context.Context, name string, clusterID proto.ClusterID) ([]cmapi.ServiceNode, error) { svrInfos, err := c.client.GetService(ctx, cmapi.GetServiceArgs{Name: name}) if err != nil { return nil, err } + nodes := make([]cmapi.ServiceNode, 0, len(svrInfos.Nodes)) for _, s := range svrInfos.Nodes { if clusterID == proto.ClusterID(s.ClusterID) { - hosts = append(hosts, s.Host) + nodes = append(nodes, s) } } - return + return nodes, nil } // AddMigrateTask adds migrate task diff --git a/blobstore/scheduler/client/clustermgr_test.go b/blobstore/scheduler/client/clustermgr_test.go index 1a830862b..104d45983 100644 --- a/blobstore/scheduler/client/clustermgr_test.go +++ b/blobstore/scheduler/client/clustermgr_test.go @@ -620,10 +620,10 @@ func TestClustermgrClientExtra(t *testing.T) { {ClusterID: 2, Host: "127.0.0.2:8080"}, }, }, nil) - hosts, err := cli.GetService(ctx, "scheduler", proto.ClusterID(1)) + nodes, err := cli.GetService(ctx, "scheduler", proto.ClusterID(1)) require.NoError(t, err) - require.Equal(t, 1, len(hosts)) - require.Equal(t, "127.0.0.1:8080", hosts[0]) + require.Equal(t, 1, len(nodes)) + require.Equal(t, "127.0.0.1:8080", nodes[0].Host) mockCli.EXPECT().GetService(any, any).Return(cmapi.ServiceInfo{}, errMock) _, err = cli.GetService(ctx, "scheduler", proto.ClusterID(1)) diff --git a/blobstore/scheduler/client_mock_test.go b/blobstore/scheduler/client_mock_test.go index 1373c242b..e99bfe486 100644 --- a/blobstore/scheduler/client_mock_test.go +++ b/blobstore/scheduler/client_mock_test.go @@ -213,10 +213,10 @@ func (mr *MockClusterMgrAPIMockRecorder) GetMigratingDisk(arg0, arg1, arg2 inter } // GetService mocks base method. -func (m *MockClusterMgrAPI) GetService(arg0 context.Context, arg1 string, arg2 proto.ClusterID) ([]string, error) { +func (m *MockClusterMgrAPI) GetService(arg0 context.Context, arg1 string, arg2 proto.ClusterID) ([]clustermgr.ServiceNode, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "GetService", arg0, arg1, arg2) - ret0, _ := ret[0].([]string) + ret0, _ := ret[0].([]clustermgr.ServiceNode) ret1, _ := ret[1].(error) return ret0, ret1 } diff --git a/blobstore/scheduler/shard_repairer.go b/blobstore/scheduler/shard_repairer.go index 271bef932..73dd66e55 100644 --- a/blobstore/scheduler/shard_repairer.go +++ b/blobstore/scheduler/shard_repairer.go @@ -19,6 +19,7 @@ import ( "encoding/json" "errors" "fmt" + "sync" "time" "github.com/Shopify/sarama" @@ -105,7 +106,7 @@ type ShardRepairMgr struct { punishTime time.Duration blobnodeCli client.BlobnodeAPI - blobnodeSelector selector.Selector + blobnodeSelector *idcSelector clusterMgrCli client.ClusterMgrAPI taskCli client.TaskAPI @@ -137,9 +138,6 @@ func NewShardRepairMgr( return nil, err } - workerSelector := selector.MakeSelector(60*1000, func() (hosts []string, err error) { - return clusterMgrCli.GetService(context.Background(), proto.ServiceNameWorker, cfg.ClusterID) - }) failMsgSender, err := base.NewMsgSender(cfg.failedProducerConfig()) if err != nil { return nil, err @@ -150,14 +148,14 @@ func NewShardRepairMgr( return nil, err } - return &ShardRepairMgr{ + mgr := &ShardRepairMgr{ blobnodeCli: blobnodeCli, clusterMgrCli: clusterMgrCli, taskCli: taskAPI, taskPool: taskpool.New(cfg.TaskPoolSize, cfg.TaskPoolSize), taskSwitch: taskSwitch, clusterTopology: clusterTopology, - blobnodeSelector: workerSelector, + blobnodeSelector: newIDCSelector(clusterMgrCli, cfg.ClusterID), kafkaConsumerClient: kafkaClient, failMsgSender: failMsgSender, @@ -174,7 +172,9 @@ func NewShardRepairMgr( cfg: cfg, Closer: closer.New(), - }, nil + } + + return mgr, nil } // Enabled returns true if shard repair task is enabled, otherwise returns false @@ -400,20 +400,21 @@ func (mgr *ShardRepairMgr) repairShard(ctx context.Context, volInfo *client.Volu span.Infof("repair shard: msg[%+v], vol info[%+v]", repairMsg, volInfo) - // update host info + // update host info and cache IDC per index + idcByVunitIdx := make([]string, len(volInfo.VunitLocations)) for idx := range volInfo.VunitLocations { location := &volInfo.VunitLocations[idx] disk, ok := mgr.clusterTopology.GetDisk(location.DiskID) if ok { location.Host = disk.Host + idcByVunitIdx[idx] = disk.Idc } } - hosts := mgr.blobnodeSelector.GetRandomN(1) - if len(hosts) == 0 { + workerHost := mgr.blobnodeSelector.get(ctx, mgr.resolveRepairIDC(ctx, repairMsg, idcByVunitIdx)) + if workerHost == "" { return volInfo, ErrBlobnodeServiceUnavailable } - workerHost := hosts[0] task := proto.ShardRepairTask{ Bid: repairMsg.Bid, @@ -519,3 +520,115 @@ func (mgr *ShardRepairMgr) send2FailQueue(ctx context.Context, msg *proto.ShardR return nil } + +// resolveRepairIDC resolves the target IDC by iterating BadIdx until it finds +// a bad index whose disk IDC is known (non-empty). Returns "" if none match. +// +// Iterating all BadIdx is strictly better than using only BadIdx[0]: +// - If BadIdx[0] maps to a disk whose topology info is stale/missing but another +// BadIdx has a valid disk, the repair can still target the correct AZ. +// - If the bad shards span multiple AZs, local repair is impossible anyway, +// so the returned AZ (whichever is found first) does not matter — the +// blobnode's own localRepairable() will fall back to global repair. +// +// In all cases, the worst outcome is an unnecessary cross-AZ request; correctness is preserved. +func (mgr *ShardRepairMgr) resolveRepairIDC(ctx context.Context, repairMsg *proto.ShardRepairMsg, idcByVunitIdx []string) string { + span := trace.SpanFromContextSafe(ctx) + span.Debugf("resolveRepairIDC: repairMsg[%+v], idcByVunitIdx[%+v]", repairMsg, idcByVunitIdx) + for _, badIdx := range repairMsg.BadIdx { + idx := int(badIdx) + if idx < len(idcByVunitIdx) && idcByVunitIdx[idx] != "" { + return idcByVunitIdx[idx] + } + } + return "" +} + +// idcSelector provides AZ-aware worker selection for shard repair. +// Each per-IDC selector and the all-worker selector are backed by +// selector.MakeSelector, which handles interval-based refresh internally. +type idcSelector struct { + clusterMgrCli client.ClusterMgrAPI + clusterID proto.ClusterID + + mu sync.RWMutex + selectors map[string]selector.Selector // key: idc ("" for all workers), created lazily +} + +func newIDCSelector(clusterMgrCli client.ClusterMgrAPI, clusterID proto.ClusterID) *idcSelector { + return &idcSelector{ + clusterMgrCli: clusterMgrCli, + clusterID: clusterID, + selectors: make(map[string]selector.Selector), + } +} + +// getSelector lazily creates and caches a per-IDC (or all) worker selector. +// The underlying MakeSelector handles its own interval-based background refresh. +func (s *idcSelector) getSelector(ctx context.Context, idc string) selector.Selector { + // fast path: read-lock check, multiple readers proceed concurrently + s.mu.RLock() + sel := s.selectors[idc] + s.mu.RUnlock() + if sel != nil { + return sel + } + + // slow path: write-lock, double-check to prevent duplicate MakeSelector + s.mu.Lock() + defer s.mu.Unlock() + if sel = s.selectors[idc]; sel != nil { + return sel + } + + sel = selector.MakeSelector(60*1000, func() ([]string, error) { + bgCtx := context.Background() + nodes, err := s.clusterMgrCli.GetService(bgCtx, proto.ServiceNameWorker, s.clusterID) + if err != nil { + span := trace.SpanFromContextSafe(bgCtx) + span.Warnf("refresh worker services failed: err[%+v]", err) + return nil, err + } + if idc == "" { + hosts := make([]string, len(nodes)) + for i, n := range nodes { + hosts[i] = n.Host + } + return hosts, nil + } + var hosts []string + for _, n := range nodes { + if n.Idc == idc { + hosts = append(hosts, n.Host) + } + } + return hosts, nil + }) + s.selectors[idc] = sel + return sel +} + +// Get returns a worker host in the given IDC. If the IDC has no workers, +// it falls back to a random worker from any IDC. Returns empty if no +// workers are available at all. +func (s *idcSelector) get(ctx context.Context, idc string) string { + span := trace.SpanFromContextSafe(ctx) + if idc != "" { + sel := s.getSelector(ctx, idc) + hosts := sel.GetRandomN(1) + if len(hosts) > 0 { + span.Debugf("selected blobnode host[%s] for idc[%s]", hosts[0], idc) + return hosts[0] + } + span.Debugf("no blobnode in idc[%s], fallback to global", idc) + } + + sel := s.getSelector(ctx, "") + hosts := sel.GetRandomN(1) + if len(hosts) > 0 { + span.Debugf("selected blobnode host[%s] from global", hosts[0]) + return hosts[0] + } + span.Debugf("no blobnode available at all") + return "" +} diff --git a/blobstore/scheduler/shard_repairer_test.go b/blobstore/scheduler/shard_repairer_test.go index 6f0b48835..268613b3e 100644 --- a/blobstore/scheduler/shard_repairer_test.go +++ b/blobstore/scheduler/shard_repairer_test.go @@ -17,7 +17,10 @@ package scheduler import ( "context" "encoding/json" + "errors" + "fmt" "os" + "sync" "testing" "time" @@ -25,6 +28,7 @@ import ( "github.com/golang/mock/gomock" "github.com/stretchr/testify/require" + clustermgr "github.com/cubefs/cubefs/blobstore/api/clustermgr" "github.com/cubefs/cubefs/blobstore/common/codemode" "github.com/cubefs/cubefs/blobstore/common/counter" errcode "github.com/cubefs/cubefs/blobstore/common/errors" @@ -35,6 +39,7 @@ import ( "github.com/cubefs/cubefs/blobstore/scheduler/client" "github.com/cubefs/cubefs/blobstore/testing/mocks" "github.com/cubefs/cubefs/blobstore/util/closer" + "github.com/cubefs/cubefs/blobstore/util/selector" "github.com/cubefs/cubefs/blobstore/util/taskpool" ) @@ -46,8 +51,8 @@ func newShardRepairMgr(t *testing.T) *ShardRepairMgr { clusterTopology.EXPECT().UpdateVolume(any).AnyTimes().Return(&client.VolumeInfoSimple{}, nil) clusterTopology.EXPECT().GetDisk(any).AnyTimes().Return(&client.DiskInfoSimple{}, true) - selector := mocks.NewMockSelector(ctr) - selector.EXPECT().GetRandomN(any).AnyTimes().Return([]string{"http://127.0.0.1:9600"}) + mockSelector := mocks.NewMockSelector(ctr) + mockSelector.EXPECT().GetRandomN(any).AnyTimes().Return([]string{"http://127.0.0.1:9600"}) blobnode := NewMockBlobnodeAPI(ctr) blobnode.EXPECT().RepairShard(any, any, any).AnyTimes().Return(nil) @@ -64,12 +69,17 @@ func newShardRepairMgr(t *testing.T) *ShardRepairMgr { clusterMgrCli := NewMockClusterMgrAPI(ctr) clusterMgrCli.EXPECT().GetConfig(any, any).AnyTimes().Return("", nil) + clusterMgrCli.EXPECT().GetService(any, any, any).AnyTimes().Return(nil, nil) switchMgr := taskswitch.NewSwitchMgr(clusterMgrCli) taskSwitch, _ := switchMgr.AddSwitch(proto.TaskTypeBlobDelete.String()) return &ShardRepairMgr{ - clusterTopology: clusterTopology, - blobnodeSelector: selector, + clusterTopology: clusterTopology, + blobnodeSelector: &idcSelector{ + clusterMgrCli: clusterMgrCli, + selectors: map[string]selector.Selector{"": mockSelector}, + }, + clusterMgrCli: clusterMgrCli, blobnodeCli: blobnode, failMsgSender: sender, kafkaConsumerClient: kafkaClient, @@ -240,7 +250,6 @@ func TestNewShardRepairMgr(t *testing.T) { blobnode.EXPECT().RepairShard(any, any, any).AnyTimes().Return(nil) clusterCli := NewMockClusterMgrAPI(ctr) - clusterCli.EXPECT().GetService(any, any, any).Return(nil, errMock) clusterCli.EXPECT().GetConsumeOffset(any, any, any).AnyTimes().Return(int64(0), nil) clusterCli.EXPECT().SetConsumeOffset(any, any, any, any).AnyTimes().Return(nil) @@ -253,6 +262,9 @@ func TestNewShardRepairMgr(t *testing.T) { require.NoError(t, err) require.False(t, mgr.Enabled()) + // verify blobnodeSelector was initialized + require.NotNil(t, mgr.blobnodeSelector) + // get stats mgr.GetErrorStats() mgr.GetTaskStats() @@ -362,6 +374,107 @@ func TestProcessDiskNotFoundErr(t *testing.T) { } } +func TestIDCSelectorInit(t *testing.T) { + t.Run("should init without calling GetService", func(t *testing.T) { + ctr := gomock.NewController(t) + cmCli := NewMockClusterMgrAPI(ctr) + sel := newIDCSelector(cmCli, proto.ClusterID(1)) + require.NotNil(t, sel) + require.Empty(t, sel.selectors) + }) + + t.Run("should build per-IDC selectors from worker GetService data", func(t *testing.T) { + ctr := gomock.NewController(t) + cmCli := NewMockClusterMgrAPI(ctr) + cmCli.EXPECT().GetService(gomock.Any(), proto.ServiceNameWorker, gomock.Any()).AnyTimes(). + Return([]clustermgr.ServiceNode{ + {ClusterID: 1, Host: "worker-a:9600", Idc: "az0"}, + {ClusterID: 1, Host: "worker-b:9600", Idc: "az1"}, + }, nil) + + sel := newIDCSelector(cmCli, proto.ClusterID(1)) + + // az0 selector must return the az0 worker only + host := sel.get(context.Background(), "az0") + require.Equal(t, "worker-a:9600", host) + // az1 selector must return the az1 worker only + host = sel.get(context.Background(), "az1") + require.Equal(t, "worker-b:9600", host) + // global selector must return all workers + host = sel.get(context.Background(), "") + require.NotEmpty(t, host) + }) + + t.Run("should handle workers only", func(t *testing.T) { + ctr := gomock.NewController(t) + cmCli := NewMockClusterMgrAPI(ctr) + cmCli.EXPECT().GetService(gomock.Any(), proto.ServiceNameWorker, gomock.Any()).AnyTimes(). + Return([]clustermgr.ServiceNode{ + {ClusterID: 1, Host: "worker-a:9600", Idc: "az0"}, + {ClusterID: 1, Host: "worker-b:9600", Idc: "az1"}, + }, nil) + + sel := newIDCSelector(cmCli, proto.ClusterID(1)) + + host := sel.get(context.Background(), "az0") + require.NotEmpty(t, host) + host = sel.get(context.Background(), "") + require.NotEmpty(t, host) + }) +} + +func TestIDCSelectorConcurrent(t *testing.T) { + ctr := gomock.NewController(t) + cmCli := NewMockClusterMgrAPI(ctr) + cmCli.EXPECT().GetService(gomock.Any(), proto.ServiceNameWorker, gomock.Any()).AnyTimes(). + Return([]clustermgr.ServiceNode{ + {ClusterID: 1, Host: "worker-a:9600", Idc: "az0"}, + {ClusterID: 1, Host: "worker-b:9600", Idc: "az1"}, + }, nil) + + sel := newIDCSelector(cmCli, proto.ClusterID(1)) + + const goroutines = 16 + var wg sync.WaitGroup + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + host := sel.get(context.Background(), "az0") + if host == "" { + t.Error("concurrent get az0 returned empty") + } + host = sel.get(context.Background(), "az1") + if host == "" { + t.Error("concurrent get az1 returned empty") + } + host = sel.get(context.Background(), "") + if host == "" { + t.Error("concurrent get global returned empty") + } + }() + } + wg.Wait() +} + +func TestIDCSelectorRefreshError(t *testing.T) { + ctr := gomock.NewController(t) + cmCli := NewMockClusterMgrAPI(ctr) + cmCli.EXPECT().GetService(gomock.Any(), proto.ServiceNameWorker, gomock.Any()).AnyTimes(). + Return(nil, errors.New("network error")) + + sel := newIDCSelector(cmCli, proto.ClusterID(1)) + + // MakeSelector calls the getter immediately and ignores the error, + // so the selector is created with empty cachedValues. + host := sel.get(context.Background(), "az0") + require.Empty(t, host, "should return empty when refresh fails") + + // Fallback to global should also fail since the same cmCli is used + host = sel.get(context.Background(), "") + require.Empty(t, host, "should return empty when global refresh also fails") +} + func TestTryRepair(t *testing.T) { ctx := context.Background() ctr := gomock.NewController(t) @@ -369,9 +482,9 @@ func TestTryRepair(t *testing.T) { { // no host for shard repair mgr := newShardRepairMgr(t) - selector := mocks.NewMockSelector(ctr) - selector.EXPECT().GetRandomN(any).Return(nil) - mgr.blobnodeSelector = selector + mockSel := mocks.NewMockSelector(ctr) + mockSel.EXPECT().GetRandomN(any).AnyTimes().Return(nil) + mgr.blobnodeSelector.selectors[""] = mockSel doneVolume, err := mgr.tryRepair(ctx, volume, &proto.ShardRepairMsg{Bid: proto.BlobID(1), Vid: proto.Vid(1), BadIdx: []uint8{0}}) require.ErrorIs(t, err, ErrBlobnodeServiceUnavailable) require.True(t, doneVolume.EqualWith(volume)) @@ -465,3 +578,327 @@ func TestTryRepair(t *testing.T) { require.ErrorIs(t, err, errcode.ErrNoSuchDisk) } } + +func TestResolveRepairIDC(t *testing.T) { + vol := MockGenVolInfo(1, codemode.EC3P3, proto.VolumeStatusActive) + for i := range vol.VunitLocations { + vol.VunitLocations[i].Host = fmt.Sprintf("http://host-%d:9600", i) + } + + t.Run("should return empty when no bad index", func(t *testing.T) { + mgr := newShardRepairMgr(t) + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1} + idc := mgr.resolveRepairIDC(context.Background(), msg, nil) + require.Equal(t, "", idc) + }) + + t.Run("should return empty when bad index out of range", func(t *testing.T) { + mgr := newShardRepairMgr(t) + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{99}} + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + idc := mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx) + require.Equal(t, "", idc) + }) + + t.Run("should return AZ when disk found", func(t *testing.T) { + mgr := newShardRepairMgr(t) + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + idcByVunitIdx[0] = "az0" + idcByVunitIdx[1] = "az1" + + { + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}} + idc := mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx) + require.Equal(t, "az0", idc) + } + { + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{1}} + idc := mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx) + require.Equal(t, "az1", idc) + } + }) + + t.Run("should return empty when no disk has IDC", func(t *testing.T) { + mgr := newShardRepairMgr(t) + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}} + idc := mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx) + require.Equal(t, "", idc) + }) + + t.Run("should iterate bad indices to find valid IDC", func(t *testing.T) { + mgr := newShardRepairMgr(t) + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + idcByVunitIdx[2] = "az2" + + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0, 1, 2}} + idc := mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx) + require.Equal(t, "az2", idc) + }) +} + +func TestPickWorkerHost(t *testing.T) { + vol := MockGenVolInfo(1, codemode.EC3P3, proto.VolumeStatusActive) + for i := range vol.VunitLocations { + vol.VunitLocations[i].Host = fmt.Sprintf("http://host-%d:9600", i) + } + + t.Run("should select same AZ worker when available", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + + mockAz0 := mocks.NewMockSelector(ctr) + mockAz0.EXPECT().GetRandomN(1).Return([]string{"az0-worker"}) + mgr.blobnodeSelector.selectors["az0"] = mockAz0 + + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + idcByVunitIdx[0] = "az0" + + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}} + host := mgr.blobnodeSelector.get(context.Background(), mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx)) + require.Equal(t, "az0-worker", host) + }) + + t.Run("should fallback to global random when target AZ has no worker", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + + mockAz0 := mocks.NewMockSelector(ctr) + mockAz0.EXPECT().GetRandomN(1).Return(nil) + mgr.blobnodeSelector.selectors["az0"] = mockAz0 + + mockGlobal := mocks.NewMockSelector(ctr) + mockGlobal.EXPECT().GetRandomN(1).Return([]string{"fallback-worker"}) + mgr.blobnodeSelector.selectors[""] = mockGlobal + + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + idcByVunitIdx[0] = "az0" + + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}} + host := mgr.blobnodeSelector.get(context.Background(), mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx)) + require.Equal(t, "fallback-worker", host) + }) + + t.Run("should route different AZ damages to different AZ workers", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + + mockAz0 := mocks.NewMockSelector(ctr) + mockAz0.EXPECT().GetRandomN(1).Return([]string{"az0-worker"}) + mgr.blobnodeSelector.selectors["az0"] = mockAz0 + + mockAz1 := mocks.NewMockSelector(ctr) + mockAz1.EXPECT().GetRandomN(1).Return([]string{"az1-worker"}) + mgr.blobnodeSelector.selectors["az1"] = mockAz1 + + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + idcByVunitIdx[0] = "az0" + idcByVunitIdx[3] = "az1" + + { + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}} + host := mgr.blobnodeSelector.get(context.Background(), mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx)) + require.Equal(t, "az0-worker", host) + } + { + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{3}} + host := mgr.blobnodeSelector.get(context.Background(), mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx)) + require.Equal(t, "az1-worker", host) + } + }) + + t.Run("should return empty when all selectors have no hosts", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + + mockAz0 := mocks.NewMockSelector(ctr) + mockAz0.EXPECT().GetRandomN(1).Return(nil) + mgr.blobnodeSelector.selectors["az0"] = mockAz0 + + mockGlobal := mocks.NewMockSelector(ctr) + mockGlobal.EXPECT().GetRandomN(1).AnyTimes().Return(nil) + mgr.blobnodeSelector.selectors[""] = mockGlobal + + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + idcByVunitIdx[0] = "az0" + + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}} + host := mgr.blobnodeSelector.get(context.Background(), mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx)) + require.Equal(t, "", host) + }) + + t.Run("should lazy refresh and find worker when fallback empty", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + mockGlobal := mocks.NewMockSelector(ctr) + mockGlobal.EXPECT().GetRandomN(any).AnyTimes().Return(nil) + mgr.blobnodeSelector.selectors[""] = mockGlobal // global fallback always returns no hosts + + // getSelector on "az0" will lazily call GetService(Worker) + cmCli := NewMockClusterMgrAPI(ctr) + cmCli.EXPECT().GetService(gomock.Any(), proto.ServiceNameWorker, gomock.Any()). + Return([]clustermgr.ServiceNode{ + {ClusterID: 1, Host: "lazy-worker:9600", Idc: "az0"}, + }, nil) + mgr.blobnodeSelector.clusterMgrCli = cmCli + + idcByVunitIdx := make([]string, len(vol.VunitLocations)) + idcByVunitIdx[0] = "az0" + + msg := &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}} + host := mgr.blobnodeSelector.get(context.Background(), mgr.resolveRepairIDC(context.Background(), msg, idcByVunitIdx)) + require.Equal(t, "lazy-worker:9600", host) + }) +} + +func TestRepairShardWithIDC(t *testing.T) { + t.Run("should repair via same AZ worker when IDC routing enabled", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + + blobnode := NewMockBlobnodeAPI(ctr) + blobnode.EXPECT().RepairShard(any, gomock.Eq("az0-worker"), any).Return(nil) + mgr.blobnodeCli = blobnode + + mockAz0 := mocks.NewMockSelector(ctr) + mockAz0.EXPECT().GetRandomN(1).Return([]string{"az0-worker"}) + mgr.blobnodeSelector.selectors["az0"] = mockAz0 + + vol := MockGenVolInfo(1, codemode.EC3P3, proto.VolumeStatusActive) + + newTopology := NewMockClusterTopology(ctr) + newTopology.EXPECT().GetVolume(any).AnyTimes().Return(&client.VolumeInfoSimple{}, nil) + newTopology.EXPECT().UpdateVolume(any).AnyTimes().Return(&client.VolumeInfoSimple{}, nil) + for i := range vol.VunitLocations { + host := fmt.Sprintf("http://host-%d:9600", i) + vol.VunitLocations[i].Host = host + diskInfo := &client.DiskInfoSimple{Host: host} + if i == 0 { + diskInfo.Idc = "az0" + } + newTopology.EXPECT().GetDisk(vol.VunitLocations[i].DiskID).Return(diskInfo, true).AnyTimes() + } + mgr.clusterTopology = newTopology + + doneVol, err := mgr.repairShard(context.Background(), vol, + &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}}) + require.NoError(t, err) + require.NotNil(t, doneVol) + }) + + t.Run("should fallback to global when unknown disk (empty IDC)", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + // GetDisk returns DiskInfoSimple without Idc, so resolveRepairIDC returns "" + // Get(ctx, "") skips per-IDC and uses the global fallback + + vol := MockGenVolInfo(1, codemode.EC3P3, proto.VolumeStatusActive) + newTopology := NewMockClusterTopology(ctr) + newTopology.EXPECT().GetVolume(any).AnyTimes().Return(&client.VolumeInfoSimple{}, nil) + newTopology.EXPECT().UpdateVolume(any).AnyTimes().Return(&client.VolumeInfoSimple{}, nil) + for i := range vol.VunitLocations { + host := fmt.Sprintf("http://host-%d:9600", i) + vol.VunitLocations[i].Host = host + newTopology.EXPECT().GetDisk(vol.VunitLocations[i].DiskID).Return( + &client.DiskInfoSimple{Host: host}, true).AnyTimes() + } + mgr.clusterTopology = newTopology + + // Replace global selector with a mock to prove it was actually called + mockGlobal := mocks.NewMockSelector(ctr) + mockGlobal.EXPECT().GetRandomN(1).Return([]string{"http://127.0.0.1:9600"}).Times(1) + mgr.blobnodeSelector.selectors[""] = mockGlobal + + doneVol, err := mgr.repairShard(context.Background(), vol, + &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}}) + require.NoError(t, err) + require.NotNil(t, doneVol) + }) + + t.Run("should fallback to global when AZ has no worker", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + // selectors["az0"] is intentionally nil — triggers fallback + + vol := MockGenVolInfo(1, codemode.EC3P3, proto.VolumeStatusActive) + newTopology := NewMockClusterTopology(ctr) + newTopology.EXPECT().GetVolume(any).AnyTimes().Return(&client.VolumeInfoSimple{}, nil) + newTopology.EXPECT().UpdateVolume(any).AnyTimes().Return(&client.VolumeInfoSimple{}, nil) + for i := range vol.VunitLocations { + host := fmt.Sprintf("http://host-%d:9600", i) + vol.VunitLocations[i].Host = host + diskInfo := &client.DiskInfoSimple{Host: host} + if i == 0 { + diskInfo.Idc = "az0" + } + newTopology.EXPECT().GetDisk(vol.VunitLocations[i].DiskID).Return(diskInfo, true).AnyTimes() + } + mgr.clusterTopology = newTopology + + // Replace global selector with a mock to prove it was hit after the nil az0 selector + mockGlobal := mocks.NewMockSelector(ctr) + mockGlobal.EXPECT().GetRandomN(1).Return([]string{"http://127.0.0.1:9600"}).Times(1) + mgr.blobnodeSelector.selectors[""] = mockGlobal + + doneVol, err := mgr.repairShard(context.Background(), vol, + &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}}) + require.NoError(t, err) + require.NotNil(t, doneVol) + }) + + t.Run("should return unavailable when no host at all", func(t *testing.T) { + ctr := gomock.NewController(t) + mgr := newShardRepairMgr(t) + mockGlobal := mocks.NewMockSelector(ctr) + mockGlobal.EXPECT().GetRandomN(any).AnyTimes().Return(nil) + mgr.blobnodeSelector.selectors[""] = mockGlobal + + vol := MockGenVolInfo(1, codemode.EC3P3, proto.VolumeStatusActive) + for i := range vol.VunitLocations { + vol.VunitLocations[i].Host = fmt.Sprintf("http://host-%d:9600", i) + } + + doneVol, err := mgr.repairShard(context.Background(), vol, + &proto.ShardRepairMsg{Bid: 1, Vid: 1, BadIdx: []uint8{0}}) + require.ErrorIs(t, err, ErrBlobnodeServiceUnavailable) + require.NotNil(t, doneVol) + }) +} + +func BenchmarkPickWorkerHostParallel(b *testing.B) { + az0Sel := selector.MakeSelector(60*1000, func() ([]string, error) { + return []string{"az0-worker:9600"}, nil + }) + az1Sel := selector.MakeSelector(60*1000, func() ([]string, error) { + return []string{"az1-worker:9600"}, nil + }) + globalSel := selector.MakeSelector(60*1000, func() ([]string, error) { + return []string{"any-worker:9600"}, nil + }) + sel := &idcSelector{ + selectors: map[string]selector.Selector{ + "az0": az0Sel, + "az1": az1Sel, + "": globalSel, + }, + } + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + host := sel.get(context.Background(), "az0") + if host == "" { + b.Error("Get az0 returned empty") + } + host = sel.get(context.Background(), "az1") + if host == "" { + b.Error("Get az1 returned empty") + } + host = sel.get(context.Background(), "") + if host == "" { + b.Error("Get global returned empty") + } + } + }) +}