orca/internal/engine/dispatcher.go

// Package engine — dispatcher.go implements the cross-node job
// dispatch logic (v0.2 P02). The dispatcher is the bridge between
// the local "should I run this?" decision (scheduler.PickNode) and
// the remote "please run this" call (transport.DispatchClient).
//
// Flow:
//
//  1. Receive a job spec (HCL bytes from the CLI).
//  2. Parse the spec into a JobSpec (cpu/mem/disk).
//  3. Check local capacity. If it fits, run locally via the local
//     executor. If not, pick a peer and dispatch.
//  4. Return the job ID and the node that actually accepted it.
package engine

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sync"

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

// Dispatcher is the public surface; constructed via NewDispatcher.
type Dispatcher struct {
	log      *slog.Logger
	capacity *store.CapacityRepo
	peers    *PeerRegistry
	executor LocalExecutor
	dedupe   *transport.IdempotencyStore
	mu       sync.Mutex
}

// LocalExecutor is the contract the dispatcher uses to run jobs on
// the local node. The engine.Executor satisfies this.
type LocalExecutor interface {
	Submit(ctx context.Context, specBytes []byte) (jobID string, err error)
	Status(ctx context.Context, jobID string) (state string, err error)
}

// NewDispatcher builds a Dispatcher.
func NewDispatcher(log *slog.Logger, capacity *store.CapacityRepo, peers *PeerRegistry, exec LocalExecutor) *Dispatcher {
	if log == nil {
		log = slog.Default()
	}
	return &Dispatcher{
		log:      log,
		capacity: capacity,
		peers:    peers,
		executor: exec,
		dedupe:   transport.NewIdempotencyStore(),
	}
}

// Dedupe exposes the in-memory dedupe store for testing.
func (d *Dispatcher) Dedupe() *transport.IdempotencyStore { return d.dedupe }

// Submit runs the spec locally if it fits, otherwise dispatches to a
// peer. Returns the (jobID, chosenNodeID) pair. If `target` is
// non-empty, it overrides bin-packing.
func (d *Dispatcher) Submit(ctx context.Context, target string, specBytes []byte, idempotencyKey string) (jobID, nodeID string, err error) {
	if len(specBytes) == 0 {
		return "", "", errors.New("Dispatcher.Submit: empty spec")
	}
	if idempotencyKey != "" {
		if jid, ok := d.dedupe.Get(idempotencyKey); ok {
			return jid, "self", nil
		}
	}

	parsed, err := parseInlineSpec(specBytes)
	if err != nil {
		return "", "", fmt.Errorf("Dispatcher.Submit: parse spec: %w", err)
	}

	// 1. Explicit target: dispatch there.
	if target != "" {
		return d.dispatchTo(ctx, target, specBytes, idempotencyKey)
	}

	// 2. Check local capacity.
	if d.capacity != nil {
		local, err := d.capacity.Get(ctx, "self")
		if err == nil && parsed.Fits(local) {
			jid, lerr := d.executor.Submit(ctx, specBytes)
			if lerr != nil {
				return "", "", fmt.Errorf("Dispatcher.Submit: local: %w", lerr)
			}
			if idempotencyKey != "" {
				d.dedupe.Put(idempotencyKey, jid)
			}
			d.log.Info("dispatch.local",
				slog.String("event", "dispatch.local"),
				slog.String("job_id", jid),
				slog.String("node_id", "self"),
			)
			return jid, "self", nil
		}
	}

	// 3. Pick a peer.
	if d.peers == nil {
		return "", "", errors.New("Dispatcher.Submit: no local capacity and no peer registry")
	}
	peers, err := d.peers.All(ctx)
	if err != nil {
		return "", "", fmt.Errorf("Dispatcher.Submit: list peers: %w", err)
	}
	if len(peers) == 0 {
		return "", "", errors.New("Dispatcher.Submit: no peers registered")
	}
	var caps []*store.NodeCapacity
	for _, p := range peers {
		caps = append(caps, p.Capacity)
	}
	best, _, err := PickNode(parsed, caps)
	if err != nil {
		return "", "", fmt.Errorf("Dispatcher.Submit: %w", err)
	}
	var chosen *Peer
	for _, p := range peers {
		if p.NodeID == best.NodeID {
			chosen = p
			break
		}
	}
	if chosen == nil {
		return "", "", fmt.Errorf("Dispatcher.Submit: chosen node %s has no peer record", best.NodeID)
	}
	return d.dispatchToPeer(ctx, chosen, specBytes, idempotencyKey)
}

