snapshot of https://github.com/seaweedfs/seaweedfs.git (master, 2026-09-16; blobs > 131072B omitted)
Some checks failed
go: build dev binaries / cleanup (push) Successful in 4s
Codespell / Check for spelling errors (push) Failing after 3s
docker: build dev containers / build-rust-binaries (amd64, x86_64-unknown-linux-musl) (push) Failing after 3s
docker: build dev containers / build-dev-containers (push) Has been skipped
docker: build dev containers / build-rust-binaries (arm64, true, aarch64-unknown-linux-musl) (push) Failing after 3s
End to End / FUSE Mount (push) Failing after 2s
EC Integration Tests / EC Integration Tests (push) Failing after 30s
FUSE DLM Integration Tests / FUSE DLM Integration Tests (push) Failing after 3s
FUSE Volume Server Failover Tests / FUSE Volume Server Failover (push) Failing after 3s
FUSE Integration Tests / FUSE Integration Testing (push) Failing after 2s
FUSE P2P Peer Chunk Sharing Integration Tests / FUSE P2P Peer Chunk Sharing (push) Failing after 3s
go: build binary / Go Vet (push) Failing after 2s
go: build binary / Go Vet 32-bit (push) Failing after 3s
go: build binary / Build (push) Failing after 4s
go: build binary / Build cross-platform (freebsd/amd64) (push) Failing after 3s
go: build binary / Build cross-platform (windows/amd64) (push) Failing after 3s
go: build binary / Build cross-platform (darwin/arm64) (push) Failing after 3s
go: build binary / Build cross-platform (windows/arm64) (push) Failing after 3s
go: build binary / Test (push) Failing after 2s
go: build binary / Test 32-bit (push) Failing after 3s
helm: lint and test charts / lint-test (push) Failing after 3s
Java Client Integration Tests / Java Integration Tests (11) (push) Failing after 3s
Java Client Integration Tests / Java Integration Tests (17) (push) Failing after 3s
go: build dev binaries / build_dev_linux_windows (amd64, linux) (push) Failing after 2s
go: build dev binaries / build_dev_linux_windows (amd64, windows) (push) Failing after 3s
go: build dev binaries / build_dev_darwin (amd64, darwin) (push) Has been skipped
go: build dev binaries / build_dev_darwin (arm64, darwin) (push) Has been skipped
Java Client Unit Tests / Java Unit Tests (11) (push) Failing after 3s
Java Client Unit Tests / Java Unit Tests (17) (push) Failing after 2s
Java Client Unit Tests / Java Unit Tests (21) (push) Failing after 3s
Java Client Unit Tests / Java Unit Tests (8) (push) Failing after 3s
Kafka Quick Test (Load Test with Schema Registry) / Kafka Client Load Test (Quick) (push) Failing after 3s
Kafka Gateway Tests / Kafka Unit Tests (unit-tests-1) (push) Failing after 11s
Kafka Gateway Tests / Kafka Integration Tests (Critical) (integration-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka End-to-End Tests (with SMQ) (e2e-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka Consumer Group Tests (Highly Isolated) (consumer-group-1) (push) Failing after 2s
Kafka Gateway Tests / Kafka Client Compatibility (with SMQ) (client-compat-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka SMQ Integration Tests (Full Stack) (smq-integration-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka Protocol Tests (Isolated) (protocol-1) (push) Failing after 3s
KMS Tests / KMS Provider Integration Tests (push) Failing after 2s
KMS Tests / S3 KMS End-to-End Tests (push) Failing after 3s
Master Cold Start Tests / Master Cold Start Tests (push) Failing after 4s
Metadata Subscribe Integration Tests / Metadata Subscribe Integration Tests (push) Failing after 3s
Multi-Master Tests / Multi-Master Failover Tests (push) Failing after 3s
Performance / CPU and Heap Profile (push) Failing after 2s
Performance / S3 Read/Write Benchmark (rust volume) (push) Failing after 3s
pjdfstest POSIX Compliance / pjdfstest (push) Failing after 2s
Plugin Worker Integration Tests / Plugin Worker: erasure_coding (push) Failing after 3s
Plugin Worker Integration Tests / Plugin Worker: vacuum (push) Failing after 3s
Rust Volume Server Tests / Rust Integration Tests (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 1) (push) Failing after 2s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 1) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 2) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 2) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 3) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Distributed Tests (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Performance Tests (push) Failing after 3s
Performance / Throughput Benchmark (go volume) (push) Failing after 3s
Performance / Throughput Benchmark (rust volume) (push) Failing after 2s
Performance / S3 Read/Write Benchmark (go volume) (push) Failing after 3s
Plugin Worker Integration Tests / Plugin Worker: volume_balance (push) Failing after 3s
PostgreSQL Gateway Tests / PostgreSQL Basic Tests (push) Failing after 3s
Rust Volume Server Tests / Rust Unit Tests (push) Failing after 2s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 3) (push) Failing after 3s
Rust Plugin Worker Tests / Rust Plugin Worker Build and Unit Tests (push) Failing after 3s
rust: build dev volume server binaries / cleanup (push) Successful in 3s
S3 ETag and ACL Tests / S3 ETag + ACL Integration Tests (push) Failing after 3s
S3 IAM Integration Tests / IAM Unit Tests (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (advanced) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (basic) (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Integration Tests (group) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (policy-enforcement) (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Integration Tests (sts) (push) Failing after 4s
S3 Keycloak Integration Tests / S3 Keycloak Integration Tests (push) Failing after 3s
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.11) (push) Failing after 3s
S3 PyArrow Parquet Tests / Go Unit Tests (Implicit Directory) (push) Has been cancelled
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.12) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-darwin (darwin-amd64, x86_64-apple-darwin) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-darwin (darwin-arm64, aarch64-apple-darwin) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-linux (linux-amd64, x86_64-unknown-linux-gnu) (push) Has been cancelled
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.9) (push) Has been cancelled
S3 Proxy Signature Tests / S3 Proxy Signature Verification Tests (push) Has been cancelled
S3 SDK V2 Route Disambiguation Tests / S3 SDK V2 Routing Tests (push) Has been cancelled
S3 SSE Tests / S3 SSE Integration Tests (quick) (push) Has been cancelled
S3 SSE Tests / S3 SSE Compatibility Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Metadata Persistence Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Copy Operations Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Multipart Upload Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Performance Test (push) Has been cancelled
S3 SSE Tests / S3 Volume Encryption Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Integration Tests (comprehensive) (push) Has been cancelled
Spark Integration Tests / Spark Integration Tests (push) Failing after 2s
Telemetry Integration Tests / telemetry-integration-test (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 1) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 1) (push) Failing after 3s
Ceph S3 tests / Basic S3 tests (KV store) (push) Has been cancelled
Ceph S3 tests / S3 Versioning & Object Lock tests (push) Has been cancelled
S3 Policy Integration Tests / S3 Policy Unit Tests (push) Failing after 3s
Ceph S3 tests / S3 CORS tests (push) Has been cancelled
Ceph S3 tests / SeaweedFS Custom S3 Copy tests (push) Has been cancelled
Ceph S3 tests / Basic S3 tests (SQL store) (push) Has been cancelled
test s3 over https using aws-cli / awscli-tests (push) Has been cancelled
TLS Rotation Integration Tests / TLS Rotation Integration Tests (push) Has been cancelled
Vacuum Integration Tests / Vacuum Integration Tests (push) Has been cancelled
S3 Policy Integration Tests / S3 Policy Variables Integration Tests (push) Failing after 3s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (advanced-policy) (push) Failing after 3s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (basic-policy) (push) Failing after 2s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (contextual-policy) (push) Failing after 4s
S3 Policy Integration Tests / Trusted Proxy Detection Tests (push) Failing after 2s
S3 Policy Integration Tests / S3 Policy Shell Integration Tests (push) Failing after 3s
Samba on FUSE Integration / samba-integration (push) Failing after 2s
SFTP Integration Tests / SFTP Integration Testing (push) Failing after 3s
terraform: validate and test modules / fmt, validate, plan-level tests (push) Failing after 3s
terraform: validate and test modules / local cluster smoke test (real weed) (push) Failing after 2s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 3) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 2) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 2) (push) Failing after 2s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 3) (push) Failing after 5s
mount: windows conformance / WinFsp conformance (push) Has been cancelled
mount: windows / Mount on Windows (push) Has been cancelled

This commit is contained in:
upstream-mirror
2026-09-16 03:23:42 +05:00
commit 89f5be005d
3997 changed files with 961970 additions and 0 deletions

View File

@@ -0,0 +1,27 @@
package pluginworker
import (
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
"google.golang.org/protobuf/types/known/timestamppb"
)
// BuildExecutorActivity creates an executor activity event. Exported for sub-packages.
func BuildExecutorActivity(stage string, message string) *plugin_pb.ActivityEvent {
return &plugin_pb.ActivityEvent{
Source: plugin_pb.ActivitySource_ACTIVITY_SOURCE_EXECUTOR,
Stage: stage,
Message: message,
CreatedAt: timestamppb.Now(),
}
}
// BuildDetectorActivity creates a detector activity event. Exported for sub-packages.
func BuildDetectorActivity(stage string, message string, details map[string]*plugin_pb.ConfigValue) *plugin_pb.ActivityEvent {
return &plugin_pb.ActivityEvent{
Source: plugin_pb.ActivitySource_ACTIVITY_SOURCE_DETECTOR,
Stage: stage,
Message: message,
Details: details,
CreatedAt: timestamppb.Now(),
}
}

View File

