orca/internal/cli/node_capacity.go

// node_capacity.go implements `orca node capacity` for v0.2 P02.
// The capacity declaration is per-node (cpu_millicores, memory_mib,
// disk_mib) and feeds the bin-packing scheduler.
//
// REQ-028: HCL/YAML schema for NodeCapacity — the CLI accepts the
// three numeric flags and writes a row to the `node_capacity` table.
// A future enhancement can read `~/.orca/node.hcl` at join time
// (out of scope for P02).
package cli

import (
	"context"
	"fmt"
	"time"

	"github.com/spf13/cobra"

	"git.cloudinit.dev/coreci/orca/internal/store"
)

var (
	capSetCPU  int64
	capSetMem  int64
	capSetDisk int64
	capNodeID  string
)

var nodeCapacityCmd = &cobra.Command{
	Use:   "capacity",
	Short: "Manage node capacity declarations (P02 bin-packing input)",
	Long:  "Read or write the per-node capacity used by the multi-node scheduler.",
}

var nodeCapacityShowCmd = &cobra.Command{
	Use:   "show [node-id]",
	Short: "Show capacity for a node (defaults to 'self')",
	Args:  cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		id := capNodeID
		if id == "" && len(args) > 0 {
			id = args[0]
		}
		if id == "" {
			id = "self"
		}
		ctx, cancel := context.WithTimeout(cmd.Context(), 5*time.Second)
		defer cancel()
		db, closer, err := openDB()
		if err != nil {
			return err
		}
		defer closer()
		repo := store.NewCapacityRepo(db)
		c, err := repo.Get(ctx, id)
		if err != nil {
			return fmt.Errorf("node %s: %w (use `orca node capacity --set` to declare)", id, err)
		}
		if jsonOutput {
			return printJSON(c)
		}
		fmt.Fprintf(cmd.OutOrStdout(), "Node:        %s\n", c.NodeID)
		fmt.Fprintf(cmd.OutOrStdout(), "CPU:         %d millicores\n", c.CPUMillicores)
		fmt.Fprintf(cmd.OutOrStdout(), "Memory:      %d MiB\n", c.MemoryMiB)
		fmt.Fprintf(cmd.OutOrStdout(), "Disk:        %d MiB\n", c.DiskMiB)
		fmt.Fprintf(cmd.OutOrStdout(), "Updated:     %s\n", c.UpdatedAt.UTC().Format(time.RFC3339))
		return nil
	},
}

var nodeCapacitySetCmd = &cobra.Command{
	Use:   "set",
	Short: "Declare capacity for a node (used by bin-packing)",
	Long:  "Write cpu_millicores, memory_mib, and disk_mib for the named node. Idempotent: subsequent calls overwrite.",
	RunE: func(cmd *cobra.Command, args []string) error {
		if capSetCPU <= 0 || capSetMem <= 0 || capSetDisk <= 0 {
			return fmt.Errorf("--cpu, --memory, and --disk must all be positive")
		}
		id := capNodeID
		if id == "" {
			id = "self"
		}
		ctx, cancel := context.WithTimeout(cmd.Context(), 5*time.Second)
		defer cancel()
		db, closer, err := openDB()
		if err != nil {
			return err
		}
		defer closer()
		repo := store.NewCapacityRepo(db)
		c := &store.NodeCapacity{
			NodeID:        id,
			CPUMillicores: capSetCPU,
			MemoryMiB:     capSetMem,
			DiskMiB:       capSetDisk,
		}
		if err := repo.Upsert(ctx, c); err != nil {
			return err
		}
		if jsonOutput {
			return printJSON(c)
		}
		fmt.Fprintf(cmd.OutOrStdout(), "✓ Capacity set for %s: cpu=%d mem=%d disk=%d\n",
			c.NodeID, c.CPUMillicores, c.MemoryMiB, c.DiskMiB)
		return nil
	},
}

var nodeCapacityListCmd = &cobra.Command{
	Use:   "list",
	Short: "List all node capacity declarations",
	RunE: func(cmd *cobra.Command, args []string) error {
		ctx, cancel := context.WithTimeout(cmd.Context(), 5*time.Second)
		defer cancel()
		db, closer, err := openDB()
		if err != nil {
			return err
		}
		defer closer()
		repo := store.NewCapacityRepo(db)
		rows, err := repo.List(ctx)
		if err != nil {
			return err
		}
		if jsonOutput {
			return printJSON(rows)
		}
		if len(rows) == 0 {
			fmt.Fprintln(cmd.OutOrStdout(), "No capacity declarations. Use `orca node capacity --set` to add one.")
			return nil
		}
		fmt.Fprintf(cmd.OutOrStdout(), "%-20s  %12s  %12s  %12s  %s\n", "NODE", "CPU(mc)", "MEM(MiB)", "DISK(MiB)", "UPDATED")
		for _, c := range rows {
			fmt.Fprintf(cmd.OutOrStdout(), "%-20s  %12d  %12d  %12d  %s\n",
				c.NodeID, c.CPUMillicores, c.MemoryMiB, c.DiskMiB, c.UpdatedAt.UTC().Format(time.RFC3339))
		}
		return nil
	},
}

func init() {
	nodeCapacitySetCmd.Flags().Int64Var(&capSetCPU, "cpu", 0, "CPU capacity in millicores (1000 = 1 vCPU)")
	nodeCapacitySetCmd.Flags().Int64Var(&capSetMem, "memory", 0, "Memory capacity in MiB")
	nodeCapacitySetCmd.Flags().Int64Var(&capSetDisk, "disk", 0, "Disk capacity in MiB")
	nodeCapacitySetCmd.Flags().StringVar(&capNodeID, "node", "", "node id (defaults to 'self')")
	nodeCapacityShowCmd.Flags().StringVar(&capNodeID, "node", "", "node id (defaults to 'self')")

	nodeCapacityCmd.AddCommand(nodeCapacityShowCmd, nodeCapacitySetCmd, nodeCapacityListCmd)
	nodeCmd.AddCommand(nodeCapacityCmd)
}