oikos/internal/execworker/worker.go

// Package execworker processes pending executions as a background daemon.
// This provides a Postgres-backed queue: executions survive restarts, and
// per-execution advisory locks prevent duplicate processing across instances.
package execworker

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"strings"
	"time"

	"github.com/dtoro/oikos/internal/actuator"
	"github.com/dtoro/oikos/internal/config"
	"github.com/dtoro/oikos/internal/db"
	"github.com/dtoro/oikos/internal/db/sqlcgen"
	"github.com/dtoro/oikos/internal/health"
	"github.com/dtoro/oikos/internal/remote"
	"github.com/google/uuid"
)

// Run starts the execution worker loop. Blocks until ctx is cancelled.
func Run(ctx context.Context, pool *db.Pool, cfg config.Config) {
	slog.Info("execworker: starting")

	// Liveness probe
	probe := health.New(2 * time.Minute)
	probe.Serve(ctx, cfg.HealthListen)

	recoverOrphaned(ctx, pool)
	probe.Bump()

	ticker := time.NewTicker(15 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			slog.Info("execworker: shutting down")
			return
		case <-ticker.C:
			processPending(ctx, pool)
			probe.Bump()
		}
	}
}

// recoverOrphaned marks executions stuck in 'running' as failed.
func recoverOrphaned(ctx context.Context, pool *db.Pool) {
	tag, err := pool.Exec(ctx, `UPDATE executions SET status = 'failed', result = '{"error":"worker restarted while execution was running"}'::jsonb, completed_at = now() WHERE status = 'running'`)
	if err != nil {
		slog.Error("execworker: recover orphaned", "error", err)
		return
	}
	if tag.RowsAffected() > 0 {
		slog.Warn("execworker: recovered orphaned executions", "count", tag.RowsAffected())
	}
}

// processPending polls for pending executions and dispatches them.
func processPending(ctx context.Context, pool *db.Pool) {
	rows, err := pool.Query(ctx, `
		SELECT e.entity_id, e.target_entity_id, e.action, e.risk_class, e.correlation_id, e.status,
		       COALESCE(t.slug, '') AS target_slug
		FROM executions e
		LEFT JOIN entities t ON t.id = e.target_entity_id
		WHERE e.status = 'proposed'
		ORDER BY e.created_at ASC
		LIMIT 10`)
	if err != nil {
		slog.Error("execworker: query pending", "error", err)
		return
	}
	defer rows.Close()

	q := sqlcgen.New(pool)

	for rows.Next() {
		var execID, targetID *uuid.UUID
		var action, riskClass, correlationID, status, targetSlug string
		if err := rows.Scan(&execID, &targetID, &action, &riskClass, &correlationID, &status, &targetSlug); err != nil {
			slog.Error("execworker: scan row", "error", err)
			continue
		}
		if execID == nil {
			continue
		}

		// At-most-once: try advisory lock on execution entity_id.
		// Acquire a dedicated connection so the session-scoped lock isn't
		// released when the transient pool connection is returned.
		lockKey := hashUUID(*execID)
		lockConn, err := pool.Acquire(ctx)
		if err != nil {
			slog.Error("execworker: acquire lock conn", "error", err)
			continue
		}
		var locked bool
		if err := lockConn.QueryRow(ctx, "SELECT pg_try_advisory_lock($1)", lockKey).Scan(&locked); err != nil || !locked {
			lockConn.Release()
			continue
		}

		dispatch(ctx, pool, q, *execID, targetID, action, targetSlug, correlationID)

		// Release the per-execution lock on the same connection.
		lockConn.Exec(context.WithoutCancel(ctx), "SELECT pg_advisory_unlock($1)", lockKey)
		lockConn.Release()
	}
}

func dispatch(ctx context.Context, pool *db.Pool, q *sqlcgen.Queries, execID uuid.UUID, targetID *uuid.UUID, action, targetSlug, correlationID string) {
	startedAt := time.Now()

	// Mark running
	_, err := pool.Exec(ctx,
		`UPDATE executions SET status = 'running', started_at = $2 WHERE entity_id = $1`,
		execID, startedAt)
	if err != nil {
		slog.Error("execworker: mark running", "error", err, "execution", execID)
		return
	}

	// Resolve SSH target. If targetSlug is available, use it; otherwise resolve from targetID.
	var host, user string
	if targetSlug == "" && targetID != nil {
		if err := pool.QueryRow(ctx, "SELECT slug FROM entities WHERE id = $1", *targetID).Scan(&targetSlug); err != nil {
			failExecution(ctx, pool, execID, fmt.Sprintf("resolve target slug: %v", err))
			return
		}
	}
	if targetSlug != "" {
		addr, sshUser, err := remote.ResolveHost(ctx, pool, targetSlug, "root")
		if err == nil {
			host, user = addr, sshUser
		}
	}
	if host == "" {
		failExecution(ctx, pool, execID, fmt.Sprintf("no reachable target: %s", targetSlug))
		return
	}

	// Determine the command to run from the action field.
	// Format: "action_name:{json_params}" or a raw command string.
	cmd := action
	if idx := strings.Index(action, ":"); idx > 0 && idx < len(action)-1 {
		rawParams := action[idx+1:]
		var params map[string]any
		if json.Unmarshal([]byte(rawParams), &params) == nil {
			if c, ok := params["command"].(string); ok && c != "" {
				cmd = c
			}
		}
	}

	signer, err := actuator.LoadSigner(os.Getenv("OIKOS_SSH_KEY_PATH"))
	if err != nil {
		signer, err = actuator.LoadSigner("/etc/oikos/ssh_key")
		if err != nil {
			failExecution(ctx, pool, execID, fmt.Sprintf("load ssh key: %v", err))
			return
		}
	}

	client, err := actuator.Dial(ctx, actuator.DialOptions{Host: host, User: user, Signer: signer})
	if err != nil {
		failExecution(ctx, pool, execID, fmt.Sprintf("ssh dial: %v", err))
		return
	}
	defer client.Close()

	out, err := actuator.RunCombinedOutput(ctx, client, cmd)
	if err != nil {
		failExecution(ctx, pool, execID, fmt.Sprintf("command: %v\noutput: %s", err, string(out)))
		return
	}

	duration := time.Since(startedAt).Milliseconds()
	resultJSON, _ := json.Marshal(map[string]any{"output": string(out), "success": true})
	_ = q.UpdateExecutionStatus(ctx, sqlcgen.UpdateExecutionStatusParams{
		EntityID:   execID,
		Status:     "completed",
		Result:     resultJSON,
		DurationMs: &[]int32{int32(duration)}[0],
		Verified:   true,
	})

	slog.Info("execworker: execution complete",
		"execution", execID, "target", targetSlug, "duration_ms", duration)
}

func failExecution(ctx context.Context, pool *db.Pool, execID uuid.UUID, reason string) {
	slog.Error("execworker: execution failed", "execution", execID, "error", reason)
	resultJSON, _ := json.Marshal(map[string]any{"error": reason, "success": false})
	pool.Exec(ctx, `UPDATE executions SET status='failed', result=$2::jsonb, completed_at=now() WHERE entity_id=$1`,
		execID, resultJSON)
}

func hashUUID(id uuid.UUID) int {
	h := 0
	for _, b := range id {
		h = (h*31 + int(b)) & 0x7fffffff
	}
	return h
}