@@ -0,0 +1,677 @@
package pluginworker
import (
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"regexp"
"strings"
"time"
"github.com/seaweedfs/seaweedfs/weed/glog"
"github.com/seaweedfs/seaweedfs/weed/pb"
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
"github.com/seaweedfs/seaweedfs/weed/shell"
"google.golang.org/grpc"
"google.golang.org/protobuf/types/known/timestamppb"
)
const (
adminScriptJobType = "admin_script"
maxAdminScriptOutputBytes = 16 * 1024
defaultAdminScriptRunMins = 17
adminScriptDetectTickMinutes = 17
)
const defaultAdminScript = `fs.log.purge -daysAgo=7
volume.deleteEmpty -quietFor=24h -apply
volume.fix.replication -apply
s3.clean.uploads -timeAgo=24h`
var adminScriptTokenRegex = regexp.MustCompile(`'.*?'|".*?"|\S+`)
func init() {
RegisterHandler(HandlerFactory{
JobType: "admin_script",
Category: CategoryDefault,
Aliases: []string{"admin-script", "admin.script", "script", "admin"},
Build: func(opts HandlerBuildOptions) (JobHandler, error) {
return NewAdminScriptHandler(opts.GrpcDialOption), nil
},
})
}
type AdminScriptHandler struct {
grpcDialOption grpc.DialOption
}
func NewAdminScriptHandler(grpcDialOption grpc.DialOption) *AdminScriptHandler {
return &AdminScriptHandler{grpcDialOption: grpcDialOption}
}
func (h *AdminScriptHandler) Capability() *plugin_pb.JobTypeCapability {
return &plugin_pb.JobTypeCapability{
JobType: adminScriptJobType,
CanDetect: true,
CanExecute: true,
MaxDetectionConcurrency: 1,
MaxExecutionConcurrency: 1,
DisplayName: "Admin Script",
Description: "Execute custom admin shell scripts",
Weight: 20,
}
}
func (h *AdminScriptHandler) Descriptor() *plugin_pb.JobTypeDescriptor {
return &plugin_pb.JobTypeDescriptor{
JobType: adminScriptJobType,
DisplayName: "Admin Script",
Description: "Run custom admin shell scripts not covered by built-in job types",
Icon: "fas fa-terminal",
DescriptorVersion: 1,
AdminConfigForm: &plugin_pb.ConfigForm{
FormId: "admin-script-admin",
Title: "Admin Script Configuration",
Description: "Define the admin shell script to execute.",
Sections: []*plugin_pb.ConfigSection{
{
SectionId: "script",
Title: "Script",
Description: "Commands run sequentially by the admin script worker.",
Fields: []*plugin_pb.ConfigField{
{
Name: "script",
Label: "Script",
Description: "Admin shell commands to execute (one per line).",
HelpText: "Lock/unlock are handled by the admin server; omit explicit lock/unlock commands.",
Placeholder: "volume.balance -apply\nvolume.fix.replication -apply",
FieldType: plugin_pb.ConfigFieldType_CONFIG_FIELD_TYPE_STRING,
Widget: plugin_pb.ConfigWidget_CONFIG_WIDGET_TEXTAREA,
Required: true,
},
{
Name: "run_interval_minutes",
Label: "Run Interval (minutes)",
Description: "Minimum interval between successful admin script runs.",
FieldType: plugin_pb.ConfigFieldType_CONFIG_FIELD_TYPE_INT64,
Widget: plugin_pb.ConfigWidget_CONFIG_WIDGET_NUMBER,
Required: true,
MinValue: &plugin_pb.ConfigValue{Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 1}},
},
},
},
},
DefaultValues: map[string]*plugin_pb.ConfigValue{
"script": {
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: defaultAdminScript},
},
"run_interval_minutes": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: defaultAdminScriptRunMins},
},
},
},
AdminRuntimeDefaults: &plugin_pb.AdminRuntimeDefaults{
Enabled: true,
DetectionIntervalMinutes: adminScriptDetectTickMinutes,
DetectionTimeoutSeconds: 300,
MaxJobsPerDetection: 1,
GlobalExecutionConcurrency: 1,
PerWorkerExecutionConcurrency: 1,
RetryLimit: 0,
RetryBackoffSeconds: 30,
JobTypeMaxRuntimeSeconds: 1800,
ExecutionTimeoutSeconds: 1800,
},
WorkerDefaultValues: map[string]*plugin_pb.ConfigValue{},
}
}
func (h *AdminScriptHandler) Detect(ctx context.Context, request *plugin_pb.RunDetectionRequest, sender DetectionSender) error {
if request == nil {
return fmt.Errorf("run detection request is nil")
}
if sender == nil {
return fmt.Errorf("detection sender is nil")
}
if request.JobType != "" && request.JobType != adminScriptJobType {
return fmt.Errorf("job type %q is not handled by admin_script worker", request.JobType)
}
script := normalizeAdminScript(ReadStringConfig(request.GetAdminConfigValues(), "script", ""))
scriptName := strings.TrimSpace(ReadStringConfig(request.GetAdminConfigValues(), "script_name", ""))
runIntervalMinutes := readAdminScriptRunIntervalMinutes(request.GetAdminConfigValues())
if ShouldSkipDetectionByInterval(request.GetLastSuccessfulRun(), runIntervalMinutes*60) {
_ = sender.SendActivity(BuildDetectorActivity(
"skipped_by_interval",
fmt.Sprintf("ADMIN SCRIPT: Detection skipped due to run interval (%dm)", runIntervalMinutes),
map[string]*plugin_pb.ConfigValue{
"run_interval_minutes": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: int64(runIntervalMinutes)},
},
},
))
if err := sender.SendProposals(&plugin_pb.DetectionProposals{
JobType: adminScriptJobType,
Proposals: []*plugin_pb.JobProposal{},
HasMore: false,
}); err != nil {
return err
}
return sender.SendComplete(&plugin_pb.DetectionComplete{
JobType: adminScriptJobType,
Success: true,
TotalProposals: 0,
})
}
commands := parseAdminScriptCommands(script)
execCount := countExecutableCommands(commands)
if execCount == 0 {
_ = sender.SendActivity(BuildDetectorActivity(
"no_script",
"ADMIN SCRIPT: No executable commands configured",
map[string]*plugin_pb.ConfigValue{
"command_count": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: int64(execCount)},
},
},
))
if err := sender.SendProposals(&plugin_pb.DetectionProposals{
JobType: adminScriptJobType,
Proposals: []*plugin_pb.JobProposal{},
HasMore: false,
}); err != nil {
return err
}
return sender.SendComplete(&plugin_pb.DetectionComplete{
JobType: adminScriptJobType,
Success: true,
TotalProposals: 0,
})
}
proposal := buildAdminScriptProposal(script, scriptName, execCount)
proposals := []*plugin_pb.JobProposal{proposal}
hasMore := false
maxResults := int(request.MaxResults)
if maxResults > 0 && len(proposals) > maxResults {
proposals = proposals[:maxResults]
hasMore = true
}
if err := sender.SendProposals(&plugin_pb.DetectionProposals{
JobType: adminScriptJobType,
Proposals: proposals,
HasMore: hasMore,
}); err != nil {
return err
}
return sender.SendComplete(&plugin_pb.DetectionComplete{
JobType: adminScriptJobType,
Success: true,
TotalProposals: 1,
})
}
func (h *AdminScriptHandler) Execute(ctx context.Context, request *plugin_pb.ExecuteJobRequest, sender ExecutionSender) error {
if request == nil || request.Job == nil {
return fmt.Errorf("execute job request is nil")
}
if sender == nil {
return fmt.Errorf("execution sender is nil")
}
if request.Job.JobType != "" && request.Job.JobType != adminScriptJobType {
return fmt.Errorf("job type %q is not handled by admin_script worker", request.Job.JobType)
}
script := normalizeAdminScript(ReadStringConfig(request.Job.Parameters, "script", ""))
scriptName := strings.TrimSpace(ReadStringConfig(request.Job.Parameters, "script_name", ""))
if script == "" {
script = normalizeAdminScript(ReadStringConfig(request.GetAdminConfigValues(), "script", ""))
}
if scriptName == "" {
scriptName = strings.TrimSpace(ReadStringConfig(request.GetAdminConfigValues(), "script_name", ""))
}
commands := parseAdminScriptCommands(script)
execCommands := filterExecutableCommands(commands)
if len(execCommands) == 0 {
return sender.SendCompleted(&plugin_pb.JobCompleted{
Success: false,
ErrorMessage: "no executable admin script commands configured",
})
}
commandEnv, cancel, err := h.buildAdminScriptCommandEnv(ctx, request.ClusterContext)
if err != nil {
return sender.SendCompleted(&plugin_pb.JobCompleted{
Success: false,
ErrorMessage: err.Error(),
})
}
defer cancel()
if err := sender.SendProgress(&plugin_pb.JobProgressUpdate{
JobId: request.Job.JobId,
JobType: request.Job.JobType,
State: plugin_pb.JobState_JOB_STATE_ASSIGNED,
ProgressPercent: 0,
Stage: "assigned",
Message: "admin script job accepted",
Activities: []*plugin_pb.ActivityEvent{
BuildExecutorActivity("assigned", "admin script job accepted"),
},
}); err != nil {
return err
}
output := &limitedBuffer{maxBytes: maxAdminScriptOutputBytes}
executed := 0
errorMessages := make([]string, 0)
executedCommands := make([]string, 0, len(execCommands))
for _, cmd := range execCommands {
if ctx.Err() != nil {
errorMessages = append(errorMessages, ctx.Err().Error())
break
}
commandLine := formatAdminScriptCommand(cmd)
executedCommands = append(executedCommands, commandLine)
_, _ = fmt.Fprintf(output, "$ %s\n", commandLine)
found := false
sendBroken := false
for _, command := range shell.Commands {
if command.Name() != cmd.Name {
continue
}
found = true
// The env is reused for every command in the script, and noLock
// belongs to the invocation that set it. ForceNoLock already exempts
// this path, so this changes nothing today -- it keeps the rule the
// same in all three dispatchers rather than resting on that.
commandEnv.SetNoLock(false)
if err := command.Do(cmd.Args, commandEnv, output); err != nil {
msg := fmt.Sprintf("%s: %v", cmd.Name, err)
errorMessages = append(errorMessages, msg)
if sendErr := sender.SendProgress(&plugin_pb.JobProgressUpdate{
State: plugin_pb.JobState_JOB_STATE_RUNNING,
ProgressPercent: percentProgress(executed+1, len(execCommands)),
Stage: "error",
Message: msg,
Activities: []*plugin_pb.ActivityEvent{
BuildExecutorActivity("error", msg),
},
}); sendErr != nil {
sendBroken = true
}
}
break
}
if sendBroken {
break
}
if !found {
msg := fmt.Sprintf("unknown admin command: %s", cmd.Name)
errorMessages = append(errorMessages, msg)
if sendErr := sender.SendProgress(&plugin_pb.JobProgressUpdate{
State: plugin_pb.JobState_JOB_STATE_RUNNING,
ProgressPercent: percentProgress(executed+1, len(execCommands)),
Stage: "error",
Message: msg,
Activities: []*plugin_pb.ActivityEvent{
BuildExecutorActivity("error", msg),
},
}); sendErr != nil {
break
}
}
executed++
progress := percentProgress(executed, len(execCommands))
if sendErr := sender.SendProgress(&plugin_pb.JobProgressUpdate{
State: plugin_pb.JobState_JOB_STATE_RUNNING,
ProgressPercent: progress,
Stage: "running",
Message: fmt.Sprintf("executed %d/%d command(s)", executed, len(execCommands)),
Activities: []*plugin_pb.ActivityEvent{
BuildExecutorActivity("running", commandLine),
},
}); sendErr != nil {
break
}
}
scriptHash := hashAdminScript(script)
resultSummary := fmt.Sprintf("admin script executed (%d command(s))", executed)
if scriptName != "" {
resultSummary = fmt.Sprintf("admin script %q executed (%d command(s))", scriptName, executed)
}
outputValues := map[string]*plugin_pb.ConfigValue{
"command_count": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: int64(executed)},
},
"error_count": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: int64(len(errorMessages))},
},
"script_hash": {
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: scriptHash},
},
}
if scriptName != "" {
outputValues["script_name"] = &plugin_pb.ConfigValue{
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: scriptName},
}
}
if len(executedCommands) > 0 {
outputValues["commands"] = &plugin_pb.ConfigValue{
Kind: &plugin_pb.ConfigValue_StringList{
StringList: &plugin_pb.StringList{Values: executedCommands},
},
}
}
if out := strings.TrimSpace(output.String()); out != "" {
outputValues["output"] = &plugin_pb.ConfigValue{
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: out},
}
}
if output.truncated {
outputValues["output_truncated"] = &plugin_pb.ConfigValue{
Kind: &plugin_pb.ConfigValue_BoolValue{BoolValue: true},
}
}
success := len(errorMessages) == 0 && ctx.Err() == nil
errorMessage := ""
if !success {
errorMessage = strings.Join(errorMessages, "; ")
if ctx.Err() != nil {
if errorMessage == "" {
errorMessage = ctx.Err().Error()
} else {
errorMessage = fmt.Sprintf("%s; %s", errorMessage, ctx.Err().Error())
}
}
}
return sender.SendCompleted(&plugin_pb.JobCompleted{
Success: success,
ErrorMessage: errorMessage,
Result: &plugin_pb.JobResult{
Summary: resultSummary,
OutputValues: outputValues,
},
Activities: []*plugin_pb.ActivityEvent{
BuildExecutorActivity("completed", resultSummary),
},
CompletedAt: timestamppb.Now(),
})
}
func readAdminScriptRunIntervalMinutes(values map[string]*plugin_pb.ConfigValue) int {
runIntervalMinutes := int(ReadInt64Config(values, "run_interval_minutes", defaultAdminScriptRunMins))
if runIntervalMinutes <= 0 {
return defaultAdminScriptRunMins
}
return runIntervalMinutes
}
type adminScriptCommand struct {
Name string
Args []string
Raw string
}
func normalizeAdminScript(script string) string {
script = strings.ReplaceAll(script, "\r\n", "\n")
return strings.TrimSpace(script)
}
func parseAdminScriptCommands(script string) []adminScriptCommand {
script = normalizeAdminScript(script)
if script == "" {
return nil
}
lines := strings.Split(script, "\n")
commands := make([]adminScriptCommand, 0)
for _, line := range lines {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
for _, chunk := range strings.Split(line, ";") {
chunk = strings.TrimSpace(chunk)
if chunk == "" {
continue
}
parts := adminScriptTokenRegex.FindAllString(chunk, -1)
if len(parts) == 0 {
continue
}
args := make([]string, 0, len(parts)-1)
for _, arg := range parts[1:] {
args = append(args, strings.Trim(arg, "\"'"))
}
commands = append(commands, adminScriptCommand{
Name: strings.TrimSpace(parts[0]),
Args: args,
Raw: chunk,
})
}
}
return commands
}
func filterExecutableCommands(commands []adminScriptCommand) []adminScriptCommand {
exec := make([]adminScriptCommand, 0, len(commands))
for _, cmd := range commands {
if cmd.Name == "" {
continue
}
if isAdminScriptLockCommand(cmd.Name) {
continue
}
exec = append(exec, cmd)
}
return exec
}
func countExecutableCommands(commands []adminScriptCommand) int {
count := 0
for _, cmd := range commands {
if cmd.Name == "" {
continue
}
if isAdminScriptLockCommand(cmd.Name) {
continue
}
count++
}
return count
}
func isAdminScriptLockCommand(name string) bool {
switch strings.ToLower(strings.TrimSpace(name)) {
case "lock", "unlock":
return true
default:
return false
}
}
func buildAdminScriptProposal(script, scriptName string, commandCount int) *plugin_pb.JobProposal {
scriptHash := hashAdminScript(script)
summary := "Run admin script"
if scriptName != "" {
summary = fmt.Sprintf("Run admin script: %s", scriptName)
}
detail := fmt.Sprintf("Admin script with %d command(s)", commandCount)
proposalID := fmt.Sprintf("admin-script-%s-%d", scriptHash[:8], time.Now().UnixNano())
labels := map[string]string{
"script_hash": scriptHash,
}
if scriptName != "" {
labels["script_name"] = scriptName
}
return &plugin_pb.JobProposal{
ProposalId: proposalID,
DedupeKey: "admin-script:" + scriptHash,
JobType: adminScriptJobType,
Priority: plugin_pb.JobPriority_JOB_PRIORITY_NORMAL,
Summary: summary,
Detail: detail,
Parameters: map[string]*plugin_pb.ConfigValue{
"script": {
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: script},
},
"script_name": {
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: scriptName},
},
"script_hash": {
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: scriptHash},
},
"command_count": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: int64(commandCount)},
},
},
Labels: labels,
}
}
func (h *AdminScriptHandler) buildAdminScriptCommandEnv(
ctx context.Context,
clusterContext *plugin_pb.ClusterContext,
) (*shell.CommandEnv, context.CancelFunc, error) {
options, err := h.buildAdminScriptShellOptions(clusterContext)
if err != nil {
return nil, nil, err
}
commandEnv := shell.NewCommandEnv(&options)
commandEnv.ForceNoLock()
ctx, cancel := context.WithCancel(ctx)
go commandEnv.MasterClient.KeepConnectedToMaster(ctx)
return commandEnv, cancel, nil
}
// buildAdminScriptShellOptions maps a cluster context onto shell.ShellOptions.
// It is split out from buildAdminScriptCommandEnv so the address handling is
// testable without standing up a master client.
func (h *AdminScriptHandler) buildAdminScriptShellOptions(clusterContext *plugin_pb.ClusterContext) (shell.ShellOptions, error) {
if clusterContext == nil {
return shell.ShellOptions{}, fmt.Errorf("cluster context is required")
}
masters := normalizeAddressList(clusterContext.MasterGrpcAddresses)
if len(masters) == 0 {
return shell.ShellOptions{}, fmt.Errorf("missing master addresses for admin script")
}
filerGroup := ""
mastersValue := strings.Join(masters, ",")
options := shell.ShellOptions{
Masters: &mastersValue,
GrpcDialOption: h.grpcDialOption,
FilerGroup: &filerGroup,
Directory: "/",
}
// FilerAddresses are pb.ServerAddress strings (host:httpPort.grpcPort).
// ShellOptions.FilerAddress is itself a ServerAddress; shell commands derive
// the gRPC or HTTP port from it as needed, so store it verbatim.
filers := normalizeAddressList(clusterContext.FilerAddresses)
if len(filers) > 0 {
options.FilerAddress = pb.ServerAddress(filers[0])
} else {
glog.V(1).Infof("admin script worker missing filer address; filer-dependent commands may fail")
}
return options, nil
}
func normalizeAddressList(addresses []string) []string {
normalized := make([]string, 0, len(addresses))
seen := make(map[string]struct{}, len(addresses))
for _, address := range addresses {
address = strings.TrimSpace(address)
if address == "" {
continue
}
if _, exists := seen[address]; exists {
continue
}
seen[address] = struct{}{}
normalized = append(normalized, address)
}
return normalized
}
func hashAdminScript(script string) string {
sum := sha256.Sum256([]byte(script))
return hex.EncodeToString(sum[:])
}
func formatAdminScriptCommand(cmd adminScriptCommand) string {
if len(cmd.Args) == 0 {
return cmd.Name
}
return fmt.Sprintf("%s %s", cmd.Name, strings.Join(cmd.Args, " "))
}
func percentProgress(done, total int) float64 {
if total <= 0 {
return 0
}
if done < 0 {
done = 0
}
if done > total {
done = total
}
return float64(done) / float64(total) * 100
}
type limitedBuffer struct {
buf bytes.Buffer
maxBytes int
truncated bool
}
func (b *limitedBuffer) Write(p []byte) (int, error) {
if b == nil {
return len(p), nil
}
if b.maxBytes <= 0 {
b.truncated = true
return len(p), nil
}
remaining := b.maxBytes - b.buf.Len()
if remaining <= 0 {
b.truncated = true
return len(p), nil
}
if len(p) > remaining {
_, _ = b.buf.Write(p[:remaining])
b.truncated = true
return len(p), nil
}
_, _ = b.buf.Write(p)
return len(p), nil
}
func (b *limitedBuffer) String() string {
if b == nil {
return ""
}
return b.buf.String()
}

