plugin worker: support job type categories (all, default, heavy) (#8547)

* plugin worker: add handler registry with job categories

Introduce a self-registration pattern for plugin worker job handlers.
Each handler can register itself via init() with a HandlerFactory that
declares its job type, category (default/heavy), CLI aliases, and a
builder function.

ResolveHandlerFactories accepts a mix of category names ("all",
"default", "heavy") and explicit job type names/aliases, returning the
matching factories. This enables workers to be configured by resource
profile rather than requiring explicit job type enumeration.

* plugin worker: register all handlers via init()

Each job handler now self-registers into the global handler registry
with its canonical job type, category, CLI aliases, and build function:

  - vacuum:              category=default
  - volume_balance:      category=default
  - admin_script:        category=default
  - erasure_coding:      category=heavy
  - iceberg_maintenance: category=heavy

Adding a new job type now only requires adding the init() call in the
handler file itself — no other files need to be touched.

* plugin worker: replace hardcoded job type switch with registry

Remove buildPluginWorkerHandler, parsePluginWorkerJobTypes, and
canonicalPluginWorkerJobType from worker_runtime.go. The simplified
buildPluginWorkerHandlers now delegates to
pluginworker.ResolveHandlerFactories, which resolves category names
("all", "default", "heavy") and explicit job type names/aliases.

The default job type is changed from an explicit list to "all", so new
handlers registered via init() are automatically picked up.

Update all tests to use the new API.

* plugin worker: update CLI help text for job categories

Update the -jobType flag description and command examples to document
category support (all, default, heavy) alongside explicit job type names.

* plugin worker: address review feedback

- Add CategoryAll constant; use typed constants in tokenAsCategory
- Pre-allocate result slice in ResolveHandlerFactories
- Add vacuum aliases (vol.vacuum, volume.vacuum)
- List alias examples (ec, balance, iceberg) in -jobType flag help
- Create handlers aggregator package for subpackage blank imports so
  new handler subpackages only need to be added in one place
- Make category tests relationship-based (subset/union checks) instead
  of asserting exact handler counts
- Add clarifying comments to worker_test.go and mini_plugin_test.go
  listing expected handler names next to count assertions

---------

Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
Chris Lu
2026-03-07 18:30:58 -08:00
committed by GitHub
parent f249fb7e63
commit 587c24ec89
12 changed files with 350 additions and 176 deletions

View File

@@ -6,7 +6,6 @@ import (
"net/http/httptest"
"os"
"path/filepath"
"slices"
"strings"
"testing"
@@ -15,68 +14,46 @@ import (
"google.golang.org/grpc/credentials/insecure"
)
func TestBuildPluginWorkerHandler(t *testing.T) {
func TestBuildPluginWorkerHandlerExplicitTypes(t *testing.T) {
dialOption := grpc.WithTransportCredentials(insecure.NewCredentials())
testMaxConcurrency := int(pluginworker.DefaultMaxExecutionConcurrency)
handler, err := buildPluginWorkerHandler("vacuum", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandler(vacuum) err = %v", err)
}
if handler == nil {
t.Fatalf("expected non-nil handler")
for _, jobType := range []string{"vacuum", "volume_balance", "erasure_coding", "admin_script", "iceberg_maintenance"} {
handlers, err := buildPluginWorkerHandlers(jobType, dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandlers(%s) err = %v", jobType, err)
}
if len(handlers) != 1 {
t.Fatalf("expected 1 handler for %s, got %d", jobType, len(handlers))
}
}
}
handler, err = buildPluginWorkerHandler("", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandler(default) err = %v", err)
}
if handler == nil {
t.Fatalf("expected non-nil default handler")
}
func TestBuildPluginWorkerHandlerAliases(t *testing.T) {
dialOption := grpc.WithTransportCredentials(insecure.NewCredentials())
testMaxConcurrency := int(pluginworker.DefaultMaxExecutionConcurrency)
handler, err = buildPluginWorkerHandler("volume_balance", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandler(volume_balance) err = %v", err)
}
if handler == nil {
t.Fatalf("expected non-nil volume_balance handler")
for _, alias := range []string{"balance", "ec", "iceberg", "admin", "script"} {
handlers, err := buildPluginWorkerHandlers(alias, dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandlers(%s) err = %v", alias, err)
}
if len(handlers) != 1 {
t.Fatalf("expected 1 handler for alias %s, got %d", alias, len(handlers))
}
}
}
handler, err = buildPluginWorkerHandler("balance", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandler(balance alias) err = %v", err)
}
if handler == nil {
t.Fatalf("expected non-nil balance alias handler")
}
handler, err = buildPluginWorkerHandler("erasure_coding", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandler(erasure_coding) err = %v", err)
}
if handler == nil {
t.Fatalf("expected non-nil erasure_coding handler")
}
handler, err = buildPluginWorkerHandler("ec", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandler(ec alias) err = %v", err)
}
if handler == nil {
t.Fatalf("expected non-nil ec alias handler")
}
_, err = buildPluginWorkerHandler("unknown", dialOption, testMaxConcurrency, "")
func TestBuildPluginWorkerHandlerUnknown(t *testing.T) {
dialOption := grpc.WithTransportCredentials(insecure.NewCredentials())
_, err := buildPluginWorkerHandlers("unknown", dialOption, 1, "")
if err == nil {
t.Fatalf("expected unsupported job type error")
t.Fatalf("expected error for unknown job type")
}
}
func TestBuildPluginWorkerHandlers(t *testing.T) {
dialOption := grpc.WithTransportCredentials(insecure.NewCredentials())
testMaxConcurrency := int(pluginworker.DefaultMaxExecutionConcurrency)
handlers, err := buildPluginWorkerHandlers("vacuum,volume_balance,erasure_coding", dialOption, testMaxConcurrency, "")
@@ -101,52 +78,99 @@ func TestBuildPluginWorkerHandlers(t *testing.T) {
}
}
func TestParsePluginWorkerJobTypes(t *testing.T) {
jobTypes, err := parsePluginWorkerJobTypes("")
func TestBuildPluginWorkerHandlersCategories(t *testing.T) {
dialOption := grpc.WithTransportCredentials(insecure.NewCredentials())
testMaxConcurrency := int(pluginworker.DefaultMaxExecutionConcurrency)
allHandlers, err := buildPluginWorkerHandlers("all", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("parsePluginWorkerJobTypes(default) err = %v", err)
t.Fatalf("buildPluginWorkerHandlers(all) err = %v", err)
}
if len(jobTypes) != 1 || jobTypes[0] != "vacuum" {
t.Fatalf("expected default [vacuum], got %v", jobTypes)
// "all" must include at least vacuum and erasure_coding (one default, one heavy)
allNames := handlerJobTypes(allHandlers)
for _, required := range []string{"vacuum", "erasure_coding", "iceberg_maintenance"} {
if !allNames[required] {
t.Fatalf("'all' missing expected job type %q, got %v", required, allNames)
}
}
jobTypes, err = parsePluginWorkerJobTypes(" volume_balance , ec , vacuum , volume_balance ")
defaultHandlers, err := buildPluginWorkerHandlers("default", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("parsePluginWorkerJobTypes(list) err = %v", err)
t.Fatalf("buildPluginWorkerHandlers(default) err = %v", err)
}
if len(jobTypes) != 3 {
t.Fatalf("expected 3 deduped job types, got %d (%v)", len(jobTypes), jobTypes)
defaultNames := handlerJobTypes(defaultHandlers)
heavyHandlers, err := buildPluginWorkerHandlers("heavy", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandlers(heavy) err = %v", err)
}
if jobTypes[0] != "volume_balance" || jobTypes[1] != "erasure_coding" || jobTypes[2] != "vacuum" {
t.Fatalf("unexpected parsed order %v", jobTypes)
heavyNames := handlerJobTypes(heavyHandlers)
// default and heavy must both be non-empty subsets of all
if len(defaultNames) == 0 {
t.Fatalf("'default' resolved no handlers")
}
if len(heavyNames) == 0 {
t.Fatalf("'heavy' resolved no handlers")
}
for name := range defaultNames {
if !allNames[name] {
t.Fatalf("default handler %q not in 'all'", name)
}
}
for name := range heavyNames {
if !allNames[name] {
t.Fatalf("heavy handler %q not in 'all'", name)
}
}
if _, err = parsePluginWorkerJobTypes(" , "); err != nil {
t.Fatalf("expected empty list to resolve to default vacuum: %v", err)
// default and heavy must be disjoint and their union must equal all
for name := range defaultNames {
if heavyNames[name] {
t.Fatalf("handler %q appears in both default and heavy", name)
}
}
if len(defaultNames)+len(heavyNames) != len(allNames) {
t.Fatalf("union(default=%d, heavy=%d) != all(%d)", len(defaultNames), len(heavyNames), len(allNames))
}
jobTypes, err = parsePluginWorkerJobTypes("admin-script,script,admin_script")
// mix category + explicit: "default,iceberg" adds one heavy to default set
mixedHandlers, err := buildPluginWorkerHandlers("default,iceberg", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("parsePluginWorkerJobTypes(admin script aliases) err = %v", err)
t.Fatalf("buildPluginWorkerHandlers(default,iceberg) err = %v", err)
}
if len(jobTypes) != 1 || jobTypes[0] != "admin_script" {
t.Fatalf("expected admin_script alias to resolve, got %v", jobTypes)
if len(mixedHandlers) != len(defaultHandlers)+1 {
t.Fatalf("expected default+1 handlers for 'default,iceberg', got %d (default=%d)", len(mixedHandlers), len(defaultHandlers))
}
}
func TestPluginWorkerDefaultJobTypes(t *testing.T) {
jobTypes, err := parsePluginWorkerJobTypes(defaultPluginWorkerJobTypes)
dialOption := grpc.WithTransportCredentials(insecure.NewCredentials())
testMaxConcurrency := int(pluginworker.DefaultMaxExecutionConcurrency)
// defaultPluginWorkerJobTypes is "all", so it should match the "all" category exactly
defaultHandlers, err := buildPluginWorkerHandlers(defaultPluginWorkerJobTypes, dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("parsePluginWorkerJobTypes(default setting) err = %v", err)
t.Fatalf("buildPluginWorkerHandlers(default setting) err = %v", err)
}
if len(jobTypes) != 5 {
t.Fatalf("expected default job types to include 5 handlers, got %v", jobTypes)
allHandlers, err := buildPluginWorkerHandlers("all", dialOption, testMaxConcurrency, "")
if err != nil {
t.Fatalf("buildPluginWorkerHandlers(all) err = %v", err)
}
if !slices.Contains(jobTypes, "iceberg_maintenance") {
t.Fatalf("expected iceberg_maintenance in default job types, got %v", jobTypes)
if len(defaultHandlers) != len(allHandlers) {
t.Fatalf("default setting resolved %d handlers, 'all' resolved %d", len(defaultHandlers), len(allHandlers))
}
}
// handlerJobTypes returns the set of job type names from a slice of handlers.
func handlerJobTypes(handlers []pluginworker.JobHandler) map[string]bool {
m := make(map[string]bool, len(handlers))
for _, h := range handlers {
m[h.Capability().JobType] = true
}
return m
}
func TestResolvePluginWorkerID(t *testing.T) {
dir := t.TempDir()