actions-runner-controller/cmd/githubrunnerscalesetlistener/autoScalerService.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"strings"

	"github.com/actions/actions-runner-controller/github/actions"
	"github.com/go-logr/logr"
)

type ScaleSettings struct {
	Namespace    string
	ResourceName string
	MinRunners   int
	MaxRunners   int
}

type Service struct {
	ctx                context.Context
	logger             logr.Logger
	rsClient           RunnerScaleSetClient
	kubeManager        KubernetesManager
	settings           *ScaleSettings
	currentRunnerCount int
}

func NewService(
	ctx context.Context,
	rsClient RunnerScaleSetClient,
	manager KubernetesManager,
	settings *ScaleSettings,
	options ...func(*Service),
) *Service {
	s := &Service{
		ctx:                ctx,
		rsClient:           rsClient,
		kubeManager:        manager,
		settings:           settings,
		currentRunnerCount: 0,
		logger:             logr.FromContextOrDiscard(ctx),
	}

	for _, option := range options {
		option(s)
	}

	return s
}

func (s *Service) Start() error {
	if s.settings.MinRunners > 0 {
		s.logger.Info("scale to match minimal runners.")
		err := s.scaleForAssignedJobCount(0)
		if err != nil {
			return fmt.Errorf("could not scale to match minimal runners. %w", err)
		}
	}

	for {
		s.logger.Info("waiting for message...")
		select {
		case <-s.ctx.Done():
			s.logger.Info("service is stopped.")
			return nil
		default:
			err := s.rsClient.GetRunnerScaleSetMessage(s.ctx, s.processMessage)
			if err != nil {
				return fmt.Errorf("could not get and process message. %w", err)
			}
		}
	}
}

func (s *Service) processMessage(message *actions.RunnerScaleSetMessage) error {
	s.logger.Info("process message.", "messageId", message.MessageId, "messageType", message.MessageType)
	if message.Statistics == nil {
		return fmt.Errorf("can't process message with empty statistics")
	}

	s.logger.Info("current runner scale set statistics.",
		"available jobs", message.Statistics.TotalAvailableJobs,
		"acquired jobs", message.Statistics.TotalAcquiredJobs,
		"assigned jobs", message.Statistics.TotalAssignedJobs,
		"running jobs", message.Statistics.TotalRunningJobs,
		"registered runners", message.Statistics.TotalRegisteredRunners,
		"busy runners", message.Statistics.TotalBusyRunners,
		"idle runners", message.Statistics.TotalIdleRunners)

	if message.MessageType != "RunnerScaleSetJobMessages" {
		s.logger.Info("skip message with unknown message type.", "messageType", message.MessageType)
		return nil
	}

	var batchedMessages []json.RawMessage
	if err := json.NewDecoder(strings.NewReader(message.Body)).Decode(&batchedMessages); err != nil {
		return fmt.Errorf("could not decode job messages. %w", err)
	}

	s.logger.Info("process batched runner scale set job messages.", "messageId", message.MessageId, "batchSize", len(batchedMessages))

	var availableJobs []int64
	for _, message := range batchedMessages {
		var messageType actions.JobMessageType
		if err := json.Unmarshal(message, &messageType); err != nil {
			return fmt.Errorf("could not decode job message type. %w", err)
		}

		switch messageType.MessageType {
		case "JobAvailable":
			var jobAvailable actions.JobAvailable
			if err := json.Unmarshal(message, &jobAvailable); err != nil {
				return fmt.Errorf("could not decode job available message. %w", err)
			}
			s.logger.Info("job available message received.", "RequestId", jobAvailable.RunnerRequestId)
			availableJobs = append(availableJobs, jobAvailable.RunnerRequestId)
		case "JobAssigned":
			var jobAssigned actions.JobAssigned
			if err := json.Unmarshal(message, &jobAssigned); err != nil {
				return fmt.Errorf("could not decode job assigned message. %w", err)
			}
			s.logger.Info("job assigned message received.", "RequestId", jobAssigned.RunnerRequestId)
		case "JobStarted":
			var jobStarted actions.JobStarted
			if err := json.Unmarshal(message, &jobStarted); err != nil {
				return fmt.Errorf("could not decode job started message. %w", err)
			}
			s.logger.Info("job started message received.", "RequestId", jobStarted.RunnerRequestId, "RunnerId", jobStarted.RunnerId)
			s.updateJobInfoForRunner(jobStarted)
		case "JobCompleted":
			var jobCompleted actions.JobCompleted
			if err := json.Unmarshal(message, &jobCompleted); err != nil {
				return fmt.Errorf("could not decode job completed message. %w", err)
			}
			s.logger.Info("job completed message received.", "RequestId", jobCompleted.RunnerRequestId, "Result", jobCompleted.Result, "RunnerId", jobCompleted.RunnerId, "RunnerName", jobCompleted.RunnerName)
		default:
			s.logger.Info("unknown job message type.", "messageType", messageType.MessageType)
		}
	}

	err := s.rsClient.AcquireJobsForRunnerScaleSet(s.ctx, availableJobs)
	if err != nil {
		return fmt.Errorf("could not acquire jobs. %w", err)
	}

	return s.scaleForAssignedJobCount(message.Statistics.TotalAssignedJobs)
}

func (s *Service) scaleForAssignedJobCount(count int) error {
	targetRunnerCount := int(math.Max(math.Min(float64(s.settings.MaxRunners), float64(count)), float64(s.settings.MinRunners)))
	if targetRunnerCount != s.currentRunnerCount {
		s.logger.Info("try scale runner request up/down base on assigned job count",
			"assigned job", count,
			"decision", targetRunnerCount,
			"min", s.settings.MinRunners,
			"max", s.settings.MaxRunners,
			"currentRunnerCount", s.currentRunnerCount)
		err := s.kubeManager.ScaleEphemeralRunnerSet(s.ctx, s.settings.Namespace, s.settings.ResourceName, targetRunnerCount)
		if err != nil {
			return fmt.Errorf("could not scale ephemeral runner set (%s/%s). %w", s.settings.Namespace, s.settings.ResourceName, err)
		}

		s.currentRunnerCount = targetRunnerCount
	}

	return nil
}

// updateJobInfoForRunner updates the ephemeral runner with the job info and this is best effort since the info is only for better telemetry
func (s *Service) updateJobInfoForRunner(jobInfo actions.JobStarted) {
	s.logger.Info("update job info for runner",
		"runnerName", jobInfo.RunnerName,
		"ownerName", jobInfo.OwnerName,
		"repoName", jobInfo.RepositoryName,
		"workflowRef", jobInfo.JobWorkflowRef,
		"workflowRunId", jobInfo.WorkflowRunId,
		"jobDisplayName", jobInfo.JobDisplayName,
		"requestId", jobInfo.RunnerRequestId)
	err := s.kubeManager.UpdateEphemeralRunnerWithJobInfo(s.ctx, s.settings.Namespace, jobInfo.RunnerName, jobInfo.OwnerName, jobInfo.RepositoryName, jobInfo.JobWorkflowRef, jobInfo.JobDisplayName, jobInfo.WorkflowRunId, jobInfo.RunnerRequestId)
	if err != nil {
		s.logger.Error(err, "could not update ephemeral runner with job info", "runnerName", jobInfo.RunnerName, "requestId", jobInfo.RunnerRequestId)
	}
}