View File

@@ -0,0 +1,130 @@
package pluginworker
import (
"context"
"strings"
"testing"
"time"
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
"google.golang.org/protobuf/types/known/timestamppb"
)
func TestAdminScriptFilerAddressKeepsGrpcPortOffConvention(t *testing.T) {
// FilerAddresses arrive as pb.ServerAddress (host:httpPort.grpcPort). The
// shell derives the gRPC port from FilerAddress via ToGrpcAddress(); it must
// land on the real gRPC port even when that port is off the httpPort+10000
// convention, not double the offset into a non-existent filer:28890.
h := NewAdminScriptHandler(nil)
options, err := h.buildAdminScriptShellOptions(&plugin_pb.ClusterContext{
MasterGrpcAddresses: []string{"master:19333"},
FilerAddresses: []string{"filer:8888.18890"},
})
if err != nil {
t.Fatalf("buildAdminScriptShellOptions returned err = %v", err)
}
if grpc := options.FilerAddress.ToGrpcAddress(); grpc != "filer:18890" {
t.Fatalf("filer gRPC address = %q, want filer:18890 (not filer:28890)", grpc)
}
if http := options.FilerAddress.ToHttpAddress(); http != "filer:8888" {
t.Fatalf("filer HTTP address = %q, want filer:8888", http)
}
}
func TestAdminScriptShellOptionsRequireMasters(t *testing.T) {
h := NewAdminScriptHandler(nil)
if _, err := h.buildAdminScriptShellOptions(&plugin_pb.ClusterContext{
FilerAddresses: []string{"filer:8888.18890"},
}); err == nil {
t.Fatalf("expected error when master addresses are missing")
}
}
func TestAdminScriptDescriptorDefaults(t *testing.T) {
descriptor := NewAdminScriptHandler(nil).Descriptor()
if descriptor == nil {
t.Fatalf("expected descriptor")
}
if descriptor.AdminRuntimeDefaults == nil {
t.Fatalf("expected admin runtime defaults")
}
if descriptor.AdminRuntimeDefaults.DetectionIntervalMinutes != adminScriptDetectTickMinutes {
t.Fatalf("unexpected detection interval seconds: got=%d want=%d",
descriptor.AdminRuntimeDefaults.DetectionIntervalMinutes, adminScriptDetectTickMinutes)
}
if descriptor.AdminConfigForm == nil {
t.Fatalf("expected admin config form")
}
runInterval := ReadInt64Config(descriptor.AdminConfigForm.DefaultValues, "run_interval_minutes", 0)
if runInterval != defaultAdminScriptRunMins {
t.Fatalf("unexpected run_interval_minutes default: got=%d want=%d", runInterval, defaultAdminScriptRunMins)
}
script := ReadStringConfig(descriptor.AdminConfigForm.DefaultValues, "script", "")
if strings.TrimSpace(script) == "" {
t.Fatalf("expected non-empty default script")
}
}
func TestAdminScriptDetectSkipsByRunInterval(t *testing.T) {
handler := NewAdminScriptHandler(nil)
sender := &recordingDetectionSender{}
err := handler.Detect(context.Background(), &plugin_pb.RunDetectionRequest{
JobType: adminScriptJobType,
LastSuccessfulRun: timestamppb.New(time.Now().Add(-2 * time.Minute)),
AdminConfigValues: map[string]*plugin_pb.ConfigValue{
"script": {
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: defaultAdminScript},
},
"run_interval_minutes": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 17},
},
},
}, sender)
if err != nil {
t.Fatalf("detect returned err = %v", err)
}
if sender.proposals == nil {
t.Fatalf("expected proposals message")
}
if len(sender.proposals.Proposals) != 0 {
t.Fatalf("expected zero proposals, got %d", len(sender.proposals.Proposals))
}
if sender.complete == nil || !sender.complete.Success {
t.Fatalf("expected successful completion message")
}
if len(sender.events) == 0 {
t.Fatalf("expected detector activity events")
}
if !strings.Contains(sender.events[0].Message, "run interval") {
t.Fatalf("unexpected skip message: %q", sender.events[0].Message)
}
}
func TestAdminScriptDetectCreatesProposalWhenIntervalElapsed(t *testing.T) {
handler := NewAdminScriptHandler(nil)
sender := &recordingDetectionSender{}
err := handler.Detect(context.Background(), &plugin_pb.RunDetectionRequest{
JobType: adminScriptJobType,
LastSuccessfulRun: timestamppb.New(time.Now().Add(-20 * time.Minute)),
AdminConfigValues: map[string]*plugin_pb.ConfigValue{
"script": {
Kind: &plugin_pb.ConfigValue_StringValue{StringValue: defaultAdminScript},
},
"run_interval_minutes": {
Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 17},
},
},
}, sender)
if err != nil {
t.Fatalf("detect returned err = %v", err)
}
if sender.proposals == nil {
t.Fatalf("expected proposals message")
}
if len(sender.proposals.Proposals) != 1 {
t.Fatalf("expected one proposal, got %d", len(sender.proposals.Proposals))
}
if sender.complete == nil || !sender.complete.Success || sender.complete.TotalProposals != 1 {
t.Fatalf("unexpected completion message: %+v", sender.complete)
}
}

View File

@@ -0,0 +1,227 @@
package pluginworker
import (
"fmt"
"math"
"strconv"
"strings"
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
workertypes "github.com/seaweedfs/seaweedfs/weed/worker/types"
)
// ReadStringConfig reads a string-valued plugin config field, returning fallback
// when the value is missing or of an unsupported kind.
func ReadStringConfig(values map[string]*plugin_pb.ConfigValue, field string, fallback string) string {
if values == nil {
return fallback
}
value := values[field]
if value == nil {
return fallback
}
switch kind := value.Kind.(type) {
case *plugin_pb.ConfigValue_StringValue:
return kind.StringValue
case *plugin_pb.ConfigValue_Int64Value:
return strconv.FormatInt(kind.Int64Value, 10)
case *plugin_pb.ConfigValue_DoubleValue:
return strconv.FormatFloat(kind.DoubleValue, 'f', -1, 64)
case *plugin_pb.ConfigValue_BoolValue:
return strconv.FormatBool(kind.BoolValue)
}
return fallback
}
// ReadDoubleConfig reads a double-valued plugin config field, returning
// fallback when the value is missing or unparseable.
func ReadDoubleConfig(values map[string]*plugin_pb.ConfigValue, field string, fallback float64) float64 {
if values == nil {
return fallback
}
value := values[field]
if value == nil {
return fallback
}
switch kind := value.Kind.(type) {
case *plugin_pb.ConfigValue_DoubleValue:
return kind.DoubleValue
case *plugin_pb.ConfigValue_Int64Value:
return float64(kind.Int64Value)
case *plugin_pb.ConfigValue_StringValue:
parsed, err := strconv.ParseFloat(strings.TrimSpace(kind.StringValue), 64)
if err == nil {
return parsed
}
case *plugin_pb.ConfigValue_BoolValue:
if kind.BoolValue {
return 1
}
return 0
}
return fallback
}
// ReadInt64Config reads an int64-valued plugin config field, returning fallback
// when the value is missing or unparseable.
func ReadInt64Config(values map[string]*plugin_pb.ConfigValue, field string, fallback int64) int64 {
if values == nil {
return fallback
}
value := values[field]
if value == nil {
return fallback
}
switch kind := value.Kind.(type) {
case *plugin_pb.ConfigValue_Int64Value:
return kind.Int64Value
case *plugin_pb.ConfigValue_DoubleValue:
return int64(kind.DoubleValue)
case *plugin_pb.ConfigValue_StringValue:
parsed, err := strconv.ParseInt(strings.TrimSpace(kind.StringValue), 10, 64)
if err == nil {
return parsed
}
case *plugin_pb.ConfigValue_BoolValue:
if kind.BoolValue {
return 1
}
return 0
}
return fallback
}
// ReadInt32Config reads an int32-valued plugin config field, returning fallback
// when the value is missing or out of int32 range. Used for protobuf int32
// fields whose admin/worker config values arrive as int64.
func ReadInt32Config(values map[string]*plugin_pb.ConfigValue, field string, fallback int32) int32 {
v := ReadInt64Config(values, field, int64(fallback))
if v < int64(math.MinInt32) || v > int64(math.MaxInt32) {
return fallback
}
return int32(v)
}
// ReadUint32Config reads a uint32-valued plugin config field, returning
// fallback when the value is missing, negative, or exceeds math.MaxUint32.
// Used for protobuf uint32 fields (volume IDs, shard counts, …).
func ReadUint32Config(values map[string]*plugin_pb.ConfigValue, field string, fallback uint32) uint32 {
v := ReadInt64Config(values, field, int64(fallback))
if v < 0 || v > int64(math.MaxUint32) {
return fallback
}
return uint32(v)
}
// ReadIntConfig reads an int-valued plugin config field, returning fallback
// when the value is missing or outside the int32 range. The int32 range is
// used as the platform-portable safe range so that the same value parses
// identically on 32-bit and 64-bit builds.
func ReadIntConfig(values map[string]*plugin_pb.ConfigValue, field string, fallback int) int {
v := ReadInt64Config(values, field, int64(fallback))
if v < int64(math.MinInt32) || v > int64(math.MaxInt32) {
return fallback
}
return int(v)
}
// ReadBytesConfig reads a bytes-valued plugin config field, returning nil when
// the value is missing or of a different kind.
func ReadBytesConfig(values map[string]*plugin_pb.ConfigValue, field string) []byte {
if values == nil {
return nil
}
value := values[field]
if value == nil {
return nil
}
if kind, ok := value.Kind.(*plugin_pb.ConfigValue_BytesValue); ok {
return kind.BytesValue
}
return nil
}
// ReadStringListConfig reads a list-of-strings plugin config field, returning
// nil when the value is missing. Accepts ConfigValue_StringList,
// ConfigValue_ListValue, or a comma-separated ConfigValue_StringValue.
func ReadStringListConfig(values map[string]*plugin_pb.ConfigValue, field string) []string {
if values == nil {
return nil
}
value := values[field]
if value == nil {
return nil
}
switch kind := value.Kind.(type) {
case *plugin_pb.ConfigValue_StringList:
return normalizeStringList(kind.StringList.GetValues())
case *plugin_pb.ConfigValue_ListValue:
out := make([]string, 0, len(kind.ListValue.GetValues()))
for _, item := range kind.ListValue.GetValues() {
itemText := readStringFromConfigValue(item)
if itemText != "" {
out = append(out, itemText)
}
}
return normalizeStringList(out)
case *plugin_pb.ConfigValue_StringValue:
return normalizeStringList(strings.Split(kind.StringValue, ","))
}
return nil
}
func readStringFromConfigValue(value *plugin_pb.ConfigValue) string {
if value == nil {
return ""
}
switch kind := value.Kind.(type) {
case *plugin_pb.ConfigValue_StringValue:
return strings.TrimSpace(kind.StringValue)
case *plugin_pb.ConfigValue_Int64Value:
return fmt.Sprintf("%d", kind.Int64Value)
case *plugin_pb.ConfigValue_DoubleValue:
return fmt.Sprintf("%g", kind.DoubleValue)
case *plugin_pb.ConfigValue_BoolValue:
if kind.BoolValue {
return "true"
}
return "false"
}
return ""
}
func normalizeStringList(values []string) []string {
normalized := make([]string, 0, len(values))
seen := make(map[string]struct{}, len(values))
for _, value := range values {
item := strings.TrimSpace(value)
if item == "" {
continue
}
if _, found := seen[item]; found {
continue
}
seen[item] = struct{}{}
normalized = append(normalized, item)
}
return normalized
}
// MapTaskPriority converts a worker-task priority into the plugin protocol's
// JobPriority enum.
func MapTaskPriority(priority workertypes.TaskPriority) plugin_pb.JobPriority {
switch strings.ToLower(string(priority)) {
case "low":
return plugin_pb.JobPriority_JOB_PRIORITY_LOW
case "medium", "normal":
return plugin_pb.JobPriority_JOB_PRIORITY_NORMAL
case "high":
return plugin_pb.JobPriority_JOB_PRIORITY_HIGH
case "critical":
return plugin_pb.JobPriority_JOB_PRIORITY_CRITICAL
default:
return plugin_pb.JobPriority_JOB_PRIORITY_NORMAL
}
}

