oto/services/core/internal/runnerregistry/registry.go

package runnerregistry

import (
	"fmt"
	"strings"
	"sync"
	"time"

	otopb "github.com/toki/oto/services/core/oto"
)

const (
	StatusAccepted         = "accepted"
	StatusOnline           = "online"
	StatusDisconnected     = "disconnected"
	StatusHeartbeatTimeout = "heartbeat_timeout"
)

type RunnerCapability struct {
	Name    string
	Version string
}

type RunnerRecord struct {
	RunnerID         string
	Alias            string
	ProtocolVersion  string
	Capability       RunnerCapability
	CommandTypes     []string
	Status           string
	AcceptedAt       time.Time
	FirstHeartbeatAt time.Time
	LastHeartbeatAt  time.Time
	FailureReason    string
}

type Registry struct {
	mu      sync.RWMutex
	now     func() time.Time
	runners map[string]RunnerRecord
}

func New() *Registry {
	return NewWithClock(time.Now)
}

func NewWithClock(now func() time.Time) *Registry {
	if now == nil {
		now = time.Now
	}
	return &Registry{
		now:     now,
		runners: map[string]RunnerRecord{},
	}
}

func (r *Registry) Register(req *otopb.RegisterRunnerRequest) *otopb.RegisterRunnerResponse {
	if req == nil {
		return rejected("missing registration request")
	}

	token := strings.TrimSpace(req.GetEnrollmentToken())
	runnerID := strings.TrimSpace(req.GetRunnerId())
	protocolVersion := strings.TrimSpace(req.GetProtocolVersion())
	if token == "" {
		return rejected("missing enrollment token")
	}
	if runnerID == "" {
		return rejected("missing runner id")
	}
	if protocolVersion == "" {
		return rejected("missing protocol version")
	}

	if err := validateProtocolCompatibility(protocolVersion, req.GetCapability()); err != nil {
		return rejectedWithCode("incompatible_runner", err.Error())
	}

	record := RunnerRecord{
		RunnerID:        runnerID,
		Alias:           strings.TrimSpace(req.GetAlias()),
		ProtocolVersion: protocolVersion,
		CommandTypes:    append([]string(nil), req.GetCommandCatalog().GetCommandTypes()...),
		Status:          StatusAccepted,
		AcceptedAt:      r.now(),
	}
	if capability := req.GetCapability(); capability != nil {
		record.Capability = RunnerCapability{
			Name:    strings.TrimSpace(capability.GetName()),
			Version: strings.TrimSpace(capability.GetVersion()),
		}
	}

	r.mu.Lock()
	if existing, ok := r.runners[runnerID]; ok && strings.TrimSpace(existing.Alias) != "" {
		record.Alias = existing.Alias
	}
	r.runners[runnerID] = record
	r.mu.Unlock()

	return &otopb.RegisterRunnerResponse{
		Accepted: true,
		RunnerId: runnerID,
		Alias:    record.Alias,
	}
}

func (r *Registry) Heartbeat(runnerID string, status otopb.HeartbeatStatus) *otopb.HeartbeatResponse {
	r.mu.Lock()
	defer r.mu.Unlock()
	record, ok := r.runners[runnerID]
	if !ok {
		return &otopb.HeartbeatResponse{
			Success:      false,
			ErrorMessage: "unknown runner",
		}
	}

	if record.Status == StatusDisconnected || record.Status == StatusHeartbeatTimeout {
		return &otopb.HeartbeatResponse{
			Success:      false,
			ErrorMessage: "runner in terminal state; re-register required",
		}
	}

	if record.Status == StatusAccepted {
		record.FirstHeartbeatAt = r.now()
	}
	record.LastHeartbeatAt = r.now()
	record.Status = StatusOnline

	r.runners[runnerID] = record

	return &otopb.HeartbeatResponse{
		Success: true,
	}
}

func (r *Registry) Disconnect(runnerID string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	record, ok := r.runners[runnerID]
	if !ok {
		return false
	}
	record.Status = StatusDisconnected
	r.runners[runnerID] = record
	return true
}

func (r *Registry) UpdateAlias(runnerID, alias string) (RunnerRecord, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	record, ok := r.runners[runnerID]
	if !ok {
		return RunnerRecord{}, false
	}
	record.Alias = strings.TrimSpace(alias)
	r.runners[runnerID] = record
	record.CommandTypes = append([]string(nil), record.CommandTypes...)
	return record, true
}

func (r *Registry) CheckTimeouts(timeoutDuration time.Duration) {
	r.mu.Lock()
	defer r.mu.Unlock()
	now := r.now()
	for id, record := range r.runners {
		if record.Status == StatusOnline {
			if now.Sub(record.LastHeartbeatAt) > timeoutDuration {
				record.Status = StatusHeartbeatTimeout
				record.FailureReason = "heartbeat timeout"
				r.runners[id] = record
			}
		} else if record.Status == StatusAccepted {
			if now.Sub(record.AcceptedAt) > timeoutDuration {
				record.Status = StatusHeartbeatTimeout
				record.FailureReason = "heartbeat timeout before first heartbeat"
				r.runners[id] = record
			}
		}
	}
}

func (r *Registry) Snapshot(runnerID string) (RunnerRecord, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	record, ok := r.runners[runnerID]
	if !ok {
		return RunnerRecord{}, false
	}
	record.CommandTypes = append([]string(nil), record.CommandTypes...)
	return record, true
}

func validateProtocolCompatibility(protocolVersion string, capability *otopb.RunnerCapability) error {
	if protocolVersion != "oto.runner.v1" {
		return fmt.Errorf("unsupported protocol version: %q", protocolVersion)
	}
	if capability == nil {
		return fmt.Errorf("missing runner capability")
	}
	if strings.TrimSpace(capability.GetName()) != "oto-runner" {
		return fmt.Errorf("unsupported capability name: %q", capability.GetName())
	}
	version := strings.TrimSpace(capability.GetVersion())
	if version == "" {
		return fmt.Errorf("missing capability version")
	}
	parts := strings.Split(version, ".")
	if parts[0] != "1" {
		return fmt.Errorf("incompatible capability major version: %q, want major version 1", version)
	}
	return nil
}

func rejectedWithCode(code, message string) *otopb.RegisterRunnerResponse {
	return &otopb.RegisterRunnerResponse{
		Accepted:     false,
		RejectReason: message,
		Error: &otopb.ProtocolError{
			Code:    code,
			Message: message,
		},
	}
}

func rejected(reason string) *otopb.RegisterRunnerResponse {
	return rejectedWithCode("registration_rejected", reason)
}