diff --git a/cmd/spindle/main.go b/cmd/spindle/main.go index 8ebbad79..bed952c5 100644 --- a/cmd/spindle/main.go +++ b/cmd/spindle/main.go @@ -6,6 +6,7 @@ import ( "encoding/hex" "fmt" "log/slog" + "math" "os" "os/signal" "strings" @@ -13,12 +14,15 @@ import ( "text/tabwriter" "time" + "github.com/bluesky-social/indigo/atproto/syntax" "github.com/urfave/cli/v3" tlog "tangled.org/core/log" "tangled.org/core/spindle" + "tangled.org/core/spindle/config" "tangled.org/core/spindle/db" "tangled.org/core/spindle/mill" + "tangled.org/core/spindle/quota" ) func main() { @@ -29,6 +33,7 @@ func main() { Command(), millCommand(), executorCommand(), + quotaCommand(), }, DefaultCommand: "run", } @@ -331,3 +336,253 @@ func millCommand() *cli.Command { }, } } + +func openQuotaStore(ctx context.Context, cmd *cli.Command) (*db.QuotaStore, error) { + d, err := openDB(ctx, cmd) + if err != nil { + return nil, err + } + return db.NewQuotaStore(d, quota.Defaults{}), nil +} + +type quotaKey struct { + did string + resource string +} + +func quotaCommand() *cli.Command { + return &cli.Command{ + Name: "quota", + Usage: "manage resource quotas for repos and users", + Commands: []*cli.Command{ + { + Name: "usage", + Usage: "show usage against limits for every repo and user", + Flags: []cli.Flag{dbFlag}, + Action: func(ctx context.Context, cmd *cli.Command) error { + if cmd.Args().Len() != 0 { + return fmt.Errorf("unexpected arguments") + } + qs, err := openQuotaStore(ctx, cmd) + if err != nil { + return err + } + defer qs.Close() + + defaults, err := config.LoadQuotaDefaults(ctx) + if err != nil { + return fmt.Errorf("invalid quota configuration: %w", err) + } + + limits, err := qs.ListLimits(ctx) + if err != nil { + return fmt.Errorf("listing limits: %w", err) + } + overrides := make(map[quotaKey]int64, len(limits)) + for _, l := range limits { + overrides[quotaKey{did: l.DID, resource: l.Resource}] = l.Limit + } + + usages, err := qs.ListUsage(ctx) + if err != nil { + return fmt.Errorf("listing usage: %w", err) + } + + w := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', 0) + fmt.Fprintln(w, "SCOPE\tDID\tRESOURCE\tUSED\tLIMIT") + + for _, u := range usages { + var limitVal int64 = -1 + if ovLimit, ok := overrides[quotaKey{did: u.DID, resource: u.Resource}]; ok { + limitVal = ovLimit + } else if scopeDefs, ok := defaults[u.Scope]; ok { + if defLimit, ok := scopeDefs[u.Resource]; ok { + limitVal = normalizeDefaultLimit(defLimit) + } + } + fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", string(u.Scope), u.DID, string(u.Resource), formatLimit(u.Resource, u.Used), formatLimit(u.Resource, limitVal)) + } + return w.Flush() + }, + }, + { + Name: "list", + Usage: "list all custom limits", + Flags: []cli.Flag{dbFlag}, + Action: func(ctx context.Context, cmd *cli.Command) error { + if cmd.Args().Len() != 0 { + return fmt.Errorf("unexpected arguments") + } + qs, err := openQuotaStore(ctx, cmd) + if err != nil { + return err + } + defer qs.Close() + limits, err := qs.ListLimits(ctx) + if err != nil { + return fmt.Errorf("listing limits: %w", err) + } + w := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', 0) + fmt.Fprintln(w, "DID\tRESOURCE\tLIMIT") + for _, l := range limits { + fmt.Fprintf(w, "%s\t%s\t%s\n", l.DID, string(l.Resource), formatLimit(l.Resource, l.Limit)) + } + return w.Flush() + }, + }, + { + Name: "set", + Usage: "give a repo or user a custom limit", + Flags: []cli.Flag{ + dbFlag, + &cli.StringFlag{ + Name: "did", + Usage: "did of the user or repo", + Required: true, + }, + &cli.StringFlag{ + Name: "resource", + Usage: "resource to limit (e.g. workflows, cache_storage_bytes)", + Required: true, + }, + &cli.IntFlag{ + Name: "limit", + Usage: "new limit (MiB for cache_storage_bytes, memory_mib, and disk_mib)", + }, + &cli.BoolFlag{ + Name: "unlimited", + Usage: "remove the limit entirely", + }, + }, + Action: func(ctx context.Context, cmd *cli.Command) error { + if cmd.Args().Len() != 0 { + return fmt.Errorf("unexpected arguments") + } + + did, err := syntax.ParseDID(cmd.String("did")) + if err != nil { + return fmt.Errorf("invalid did: %w", err) + } + + resource := cmd.String("resource") + + if err := quota.ValidateOverride(did.String(), resource); err != nil { + return err + } + + hasLimit := cmd.IsSet("limit") + hasUnlimited := cmd.Bool("unlimited") + + if !hasLimit && !hasUnlimited { + return fmt.Errorf("must specify either --limit or --unlimited") + } + if hasLimit && hasUnlimited { + return fmt.Errorf("cannot specify both --limit and --unlimited") + } + + var limitVal int64 = -1 + if hasLimit { + limitInt := cmd.Int("limit") + if limitInt < 0 { + return fmt.Errorf("limit cannot be negative: %d", limitInt) + } + + if resource == quota.ResourceCacheStorageBytes { + if limitInt > math.MaxInt64/1048576 { + return fmt.Errorf("limit %d MiB would overflow bytes", limitInt) + } + limitVal = int64(limitInt) * 1048576 + } else { + limitVal = int64(limitInt) + } + } + + qs, err := openQuotaStore(ctx, cmd) + if err != nil { + return err + } + defer qs.Close() + if err := qs.SetLimit(ctx, did.String(), resource, limitVal); err != nil { + return fmt.Errorf("setting limit: %w", err) + } + fmt.Printf("quota override set for %s %s\n", did.String(), resource) + return nil + }, + }, + { + Name: "unset", + Usage: "remove a custom limit and go back to the defaults", + Flags: []cli.Flag{ + dbFlag, + &cli.StringFlag{ + Name: "did", + Usage: "did of the user or repo", + Required: true, + }, + &cli.StringFlag{ + Name: "resource", + Usage: "resource to limit (e.g. workflows, cache_storage_bytes)", + Required: true, + }, + }, + Action: func(ctx context.Context, cmd *cli.Command) error { + if cmd.Args().Len() != 0 { + return fmt.Errorf("unexpected arguments") + } + + did, err := syntax.ParseDID(cmd.String("did")) + if err != nil { + return fmt.Errorf("invalid did: %w", err) + } + + resource := cmd.String("resource") + + if err := quota.ValidateOverride(did.String(), resource); err != nil { + return err + } + + qs, err := openQuotaStore(ctx, cmd) + if err != nil { + return err + } + defer qs.Close() + if err := qs.UnsetLimit(ctx, did.String(), resource); err != nil { + return fmt.Errorf("unsetting limit: %w", err) + } + fmt.Printf("quota override removed for %s %s\n", did.String(), resource) + return nil + }, + }, + }, + } +} + +func normalizeDefaultLimit(limit int64) int64 { + if limit == 0 { + return -1 + } + return limit +} + +func formatLimit(resource string, limit int64) string { + if limit < 0 { + return "unlimited" + } + switch resource { + case quota.ResourceCacheStorageBytes: + if limit%(1024*1024) == 0 { + return fmt.Sprintf("%d MiB", limit/(1024*1024)) + } + return fmt.Sprintf("%d B", limit) + case quota.ResourceMemoryMiB: + return fmt.Sprintf("%d MiB", limit) + case quota.ResourceDiskMiB: + return fmt.Sprintf("%d MiB", limit) + case quota.ResourceWorkflows: + return fmt.Sprintf("%d workflows", limit) + case quota.ResourceVCPUs: + return fmt.Sprintf("%d vCPUs", limit) + default: + return fmt.Sprintf("%d", limit) + } +} diff --git a/cmd/spindle/main_test.go b/cmd/spindle/main_test.go index 6e4c4fe0..739b4e16 100644 --- a/cmd/spindle/main_test.go +++ b/cmd/spindle/main_test.go @@ -4,8 +4,11 @@ import ( "bytes" "os" "os/exec" + "path/filepath" "strings" "testing" + + "tangled.org/core/spindle/quota" ) func TestCLIInvalidLogFormat(t *testing.T) { @@ -46,3 +49,165 @@ func TestHelperProcess(t *testing.T) { } main() } + +func TestCLIQuotaCommands(t *testing.T) { + tmpDb := filepath.Join(t.TempDir(), "spindle.db") + + tests := []struct { + name string + args []string + wantExitCode int + wantError string + wantOutput string + }{ + { + name: "invalid DID", + args: []string{"quota", "set", "--did", "alice", "--resource", "cache_storage_bytes", "--limit", "10"}, + wantExitCode: 255, + wantError: "invalid did", + }, + { + name: "invalid resource", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", strings.Repeat("a", 65), "--limit", "10"}, + wantExitCode: 255, + wantError: "invalid resource", + }, + { + name: "conflicting flags limit and unlimited", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "cache_storage_bytes", "--limit", "10", "--unlimited"}, + wantExitCode: 255, + wantError: "cannot specify both --limit and --unlimited", + }, + { + name: "missing flags limit and unlimited", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "cache_storage_bytes"}, + wantExitCode: 255, + wantError: "must specify either --limit or --unlimited", + }, + { + name: "cache storage limit overflow", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "cache_storage_bytes", "--limit", "9223372036854775807"}, + wantExitCode: 255, + wantError: "would overflow bytes", + }, + { + name: "set user cache_storage_bytes limit", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "cache_storage_bytes", "--limit", "10"}, + wantExitCode: 0, + wantOutput: "quota override set", + }, + { + name: "set user workflows limit", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "workflows", "--limit", "5"}, + wantExitCode: 0, + wantOutput: "quota override set", + }, + { + name: "set user vcpus limit", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "vcpus", "--limit", "4"}, + wantExitCode: 0, + wantOutput: "quota override set", + }, + { + name: "set user memory_mib limit", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "memory_mib", "--limit", "2048"}, + wantExitCode: 0, + wantOutput: "quota override set", + }, + { + name: "set user disk_mib limit", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "disk_mib", "--limit", "4096"}, + wantExitCode: 0, + wantOutput: "quota override set", + }, + { + name: "set user cache_storage_bytes unlimited", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "cache_storage_bytes", "--unlimited"}, + wantExitCode: 0, + wantOutput: "quota override set", + }, + { + name: "list quota limits", + args: []string{"quota", "list"}, + wantExitCode: 0, + wantOutput: "LIMIT", + }, + { + name: "usage list", + args: []string{"quota", "usage"}, + wantExitCode: 0, + wantOutput: "LIMIT", + }, + { + name: "unset user cache_storage_bytes limit", + args: []string{"quota", "unset", "--did", "did:web:alice", "--resource", "cache_storage_bytes"}, + wantExitCode: 0, + wantOutput: "quota override removed", + }, + { + name: "unset user workflows limit", + args: []string{"quota", "unset", "--did", "did:web:alice", "--resource", "workflows"}, + wantExitCode: 0, + wantOutput: "quota override removed", + }, + { + name: "set user generic resource limit", + args: []string{"quota", "set", "--did", "did:web:alice", "--resource", "custom_gpu", "--limit", "8"}, + wantExitCode: 0, + wantOutput: "quota override set", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + fullArgs := append([]string{"-test.run=TestHelperProcess", "--"}, tt.args...) + fullArgs = append(fullArgs, "--db", tmpDb) + cmd := exec.Command(os.Args[0], fullArgs...) + cmd.Env = append(os.Environ(), + "WANT_HELPER_PROCESS=1", + "SPINDLE_SERVER_HOSTNAME=spindle.example.com", + "SPINDLE_SERVER_OWNER=did:web:spindle.example.com", + ) + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + exitCode := 0 + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + exitCode = exitErr.ExitCode() + } else { + t.Fatalf("unexpected run error: %v", err) + } + } + + expectedCode := tt.wantExitCode + if expectedCode == -1 { + expectedCode = 255 + } + if exitCode != expectedCode { + t.Errorf("exitCode = %d, want %d. Stderr: %q, Stdout: %q", exitCode, expectedCode, stderr.String(), stdout.String()) + } + + combined := stdout.String() + "\n" + stderr.String() + if tt.wantError != "" && !strings.Contains(combined, tt.wantError) { + t.Errorf("expected error/output to contain %q, got combined: %q", tt.wantError, combined) + } + if tt.wantOutput != "" && !strings.Contains(combined, tt.wantOutput) { + t.Errorf("expected output to contain %q, got combined: %q", tt.wantOutput, combined) + } + }) + } +} + +func TestNormalizeDefaultLimit(t *testing.T) { + if got := normalizeDefaultLimit(0); got != -1 { + t.Fatalf("zero configured limit = %d, want -1", got) + } + if got := normalizeDefaultLimit(2); got != 2 { + t.Fatalf("finite configured limit = %d, want 2", got) + } + if got := formatLimit(quota.ResourceWorkflows, 0); got != "0 workflows" { + t.Fatalf("zero override = %q, want 0 workflows", got) + } +}