View File

@@ -0,0 +1,130 @@
package pluginworker
import (
"fmt"
"strings"
"sync"
"google.golang.org/grpc"
)
// JobCategory groups job types by resource profile so that workers can be
// configured with a category name instead of an explicit list of job types.
type JobCategory string
const (
CategoryAll JobCategory = "all" // pseudo-category matching every handler
CategoryDefault JobCategory = "default" // lightweight, safe for any worker
CategoryHeavy JobCategory = "heavy" // resource-intensive jobs
)
// HandlerFactory describes how to build a JobHandler for a single job type.
type HandlerFactory struct {
// JobType is the canonical job type string (e.g. "vacuum").
JobType string
// Category controls which category label selects this handler.
Category JobCategory
// Aliases are alternative CLI names that resolve to this job type
// (e.g. "ec" for "erasure_coding").
Aliases []string
// Build constructs the JobHandler.
Build func(opts HandlerBuildOptions) (JobHandler, error)
}
// HandlerBuildOptions carries parameters forwarded from the CLI to handler
// constructors.
type HandlerBuildOptions struct {
GrpcDialOption grpc.DialOption
MaxExecute int
WorkingDir string
}
var (
registryMu sync.Mutex
registry []HandlerFactory
)
// RegisterHandler adds a handler factory to the global registry.
// It is intended to be called from handler init() functions.
func RegisterHandler(f HandlerFactory) {
registryMu.Lock()
defer registryMu.Unlock()
registry = append(registry, f)
}
// ResolveHandlerFactories takes a comma-separated token list that can contain
// category names ("all", "default", "heavy") and/or explicit job type names
// (including aliases). It returns a deduplicated, ordered slice of factories.
func ResolveHandlerFactories(tokens string) ([]HandlerFactory, error) {
registryMu.Lock()
snapshot := make([]HandlerFactory, len(registry))
copy(snapshot, registry)
registryMu.Unlock()
parts := strings.Split(tokens, ",")
result := make([]HandlerFactory, 0, len(snapshot))
seen := make(map[string]bool)
for _, raw := range parts {
tok := strings.ToLower(strings.TrimSpace(raw))
if tok == "" {
continue
}
if cat, ok := tokenAsCategory(tok); ok {
for _, f := range snapshot {
if cat == CategoryAll || f.Category == cat {
if !seen[f.JobType] {
seen[f.JobType] = true
result = append(result, f)
}
}
}
continue
}
f, err := findFactory(snapshot, tok)
if err != nil {
return nil, err
}
if !seen[f.JobType] {
seen[f.JobType] = true
result = append(result, f)
}
}
if len(result) == 0 {
return nil, fmt.Errorf("no job types resolved from %q", tokens)
}
return result, nil
}
// tokenAsCategory returns the category and true when tok is a known category
// keyword. "all" is treated as a special pseudo-category that matches every
// registered handler.
func tokenAsCategory(tok string) (JobCategory, bool) {
switch tok {
case string(CategoryAll):
return CategoryAll, true
case string(CategoryDefault):
return CategoryDefault, true
case string(CategoryHeavy):
return CategoryHeavy, true
default:
return "", false
}
}
func findFactory(factories []HandlerFactory, tok string) (HandlerFactory, error) {
for _, f := range factories {
if strings.EqualFold(f.JobType, tok) {
return f, nil
}
for _, alias := range f.Aliases {
if strings.EqualFold(alias, tok) {
return f, nil
}
}
}
return HandlerFactory{}, fmt.Errorf("unknown job type %q", tok)
}

View File

@@ -0,0 +1,14 @@
// Package handlers is an aggregator that blank-imports every plugin worker
// handler subpackage so their init() functions register with the handler
// registry. Import this package instead of individual subpackages when you
// need all handlers available.
package handlers
import (
_ "github.com/seaweedfs/seaweedfs/weed/worker/tasks/balance" // register volume_balance handler
_ "github.com/seaweedfs/seaweedfs/weed/worker/tasks/ec_balance" // register ec_balance handler
_ "github.com/seaweedfs/seaweedfs/weed/worker/tasks/erasure_coding" // register erasure_coding handler
_ "github.com/seaweedfs/seaweedfs/weed/worker/tasks/iceberg" // register iceberg_maintenance handler
_ "github.com/seaweedfs/seaweedfs/weed/worker/tasks/s3_lifecycle" // register s3_lifecycle handler
_ "github.com/seaweedfs/seaweedfs/weed/worker/tasks/vacuum" // register vacuum handler
)

View File

@@ -0,0 +1,20 @@
package pluginworker
import (
"time"
"google.golang.org/protobuf/types/known/timestamppb"
)
// ShouldSkipDetectionByInterval returns true when less than minIntervalSeconds
// have elapsed since lastSuccessfulRun. Exported so sub-packages can reuse it.
func ShouldSkipDetectionByInterval(lastSuccessfulRun *timestamppb.Timestamp, minIntervalSeconds int) bool {
if lastSuccessfulRun == nil || minIntervalSeconds <= 0 {
return false
}
lastRun := lastSuccessfulRun.AsTime()
if lastRun.IsZero() {
return false
}
return time.Since(lastRun) < time.Duration(minIntervalSeconds)*time.Second
}

View File

@@ -0,0 +1,39 @@
package pluginworker
import (
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
)
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
}

View File

@@ -0,0 +1,75 @@
package pluginworker
import (
"context"
"fmt"
"strconv"
"time"
"github.com/seaweedfs/seaweedfs/weed/pb"
"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
"google.golang.org/grpc"
)
// LookupVolumeLocations asks the masters for the current registered locations
// of a volume, bypassing any client-side caches. Execution handlers use it to
// re-check a proposal against the live topology, since a proposal can go
// stale between detection and execution. Returns the location URLs
// (host:port) the master reports; an empty slice means the volume is not
// registered anywhere.
func LookupVolumeLocations(ctx context.Context, masterAddresses []string, grpcDialOption grpc.DialOption, volumeID uint32) ([]string, error) {
if grpcDialOption == nil {
return nil, fmt.Errorf("grpc dial option is not configured")
}
if len(masterAddresses) == 0 {
return nil, fmt.Errorf("no master addresses provided in cluster context")
}
vid := strconv.FormatUint(uint64(volumeID), 10)
var lastErr error
for _, address := range masterAddresses {
for _, candidate := range MasterAddressCandidates(address) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
dialCtx, cancelDial := context.WithTimeout(ctx, 5*time.Second)
conn, err := pb.GrpcDial(dialCtx, candidate, false, grpcDialOption)
cancelDial()
if err != nil {
lastErr = err
continue
}
client := master_pb.NewSeaweedClient(conn)
callCtx, cancelCall := context.WithTimeout(ctx, 10*time.Second)
resp, callErr := client.LookupVolume(callCtx, &master_pb.LookupVolumeRequest{
VolumeOrFileIds: []string{vid},
})
cancelCall()
_ = conn.Close()
if callErr != nil {
lastErr = callErr
continue
}
for _, vidLocation := range resp.VolumeIdLocations {
if vidLocation.VolumeOrFileId != vid {
continue
}
locations := make([]string, 0, len(vidLocation.Locations))
for _, location := range vidLocation.Locations {
locations = append(locations, location.Url)
}
return locations, nil
}
// The master answered but does not know the volume.
return nil, nil
}
}
if lastErr == nil {
lastErr = fmt.Errorf("no valid master address candidate")
}
return nil, lastErr
}

View File

@@ -0,0 +1,243 @@
package pluginworker
import (
"context"
"errors"
"fmt"
"sort"
"strings"
"time"
"github.com/seaweedfs/seaweedfs/weed/admin/topology"
"github.com/seaweedfs/seaweedfs/weed/glog"
"github.com/seaweedfs/seaweedfs/weed/pb"
"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
"github.com/seaweedfs/seaweedfs/weed/util/wildcard"
workertypes "github.com/seaweedfs/seaweedfs/weed/worker/types"
"google.golang.org/grpc"
)
// CollectVolumeMetricsFromMasters dials the provided master addresses in order
// until one returns a usable volume list, then converts that into per-volume
// health metrics, an active-topology view, and a replica-location map.
func CollectVolumeMetricsFromMasters(
ctx context.Context,
masterAddresses []string,
collectionFilter string,
grpcDialOption grpc.DialOption,
) ([]*workertypes.VolumeHealthMetrics, *topology.ActiveTopology, map[uint32][]workertypes.ReplicaLocation, error) {
if grpcDialOption == nil {
return nil, nil, nil, fmt.Errorf("grpc dial option is not configured")
}
if len(masterAddresses) == 0 {
return nil, nil, nil, fmt.Errorf("no master addresses provided in cluster context")
}
for _, masterAddress := range masterAddresses {
response, err := FetchVolumeList(ctx, masterAddress, grpcDialOption)
if err != nil {
glog.Warningf("Plugin worker failed master volume list at %s: %v", masterAddress, err)
continue
}
metrics, activeTopology, replicaMap, buildErr := buildVolumeMetrics(response, collectionFilter)
if buildErr != nil {
// Configuration errors (e.g. invalid regex) will fail on every master,
// so return immediately instead of masking them with retries.
if isConfigError(buildErr) {
return nil, nil, nil, buildErr
}
glog.Warningf("Plugin worker failed to build metrics from master %s: %v", masterAddress, buildErr)
continue
}
return metrics, activeTopology, replicaMap, nil
}
return nil, nil, nil, fmt.Errorf("failed to load topology from all provided masters")
}
// FetchVolumeList dials the given master address (trying both the address as
// given and the gRPC port variant) and returns the master's volume list. Used
// by detection helpers that already know which master address to talk to.
// FetchDefaultReplicaPlacement returns the master's configured default replication
// (GetMasterConfiguration), used by detectors as the replica-placement fallback so
// the plugin path matches the shell. Returns "" if it cannot be fetched, so callers
// fall back to even spread rather than failing detection.
func FetchDefaultReplicaPlacement(ctx context.Context, masterAddresses []string, grpcDialOption grpc.DialOption) string {
if grpcDialOption == nil {
return ""
}
for _, address := range masterAddresses {
for _, candidate := range MasterAddressCandidates(address) {
if ctx.Err() != nil {
return ""
}
dialCtx, cancelDial := context.WithTimeout(ctx, 5*time.Second)
conn, err := pb.GrpcDial(dialCtx, candidate, false, grpcDialOption)
cancelDial()
if err != nil {
continue
}
client := master_pb.NewSeaweedClient(conn)
callCtx, cancelCall := context.WithTimeout(ctx, 10*time.Second)
resp, callErr := client.GetMasterConfiguration(callCtx, &master_pb.GetMasterConfigurationRequest{})
cancelCall()
_ = conn.Close()
if callErr == nil {
return resp.DefaultReplication
}
}
}
return ""
}
func FetchVolumeList(ctx context.Context, address string, grpcDialOption grpc.DialOption) (*master_pb.VolumeListResponse, error) {
var lastErr error
for _, candidate := range MasterAddressCandidates(address) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
dialCtx, cancelDial := context.WithTimeout(ctx, 5*time.Second)
conn, err := pb.GrpcDial(dialCtx, candidate, false, grpcDialOption)
cancelDial()
if err != nil {
lastErr = err
continue
}
client := master_pb.NewSeaweedClient(conn)
callCtx, cancelCall := context.WithTimeout(ctx, 10*time.Second)
response, callErr := pb.CollectVolumeList(callCtx, client, &master_pb.VolumeListRequest{})
cancelCall()
_ = conn.Close()
if callErr == nil {
return response, nil
}
lastErr = callErr
}
if lastErr == nil {
lastErr = fmt.Errorf("no valid master address candidate")
}
return nil, lastErr
}
func buildVolumeMetrics(
response *master_pb.VolumeListResponse,
collectionFilter string,
) ([]*workertypes.VolumeHealthMetrics, *topology.ActiveTopology, map[uint32][]workertypes.ReplicaLocation, error) {
if response == nil || response.TopologyInfo == nil {
return nil, nil, nil, fmt.Errorf("volume list response has no topology info")
}
activeTopology := topology.NewActiveTopology(10)
if err := activeTopology.UpdateTopology(response.TopologyInfo); err != nil {
return nil, nil, nil, err
}
collectionMatcher, err := wildcard.CompileCollectionMatcher(collectionFilter)
if err != nil {
return nil, nil, nil, &configError{err: err}
}
volumeSizeLimitBytes := uint64(response.VolumeSizeLimitMb) * 1024 * 1024
now := time.Now()
metrics := make([]*workertypes.VolumeHealthMetrics, 0, 256)
replicaMap := make(map[uint32][]workertypes.ReplicaLocation)
for _, dc := range response.TopologyInfo.DataCenterInfos {
for _, rack := range dc.RackInfos {
for _, node := range rack.DataNodeInfos {
for diskType, diskInfo := range node.DiskInfos {
for _, volume := range diskInfo.VolumeInfos {
// Build replica map from ALL volumes BEFORE collection filtering,
// since replicas may span filtered/unfiltered nodes.
replicaMap[volume.Id] = append(replicaMap[volume.Id], workertypes.ReplicaLocation{
DataCenter: dc.Id,
Rack: rack.Id,
NodeID: node.Id,
Host: pb.NewServerAddressFromDataNode(node).ToHost(),
})
if !collectionMatcher.Matches(volume.Collection) {
continue
}
metric := &workertypes.VolumeHealthMetrics{
VolumeID: volume.Id,
Server: node.Id,
ServerAddress: string(pb.NewServerAddressFromDataNode(node)),
DiskType: diskType,
DiskId: volume.DiskId,
DataCenter: dc.Id,
Rack: rack.Id,
Collection: volume.Collection,
Size: volume.Size,
DeletedBytes: volume.DeletedByteCount,
LastModified: time.Unix(volume.ModifiedAtSecond, 0),
ReplicaCount: 1,
ExpectedReplicas: int(volume.ReplicaPlacement),
IsReadOnly: volume.ReadOnly,
HasRemoteCopy: volume.RemoteStorageName != "",
}
if metric.Size > 0 {
metric.GarbageRatio = float64(metric.DeletedBytes) / float64(metric.Size)
}
if volumeSizeLimitBytes > 0 {
metric.FullnessRatio = float64(metric.Size) / float64(volumeSizeLimitBytes)
}
metric.Age = now.Sub(metric.LastModified)
metrics = append(metrics, metric)
}
}
}
}
}
replicaCounts := make(map[uint32]int)
for _, metric := range metrics {
replicaCounts[metric.VolumeID]++
}
for _, metric := range metrics {
metric.ReplicaCount = replicaCounts[metric.VolumeID]
}
return metrics, activeTopology, replicaMap, nil
}
// configError wraps configuration errors that should not be retried across masters.
type configError struct {
err error
}
func (e *configError) Error() string { return e.err.Error() }
func (e *configError) Unwrap() error { return e.err }
func isConfigError(err error) bool {
var ce *configError
return errors.As(err, &ce)
}
// MasterAddressCandidates returns address forms to try when dialing a master:
// the address as given plus the gRPC variant (port + 10000). Both are tried
// because callers may pass either an HTTP-port or gRPC-port address.
func MasterAddressCandidates(address string) []string {
trimmed := strings.TrimSpace(address)
if trimmed == "" {
return nil
}
candidateSet := map[string]struct{}{
trimmed: {},
}
converted := pb.ServerToGrpcAddress(trimmed)
candidateSet[converted] = struct{}{}
candidates := make([]string, 0, len(candidateSet))
for candidate := range candidateSet {
candidates = append(candidates, candidate)
}
sort.Strings(candidates)
return candidates
}

