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

package httpserver

import (
	"context"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/toki/oto/services/core/internal/cicdstate"
	"github.com/toki/oto/services/core/internal/runnerregistry"
	"github.com/toki/oto/services/core/internal/runnersocket"
)

const (
	DefaultHeartbeatTimeout = 90 * time.Second
	DefaultScanInterval     = 30 * time.Second
)

// ServerConfig holds timeout loop configuration for the server's runner maintenance loop.
type ServerConfig struct {
	HeartbeatTimeout time.Duration
	ScanInterval     time.Duration
}

// Server wraps the HTTP server for the OTO Core service.
type Server struct {
	httpServer       *http.Server
	registry         *runnerregistry.Registry
	store            *cicdstate.Store
	heartbeatTimeout time.Duration
	scanInterval     time.Duration
	runnerSocket     *runnersocket.Server

	mu           sync.Mutex
	loopCancel   context.CancelFunc
	loopDone     <-chan struct{}
	socketCancel context.CancelFunc
	startOnce    sync.Once
}

// NewServer creates a new instance of Server.
func NewServer(addr string) *Server {
	return newServer(addr, runnerregistry.New(), cicdstate.NewStore(), DefaultHeartbeatTimeout, DefaultScanInterval)
}

// NewServerWithRegistry creates a server using an injected runner registry.
func NewServerWithRegistry(addr string, registry *runnerregistry.Registry) *Server {
	return newServer(addr, registry, cicdstate.NewStore(), DefaultHeartbeatTimeout, DefaultScanInterval)
}

// NewServerWithRegistryAndStore creates a server with injected runner registry and CICD store.
func NewServerWithRegistryAndStore(addr string, registry *runnerregistry.Registry, store *cicdstate.Store) *Server {
	return newServer(addr, registry, store, DefaultHeartbeatTimeout, DefaultScanInterval)
}

// NewServerWithConfig creates a server with custom timeout loop configuration.
func NewServerWithConfig(addr string, registry *runnerregistry.Registry, store *cicdstate.Store, cfg ServerConfig) *Server {
	return newServer(addr, registry, store, cfg.HeartbeatTimeout, cfg.ScanInterval)
}

func newServer(addr string, registry *runnerregistry.Registry, store *cicdstate.Store, heartbeatTimeout, scanInterval time.Duration) *Server {
	mux := http.NewServeMux()

	s := &Server{
		httpServer: &http.Server{
			Addr:    addr,
			Handler: mux,
		},
		registry:         registry,
		store:            store,
		heartbeatTimeout: heartbeatTimeout,
		scanInterval:     scanInterval,
	}
	registerRoutes(mux, registry, store, s)
	return s
}

// EnableRunnerSocket starts a runner-facing proto-socket listener together with
// the HTTP server. The listener shares the same registry and CICD store.
func (s *Server) EnableRunnerSocket(addr string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if addr == "" {
		s.runnerSocket = nil
		return
	}
	s.runnerSocket = runnersocket.New(addr, s.registry, s.store)
}

func (s *Server) startTimeoutLoop() {
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})

	s.mu.Lock()
	s.loopCancel = cancel
	s.loopDone = done
	s.mu.Unlock()

	go func() {
		defer close(done)
		if s.registry == nil || s.scanInterval <= 0 {
			return
		}
		ticker := time.NewTicker(s.scanInterval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				s.registry.CheckTimeouts(s.heartbeatTimeout)
			case <-ctx.Done():
				return
			}
		}
	}()
}

func (s *Server) stopTimeoutLoop() {
	s.mu.Lock()
	cancel := s.loopCancel
	done := s.loopDone
	s.mu.Unlock()
	if cancel != nil {
		cancel()
		<-done
	}
}

func (s *Server) startRunnerSocket() error {
	s.mu.Lock()
	runnerSocket := s.runnerSocket
	if runnerSocket == nil {
		s.mu.Unlock()
		return nil
	}
	ctx, cancel := context.WithCancel(context.Background())
	s.socketCancel = cancel
	s.mu.Unlock()
	return runnerSocket.Start(ctx)
}

func (s *Server) stopRunnerSocket() {
	s.mu.Lock()
	cancel := s.socketCancel
	runnerSocket := s.runnerSocket
	s.socketCancel = nil
	s.mu.Unlock()
	if cancel != nil {
		cancel()
	}
	if runnerSocket != nil {
		_ = runnerSocket.Stop()
	}
}

// DispatchQueuedJobs lets HTTP handlers notify the runner socket dispatcher.
func (s *Server) DispatchQueuedJobs() {
	s.mu.Lock()
	runnerSocket := s.runnerSocket
	s.mu.Unlock()
	if runnerSocket != nil {
		runnerSocket.DispatchQueuedJobs()
	}
}

// CancelExecution forwards a cancellation request to a connected runner socket.
func (s *Server) CancelExecution(runnerID, execID, reason string) error {
	s.mu.Lock()
	runnerSocket := s.runnerSocket
	s.mu.Unlock()
	if runnerSocket == nil {
		return nil
	}
	return runnerSocket.CancelExecution(runnerID, execID, reason)
}

// Start starts the HTTP server.
func (s *Server) Start() error {
	s.startOnce.Do(s.startTimeoutLoop)
	if err := s.startRunnerSocket(); err != nil {
		s.stopTimeoutLoop()
		return err
	}
	err := s.httpServer.ListenAndServe()
	s.stopRunnerSocket()
	s.stopTimeoutLoop()
	return err
}

// StartListener starts the HTTP server using a custom net.Listener.
// Useful for testing with dynamic ports.
func (s *Server) StartListener(ln net.Listener) error {
	s.startOnce.Do(s.startTimeoutLoop)
	if err := s.startRunnerSocket(); err != nil {
		s.stopTimeoutLoop()
		return err
	}
	err := s.httpServer.Serve(ln)
	s.stopRunnerSocket()
	s.stopTimeoutLoop()
	return err
}

// Shutdown gracefully shuts down the server.
func (s *Server) Shutdown(ctx context.Context) error {
	s.stopRunnerSocket()
	s.stopTimeoutLoop()
	return s.httpServer.Shutdown(ctx)
}