oto/services/core/internal/runnersocket/server.go

package runnersocket

import (
	"context"
	"fmt"
	"net"
	"sort"
	"strconv"
	"strings"
	"sync"

	protoSocket "git.toki-labs.com/toki/proto-socket/go"
	"github.com/toki/oto/services/core/internal/cicdstate"
	"github.com/toki/oto/services/core/internal/runnerregistry"
	otopb "github.com/toki/oto/services/core/oto"
	"google.golang.org/protobuf/proto"
)

const (
	heartbeatIntervalSeconds = 30
	heartbeatWaitSeconds     = 45
)

// Server owns the runner-facing proto-socket listener.
type Server struct {
	addr     string
	registry *runnerregistry.Registry
	store    *cicdstate.Store

	mu      sync.RWMutex
	server  *protoSocket.TcpServer
	clients map[string]*protoSocket.TcpClient
}

// New creates a runner socket server sharing the core registry and CICD store.
func New(addr string, registry *runnerregistry.Registry, store *cicdstate.Store) *Server {
	return &Server{
		addr:     addr,
		registry: registry,
		store:    store,
		clients:  make(map[string]*protoSocket.TcpClient),
	}
}

// Start binds the proto-socket TCP listener.
func (s *Server) Start(ctx context.Context) error {
	host, port, err := splitAddr(s.addr)
	if err != nil {
		return err
	}
	server := protoSocket.NewTcpServer(host, port, func(conn net.Conn) *protoSocket.TcpClient {
		client := protoSocket.NewTcpClient(conn, heartbeatIntervalSeconds, heartbeatWaitSeconds, parserMap())
		s.attachClientHandlers(client)
		return client
	})

	s.mu.Lock()
	s.server = server
	s.mu.Unlock()

	return server.Start(ctx)
}

// Stop closes the listener and connected runner sockets.
func (s *Server) Stop() error {
	s.mu.RLock()
	server := s.server
	s.mu.RUnlock()
	if server == nil {
		return nil
	}
	return server.Stop()
}

// DispatchQueuedJobs pushes queued jobs to currently connected idle runners.
func (s *Server) DispatchQueuedJobs() {
	for _, runnerID := range s.connectedRunnerIDs() {
		if s.runnerHasActiveExecution(runnerID) {
			continue
		}
		client := s.clientFor(runnerID)
		if client == nil {
			continue
		}
		runRequest, jobID, execID, err := s.claimNextQueuedJob(runnerID)
		if err != nil {
			if jobID != "" {
				if execID != "" {
					_ = s.store.AppendLog(execID, fmt.Sprintf("failed to prepare run request over proto-socket: %v", err))
					_ = s.store.TransitionExecution(execID, cicdstate.StateFailed)
				}
				_ = s.store.TransitionJob(jobID, cicdstate.StateFailed)
			}
			continue
		}
		if runRequest == nil {
			continue
		}
		if err := client.Send(runRequest); err != nil {
			_ = s.store.AppendLog(execID, fmt.Sprintf("failed to dispatch run request over proto-socket: %v", err))
			_ = s.store.TransitionExecution(execID, cicdstate.StateFailed)
			_ = s.store.TransitionJob(jobID, cicdstate.StateFailed)
			continue
		}
	}
}

// CancelExecution asks a connected runner to cancel an execution.
func (s *Server) CancelExecution(runnerID, execID, reason string) error {
	client := s.clientFor(runnerID)
	if client == nil {
		return fmt.Errorf("runner socket is not connected")
	}
	return client.Send(&otopb.CancelRunRequest{
		RunnerId:    runnerID,
		ExecutionId: execID,
		Reason:      reason,
	})
}