View File

@@ -0,0 +1,130 @@
package pluginworker
import (
"testing"
"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
"github.com/seaweedfs/seaweedfs/weed/util/wildcard"
)
func makeTestVolumeListResponse(volumes ...*master_pb.VolumeInformationMessage) *master_pb.VolumeListResponse {
return &master_pb.VolumeListResponse{
VolumeSizeLimitMb: 30000,
TopologyInfo: &master_pb.TopologyInfo{
DataCenterInfos: []*master_pb.DataCenterInfo{
{
Id: "dc1",
RackInfos: []*master_pb.RackInfo{
{
Id: "rack1",
DataNodeInfos: []*master_pb.DataNodeInfo{
{
Id: "server1:8080",
DiskInfos: map[string]*master_pb.DiskInfo{
"hdd": {
VolumeInfos: volumes,
},
},
},
},
},
},
},
},
},
}
}
func TestBuildVolumeMetricsEmptyFilter(t *testing.T) {
resp := makeTestVolumeListResponse(
&master_pb.VolumeInformationMessage{Id: 1, Collection: "photos", Size: 100},
&master_pb.VolumeInformationMessage{Id: 2, Collection: "videos", Size: 200},
)
metrics, _, _, err := buildVolumeMetrics(resp, "")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(metrics) != 2 {
t.Fatalf("expected 2 metrics, got %d", len(metrics))
}
}
func TestBuildVolumeMetricsAllCollections(t *testing.T) {
resp := makeTestVolumeListResponse(
&master_pb.VolumeInformationMessage{Id: 1, Collection: "photos", Size: 100},
&master_pb.VolumeInformationMessage{Id: 2, Collection: "videos", Size: 200},
)
metrics, _, _, err := buildVolumeMetrics(resp, string(wildcard.CollectionFilterAll))
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(metrics) != 2 {
t.Fatalf("expected 2 metrics, got %d", len(metrics))
}
}
func TestBuildVolumeMetricsEachCollection(t *testing.T) {
resp := makeTestVolumeListResponse(
&master_pb.VolumeInformationMessage{Id: 1, Collection: "photos", Size: 100},
&master_pb.VolumeInformationMessage{Id: 2, Collection: "videos", Size: 200},
)
// EACH_COLLECTION passes all volumes through; filtering happens in the handler
metrics, _, _, err := buildVolumeMetrics(resp, string(wildcard.CollectionFilterEach))
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(metrics) != 2 {
t.Fatalf("expected 2 metrics, got %d", len(metrics))
}
}
func TestBuildVolumeMetricsRegexFilter(t *testing.T) {
resp := makeTestVolumeListResponse(
&master_pb.VolumeInformationMessage{Id: 1, Collection: "photos", Size: 100},
&master_pb.VolumeInformationMessage{Id: 2, Collection: "videos", Size: 200},
&master_pb.VolumeInformationMessage{Id: 3, Collection: "photos-backup", Size: 300},
)
metrics, _, _, err := buildVolumeMetrics(resp, "^photos$")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(metrics) != 1 {
t.Fatalf("expected 1 metric, got %d", len(metrics))
}
if metrics[0].Collection != "photos" {
t.Fatalf("expected collection 'photos', got %q", metrics[0].Collection)
}
}
func TestBuildVolumeMetricsCollectionList(t *testing.T) {
resp := makeTestVolumeListResponse(
&master_pb.VolumeInformationMessage{Id: 1, Collection: "photos", Size: 100},
&master_pb.VolumeInformationMessage{Id: 2, Collection: "videos", Size: 200},
&master_pb.VolumeInformationMessage{Id: 3, Collection: "clips", Size: 300},
)
metrics, _, _, err := buildVolumeMetrics(resp, "photos,videos")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(metrics) != 2 {
t.Fatalf("expected 2 metrics, got %d", len(metrics))
}
for _, metric := range metrics {
if metric.Collection != "photos" && metric.Collection != "videos" {
t.Fatalf("unexpected collection %q", metric.Collection)
}
}
}
func TestBuildVolumeMetricsInvalidRegex(t *testing.T) {
resp := makeTestVolumeListResponse(
&master_pb.VolumeInformationMessage{Id: 1, Collection: "photos", Size: 100},
)
_, _, _, err := buildVolumeMetrics(resp, "[invalid")
if err == nil {
t.Fatal("expected error for invalid regex")
}
if !isConfigError(err) {
t.Fatalf("expected config error for invalid regex, got: %v", err)
}
}

View File

@@ -0,0 +1,76 @@
package pluginworker
import (
"github.com/seaweedfs/seaweedfs/weed/util/wildcard"
workertypes "github.com/seaweedfs/seaweedfs/weed/worker/types"
)
// VolumeState controls which volumes participate in detection (e.g. balance,
// vacuum). It is passed to FilterMetricsByVolumeState by the various plugin
// handlers.
type VolumeState string
const (
VolumeStateAll VolumeState = "ALL"
VolumeStateActive VolumeState = "ACTIVE"
VolumeStateFull VolumeState = "FULL"
)
// FilterMetricsByVolumeState filters volume metrics by state.
// "ACTIVE" keeps volumes with FullnessRatio < 1.01 (writable, below size limit).
// "FULL" keeps volumes with FullnessRatio >= 1.01 (read-only, above size limit).
// "ALL" or any other value returns all metrics unfiltered.
func FilterMetricsByVolumeState(metrics []*workertypes.VolumeHealthMetrics, state VolumeState) []*workertypes.VolumeHealthMetrics {
const fullnessThreshold = 1.01
var predicate func(m *workertypes.VolumeHealthMetrics) bool
switch state {
case VolumeStateActive:
predicate = func(m *workertypes.VolumeHealthMetrics) bool {
return m.FullnessRatio < fullnessThreshold
}
case VolumeStateFull:
predicate = func(m *workertypes.VolumeHealthMetrics) bool {
return m.FullnessRatio >= fullnessThreshold
}
default:
return metrics
}
filtered := make([]*workertypes.VolumeHealthMetrics, 0, len(metrics))
for _, m := range metrics {
if m == nil {
continue
}
if predicate(m) {
filtered = append(filtered, m)
}
}
return filtered
}
// FilterMetricsByLocation filters volume metrics by data center, rack, and node
// wildcards. Empty filters match anything in their respective dimension.
func FilterMetricsByLocation(metrics []*workertypes.VolumeHealthMetrics, dcFilter, rackFilter, nodeFilter string) []*workertypes.VolumeHealthMetrics {
dcMatchers := wildcard.CompileWildcardMatchers(dcFilter)
rackMatchers := wildcard.CompileWildcardMatchers(rackFilter)
nodeMatchers := wildcard.CompileWildcardMatchers(nodeFilter)
filtered := make([]*workertypes.VolumeHealthMetrics, 0, len(metrics))
for _, m := range metrics {
if m == nil {
continue
}
if !wildcard.MatchesAnyWildcard(dcMatchers, m.DataCenter) {
continue
}
if !wildcard.MatchesAnyWildcard(rackMatchers, m.Rack) {
continue
}
if !wildcard.MatchesAnyWildcard(nodeMatchers, m.Server) {
continue
}
filtered = append(filtered, m)
}
return filtered
}

View File

