Files
seaweedfs/weed/worker/tasks/balance/plugin_handler_test.go
Chris LuandGitHub 1f6f473995 refactor(worker): co-locate plugin handlers with their task packages (#9301)
* refactor(worker): co-locate plugin handlers with their task packages

Move every per-task plugin handler from weed/plugin/worker/ into the
matching weed/worker/tasks/<name>/ package, so each task owns its
detection, scheduling, execution, and plugin handler in one place.

Step 0 (within pluginworker, no behavior change): extract shared helpers
that previously lived inside individual handler files into dedicated
files and export the ones now consumed across packages.

  - activity.go: BuildExecutorActivity, BuildDetectorActivity
  - config.go: ReadStringConfig/Double/Int64/Bytes/StringList, MapTaskPriority
  - interval.go: ShouldSkipDetectionByInterval
  - volume_state.go: VolumeState + consts, FilterMetricsByVolumeState/Location
  - collection_filter.go: CollectionFilterMode + consts
  - volume_metrics.go: export CollectVolumeMetricsFromMasters,
    MasterAddressCandidates, FetchVolumeList
  - testing_senders_test.go: shared test stubs

Phase 1: move the per-task plugin handlers (and the iceberg subpackage)
into their task packages.

  weed/plugin/worker/vacuum_handler.go         -> weed/worker/tasks/vacuum/plugin_handler.go
  weed/plugin/worker/ec_balance_handler.go     -> weed/worker/tasks/ec_balance/plugin_handler.go
  weed/plugin/worker/erasure_coding_handler.go -> weed/worker/tasks/erasure_coding/plugin_handler.go
  weed/plugin/worker/volume_balance_handler.go -> weed/worker/tasks/balance/plugin_handler.go
  weed/plugin/worker/iceberg/                   -> weed/worker/tasks/iceberg/

  weed/plugin/worker/handlers/handlers.go now blank-imports all five
  task subpackages so their init() registrations fire.

  weed/command/mini.go and the worker tests construct the handler with
  vacuum.DefaultMaxExecutionConcurrency (the constant moved with the
  vacuum handler).

admin_script remains in weed/plugin/worker/ because there is no
underlying weed/worker/tasks/admin_script/ package to merge with.

* refactor(worker): update test/plugin_workers imports for moved handlers

Three handler constructors moved out of pluginworker into their task
packages — update the integration test files in test/plugin_workers/
to import from the new locations:

  pluginworker.NewVacuumHandler        -> vacuum.NewVacuumHandler
  pluginworker.NewVolumeBalanceHandler -> balance.NewVolumeBalanceHandler
  pluginworker.NewErasureCodingHandler -> erasure_coding.NewErasureCodingHandler

The pluginworker import is kept where the file still uses
pluginworker.WorkerOptions / pluginworker.JobHandler.

* refactor(worker): update test/s3tables iceberg import path

The iceberg subpackage moved from weed/plugin/worker/iceberg/ to
weed/worker/tasks/iceberg/. test/s3tables/maintenance/maintenance_integration_test.go
still imported the old path, breaking S3 Tables / RisingWave / Trino /
Spark / Iceberg-catalog / STS integration test builds.

Mirrors the OSS-side fix needed by every job in the run that
transitively imports test/s3tables/maintenance.

* chore: gofmt PR-touched files

The S3 Tables Format Check job runs `gofmt -l` over weed/s3api/s3tables
and test/s3tables, then fails if anything is unformatted. Files this
PR moved or modified had import-grouping and trailing-spacing issues
introduced by perl-based renames; reformat them with gofmt -w.

Touched files:
  test/plugin_workers/erasure_coding/{detection,execution}_test.go
  test/s3tables/maintenance/maintenance_integration_test.go
  weed/plugin/worker/handlers/handlers.go
  weed/worker/tasks/{balance,ec_balance,erasure_coding,vacuum}/plugin_handler*.go

* refactor(worker): bounds-checked int conversions for plugin config values

CodeQL flagged 18 go/incorrect-integer-conversion warnings on the moved
plugin handler files: results of pluginworker.ReadInt64Config (which
ultimately calls strconv.ParseInt with bit size 64) were being narrowed
to int32/uint32/int without an upper-bound check, so a malicious or
malformed admin/worker config value could overflow the target type.

Add three helpers in weed/plugin/worker/config.go that wrap
ReadInt64Config and clamp out-of-range values back to the caller's
fallback:

  ReadInt32Config (math.MinInt32 .. math.MaxInt32)
  ReadUint32Config (0 .. math.MaxUint32)
  ReadIntConfig    (math.MinInt32 .. math.MaxInt32, platform-portable)

Update each flagged call site in the four moved task packages to use
the bounds-checked helper. For protobuf uint32 fields (volume IDs)
the variable type also becomes uint32, removing the trailing
uint32(volumeID) casts and changing the "missing volume_id" check
from `<= 0` to `== 0`.

Touched files:
  weed/plugin/worker/config.go
  weed/worker/tasks/balance/plugin_handler.go
  weed/worker/tasks/erasure_coding/plugin_handler.go
  weed/worker/tasks/vacuum/plugin_handler.go

* refactor(worker): use ReadIntConfig for clamped derive-worker-config helpers

CodeQL still flagged three call sites where ReadInt64Config was being
narrowed to int after a value-range clamp (max_concurrent_moves <= 50,
batch_size <= 100, min_server_count >= 2). The clamp is correct but
CodeQL's flow analysis didn't recognize the bound, so it flagged them
as unbounded narrowing.

Switch to ReadIntConfig (already int32-bounded by the helper) for
those three sites, drop the now-redundant int64 intermediate variables.

Also drops the now-unused `> math.MaxInt32` clamp in
ec_balance.deriveECBalanceWorkerConfig (the helper covers it).
2026-05-02 18:03:13 -07:00

795 lines
24 KiB
Go

package balance
import (
"context"
"fmt"
"strings"
"sync"
"testing"
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
"github.com/seaweedfs/seaweedfs/weed/pb/worker_pb"
pluginworker "github.com/seaweedfs/seaweedfs/weed/plugin/worker"
workertypes "github.com/seaweedfs/seaweedfs/weed/worker/types"
"google.golang.org/protobuf/proto"
)
func TestDecodeVolumeBalanceTaskParamsFromPayload(t *testing.T) {
expected := &worker_pb.TaskParams{
TaskId: "task-1",
VolumeId: 42,
Collection: "photos",
Sources: []*worker_pb.TaskSource{
{
Node: "10.0.0.1:8080",
VolumeId: 42,
},
},
Targets: []*worker_pb.TaskTarget{
{
Node: "10.0.0.2:8080",
VolumeId: 42,
},
},
TaskParams: &worker_pb.TaskParams_BalanceParams{
BalanceParams: &worker_pb.BalanceTaskParams{
ForceMove: true,
TimeoutSeconds: 1200,
},
},
}
payload, err := proto.Marshal(expected)
if err != nil {
t.Fatalf("marshal payload: %v", err)
}
job := &plugin_pb.JobSpec{
JobId: "job-from-admin",
Parameters: map[string]*plugin_pb.ConfigValue{
"task_params_pb": {Kind: &plugin_pb.ConfigValue_BytesValue{BytesValue: payload}},
},
}
actual, err := decodeVolumeBalanceTaskParams(job)
if err != nil {
t.Fatalf("decodeVolumeBalanceTaskParams() err = %v", err)
}
if !proto.Equal(expected, actual) {
t.Fatalf("decoded params mismatch\nexpected: %+v\nactual: %+v", expected, actual)
}
}
func TestDecodeVolumeBalanceTaskParamsFallback(t *testing.T) {
job := &plugin_pb.JobSpec{
JobId: "job-2",
Parameters: map[string]*plugin_pb.ConfigValue{
"volume_id": {Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 7}},
"source_server": {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: "127.0.0.1:8080"}},
"target_server": {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: "127.0.0.2:8080"}},
"collection": {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: "videos"}},
},
}
params, err := decodeVolumeBalanceTaskParams(job)
if err != nil {
t.Fatalf("decodeVolumeBalanceTaskParams() err = %v", err)
}
if params.TaskId != "job-2" || params.VolumeId != 7 || params.Collection != "videos" {
t.Fatalf("unexpected basic params: %+v", params)
}
if len(params.Sources) != 1 || params.Sources[0].Node != "127.0.0.1:8080" {
t.Fatalf("unexpected sources: %+v", params.Sources)
}
if len(params.Targets) != 1 || params.Targets[0].Node != "127.0.0.2:8080" {
t.Fatalf("unexpected targets: %+v", params.Targets)
}
if params.GetBalanceParams() == nil {
t.Fatalf("expected fallback balance params")
}
}
func TestDeriveBalanceWorkerConfig(t *testing.T) {
values := map[string]*plugin_pb.ConfigValue{
"imbalance_threshold": {
Kind: &plugin_pb.ConfigValue_DoubleValue{DoubleValue: 0.45},
},
"min_server_count": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 5},
},
}
cfg := deriveBalanceWorkerConfig(values)
if cfg.TaskConfig.ImbalanceThreshold != 0.45 {
t.Fatalf("expected imbalance_threshold 0.45, got %v", cfg.TaskConfig.ImbalanceThreshold)
}
if cfg.TaskConfig.MinServerCount != 5 {
t.Fatalf("expected min_server_count 5, got %d", cfg.TaskConfig.MinServerCount)
}
// Defaults for batch config when not specified
if cfg.MaxConcurrentMoves != defaultMaxConcurrentMoves {
t.Fatalf("expected default max_concurrent_moves %d, got %d", defaultMaxConcurrentMoves, cfg.MaxConcurrentMoves)
}
if cfg.BatchSize != 20 {
t.Fatalf("expected default batch_size 20, got %d", cfg.BatchSize)
}
}
func TestDeriveBalanceWorkerConfigBatchFields(t *testing.T) {
values := map[string]*plugin_pb.ConfigValue{
"max_concurrent_moves": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 10},
},
"batch_size": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 50},
},
}
cfg := deriveBalanceWorkerConfig(values)
if cfg.MaxConcurrentMoves != 10 {
t.Fatalf("expected max_concurrent_moves 10, got %d", cfg.MaxConcurrentMoves)
}
if cfg.BatchSize != 50 {
t.Fatalf("expected batch_size 50, got %d", cfg.BatchSize)
}
}
func TestDeriveBalanceWorkerConfigBatchClamping(t *testing.T) {
values := map[string]*plugin_pb.ConfigValue{
"max_concurrent_moves": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 999},
},
"batch_size": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 0},
},
}
cfg := deriveBalanceWorkerConfig(values)
if cfg.MaxConcurrentMoves != 50 {
t.Fatalf("expected max_concurrent_moves clamped to 50, got %d", cfg.MaxConcurrentMoves)
}
if cfg.BatchSize != 1 {
t.Fatalf("expected batch_size clamped to 1, got %d", cfg.BatchSize)
}
}
func makeDetectionResult(volumeID uint32, source, target, collection string) *workertypes.TaskDetectionResult {
return &workertypes.TaskDetectionResult{
TaskID: fmt.Sprintf("balance-%d", volumeID),
TaskType: workertypes.TaskTypeBalance,
VolumeID: volumeID,
Server: source,
Collection: collection,
Priority: workertypes.TaskPriorityNormal,
Reason: "imbalanced",
TypedParams: &worker_pb.TaskParams{
VolumeId: volumeID,
Collection: collection,
VolumeSize: 1024,
Sources: []*worker_pb.TaskSource{
{Node: source, VolumeId: volumeID},
},
Targets: []*worker_pb.TaskTarget{
{Node: target, VolumeId: volumeID},
},
TaskParams: &worker_pb.TaskParams_BalanceParams{
BalanceParams: &worker_pb.BalanceTaskParams{TimeoutSeconds: 600},
},
},
}
}
func TestBuildBatchVolumeBalanceProposals_SingleBatch(t *testing.T) {
results := []*workertypes.TaskDetectionResult{
makeDetectionResult(1, "s1:8080", "t1:8080", "c1"),
makeDetectionResult(2, "s2:8080", "t2:8080", "c1"),
makeDetectionResult(3, "s1:8080", "t2:8080", "c1"),
}
proposals := buildBatchVolumeBalanceProposals(results, 10, 5)
if len(proposals) != 1 {
t.Fatalf("expected 1 batch proposal, got %d", len(proposals))
}
p := proposals[0]
if p.Labels["batch"] != "true" {
t.Fatalf("expected batch label")
}
if p.Labels["batch_size"] != "3" {
t.Fatalf("expected batch_size label '3', got %q", p.Labels["batch_size"])
}
// Decode and verify moves
payload := p.Parameters["task_params_pb"].GetBytesValue()
if len(payload) == 0 {
t.Fatalf("expected task_params_pb payload")
}
decoded := &worker_pb.TaskParams{}
if err := proto.Unmarshal(payload, decoded); err != nil {
t.Fatalf("unmarshal: %v", err)
}
moves := decoded.GetBalanceParams().GetMoves()
if len(moves) != 3 {
t.Fatalf("expected 3 moves, got %d", len(moves))
}
if moves[0].VolumeId != 1 || moves[1].VolumeId != 2 || moves[2].VolumeId != 3 {
t.Fatalf("unexpected volume IDs: %v", moves)
}
if decoded.GetBalanceParams().MaxConcurrentMoves != 5 {
t.Fatalf("expected MaxConcurrentMoves 5, got %d", decoded.GetBalanceParams().MaxConcurrentMoves)
}
}
func TestBuildBatchVolumeBalanceProposals_MultipleBatches(t *testing.T) {
results := make([]*workertypes.TaskDetectionResult, 5)
for i := range results {
results[i] = makeDetectionResult(uint32(i+1), "s1:8080", "t1:8080", "c1")
}
proposals := buildBatchVolumeBalanceProposals(results, 2, 3)
// 5 results / batch_size 2 = 3 proposals (2, 2, 1)
if len(proposals) != 3 {
t.Fatalf("expected 3 proposals, got %d", len(proposals))
}
// First two should be batch proposals
if proposals[0].Labels["batch"] != "true" {
t.Fatalf("first proposal should be batch")
}
if proposals[1].Labels["batch"] != "true" {
t.Fatalf("second proposal should be batch")
}
// Last one has only 1 result, should fall back to single-move proposal
if proposals[2].Labels["batch"] == "true" {
t.Fatalf("last proposal with 1 result should be single-move, not batch")
}
}
func TestBuildBatchVolumeBalanceProposals_BatchSizeOne(t *testing.T) {
results := []*workertypes.TaskDetectionResult{
makeDetectionResult(1, "s1:8080", "t1:8080", "c1"),
makeDetectionResult(2, "s2:8080", "t2:8080", "c1"),
}
// batch_size=1 should not be called (Detect guards this), but test the function directly
proposals := buildBatchVolumeBalanceProposals(results, 1, 5)
// Each result becomes its own single-move proposal
if len(proposals) != 2 {
t.Fatalf("expected 2 proposals, got %d", len(proposals))
}
}
func TestVolumeBalanceDescriptorHasBatchFields(t *testing.T) {
descriptor := NewVolumeBalanceHandler(nil).Descriptor()
if !workerConfigFormHasField(descriptor.WorkerConfigForm, "max_concurrent_moves") {
t.Fatalf("expected max_concurrent_moves in worker config form")
}
if !workerConfigFormHasField(descriptor.WorkerConfigForm, "batch_size") {
t.Fatalf("expected batch_size in worker config form")
}
}
func TestBuildVolumeBalanceProposal(t *testing.T) {
params := &worker_pb.TaskParams{
TaskId: "balance-task-1",
VolumeId: 55,
Collection: "images",
Sources: []*worker_pb.TaskSource{
{
Node: "source-a:8080",
VolumeId: 55,
},
},
Targets: []*worker_pb.TaskTarget{
{
Node: "target-b:8080",
VolumeId: 55,
},
},
TaskParams: &worker_pb.TaskParams_BalanceParams{
BalanceParams: &worker_pb.BalanceTaskParams{
TimeoutSeconds: 600,
},
},
}
result := &workertypes.TaskDetectionResult{
TaskID: "balance-task-1",
TaskType: workertypes.TaskTypeBalance,
VolumeID: 55,
Server: "source-a",
Collection: "images",
Priority: workertypes.TaskPriorityHigh,
Reason: "imbalanced load",
TypedParams: params,
}
proposal, err := buildVolumeBalanceProposal(result)
if err != nil {
t.Fatalf("buildVolumeBalanceProposal() err = %v", err)
}
if proposal.JobType != "volume_balance" {
t.Fatalf("unexpected job type %q", proposal.JobType)
}
if proposal.DedupeKey == "" {
t.Fatalf("expected dedupe key")
}
if proposal.Parameters["task_params_pb"] == nil {
t.Fatalf("expected serialized task params")
}
if proposal.Labels["source_node"] != "source-a:8080" {
t.Fatalf("unexpected source label %q", proposal.Labels["source_node"])
}
if proposal.Labels["target_node"] != "target-b:8080" {
t.Fatalf("unexpected target label %q", proposal.Labels["target_node"])
}
}
func TestVolumeBalanceHandlerRejectsUnsupportedJobType(t *testing.T) {
handler := NewVolumeBalanceHandler(nil)
err := handler.Detect(context.Background(), &plugin_pb.RunDetectionRequest{
JobType: "vacuum",
}, noopDetectionSender{})
if err == nil {
t.Fatalf("expected detect job type mismatch error")
}
err = handler.Execute(context.Background(), &plugin_pb.ExecuteJobRequest{
Job: &plugin_pb.JobSpec{JobId: "job-1", JobType: "vacuum"},
}, noopExecutionSender{})
if err == nil {
t.Fatalf("expected execute job type mismatch error")
}
}
func TestEmitVolumeBalanceDetectionDecisionTraceNoTasks(t *testing.T) {
sender := &recordingDetectionSender{}
config := NewDefaultConfig()
config.ImbalanceThreshold = 0.2
config.MinServerCount = 2
metrics := []*workertypes.VolumeHealthMetrics{
{VolumeID: 1, Server: "server-a", DiskType: "hdd"},
{VolumeID: 2, Server: "server-a", DiskType: "hdd"},
{VolumeID: 3, Server: "server-b", DiskType: "hdd"},
{VolumeID: 4, Server: "server-b", DiskType: "hdd"},
}
if err := emitVolumeBalanceDetectionDecisionTrace(sender, metrics, nil, config, nil); err != nil {
t.Fatalf("emitVolumeBalanceDetectionDecisionTrace error: %v", err)
}
if len(sender.events) < 2 {
t.Fatalf("expected at least 2 detection events, got %d", len(sender.events))
}
if sender.events[0].Source != plugin_pb.ActivitySource_ACTIVITY_SOURCE_DETECTOR {
t.Fatalf("expected detector source, got %v", sender.events[0].Source)
}
if !strings.Contains(sender.events[0].Message, "BALANCE: No tasks created for 4 volumes") {
t.Fatalf("unexpected summary message: %q", sender.events[0].Message)
}
foundDiskTypeDecision := false
for _, event := range sender.events {
if strings.Contains(event.Message, "BALANCE [hdd]: No tasks created - cluster well balanced") {
foundDiskTypeDecision = true
break
}
}
if !foundDiskTypeDecision {
t.Fatalf("expected per-disk-type decision message")
}
}
func TestVolumeBalanceDescriptorOmitsExecutionTuningFields(t *testing.T) {
descriptor := NewVolumeBalanceHandler(nil).Descriptor()
if descriptor == nil || descriptor.WorkerConfigForm == nil {
t.Fatalf("expected worker config form in descriptor")
}
if workerConfigFormHasField(descriptor.WorkerConfigForm, "timeout_seconds") {
t.Fatalf("unexpected timeout_seconds in volume balance worker config form")
}
if workerConfigFormHasField(descriptor.WorkerConfigForm, "force_move") {
t.Fatalf("unexpected force_move in volume balance worker config form")
}
}
type recordingExecutionSender struct {
mu sync.Mutex
progress []*plugin_pb.JobProgressUpdate
completed *plugin_pb.JobCompleted
}
func (r *recordingExecutionSender) SendProgress(p *plugin_pb.JobProgressUpdate) error {
r.mu.Lock()
defer r.mu.Unlock()
r.progress = append(r.progress, proto.Clone(p).(*plugin_pb.JobProgressUpdate))
return nil
}
func (r *recordingExecutionSender) SendCompleted(c *plugin_pb.JobCompleted) error {
r.mu.Lock()
defer r.mu.Unlock()
r.completed = proto.Clone(c).(*plugin_pb.JobCompleted)
return nil
}
func TestBuildMoveTaskParams(t *testing.T) {
move := &worker_pb.BalanceMoveSpec{
VolumeId: 42,
SourceNode: "10.0.0.1:8080",
TargetNode: "10.0.0.2:8080",
Collection: "photos",
VolumeSize: 1024 * 1024,
}
outerParams := &worker_pb.BalanceTaskParams{
ForceMove: true,
TimeoutSeconds: 300,
}
params := buildMoveTaskParams(move, outerParams)
if params.VolumeId != 42 {
t.Fatalf("expected volume_id 42, got %d", params.VolumeId)
}
if params.Collection != "photos" {
t.Fatalf("expected collection photos, got %s", params.Collection)
}
if params.VolumeSize != 1024*1024 {
t.Fatalf("expected volume_size %d, got %d", 1024*1024, params.VolumeSize)
}
if len(params.Sources) != 1 || params.Sources[0].Node != "10.0.0.1:8080" {
t.Fatalf("unexpected sources: %+v", params.Sources)
}
if len(params.Targets) != 1 || params.Targets[0].Node != "10.0.0.2:8080" {
t.Fatalf("unexpected targets: %+v", params.Targets)
}
bp := params.GetBalanceParams()
if bp == nil {
t.Fatalf("expected balance params")
}
if bp.TimeoutSeconds != 300 {
t.Fatalf("expected timeout 300, got %d", bp.TimeoutSeconds)
}
if !bp.ForceMove {
t.Fatalf("expected force_move to be propagated from outer params")
}
}
func TestBuildMoveTaskParamsDefaultTimeout(t *testing.T) {
move := &worker_pb.BalanceMoveSpec{
VolumeId: 1,
SourceNode: "a:8080",
TargetNode: "b:8080",
}
params := buildMoveTaskParams(move, nil)
if params.GetBalanceParams().TimeoutSeconds != defaultBalanceTimeoutSeconds {
t.Fatalf("expected default timeout %d, got %d", defaultBalanceTimeoutSeconds, params.GetBalanceParams().TimeoutSeconds)
}
if params.GetBalanceParams().ForceMove {
t.Fatalf("expected force_move to default to false with nil outer params")
}
}
func TestExecuteDispatchesBatchPath(t *testing.T) {
// Build a job with batch moves in BalanceTaskParams
bp := &worker_pb.BalanceTaskParams{
TimeoutSeconds: 60,
MaxConcurrentMoves: 2,
Moves: []*worker_pb.BalanceMoveSpec{
{VolumeId: 1, SourceNode: "s1:8080", TargetNode: "t1:8080", Collection: "c1"},
{VolumeId: 2, SourceNode: "s2:8080", TargetNode: "t2:8080", Collection: "c1"},
},
}
taskParams := &worker_pb.TaskParams{
TaskId: "batch-1",
TaskParams: &worker_pb.TaskParams_BalanceParams{
BalanceParams: bp,
},
}
payload, err := proto.Marshal(taskParams)
if err != nil {
t.Fatalf("marshal: %v", err)
}
job := &plugin_pb.JobSpec{
JobId: "batch-job-1",
JobType: "volume_balance",
Parameters: map[string]*plugin_pb.ConfigValue{
"task_params_pb": {Kind: &plugin_pb.ConfigValue_BytesValue{BytesValue: payload}},
},
}
handler := NewVolumeBalanceHandler(nil)
sender := &recordingExecutionSender{}
// Execute will enter the batch path. It will fail because there's no real gRPC server,
// but we verify it sends the assigned progress and eventually a completion.
err = handler.Execute(context.Background(), &plugin_pb.ExecuteJobRequest{
Job: job,
}, sender)
// Expect an error since no real volume servers exist
// But verify the batch path was taken by checking the assigned message
sender.mu.Lock()
defer sender.mu.Unlock()
if len(sender.progress) == 0 {
t.Fatalf("expected progress messages from batch path")
}
// First progress should be "assigned" with batch info
first := sender.progress[0]
if first.Stage != "assigned" {
t.Fatalf("expected first stage 'assigned', got %q", first.Stage)
}
if !strings.Contains(first.Message, "batch") || !strings.Contains(first.Message, "2 moves") {
t.Fatalf("expected batch assigned message, got %q", first.Message)
}
// Should have a completion with failure details (since no servers)
if sender.completed == nil {
t.Fatalf("expected completion message")
}
if sender.completed.Success {
t.Fatalf("expected failure since no real gRPC servers")
}
// Should report 0 succeeded, 2 failed
if v, ok := sender.completed.Result.OutputValues["failed"]; !ok || v.GetInt64Value() != 2 {
t.Fatalf("expected 2 failed moves, got %+v", sender.completed.Result.OutputValues)
}
}
func TestExecuteSingleMovePathUnchanged(t *testing.T) {
// Build a single-move job (no batch moves)
taskParams := &worker_pb.TaskParams{
TaskId: "single-1",
VolumeId: 99,
Collection: "videos",
Sources: []*worker_pb.TaskSource{
{Node: "src:8080", VolumeId: 99},
},
Targets: []*worker_pb.TaskTarget{
{Node: "dst:8080", VolumeId: 99},
},
TaskParams: &worker_pb.TaskParams_BalanceParams{
BalanceParams: &worker_pb.BalanceTaskParams{
TimeoutSeconds: 60,
},
},
}
payload, err := proto.Marshal(taskParams)
if err != nil {
t.Fatalf("marshal: %v", err)
}
job := &plugin_pb.JobSpec{
JobId: "single-job-1",
JobType: "volume_balance",
Parameters: map[string]*plugin_pb.ConfigValue{
"task_params_pb": {Kind: &plugin_pb.ConfigValue_BytesValue{BytesValue: payload}},
},
}
handler := NewVolumeBalanceHandler(nil)
sender := &recordingExecutionSender{}
// Execute single-move path. Will fail on gRPC but verify it takes the single-move path.
_ = handler.Execute(context.Background(), &plugin_pb.ExecuteJobRequest{
Job: job,
}, sender)
sender.mu.Lock()
defer sender.mu.Unlock()
if len(sender.progress) == 0 {
t.Fatalf("expected progress messages")
}
// Single-move path sends "volume balance job accepted" not "batch volume balance"
first := sender.progress[0]
if first.Stage != "assigned" {
t.Fatalf("expected first stage 'assigned', got %q", first.Stage)
}
if strings.Contains(first.Message, "batch") {
t.Fatalf("single-move path should not mention batch, got %q", first.Message)
}
}
func TestFilterMetricsByLocation(t *testing.T) {
metrics := []*workertypes.VolumeHealthMetrics{
{VolumeID: 1, Server: "node-a", DataCenter: "dc1", Rack: "rack1"},
{VolumeID: 2, Server: "node-b", DataCenter: "dc1", Rack: "rack2"},
{VolumeID: 3, Server: "node-c", DataCenter: "dc2", Rack: "rack1"},
{VolumeID: 4, Server: "node-d", DataCenter: "dc2", Rack: "rack3"},
}
// Filter by DC
filtered := pluginworker.FilterMetricsByLocation(metrics, "dc1", "", "")
if len(filtered) != 2 {
t.Fatalf("DC filter: expected 2, got %d", len(filtered))
}
// Filter by rack
filtered = pluginworker.FilterMetricsByLocation(metrics, "", "rack1,rack2", "")
if len(filtered) != 3 {
t.Fatalf("rack filter: expected 3, got %d", len(filtered))
}
// Filter by node
filtered = pluginworker.FilterMetricsByLocation(metrics, "", "", "node-a,node-c")
if len(filtered) != 2 {
t.Fatalf("node filter: expected 2, got %d", len(filtered))
}
// Combined DC + rack
filtered = pluginworker.FilterMetricsByLocation(metrics, "dc2", "rack3", "")
if len(filtered) != 1 {
t.Fatalf("DC+rack filter: expected 1, got %d", len(filtered))
}
// Empty filters pass all
filtered = pluginworker.FilterMetricsByLocation(metrics, "", "", "")
if len(filtered) != 4 {
t.Fatalf("no filter: expected 4, got %d", len(filtered))
}
}
func TestFilterMetricsByVolumeState(t *testing.T) {
metrics := []*workertypes.VolumeHealthMetrics{
{VolumeID: 1, FullnessRatio: 0.5}, // active
{VolumeID: 2, FullnessRatio: 1.0}, // active (below 1.01)
{VolumeID: 3, FullnessRatio: 1.009}, // active (below 1.01)
{VolumeID: 4, FullnessRatio: 1.01}, // full (exactly at threshold)
{VolumeID: 5, FullnessRatio: 1.5}, // full
{VolumeID: 6, FullnessRatio: 2.0}, // full
}
tests := []struct {
name string
state string
expectedIDs []uint32
}{
{
name: "ALL returns everything",
state: "ALL",
expectedIDs: []uint32{1, 2, 3, 4, 5, 6},
},
{
name: "empty string returns everything",
state: "",
expectedIDs: []uint32{1, 2, 3, 4, 5, 6},
},
{
name: "ACTIVE keeps FullnessRatio below 1.01",
state: "ACTIVE",
expectedIDs: []uint32{1, 2, 3},
},
{
name: "FULL keeps FullnessRatio at or above 1.01",
state: "FULL",
expectedIDs: []uint32{4, 5, 6},
},
{
name: "unknown value returns everything",
state: "INVALID",
expectedIDs: []uint32{1, 2, 3, 4, 5, 6},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := pluginworker.FilterMetricsByVolumeState(metrics, pluginworker.VolumeState(tt.state))
if len(result) != len(tt.expectedIDs) {
t.Fatalf("expected %d metrics, got %d", len(tt.expectedIDs), len(result))
}
for i, m := range result {
if m.VolumeID != tt.expectedIDs[i] {
t.Errorf("result[%d].VolumeID = %d, want %d", i, m.VolumeID, tt.expectedIDs[i])
}
}
})
}
}
func TestFilterMetricsByVolumeState_NilElement(t *testing.T) {
metrics := []*workertypes.VolumeHealthMetrics{
nil,
{VolumeID: 1, FullnessRatio: 0.5},
nil,
{VolumeID: 2, FullnessRatio: 1.5},
}
result := pluginworker.FilterMetricsByVolumeState(metrics, pluginworker.VolumeStateActive)
if len(result) != 1 || result[0].VolumeID != 1 {
t.Fatalf("expected [vol 1] for ACTIVE with nil elements, got %d results", len(result))
}
result = pluginworker.FilterMetricsByVolumeState(metrics, pluginworker.VolumeStateFull)
if len(result) != 1 || result[0].VolumeID != 2 {
t.Fatalf("expected [vol 2] for FULL with nil elements, got %d results", len(result))
}
}
func TestFilterMetricsByVolumeState_EmptyInput(t *testing.T) {
result := pluginworker.FilterMetricsByVolumeState(nil, pluginworker.VolumeStateActive)
if len(result) != 0 {
t.Fatalf("expected 0 metrics for nil input, got %d", len(result))
}
result = pluginworker.FilterMetricsByVolumeState([]*workertypes.VolumeHealthMetrics{}, pluginworker.VolumeStateFull)
if len(result) != 0 {
t.Fatalf("expected 0 metrics for empty input, got %d", len(result))
}
}
func TestVolumeBalanceDescriptorHasVolumeStateField(t *testing.T) {
descriptor := NewVolumeBalanceHandler(nil).Descriptor()
if descriptor == nil || descriptor.AdminConfigForm == nil {
t.Fatalf("expected admin config form in descriptor")
}
found := false
for _, section := range descriptor.AdminConfigForm.Sections {
for _, field := range section.Fields {
if field.Name == "volume_state" {
found = true
break
}
}
}
if !found {
t.Fatalf("expected volume_state field in admin config form")
}
defaultVal, ok := descriptor.AdminConfigForm.DefaultValues["volume_state"]
if !ok {
t.Fatalf("expected volume_state default value")
}
if defaultVal.GetStringValue() != "ALL" {
t.Fatalf("expected volume_state default 'ALL', got %q", defaultVal.GetStringValue())
}
}
func workerConfigFormHasField(form *plugin_pb.ConfigForm, fieldName string) bool {
if form == nil {
return false
}
for _, section := range form.Sections {
if section == nil {
continue
}
for _, field := range section.Fields {
if field != nil && field.Name == fieldName {
return true
}
}
}
return false
}
type noopDetectionSender struct{}
func (noopDetectionSender) SendProposals(*plugin_pb.DetectionProposals) error { return nil }
func (noopDetectionSender) SendComplete(*plugin_pb.DetectionComplete) error { return nil }
func (noopDetectionSender) SendActivity(*plugin_pb.ActivityEvent) error { return nil }
type noopExecutionSender struct{}
func (noopExecutionSender) SendProgress(*plugin_pb.JobProgressUpdate) error { return nil }
func (noopExecutionSender) SendCompleted(*plugin_pb.JobCompleted) error { return nil }
type recordingDetectionSender struct {
proposals *plugin_pb.DetectionProposals
complete *plugin_pb.DetectionComplete
events []*plugin_pb.ActivityEvent
}
func (r *recordingDetectionSender) SendProposals(proposals *plugin_pb.DetectionProposals) error {
r.proposals = proposals
return nil
}
func (r *recordingDetectionSender) SendComplete(complete *plugin_pb.DetectionComplete) error {
r.complete = complete
return nil
}
func (r *recordingDetectionSender) SendActivity(event *plugin_pb.ActivityEvent) error {
if event != nil {
r.events = append(r.events, event)
}
return nil
}