gian/cubefs
1
0
Fork 0
mirror of https://github.com/cubefs/cubefs.git synced 2026-08-02 02:00:56 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat(blobstore): add AZ-aware blobnode selection for shard repair

Browse Source 
The scheduler now routes shard repair tasks to blobnode workers
in the same availability zone (IDC) as the damaged shard,
reducing cross-AZ data transfer during local stripe repair.

. #1001113674

Signed-off-by: yujrchyang <yujrchyang@gmail.com>
This commit is contained in:
yujrchyang 2026-06-22 15:04:36 +08:00 committed by slasher
parent 898b19eb44
commit ff1b231007
5 changed files with 580 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
blobstore/scheduler/client/clustermgr.go
View File

@ -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

6
blobstore/scheduler/client/clustermgr_test.go
View File

@ -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))

4
blobstore/scheduler/client_mock_test.go
View File

@ -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
}

135
blobstore/scheduler/shard_repairer.go
View File

@ -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 ""
}

453
blobstore/scheduler/shard_repairer_test.go
View File

@ -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")
}
}
})
}