anchorage/internal/pkg/scheduler/consumer.go

// Package scheduler owns a node's job consumer, rebalancer, requeue
// sweeper, and reconciler loops.
//
// Job consumer: one pull subscription per node on pin.jobs.<myNodeID>.
// Dispatches to the local Kubo backend, updates the placement row with
// WHERE fence = $n, bumps or decrements pin_refcount, publishes an event.
package scheduler

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

	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"

	"anchorage/internal/pkg/ids"
	"anchorage/internal/pkg/ipfs"
	"anchorage/internal/pkg/metrics"
	"anchorage/internal/pkg/pin"
	"anchorage/internal/pkg/store"
)

// PinJobsStream is the JetStream name for the pin work queue.
const PinJobsStream = "PIN_JOBS"

// EnsureStreams creates (or updates) the PIN_JOBS and PIN_EVENTS streams.
// Idempotent; safe to call on every boot.
//
// The replicas argument is the operator's requested replication factor;
// if the running NATS deployment is smaller (e.g., single-node embedded
// NATS with replicas=3 in config) we retry with replicas=1. This matches
// the plan's "Capped at clusterSize at runtime" note without poking
// brittle internals.
func EnsureStreams(ctx context.Context, js jetstream.JetStream, replicas int) error {
	if replicas <= 0 {
		replicas = 1
	}
	mk := func(cfg jetstream.StreamConfig) error {
		_, err := js.CreateOrUpdateStream(ctx, cfg)
		if err == nil {
			return nil
		}
		// JetStream reports "insufficient resources" / "replicas > cluster size"
		// variants; rather than pattern-matching their error strings, retry
		// once with replicas=1 which every deployment accepts.
		if replicas > 1 {
			cfg.Replicas = 1
			if _, err2 := js.CreateOrUpdateStream(ctx, cfg); err2 == nil {
				return nil
			}
		}
		return err
	}

	if err := mk(jetstream.StreamConfig{
		Name:       PinJobsStream,
		Subjects:   []string{pin.JobsSubjectPrefix + ".*"},
		Retention:  jetstream.WorkQueuePolicy,
		Duplicates: 5 * time.Minute,
		Replicas:   replicas,
	}); err != nil {
		return fmt.Errorf("ensure PIN_JOBS: %w", err)
	}
	if err := mk(jetstream.StreamConfig{
		Name:      "PIN_EVENTS",
		Subjects:  []string{pin.EventsSubjectPrefix + ".>"},
		Retention: jetstream.LimitsPolicy,
		MaxAge:    1 * time.Hour,
		Replicas:  replicas,
	}); err != nil {
		return fmt.Errorf("ensure PIN_EVENTS: %w", err)
	}
	return nil
}

// Consumer is a single node's job worker.
type Consumer struct {
	NodeID  ids.NodeID
	Store   store.Store
	Backend ipfs.Backend
	NC      *nats.Conn
	JS      jetstream.JetStream
}

// Run pulls pin.jobs.<NodeID> messages until ctx is cancelled. Honors
// nodes.status: on transitions to drained, the pull loop stops; the
// caller is responsible for restarting it when the node is uncordoned.
func (c *Consumer) Run(ctx context.Context) error {
	consName := "worker-" + c.NodeID.String()
	cons, err := c.JS.CreateOrUpdateConsumer(ctx, PinJobsStream, jetstream.ConsumerConfig{
		Name:          consName,
		Durable:       consName,
		FilterSubject: fmt.Sprintf("%s.%s", pin.JobsSubjectPrefix, c.NodeID.String()),
		AckPolicy:     jetstream.AckExplicitPolicy,
		AckWait:       30 * time.Second,
		MaxDeliver:    5,
	})
	if err != nil {
		return fmt.Errorf("create consumer: %w", err)
	}

	for {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		// Check for drain: if nodes.status='drained', pause the loop.
		// One-shot check per outer iteration is cheap.
		if n, err := c.Store.Nodes().Get(ctx, c.NodeID); err == nil && n.Status == store.NodeStatusDrained {
			slog.Info("scheduler: consumer paused (drained)", "node_id", c.NodeID)
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(2 * time.Second):
			}
			continue
		}

		msgs, err := cons.Fetch(10, jetstream.FetchMaxWait(2*time.Second))
		if err != nil {
			if errors.Is(err, nats.ErrTimeout) || errors.Is(err, context.DeadlineExceeded) {
				metrics.SchedulerFetch.WithLabelValues(c.NodeID.String(), "timeout").Inc()
				continue
			}
			metrics.SchedulerFetch.WithLabelValues(c.NodeID.String(), "err").Inc()
			slog.Warn("scheduler: fetch", "err", err)
			time.Sleep(500 * time.Millisecond)
			continue
		}
		for m := range msgs.Messages() {
			metrics.SchedulerFetch.WithLabelValues(c.NodeID.String(), "ok").Inc()
			c.handle(ctx, m)
		}
	}
}