@@ -0,0 +1,939 @@
package pluginworker
import (
"context"
"crypto/rand"
"encoding/hex"
"fmt"
"os"
"sort"
"strings"
"sync"
"time"
"github.com/seaweedfs/seaweedfs/weed/glog"
"github.com/seaweedfs/seaweedfs/weed/pb"
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
"google.golang.org/grpc"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/timestamppb"
)
const (
defaultHeartbeatInterval = 15 * time.Second
defaultReconnectDelay = 5 * time.Second
defaultSendBufferSize = 256
)
// DetectionSender sends detection responses for one request.
type DetectionSender interface {
SendProposals(*plugin_pb.DetectionProposals) error
SendComplete(*plugin_pb.DetectionComplete) error
SendActivity(*plugin_pb.ActivityEvent) error
}
// ExecutionSender sends execution progress/completion responses for one request.
type ExecutionSender interface {
SendProgress(*plugin_pb.JobProgressUpdate) error
SendCompleted(*plugin_pb.JobCompleted) error
}
// JobHandler implements one plugin job type on the worker side.
type JobHandler interface {
Capability() *plugin_pb.JobTypeCapability
Descriptor() *plugin_pb.JobTypeDescriptor
Detect(context.Context, *plugin_pb.RunDetectionRequest, DetectionSender) error
Execute(context.Context, *plugin_pb.ExecuteJobRequest, ExecutionSender) error
}
// WorkerOptions configures one plugin worker process.
type WorkerOptions struct {
AdminServer string
WorkerID string
WorkerVersion string
WorkerAddress string
HeartbeatInterval time.Duration
ReconnectDelay time.Duration
MaxDetectionConcurrency int
MaxExecutionConcurrency int
GrpcDialOption grpc.DialOption
Handlers []JobHandler
Handler JobHandler
}
// Worker runs one plugin job handler over plugin.proto stream.
type Worker struct {
opts WorkerOptions
detectSlots chan struct{}
execSlots chan struct{}
handlers map[string]JobHandler
runningMu sync.RWMutex
runningWork map[string]*plugin_pb.RunningWork
workCancelMu sync.Mutex
workCancel map[string]context.CancelFunc
workerID string
connectionMu sync.RWMutex
connected bool
}
// NewWorker creates a plugin worker instance.
func NewWorker(options WorkerOptions) (*Worker, error) {
if strings.TrimSpace(options.AdminServer) == "" {
return nil, fmt.Errorf("admin server is required")
}
if options.GrpcDialOption == nil {
return nil, fmt.Errorf("grpc dial option is required")
}
if options.HeartbeatInterval <= 0 {
options.HeartbeatInterval = defaultHeartbeatInterval
}
if options.ReconnectDelay <= 0 {
options.ReconnectDelay = defaultReconnectDelay
}
if options.MaxDetectionConcurrency <= 0 {
options.MaxDetectionConcurrency = 1
}
if options.MaxExecutionConcurrency <= 0 {
options.MaxExecutionConcurrency = 1
}
if strings.TrimSpace(options.WorkerVersion) == "" {
options.WorkerVersion = "dev"
}
workerID := strings.TrimSpace(options.WorkerID)
if workerID == "" {
workerID = generateWorkerID()
}
workerAddress := strings.TrimSpace(options.WorkerAddress)
if workerAddress == "" {
hostname, _ := os.Hostname()
workerAddress = hostname
}
opts := options
opts.WorkerAddress = workerAddress
allHandlers := make([]JobHandler, 0, len(opts.Handlers)+1)
if opts.Handler != nil {
allHandlers = append(allHandlers, opts.Handler)
}
allHandlers = append(allHandlers, opts.Handlers...)
if len(allHandlers) == 0 {
return nil, fmt.Errorf("at least one job handler is required")
}
handlers := make(map[string]JobHandler, len(allHandlers))
for i, handler := range allHandlers {
if handler == nil {
return nil, fmt.Errorf("job handler at index %d is nil", i)
}
handlerJobType, err := resolveHandlerJobType(handler)
if err != nil {
return nil, fmt.Errorf("resolve job handler at index %d: %w", i, err)
}
key := normalizeJobTypeKey(handlerJobType)
if key == "" {
return nil, fmt.Errorf("job handler at index %d has empty job type", i)
}
if _, found := handlers[key]; found {
return nil, fmt.Errorf("duplicate job handler for job type %q", handlerJobType)
}
handlers[key] = handler
}
if opts.Handler == nil {
opts.Handler = allHandlers[0]
}
w := &Worker{
opts: opts,
detectSlots: make(chan struct{}, opts.MaxDetectionConcurrency),
execSlots: make(chan struct{}, opts.MaxExecutionConcurrency),
handlers: handlers,
runningWork: make(map[string]*plugin_pb.RunningWork),
workCancel: make(map[string]context.CancelFunc),
workerID: workerID,
}
return w, nil
}
// Run keeps the plugin worker connected and reconnects on stream failures.
func (w *Worker) Run(ctx context.Context) error {
adminAddress := pb.ServerToGrpcAddress(w.opts.AdminServer)
for {
select {
case <-ctx.Done():
return nil
default:
}
if err := w.runOnce(ctx, adminAddress); err != nil {
if ctx.Err() != nil {
return nil
}
glog.Warningf("Plugin worker %s stream ended: %v", w.workerID, err)
}
select {
case <-ctx.Done():
return nil
case <-time.After(w.opts.ReconnectDelay):
}
}
}
func (w *Worker) runOnce(ctx context.Context, adminAddress string) error {
defer w.setConnected(false)
dialCtx, cancelDial := context.WithTimeout(ctx, 5*time.Second)
defer cancelDial()
conn, err := pb.GrpcDial(dialCtx, adminAddress, false, w.opts.GrpcDialOption)
if err != nil {
return fmt.Errorf("dial admin %s: %w", adminAddress, err)
}
defer conn.Close()
client := plugin_pb.NewPluginControlServiceClient(conn)
connCtx, cancelConn := context.WithCancel(ctx)
defer cancelConn()
stream, err := client.WorkerStream(connCtx)
if err != nil {
return fmt.Errorf("open worker stream: %w", err)
}
w.setConnected(true)
sendCh := make(chan *plugin_pb.WorkerToAdminMessage, defaultSendBufferSize)
sendErrCh := make(chan error, 1)
send := func(msg *plugin_pb.WorkerToAdminMessage) bool {
if msg == nil {
return false
}
msg.WorkerId = w.workerID
if msg.SentAt == nil {
msg.SentAt = timestamppb.Now()
}
select {
case <-connCtx.Done():
return false
case sendCh <- msg:
return true
}
}
go func() {
for {
select {
case <-connCtx.Done():
return
case msg := <-sendCh:
if msg == nil {
continue
}
if err := stream.Send(msg); err != nil {
select {
case sendErrCh <- err:
default:
}
cancelConn()
return
}
}
}
}()
if !send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_Hello{Hello: w.buildHello()},
}) {
return fmt.Errorf("send worker hello: stream closed")
}
heartbeatTicker := time.NewTicker(w.opts.HeartbeatInterval)
defer heartbeatTicker.Stop()
go func() {
for {
select {
case <-connCtx.Done():
return
case <-heartbeatTicker.C:
w.sendHeartbeat(send)
}
}
}()
for {
select {
case <-connCtx.Done():
return connCtx.Err()
case err := <-sendErrCh:
return fmt.Errorf("send to admin stream: %w", err)
default:
}
message, err := stream.Recv()
if err != nil {
return fmt.Errorf("recv admin message: %w", err)
}
w.handleAdminMessage(connCtx, message, send)
}
}
// IsConnected reports whether the worker currently has an active stream to admin.
func (w *Worker) IsConnected() bool {
w.connectionMu.RLock()
defer w.connectionMu.RUnlock()
return w.connected
}
func (w *Worker) setConnected(connected bool) {
w.connectionMu.Lock()
w.connected = connected
w.connectionMu.Unlock()
}
func (w *Worker) handleAdminMessage(
ctx context.Context,
message *plugin_pb.AdminToWorkerMessage,
send func(*plugin_pb.WorkerToAdminMessage) bool,
) {
if message == nil {
return
}
switch body := message.Body.(type) {
case *plugin_pb.AdminToWorkerMessage_Hello:
_ = body
case *plugin_pb.AdminToWorkerMessage_RequestConfigSchema:
w.handleSchemaRequest(message.GetRequestId(), body.RequestConfigSchema, send)
case *plugin_pb.AdminToWorkerMessage_RunDetectionRequest:
w.handleDetectionRequest(ctx, message.GetRequestId(), body.RunDetectionRequest, send)
case *plugin_pb.AdminToWorkerMessage_ExecuteJobRequest:
w.handleExecuteRequest(ctx, message.GetRequestId(), body.ExecuteJobRequest, send)
case *plugin_pb.AdminToWorkerMessage_CancelRequest:
cancel := body.CancelRequest
targetID := ""
if cancel != nil {
targetID = strings.TrimSpace(cancel.TargetId)
}
accepted := false
ackMessage := "cancel target is required"
if targetID != "" {
if w.cancelWork(targetID) {
accepted = true
ackMessage = "cancel request accepted"
} else {
ackMessage = "cancel target not found"
}
}
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_Acknowledge{Acknowledge: &plugin_pb.WorkerAcknowledge{
RequestId: message.GetRequestId(),
Accepted: accepted,
Message: ackMessage,
}},
})
case *plugin_pb.AdminToWorkerMessage_Shutdown:
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_Acknowledge{Acknowledge: &plugin_pb.WorkerAcknowledge{
RequestId: message.GetRequestId(),
Accepted: true,
Message: "shutdown acknowledged",
}},
})
default:
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_Acknowledge{Acknowledge: &plugin_pb.WorkerAcknowledge{
RequestId: message.GetRequestId(),
Accepted: false,
Message: "unsupported request body",
}},
})
}
}
func (w *Worker) handleSchemaRequest(requestID string, request *plugin_pb.RequestConfigSchema, send func(*plugin_pb.WorkerToAdminMessage) bool) {
jobType := ""
if request != nil {
jobType = strings.TrimSpace(request.JobType)
}
handler, resolvedJobType, err := w.findHandler(jobType)
if err != nil {
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_ConfigSchemaResponse{ConfigSchemaResponse: &plugin_pb.ConfigSchemaResponse{
RequestId: requestID,
JobType: jobType,
Success: false,
ErrorMessage: err.Error(),
}},
})
return
}
descriptor := handler.Descriptor()
if descriptor == nil || descriptor.JobType == "" {
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_ConfigSchemaResponse{ConfigSchemaResponse: &plugin_pb.ConfigSchemaResponse{
RequestId: requestID,
JobType: resolvedJobType,
Success: false,
ErrorMessage: "handler descriptor is not configured",
}},
})
return
}
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_ConfigSchemaResponse{ConfigSchemaResponse: &plugin_pb.ConfigSchemaResponse{
RequestId: requestID,
JobType: descriptor.JobType,
Success: true,
JobTypeDescriptor: descriptor,
}},
})
}
func (w *Worker) handleDetectionRequest(
ctx context.Context,
requestID string,
request *plugin_pb.RunDetectionRequest,
send func(*plugin_pb.WorkerToAdminMessage) bool,
) {
if request == nil {
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_DetectionComplete{DetectionComplete: &plugin_pb.DetectionComplete{
RequestId: requestID,
Success: false,
ErrorMessage: "run detection request is nil",
}},
})
return
}
handler, resolvedJobType, err := w.findHandler(request.JobType)
if err != nil {
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_DetectionComplete{DetectionComplete: &plugin_pb.DetectionComplete{
RequestId: requestID,
JobType: request.JobType,
Success: false,
ErrorMessage: err.Error(),
}},
})
return
}
workKey := "detect:" + requestID
w.setRunningWork(workKey, &plugin_pb.RunningWork{
WorkId: requestID,
Kind: plugin_pb.WorkKind_WORK_KIND_DETECTION,
JobType: resolvedJobType,
State: plugin_pb.JobState_JOB_STATE_ASSIGNED,
ProgressPercent: 0,
Stage: "queued",
})
w.sendHeartbeat(send)
requestCtx, cancelRequest := context.WithCancel(ctx)
w.setWorkCancel(cancelRequest, requestID)
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_Acknowledge{Acknowledge: &plugin_pb.WorkerAcknowledge{
RequestId: requestID,
Accepted: true,
Message: "detection request accepted",
}},
})
go func() {
detectionSender := &detectionSender{
requestID: requestID,
jobType: resolvedJobType,
send: send,
}
defer func() {
w.clearWorkCancel(requestID)
cancelRequest()
w.clearRunningWork(workKey)
w.sendHeartbeat(send)
}()
select {
case <-requestCtx.Done():
detectionSender.SendComplete(&plugin_pb.DetectionComplete{
Success: false,
ErrorMessage: requestCtx.Err().Error(),
})
return
case w.detectSlots <- struct{}{}:
}
defer func() {
<-w.detectSlots
w.sendHeartbeat(send)
}()
w.setRunningWork(workKey, &plugin_pb.RunningWork{
WorkId: requestID,
Kind: plugin_pb.WorkKind_WORK_KIND_DETECTION,
JobType: resolvedJobType,
State: plugin_pb.JobState_JOB_STATE_RUNNING,
ProgressPercent: 0,
Stage: "detecting",
})
w.sendHeartbeat(send)
if err := handler.Detect(requestCtx, request, detectionSender); err != nil {
detectionSender.SendComplete(&plugin_pb.DetectionComplete{
Success: false,
ErrorMessage: err.Error(),
})
}
}()
}
func (w *Worker) handleExecuteRequest(
ctx context.Context,
requestID string,
request *plugin_pb.ExecuteJobRequest,
send func(*plugin_pb.WorkerToAdminMessage) bool,
) {
if request == nil || request.Job == nil {
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_JobCompleted{JobCompleted: &plugin_pb.JobCompleted{
RequestId: requestID,
Success: false,
ErrorMessage: "execute request/job is nil",
}},
})
return
}
handler, resolvedJobType, err := w.findHandler(request.Job.JobType)
if err != nil {
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_JobCompleted{JobCompleted: &plugin_pb.JobCompleted{
RequestId: requestID,
JobId: request.Job.JobId,
JobType: request.Job.JobType,
Success: false,
ErrorMessage: err.Error(),
}},
})
return
}
select {
case w.execSlots <- struct{}{}:
default:
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_JobCompleted{JobCompleted: &plugin_pb.JobCompleted{
RequestId: requestID,
JobId: request.Job.JobId,
JobType: resolvedJobType,
Success: false,
ErrorMessage: "executor is at capacity",
}},
})
return
}
w.sendHeartbeat(send)
workKey := "exec:" + requestID
w.setRunningWork(workKey, &plugin_pb.RunningWork{
WorkId: request.Job.JobId,
Kind: plugin_pb.WorkKind_WORK_KIND_EXECUTION,
JobType: resolvedJobType,
State: plugin_pb.JobState_JOB_STATE_RUNNING,
ProgressPercent: 0,
Stage: "starting",
})
w.sendHeartbeat(send)
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_Acknowledge{Acknowledge: &plugin_pb.WorkerAcknowledge{
RequestId: requestID,
Accepted: true,
Message: "execute request accepted",
}},
})
go func() {
requestCtx, cancelRequest := context.WithCancel(ctx)
w.setWorkCancel(cancelRequest, requestID, request.Job.JobId)
defer func() {
w.clearWorkCancel(requestID, request.Job.JobId)
cancelRequest()
<-w.execSlots
w.clearRunningWork(workKey)
w.sendHeartbeat(send)
}()
executionSender := &executionSender{
requestID: requestID,
jobID: request.Job.JobId,
jobType: resolvedJobType,
send: send,
onProgress: func(progress float64, stage string) {
w.updateRunningExecution(workKey, progress, stage)
},
}
if err := handler.Execute(requestCtx, request, executionSender); err != nil {
executionSender.SendCompleted(&plugin_pb.JobCompleted{
Success: false,
ErrorMessage: err.Error(),
})
}
}()
}
func (w *Worker) buildHello() *plugin_pb.WorkerHello {
jobTypeKeys := make([]string, 0, len(w.handlers))
for key := range w.handlers {
jobTypeKeys = append(jobTypeKeys, key)
}
sort.Strings(jobTypeKeys)
capabilities := make([]*plugin_pb.JobTypeCapability, 0, len(jobTypeKeys))
jobTypes := make([]string, 0, len(jobTypeKeys))
for _, key := range jobTypeKeys {
handler := w.handlers[key]
if handler == nil {
continue
}
jobType, _ := resolveHandlerJobType(handler)
capability := handler.Capability()
if capability == nil {
capability = &plugin_pb.JobTypeCapability{}
} else {
capability = proto.Clone(capability).(*plugin_pb.JobTypeCapability)
}
if strings.TrimSpace(capability.JobType) == "" {
capability.JobType = jobType
}
capability.MaxDetectionConcurrency = int32(cap(w.detectSlots))
capability.MaxExecutionConcurrency = int32(cap(w.execSlots))
capabilities = append(capabilities, capability)
if capability.JobType != "" {
jobTypes = append(jobTypes, capability.JobType)
}
}
instanceID := generateWorkerID()
return &plugin_pb.WorkerHello{
WorkerId: w.workerID,
WorkerInstanceId: "inst-" + instanceID,
Address: w.opts.WorkerAddress,
WorkerVersion: w.opts.WorkerVersion,
ProtocolVersion: "plugin.v1",
Capabilities: capabilities,
Metadata: map[string]string{
"runtime": "plugin",
"job_types": strings.Join(jobTypes, ","),
},
}
}
func (w *Worker) buildHeartbeat() *plugin_pb.WorkerHeartbeat {
w.runningMu.RLock()
running := make([]*plugin_pb.RunningWork, 0, len(w.runningWork))
for _, work := range w.runningWork {
if work == nil {
continue
}
cloned := *work
running = append(running, &cloned)
}
w.runningMu.RUnlock()
detectUsed := len(w.detectSlots)
execUsed := len(w.execSlots)
return &plugin_pb.WorkerHeartbeat{
WorkerId: w.workerID,
RunningWork: running,
DetectionSlotsUsed: int32(detectUsed),
DetectionSlotsTotal: int32(cap(w.detectSlots)),
ExecutionSlotsUsed: int32(execUsed),
ExecutionSlotsTotal: int32(cap(w.execSlots)),
QueuedJobsByType: map[string]int32{},
Metadata: map[string]string{
"runtime": "plugin",
},
}
}
func (w *Worker) sendHeartbeat(send func(*plugin_pb.WorkerToAdminMessage) bool) {
if send == nil {
return
}
send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_Heartbeat{
Heartbeat: w.buildHeartbeat(),
},
})
}
func (w *Worker) setRunningWork(key string, work *plugin_pb.RunningWork) {
if strings.TrimSpace(key) == "" || work == nil {
return
}
w.runningMu.Lock()
w.runningWork[key] = work
w.runningMu.Unlock()
}
func (w *Worker) clearRunningWork(key string) {
w.runningMu.Lock()
delete(w.runningWork, key)
w.runningMu.Unlock()
}
func (w *Worker) updateRunningExecution(key string, progress float64, stage string) {
w.runningMu.Lock()
if running := w.runningWork[key]; running != nil {
running.ProgressPercent = progress
if strings.TrimSpace(stage) != "" {
running.Stage = stage
}
running.State = plugin_pb.JobState_JOB_STATE_RUNNING
}
w.runningMu.Unlock()
}
type detectionSender struct {
requestID string
jobType string
send func(*plugin_pb.WorkerToAdminMessage) bool
}
func (s *detectionSender) SendProposals(proposals *plugin_pb.DetectionProposals) error {
if proposals == nil {
return fmt.Errorf("detection proposals are nil")
}
if proposals.RequestId == "" {
proposals.RequestId = s.requestID
}
if proposals.JobType == "" {
proposals.JobType = s.jobType
}
if !s.send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_DetectionProposals{DetectionProposals: proposals},
}) {
return fmt.Errorf("stream closed")
}
return nil
}
func (s *detectionSender) SendComplete(complete *plugin_pb.DetectionComplete) error {
if complete == nil {
return fmt.Errorf("detection complete is nil")
}
if complete.RequestId == "" {
complete.RequestId = s.requestID
}
if complete.JobType == "" {
complete.JobType = s.jobType
}
if !s.send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_DetectionComplete{DetectionComplete: complete},
}) {
return fmt.Errorf("stream closed")
}
return nil
}
func (s *detectionSender) SendActivity(activity *plugin_pb.ActivityEvent) error {
if activity == nil {
return fmt.Errorf("detection activity is nil")
}
if activity.CreatedAt == nil {
activity.CreatedAt = timestamppb.Now()
}
if activity.Source == plugin_pb.ActivitySource_ACTIVITY_SOURCE_UNSPECIFIED {
activity.Source = plugin_pb.ActivitySource_ACTIVITY_SOURCE_DETECTOR
}
update := &plugin_pb.JobProgressUpdate{
RequestId: s.requestID,
JobType: s.jobType,
State: plugin_pb.JobState_JOB_STATE_RUNNING,
ProgressPercent: 0,
Stage: activity.Stage,
Message: activity.Message,
Activities: []*plugin_pb.ActivityEvent{activity},
UpdatedAt: timestamppb.Now(),
}
if !s.send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_JobProgressUpdate{JobProgressUpdate: update},
}) {
return fmt.Errorf("stream closed")
}
return nil
}
type executionSender struct {
requestID string
jobID string
jobType string
send func(*plugin_pb.WorkerToAdminMessage) bool
onProgress func(progress float64, stage string)
}
func (s *executionSender) SendProgress(progress *plugin_pb.JobProgressUpdate) error {
if progress == nil {
return fmt.Errorf("job progress is nil")
}
if progress.RequestId == "" {
progress.RequestId = s.requestID
}
if progress.JobId == "" {
progress.JobId = s.jobID
}
if progress.JobType == "" {
progress.JobType = s.jobType
}
if progress.UpdatedAt == nil {
progress.UpdatedAt = timestamppb.Now()
}
if s.onProgress != nil {
s.onProgress(progress.ProgressPercent, progress.Stage)
}
if !s.send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_JobProgressUpdate{JobProgressUpdate: progress},
}) {
return fmt.Errorf("stream closed")
}
return nil
}
func (s *executionSender) SendCompleted(completed *plugin_pb.JobCompleted) error {
if completed == nil {
return fmt.Errorf("job completed is nil")
}
if completed.RequestId == "" {
completed.RequestId = s.requestID
}
if completed.JobId == "" {
completed.JobId = s.jobID
}
if completed.JobType == "" {
completed.JobType = s.jobType
}
if completed.CompletedAt == nil {
completed.CompletedAt = timestamppb.Now()
}
if !s.send(&plugin_pb.WorkerToAdminMessage{
Body: &plugin_pb.WorkerToAdminMessage_JobCompleted{JobCompleted: completed},
}) {
return fmt.Errorf("stream closed")
}
return nil
}
func generateWorkerID() string {
random := make([]byte, 3)
if _, err := rand.Read(random); err != nil {
return fmt.Sprintf("plugin-%d", time.Now().UnixNano())
}
return "plugin-" + hex.EncodeToString(random)
}
func (w *Worker) setWorkCancel(cancel context.CancelFunc, keys ...string) {
if cancel == nil {
return
}
w.workCancelMu.Lock()
defer w.workCancelMu.Unlock()
for _, key := range keys {
key = strings.TrimSpace(key)
if key == "" {
continue
}
w.workCancel[key] = cancel
}
}
func (w *Worker) clearWorkCancel(keys ...string) {
w.workCancelMu.Lock()
defer w.workCancelMu.Unlock()
for _, key := range keys {
key = strings.TrimSpace(key)
if key == "" {
continue
}
delete(w.workCancel, key)
}
}
func (w *Worker) cancelWork(targetID string) bool {
targetID = strings.TrimSpace(targetID)
if targetID == "" {
return false
}
w.workCancelMu.Lock()
cancel := w.workCancel[targetID]
w.workCancelMu.Unlock()
if cancel == nil {
return false
}
cancel()
return true
}
func (w *Worker) findHandler(jobType string) (JobHandler, string, error) {
trimmed := strings.TrimSpace(jobType)
if trimmed == "" {
if len(w.handlers) == 1 {
for _, handler := range w.handlers {
resolvedJobType, err := resolveHandlerJobType(handler)
return handler, resolvedJobType, err
}
}
return nil, "", fmt.Errorf("job type is required when worker serves multiple job types")
}
key := normalizeJobTypeKey(trimmed)
handler := w.handlers[key]
if handler == nil {
return nil, "", fmt.Errorf("job type %q is not handled by this worker", trimmed)
}
resolvedJobType, err := resolveHandlerJobType(handler)
if err != nil {
return nil, "", err
}
return handler, resolvedJobType, nil
}
func resolveHandlerJobType(handler JobHandler) (string, error) {
if handler == nil {
return "", fmt.Errorf("job handler is nil")
}
if descriptor := handler.Descriptor(); descriptor != nil {
if jobType := strings.TrimSpace(descriptor.JobType); jobType != "" {
return jobType, nil
}
}
if capability := handler.Capability(); capability != nil {
if jobType := strings.TrimSpace(capability.JobType); jobType != "" {
return jobType, nil
}
}
return "", fmt.Errorf("handler job type is not configured")
}
func normalizeJobTypeKey(jobType string) string {
return strings.ToLower(strings.TrimSpace(jobType))
}