func (s *Server) attachClientHandlers(client *protoSocket.TcpClient) {
	var runnerIDMu sync.RWMutex
	runnerID := ""

	setRunnerID := func(id string) {
		runnerIDMu.Lock()
		runnerID = id
		runnerIDMu.Unlock()
	}
	getRunnerID := func() string {
		runnerIDMu.RLock()
		defer runnerIDMu.RUnlock()
		return runnerID
	}

	protoSocket.AddRequestListenerTyped[*otopb.RegisterRunnerRequest, *otopb.RegisterRunnerResponse](
		&client.Communicator,
		func(req *otopb.RegisterRunnerRequest) (*otopb.RegisterRunnerResponse, error) {
			res := s.registry.Register(req)
			if res.GetAccepted() {
				id := res.GetRunnerId()
				setRunnerID(id)
				s.setClient(id, client)
				s.registry.Heartbeat(id, otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
				go s.DispatchQueuedJobs()
			}
			return res, nil
		},
	)

	protoSocket.AddRequestListenerTyped[*otopb.HeartbeatRequest, *otopb.HeartbeatResponse](
		&client.Communicator,
		func(req *otopb.HeartbeatRequest) (*otopb.HeartbeatResponse, error) {
			runnerID := strings.TrimSpace(req.GetRunnerId())
			if runnerID == "" {
				runnerID = getRunnerID()
			}
			return s.registry.Heartbeat(runnerID, req.GetStatus()), nil
		},
	)

	protoSocket.AddRequestListenerTyped[*otopb.ExecutionReportRequest, *otopb.ExecutionReportResponse](
		&client.Communicator,
		func(req *otopb.ExecutionReportRequest) (*otopb.ExecutionReportResponse, error) {
			res := s.handleExecutionReport(req)
			if res.GetAccepted() {
				go s.DispatchQueuedJobs()
			}
			return res, nil
		},
	)

	protoSocket.AddRequestListenerTyped[*otopb.LogAppendRequest, *otopb.LogAppendResponse](
		&client.Communicator,
		func(req *otopb.LogAppendRequest) (*otopb.LogAppendResponse, error) {
			return s.handleLogAppend(req), nil
		},
	)

	protoSocket.AddRequestListenerTyped[*otopb.ArtifactReportRequest, *otopb.ArtifactReportResponse](
		&client.Communicator,
		func(req *otopb.ArtifactReportRequest) (*otopb.ArtifactReportResponse, error) {
			return s.handleArtifactReport(req), nil
		},
	)

	client.AddDisconnectListener(func(c *protoSocket.TcpClient) {
		id := getRunnerID()
		if id == "" {
			return
		}
		if s.removeClient(id, c) {
			s.registry.Disconnect(id)
		}
	})
}

func (s *Server) handleExecutionReport(req *otopb.ExecutionReportRequest) *otopb.ExecutionReportResponse {
	runnerID := strings.TrimSpace(req.GetRunnerId())
	execID := strings.TrimSpace(req.GetExecutionId())
	if err := s.ensureRunnerKnown(runnerID); err != nil {
		return &otopb.ExecutionReportResponse{RunnerId: runnerID, ExecutionId: execID, ErrorMessage: err.Error()}
	}
	exec, err := s.store.GetExecution(execID)
	if err != nil {
		return &otopb.ExecutionReportResponse{RunnerId: runnerID, ExecutionId: execID, ErrorMessage: "execution not found"}
	}
	if exec.RunnerID != runnerID {
		return &otopb.ExecutionReportResponse{RunnerId: runnerID, ExecutionId: execID, ErrorMessage: "execution owner mismatch"}
	}
	jobID := strings.TrimSpace(req.GetJobId())
	if jobID == "" {
		jobID = exec.JobID
	} else if jobID != exec.JobID {
		return &otopb.ExecutionReportResponse{RunnerId: runnerID, JobId: jobID, ExecutionId: execID, ErrorMessage: "job id mismatch for execution"}
	}

	targetState := cicdstate.StateFailed
	if req.GetSuccess() {
		targetState = cicdstate.StateSucceeded
	}
	if err := s.store.TransitionExecution(execID, targetState); err != nil {
		return &otopb.ExecutionReportResponse{RunnerId: runnerID, JobId: jobID, ExecutionId: execID, ErrorMessage: err.Error()}
	}
	if err := s.store.TransitionJob(jobID, targetState); err != nil {
		return &otopb.ExecutionReportResponse{RunnerId: runnerID, JobId: jobID, ExecutionId: execID, ErrorMessage: err.Error()}
	}
	if strings.TrimSpace(req.GetMessage()) != "" {
		_ = s.store.AppendLog(execID, req.GetMessage())
	}
	return &otopb.ExecutionReportResponse{
		Accepted:    true,
		RunnerId:    runnerID,
		JobId:       jobID,
		ExecutionId: execID,
		State:       targetState,
	}
}

func (s *Server) handleLogAppend(req *otopb.LogAppendRequest) *otopb.LogAppendResponse {
	runnerID := strings.TrimSpace(req.GetRunnerId())
	execID := strings.TrimSpace(req.GetExecutionId())
	if err := s.ensureRunnerKnown(runnerID); err != nil {
		return &otopb.LogAppendResponse{ErrorMessage: err.Error()}
	}
	if err := s.ensureExecutionOwner(runnerID, execID); err != nil {
		return &otopb.LogAppendResponse{ErrorMessage: err.Error()}
	}
	line := strings.TrimSpace(req.GetLine())
	if line == "" {
		return &otopb.LogAppendResponse{ErrorMessage: "line is required"}
	}
	if err := s.store.AppendLog(execID, line); err != nil {
		return &otopb.LogAppendResponse{ErrorMessage: "execution not found"}
	}
	return &otopb.LogAppendResponse{Accepted: true}
}

func (s *Server) handleArtifactReport(req *otopb.ArtifactReportRequest) *otopb.ArtifactReportResponse {
	runnerID := strings.TrimSpace(req.GetRunnerId())
	execID := strings.TrimSpace(req.GetExecutionId())
	if err := s.ensureRunnerKnown(runnerID); err != nil {
		return &otopb.ArtifactReportResponse{ErrorMessage: err.Error()}
	}
	if err := s.ensureExecutionOwner(runnerID, execID); err != nil {
		return &otopb.ArtifactReportResponse{ErrorMessage: err.Error()}
	}
	name := strings.TrimSpace(req.GetName())
	path := strings.TrimSpace(req.GetPath())
	if name == "" || path == "" {
		return &otopb.ArtifactReportResponse{ErrorMessage: "name and path are required"}
	}
	if err := s.store.AppendArtifact(execID, name, path); err != nil {
		return &otopb.ArtifactReportResponse{ErrorMessage: "execution not found"}
	}
	return &otopb.ArtifactReportResponse{Accepted: true}
}

func (s *Server) claimNextQueuedJob(runnerID string) (*otopb.RunRequest, string, string, error) {
	if err := s.ensureRunnerKnown(runnerID); err != nil {
		return nil, "", "", err
	}
	var foundJob *cicdstate.Job
	for _, job := range s.store.SnapshotJobs() {
		if job.State != cicdstate.StateQueued {
			continue
		}
		if foundJob == nil || job.CreatedAt.Before(foundJob.CreatedAt) {
			copyJob := job
			foundJob = &copyJob
		}
	}
	if foundJob == nil {
		return nil, "", "", nil
	}
	if foundJob.RunInput == nil {
		return nil, foundJob.ID, "", fmt.Errorf("job has no run request")
	}
	if err := validateRunInput(foundJob.RunInput); err != nil {
		return nil, foundJob.ID, "", err
	}

	execID := s.store.NextExecutionID()
	if _, err := s.store.CreateExecution(foundJob.ID, execID); err != nil {
		return nil, foundJob.ID, execID, err
	}
	_ = s.store.SetExecutionRunnerID(execID, runnerID)
	if err := s.store.TransitionJob(foundJob.ID, cicdstate.StateRunning); err != nil {
		return nil, foundJob.ID, execID, err
	}
	if err := s.store.TransitionExecution(execID, cicdstate.StateRunning); err != nil {
		return nil, foundJob.ID, execID, err
	}
	return runRequestToProto(foundJob.RunInput, runnerID, foundJob.ID, execID), foundJob.ID, execID, nil
}

func (s *Server) runnerHasActiveExecution(runnerID string) bool {
	for _, exec := range s.store.SnapshotExecutions() {
		if exec.RunnerID == runnerID && exec.State == cicdstate.StateRunning {
			return true
		}
	}
	return false
}

func (s *Server) ensureRunnerKnown(runnerID string) error {
	if strings.TrimSpace(runnerID) == "" {
		return fmt.Errorf("missing runner id")
	}
	if s.registry == nil {
		return nil
	}
	if _, ok := s.registry.Snapshot(runnerID); !ok {
		return fmt.Errorf("runner not found")
	}
	return nil
}

func (s *Server) ensureExecutionOwner(runnerID, execID string) error {
	exec, err := s.store.GetExecution(execID)
	if err != nil {
		return fmt.Errorf("execution not found")
	}
	if exec.RunnerID != runnerID {
		return fmt.Errorf("execution owner mismatch")
	}
	return nil
}

func (s *Server) setClient(runnerID string, client *protoSocket.TcpClient) {
	s.mu.Lock()
	old := s.clients[runnerID]
	s.clients[runnerID] = client
	s.mu.Unlock()

	// Close outside the lock: Close() fires disconnect listeners synchronously,
	// and those listeners call removeClient which also acquires s.mu. Closing
	// inside the lock would cause a deadlock on re-entry.
	if old != nil && old != client {
		_ = old.Close()
	}
}

func (s *Server) removeClient(runnerID string, client *protoSocket.TcpClient) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.clients[runnerID] == client {
		delete(s.clients, runnerID)
		return true
	}
	return false
}