// dispatchTo sends a Submit to a specific node id (looked up in the peer registry).
func (d *Dispatcher) dispatchTo(ctx context.Context, targetNode string, specBytes []byte, idempotencyKey string) (string, string, error) {
	if d.peers == nil {
		return "", "", errors.New("dispatchTo: no peer registry")
	}
	peers, err := d.peers.All(ctx)
	if err != nil {
		return "", "", fmt.Errorf("dispatchTo: list peers: %w", err)
	}
	for _, p := range peers {
		if p.NodeID == targetNode {
			return d.dispatchToPeer(ctx, p, specBytes, idempotencyKey)
		}
	}
	return "", "", fmt.Errorf("dispatchTo: target node %q not found in peer registry", targetNode)
}

// dispatchToPeer opens an mTLS client and calls Submit on the peer.
func (d *Dispatcher) dispatchToPeer(ctx context.Context, p *Peer, specBytes []byte, idempotencyKey string) (string, string, error) {
	if p.CAPath == "" || p.ServerName == "" {
		return "", "", fmt.Errorf("dispatchToPeer: peer %s missing CA or server name", p.NodeID)
	}
	client, err := transport.NewDispatchClient(p.CAPath, p.ServerName, "https://"+p.Address)
	if err != nil {
		return "", "", fmt.Errorf("dispatchToPeer: %w", err)
	}
	resp, err := client.Submit(ctx, specBytes, idempotencyKey)
	if err != nil {
		return "", "", fmt.Errorf("dispatchToPeer: %w", err)
	}
	if idempotencyKey != "" {
		d.dedupe.Put(idempotencyKey, resp.JobID)
	}
	d.log.Info("dispatch.peer",
		slog.String("event", "dispatch.peer"),
		slog.String("job_id", resp.JobID),
		slog.String("node_id", p.NodeID),
	)
	return resp.JobID, p.NodeID, nil
}

// LocalSubmit / LocalStatus satisfy the transport.Dispatcher
// interface (the server-side counterpart of DispatchClient).
func (d *Dispatcher) LocalSubmit(ctx context.Context, specBytes []byte) (string, error) {
	if d.executor == nil {
		return "", errors.New("Dispatcher.LocalSubmit: no local executor")
	}
	return d.executor.Submit(ctx, specBytes)
}

func (d *Dispatcher) LocalStatus(ctx context.Context, jobID string) (string, error) {
	if d.executor == nil {
		return "", errors.New("Dispatcher.LocalStatus: no local executor")
	}
	return d.executor.Status(ctx, jobID)
}

// parseInlineSpec parses a minimal JSON spec with cpu_millicores,
// memory_mib, disk_mib fields. The CLI uses this as the wire format
// for cross-node dispatch; full HCL parsing is in internal/jobspec.
func parseInlineSpec(b []byte) (JobSpec, error) {
	type wire struct {
		CPUMillicores int64 `json:"cpu_millicores"`
		MemoryMiB     int64 `json:"memory_mib"`
		DiskMiB       int64 `json:"disk_mib"`
	}
	var w wire
	if err := json.Unmarshal(b, &w); err != nil {
		return JobSpec{}, fmt.Errorf("parseInlineSpec: %w", err)
	}
	return JobSpec{
		CPUMillicores: w.CPUMillicores,
		MemoryMiB:     w.MemoryMiB,
		DiskMiB:       w.DiskMiB,
	}, nil
}