View File

@@ -0,0 +1,599 @@
package pluginworker
import (
"context"
"testing"
"time"
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
)
func TestWorkerBuildHelloUsesConfiguredConcurrency(t *testing.T) {
handler := &testJobHandler{
capability: &plugin_pb.JobTypeCapability{
JobType: "vacuum",
CanDetect: true,
CanExecute: true,
MaxDetectionConcurrency: 99,
MaxExecutionConcurrency: 88,
},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "vacuum"},
}
worker, err := NewWorker(WorkerOptions{
AdminServer: "localhost:23646",
GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
Handler: handler,
MaxDetectionConcurrency: 3,
MaxExecutionConcurrency: 4,
})
if err != nil {
t.Fatalf("NewWorker error = %v", err)
}
hello := worker.buildHello()
if hello == nil || len(hello.Capabilities) != 1 {
t.Fatalf("expected one capability in hello")
}
capability := hello.Capabilities[0]
if capability.MaxDetectionConcurrency != 3 {
t.Fatalf("expected max_detection_concurrency=3, got=%d", capability.MaxDetectionConcurrency)
}
if capability.MaxExecutionConcurrency != 4 {
t.Fatalf("expected max_execution_concurrency=4, got=%d", capability.MaxExecutionConcurrency)
}
if capability.JobType != "vacuum" {
t.Fatalf("expected job type vacuum, got=%q", capability.JobType)
}
}
func TestWorkerBuildHelloIncludesMultipleCapabilities(t *testing.T) {
worker, err := NewWorker(WorkerOptions{
AdminServer: "localhost:23646",
GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
Handlers: []JobHandler{
&testJobHandler{
capability: &plugin_pb.JobTypeCapability{JobType: "vacuum", CanDetect: true, CanExecute: true},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "vacuum"},
},
&testJobHandler{
capability: &plugin_pb.JobTypeCapability{JobType: "volume_balance", CanDetect: true, CanExecute: true},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "volume_balance"},
},
},
MaxDetectionConcurrency: 2,
MaxExecutionConcurrency: 3,
})
if err != nil {
t.Fatalf("NewWorker error = %v", err)
}
hello := worker.buildHello()
if hello == nil || len(hello.Capabilities) != 2 {
t.Fatalf("expected two capabilities in hello")
}
found := map[string]bool{}
for _, capability := range hello.Capabilities {
found[capability.JobType] = true
if capability.MaxDetectionConcurrency != 2 {
t.Fatalf("expected max_detection_concurrency=2, got=%d", capability.MaxDetectionConcurrency)
}
if capability.MaxExecutionConcurrency != 3 {
t.Fatalf("expected max_execution_concurrency=3, got=%d", capability.MaxExecutionConcurrency)
}
}
if !found["vacuum"] || !found["volume_balance"] {
t.Fatalf("expected capabilities for vacuum and volume_balance, got=%v", found)
}
}
func TestWorkerCancelWorkByTargetID(t *testing.T) {
worker, err := NewWorker(WorkerOptions{
AdminServer: "localhost:23646",
GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
Handler: &testJobHandler{
capability: &plugin_pb.JobTypeCapability{JobType: "vacuum"},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "vacuum"},
},
})
if err != nil {
t.Fatalf("NewWorker error = %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
worker.setWorkCancel(cancel, "request-1", "job-1")
if !worker.cancelWork("request-1") {
t.Fatalf("expected cancel by request id to succeed")
}
select {
case <-ctx.Done():
case <-time.After(100 * time.Millisecond):
t.Fatalf("expected context to be canceled")
}
if !worker.cancelWork("job-1") {
t.Fatalf("expected cancel by job id to succeed")
}
if worker.cancelWork("unknown-target") {
t.Fatalf("expected cancel unknown target to fail")
}
}
func TestWorkerHandleCancelRequestAck(t *testing.T) {
worker, err := NewWorker(WorkerOptions{
AdminServer: "localhost:23646",
GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
Handler: &testJobHandler{
capability: &plugin_pb.JobTypeCapability{JobType: "vacuum"},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "vacuum"},
},
})
if err != nil {
t.Fatalf("NewWorker error = %v", err)
}
canceled := false
worker.setWorkCancel(func() { canceled = true }, "job-42")
var response *plugin_pb.WorkerToAdminMessage
ok := worker.handleAdminMessageForTest(&plugin_pb.AdminToWorkerMessage{
RequestId: "cancel-req-1",
Body: &plugin_pb.AdminToWorkerMessage_CancelRequest{
CancelRequest: &plugin_pb.CancelRequest{TargetId: "job-42"},
},
}, func(msg *plugin_pb.WorkerToAdminMessage) bool {
response = msg
return true
})
if !ok {
t.Fatalf("expected send callback to be invoked")
}
if !canceled {
t.Fatalf("expected registered work cancel function to be called")
}
if response == nil || response.GetAcknowledge() == nil || !response.GetAcknowledge().Accepted {
t.Fatalf("expected accepted acknowledge response, got=%+v", response)
}
response = nil
ok = worker.handleAdminMessageForTest(&plugin_pb.AdminToWorkerMessage{
RequestId: "cancel-req-2",
Body: &plugin_pb.AdminToWorkerMessage_CancelRequest{
CancelRequest: &plugin_pb.CancelRequest{TargetId: "missing"},
},
}, func(msg *plugin_pb.WorkerToAdminMessage) bool {
response = msg
return true
})
if !ok {
t.Fatalf("expected send callback to be invoked")
}
if response == nil || response.GetAcknowledge() == nil || response.GetAcknowledge().Accepted {
t.Fatalf("expected rejected acknowledge for missing target, got=%+v", response)
}
}
func TestWorkerSchemaRequestRequiresJobTypeWhenMultipleHandlers(t *testing.T) {
worker, err := NewWorker(WorkerOptions{
AdminServer: "localhost:23646",
GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
Handlers: []JobHandler{
&testJobHandler{
capability: &plugin_pb.JobTypeCapability{JobType: "vacuum"},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "vacuum"},
},
&testJobHandler{
capability: &plugin_pb.JobTypeCapability{JobType: "erasure_coding"},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "erasure_coding"},
},
},
})
if err != nil {
t.Fatalf("NewWorker error = %v", err)
}
var response *plugin_pb.WorkerToAdminMessage
ok := worker.handleAdminMessageForTest(&plugin_pb.AdminToWorkerMessage{
RequestId: "schema-req-1",
Body: &plugin_pb.AdminToWorkerMessage_RequestConfigSchema{
RequestConfigSchema: &plugin_pb.RequestConfigSchema{},
},
}, func(msg *plugin_pb.WorkerToAdminMessage) bool {
response = msg
return true
})
if !ok {
t.Fatalf("expected send callback to be invoked")
}
schema := response.GetConfigSchemaResponse()
if schema == nil || schema.Success {
t.Fatalf("expected schema error response, got=%+v", response)
}
}
func TestWorkerHandleDetectionQueuesWhenAtCapacity(t *testing.T) {
handler := &detectionQueueTestHandler{
capability: &plugin_pb.JobTypeCapability{
JobType: "vacuum",
CanDetect: true,
CanExecute: false,
},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "vacuum"},
detectEntered: make(chan struct{}, 2),
detectContinue: make(chan struct{}, 2),
}
worker, err := NewWorker(WorkerOptions{
AdminServer: "localhost:23646",
GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
Handler: handler,
MaxDetectionConcurrency: 1,
})
if err != nil {
t.Fatalf("NewWorker error = %v", err)
}
msgCh := make(chan *plugin_pb.WorkerToAdminMessage, 8)
send := func(msg *plugin_pb.WorkerToAdminMessage) bool {
msgCh <- msg
return true
}
sendDetection := func(requestID string) {
worker.handleAdminMessage(context.Background(), &plugin_pb.AdminToWorkerMessage{
RequestId: requestID,
Body: &plugin_pb.AdminToWorkerMessage_RunDetectionRequest{
RunDetectionRequest: &plugin_pb.RunDetectionRequest{
JobType: "vacuum",
},
},
}, send)
}
sendDetection("detect-1")
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
ack := message.GetAcknowledge()
return ack != nil && ack.RequestId == "detect-1" && ack.Accepted
}, "detection acknowledge detect-1")
<-handler.detectEntered
sendDetection("detect-2")
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
ack := message.GetAcknowledge()
return ack != nil && ack.RequestId == "detect-2" && ack.Accepted
}, "detection acknowledge detect-2")
select {
case unexpected := <-msgCh:
t.Fatalf("did not expect detection completion before slot is available, got=%+v", unexpected)
case <-time.After(100 * time.Millisecond):
}
handler.detectContinue <- struct{}{}
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
complete := message.GetDetectionComplete()
return complete != nil && complete.RequestId == "detect-1" && complete.Success
}, "detection complete detect-1")
<-handler.detectEntered
handler.detectContinue <- struct{}{}
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
complete := message.GetDetectionComplete()
return complete != nil && complete.RequestId == "detect-2" && complete.Success
}, "detection complete detect-2")
}
func TestWorkerHeartbeatReflectsActiveDetectionLoad(t *testing.T) {
handler := &detectionQueueTestHandler{
capability: &plugin_pb.JobTypeCapability{
JobType: "vacuum",
CanDetect: true,
CanExecute: false,
},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "vacuum"},
detectEntered: make(chan struct{}, 1),
detectContinue: make(chan struct{}, 1),
}
worker, err := NewWorker(WorkerOptions{
AdminServer: "localhost:23646",
GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
Handler: handler,
MaxDetectionConcurrency: 1,
})
if err != nil {
t.Fatalf("NewWorker error = %v", err)
}
msgCh := make(chan *plugin_pb.WorkerToAdminMessage, 16)
send := func(msg *plugin_pb.WorkerToAdminMessage) bool {
msgCh <- msg
return true
}
requestID := "detect-heartbeat-1"
worker.handleAdminMessage(context.Background(), &plugin_pb.AdminToWorkerMessage{
RequestId: requestID,
Body: &plugin_pb.AdminToWorkerMessage_RunDetectionRequest{
RunDetectionRequest: &plugin_pb.RunDetectionRequest{
JobType: "vacuum",
},
},
}, send)
<-handler.detectEntered
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
heartbeat := message.GetHeartbeat()
return heartbeat != nil &&
heartbeat.DetectionSlotsUsed > 0 &&
heartbeatHasRunningWork(heartbeat, requestID, plugin_pb.WorkKind_WORK_KIND_DETECTION)
}, "active detection heartbeat")
handler.detectContinue <- struct{}{}
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
complete := message.GetDetectionComplete()
return complete != nil && complete.RequestId == requestID && complete.Success
}, "detection complete")
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
heartbeat := message.GetHeartbeat()
return heartbeat != nil && heartbeat.DetectionSlotsUsed == 0 &&
!heartbeatHasRunningWork(heartbeat, requestID, plugin_pb.WorkKind_WORK_KIND_DETECTION)
}, "idle detection heartbeat")
}
func TestWorkerHeartbeatReflectsActiveExecutionLoad(t *testing.T) {
handler := &executionHeartbeatTestHandler{
capability: &plugin_pb.JobTypeCapability{
JobType: "vacuum",
CanDetect: false,
CanExecute: true,
},
descriptor: &plugin_pb.JobTypeDescriptor{JobType: "vacuum"},
executeEntered: make(chan struct{}, 1),
executeDone: make(chan struct{}, 1),
}
worker, err := NewWorker(WorkerOptions{
AdminServer: "localhost:23646",
GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
Handler: handler,
MaxExecutionConcurrency: 1,
})
if err != nil {
t.Fatalf("NewWorker error = %v", err)
}
msgCh := make(chan *plugin_pb.WorkerToAdminMessage, 16)
send := func(msg *plugin_pb.WorkerToAdminMessage) bool {
msgCh <- msg
return true
}
requestID := "exec-heartbeat-1"
jobID := "job-heartbeat-1"
worker.handleAdminMessage(context.Background(), &plugin_pb.AdminToWorkerMessage{
RequestId: requestID,
Body: &plugin_pb.AdminToWorkerMessage_ExecuteJobRequest{
ExecuteJobRequest: &plugin_pb.ExecuteJobRequest{
Job: &plugin_pb.JobSpec{
JobId: jobID,
JobType: "vacuum",
},
},
},
}, send)
<-handler.executeEntered
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
heartbeat := message.GetHeartbeat()
return heartbeat != nil &&
heartbeat.ExecutionSlotsUsed > 0 &&
heartbeatHasRunningWork(heartbeat, jobID, plugin_pb.WorkKind_WORK_KIND_EXECUTION)
}, "active execution heartbeat")
handler.executeDone <- struct{}{}
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
completed := message.GetJobCompleted()
return completed != nil && completed.RequestId == requestID && completed.Success
}, "execution complete")
waitForWorkerMessage(t, msgCh, func(message *plugin_pb.WorkerToAdminMessage) bool {
heartbeat := message.GetHeartbeat()
return heartbeat != nil && heartbeat.ExecutionSlotsUsed == 0 &&
!heartbeatHasRunningWork(heartbeat, jobID, plugin_pb.WorkKind_WORK_KIND_EXECUTION)
}, "idle execution heartbeat")
}
type testJobHandler struct {
capability *plugin_pb.JobTypeCapability
descriptor *plugin_pb.JobTypeDescriptor
}
func (h *testJobHandler) Capability() *plugin_pb.JobTypeCapability {
return h.capability
}
func (h *testJobHandler) Descriptor() *plugin_pb.JobTypeDescriptor {
return h.descriptor
}
func (h *testJobHandler) Detect(context.Context, *plugin_pb.RunDetectionRequest, DetectionSender) error {
return nil
}
func (h *testJobHandler) Execute(context.Context, *plugin_pb.ExecuteJobRequest, ExecutionSender) error {
return nil
}
type detectionQueueTestHandler struct {
capability *plugin_pb.JobTypeCapability
descriptor *plugin_pb.JobTypeDescriptor
detectEntered chan struct{}
detectContinue chan struct{}
}
func (h *detectionQueueTestHandler) Capability() *plugin_pb.JobTypeCapability {
return h.capability
}
func (h *detectionQueueTestHandler) Descriptor() *plugin_pb.JobTypeDescriptor {
return h.descriptor
}
func (h *detectionQueueTestHandler) Detect(ctx context.Context, _ *plugin_pb.RunDetectionRequest, sender DetectionSender) error {
select {
case h.detectEntered <- struct{}{}:
default:
}
select {
case <-ctx.Done():
return ctx.Err()
case <-h.detectContinue:
}
return sender.SendComplete(&plugin_pb.DetectionComplete{
Success: true,
})
}
func (h *detectionQueueTestHandler) Execute(context.Context, *plugin_pb.ExecuteJobRequest, ExecutionSender) error {
return nil
}
type executionHeartbeatTestHandler struct {
capability *plugin_pb.JobTypeCapability
descriptor *plugin_pb.JobTypeDescriptor
executeEntered chan struct{}
executeDone chan struct{}
}
func (h *executionHeartbeatTestHandler) Capability() *plugin_pb.JobTypeCapability {
return h.capability
}
func (h *executionHeartbeatTestHandler) Descriptor() *plugin_pb.JobTypeDescriptor {
return h.descriptor
}
func (h *executionHeartbeatTestHandler) Detect(context.Context, *plugin_pb.RunDetectionRequest, DetectionSender) error {
return nil
}
func (h *executionHeartbeatTestHandler) Execute(ctx context.Context, request *plugin_pb.ExecuteJobRequest, sender ExecutionSender) error {
select {
case h.executeEntered <- struct{}{}:
default:
}
select {
case <-ctx.Done():
return ctx.Err()
case <-h.executeDone:
}
return sender.SendCompleted(&plugin_pb.JobCompleted{
JobId: request.Job.JobId,
JobType: request.Job.JobType,
Success: true,
})
}
func recvWorkerMessage(t *testing.T, msgCh <-chan *plugin_pb.WorkerToAdminMessage) *plugin_pb.WorkerToAdminMessage {
t.Helper()
select {
case msg := <-msgCh:
return msg
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for worker message")
return nil
}
}
func expectDetectionAck(t *testing.T, message *plugin_pb.WorkerToAdminMessage, requestID string) {
t.Helper()
ack := message.GetAcknowledge()
if ack == nil {
t.Fatalf("expected acknowledge for request %q, got=%+v", requestID, message)
}
if ack.RequestId != requestID {
t.Fatalf("expected acknowledge request_id=%q, got=%q", requestID, ack.RequestId)
}
if !ack.Accepted {
t.Fatalf("expected acknowledge accepted for request %q, got=%+v", requestID, ack)
}
}
func expectDetectionCompleteSuccess(t *testing.T, message *plugin_pb.WorkerToAdminMessage, requestID string) {
t.Helper()
complete := message.GetDetectionComplete()
if complete == nil {
t.Fatalf("expected detection complete for request %q, got=%+v", requestID, message)
}
if complete.RequestId != requestID {
t.Fatalf("expected detection complete request_id=%q, got=%q", requestID, complete.RequestId)
}
if !complete.Success {
t.Fatalf("expected successful detection complete for request %q, got=%+v", requestID, complete)
}
}
func waitForWorkerMessage(
t *testing.T,
msgCh <-chan *plugin_pb.WorkerToAdminMessage,
predicate func(*plugin_pb.WorkerToAdminMessage) bool,
description string,
) *plugin_pb.WorkerToAdminMessage {
t.Helper()
timeout := time.NewTimer(3 * time.Second)
defer timeout.Stop()
for {
select {
case message := <-msgCh:
if predicate(message) {
return message
}
case <-timeout.C:
t.Fatalf("timed out waiting for %s", description)
return nil
}
}
}
func heartbeatHasRunningWork(heartbeat *plugin_pb.WorkerHeartbeat, workID string, kind plugin_pb.WorkKind) bool {
if heartbeat == nil || workID == "" {
return false
}
for _, work := range heartbeat.RunningWork {
if work == nil {
continue
}
if work.WorkId == workID && work.Kind == kind {
return true
}
}
return false
}
func (w *Worker) handleAdminMessageForTest(
message *plugin_pb.AdminToWorkerMessage,
send func(*plugin_pb.WorkerToAdminMessage) bool,
) bool {
called := false
w.handleAdminMessage(context.Background(), message, func(msg *plugin_pb.WorkerToAdminMessage) bool {
called = true
return send(msg)
})
return called
}