func (s *Server) clientFor(runnerID string) *protoSocket.TcpClient {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.clients[runnerID]
}

func (s *Server) connectedRunnerIDs() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	ids := make([]string, 0, len(s.clients))
	for id := range s.clients {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	return ids
}

func parserMap() protoSocket.ParserMap {
	return protoSocket.ParserMap{
		protoSocket.TypeNameOf(&otopb.RegisterRunnerRequest{}):  parserFor(func() *otopb.RegisterRunnerRequest { return &otopb.RegisterRunnerRequest{} }),
		protoSocket.TypeNameOf(&otopb.HeartbeatRequest{}):       parserFor(func() *otopb.HeartbeatRequest { return &otopb.HeartbeatRequest{} }),
		protoSocket.TypeNameOf(&otopb.ExecutionReportRequest{}): parserFor(func() *otopb.ExecutionReportRequest { return &otopb.ExecutionReportRequest{} }),
		protoSocket.TypeNameOf(&otopb.LogAppendRequest{}):       parserFor(func() *otopb.LogAppendRequest { return &otopb.LogAppendRequest{} }),
		protoSocket.TypeNameOf(&otopb.ArtifactReportRequest{}):  parserFor(func() *otopb.ArtifactReportRequest { return &otopb.ArtifactReportRequest{} }),
	}
}

