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

package runnerregistry

import (
	"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")
	}

	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()
	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) 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 rejected(reason string) *otopb.RegisterRunnerResponse {
	return &otopb.RegisterRunnerResponse{
		Accepted:     false,
		RejectReason: reason,
	}
}