func (c *Consumer) handle(ctx context.Context, m jetstream.Msg) {
	var job pin.Job
	if err := json.Unmarshal(m.Data(), &job); err != nil {
		slog.Warn("scheduler: bad job payload", "err", err)
		_ = m.Term()
		return
	}
	if err := c.execute(ctx, job); err != nil {
		slog.Warn("scheduler: execute", "err", err, "request_id", job.RequestID)
		_ = m.Nak()
		return
	}
	_ = m.Ack()
}

// execute dispatches the job to Kubo and updates state.
//
// The fence check in UpdatePlacementFenced is what keeps zombie writes
// from clobbering a reassigned placement: if the fence has been bumped
// by the rebalancer, the UPDATE affects 0 rows and we silently drop.
func (c *Consumer) execute(ctx context.Context, j pin.Job) error {
	orgID, err := ids.ParseOrg(j.OrgID)
	if err != nil {
		return err
	}
	rid, err := ids.ParsePin(j.RequestID)
	if err != nil {
		return err
	}

	switch j.Action {
	case "pin":
		if err := c.Backend.Pin(ctx, j.CID, j.Origins); err != nil {
			return c.markPlacement(ctx, rid, j, store.PinStatusFailed, err.Error())
		}
		return c.markPlacement(ctx, rid, j, store.PinStatusPinned, "")
	case "unpin":
		// Decrement refcount; unpin Kubo only on 0.
		remaining, err := c.Store.Pins().DecRefcount(ctx, c.NodeID, j.CID)
		if err != nil && !errors.Is(err, store.ErrNotFound) {
			return err
		}
		if remaining <= 0 {
			if err := c.Backend.Unpin(ctx, j.CID); err != nil && !errors.Is(err, ipfs.ErrNotFound) {
				return err
			}
			_ = c.Store.Pins().DeleteRefcountIfZero(ctx, c.NodeID, j.CID)
		}
		pin.PublishEvent(c.NC, pin.Event{
			RequestID: j.RequestID,
			OrgID:     orgID.String(),
			Status:    "unpinned",
			NodeID:    c.NodeID.String(),
		})
		metrics.SchedulerAcks.WithLabelValues(c.NodeID.String(), "unpinned").Inc()
		return nil
	default:
		return fmt.Errorf("scheduler: unknown action %q", j.Action)
	}
}

func (c *Consumer) markPlacement(ctx context.Context, rid ids.PinID, j pin.Job, status, reason string) error {
	var fr *string
	if reason != "" {
		fr = &reason
	}
	n, err := c.Store.Pins().UpdatePlacementFenced(ctx, rid, c.NodeID, status, fr, j.Fence)
	if err != nil {
		return err
	}
	if n == 0 {
		// Fence mismatch — reassigned already. Silent no-op.
		slog.Info("scheduler: fence mismatch, dropping",
			"request_id", j.RequestID, "node_id", c.NodeID, "fence", j.Fence)
		metrics.SchedulerAcks.WithLabelValues(c.NodeID.String(), "fence-mismatch").Inc()
		return nil
	}
	pin.PublishEvent(c.NC, pin.Event{
		RequestID: j.RequestID,
		OrgID:     j.OrgID,
		Status:    status,
		NodeID:    c.NodeID.String(),
		Reason:    reason,
	})
	metrics.SchedulerAcks.WithLabelValues(c.NodeID.String(), status).Inc()
	return nil
}