func parserFor[T proto.Message](newMessage func() T) func([]byte) (proto.Message, error) {
	return func(b []byte) (proto.Message, error) {
		m := newMessage()
		return m, proto.Unmarshal(b, m)
	}
}

func splitAddr(addr string) (string, int, error) {
	host, portText, err := net.SplitHostPort(addr)
	if err != nil {
		return "", 0, err
	}
	port, err := strconv.Atoi(portText)
	if err != nil {
		return "", 0, err
	}
	return host, port, nil
}

func validateRunInput(in *cicdstate.RunInput) error {
	hasYAML := strings.TrimSpace(in.PipelineYAML) != ""
	hasPath := strings.TrimSpace(in.PipelineYAMLPath) != ""
	if hasYAML == hasPath {
		return fmt.Errorf("run request must set exactly one of pipeline_yaml or pipeline_yaml_path")
	}
	return nil
}

func runRequestToProto(in *cicdstate.RunInput, runnerID, jobID, execID string) *otopb.RunRequest {
	if in == nil {
		return nil
	}
	rr := &otopb.RunRequest{
		RunnerId:         runnerID,
		JobId:            jobID,
		ExecutionId:      execID,
		PipelineYamlPath: in.PipelineYAMLPath,
		PipelineYaml:     in.PipelineYAML,
		CommandTypes:     append([]string(nil), in.CommandTypes...),
	}
	if len(in.Variables) > 0 {
		rr.Variables = make(map[string]string, len(in.Variables))
		for k, v := range in.Variables {
			rr.Variables[k] = v
		}
	}
	return rr
}