Refactoring listener app with configurable fallback (#3096)

cmd/ghalistener/listener/listener.go

@@ -0,0 +1,388 @@
+package listener
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/actions/actions-runner-controller/cmd/ghalistener/metrics"
+	"github.com/actions/actions-runner-controller/github/actions"
+	"github.com/go-logr/logr"
+	"github.com/google/uuid"
+)
+
+const (
+	sessionCreationMaxRetries = 10
+)
+
+// message types
+const (
+	messageTypeJobAvailable = "JobAvailable"
+	messageTypeJobAssigned  = "JobAssigned"
+	messageTypeJobStarted   = "JobStarted"
+	messageTypeJobCompleted = "JobCompleted"
+)
+
+//go:generate mockery --name Client --output ./mocks --outpkg mocks --case underscore
+type Client interface {
+	GetAcquirableJobs(ctx context.Context, runnerScaleSetId int) (*actions.AcquirableJobList, error)
+	CreateMessageSession(ctx context.Context, runnerScaleSetId int, owner string) (*actions.RunnerScaleSetSession, error)
+	GetMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, lastMessageId int64) (*actions.RunnerScaleSetMessage, error)
+	DeleteMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, messageId int64) error
+	AcquireJobs(ctx context.Context, runnerScaleSetId int, messageQueueAccessToken string, requestIds []int64) ([]int64, error)
+	RefreshMessageSession(ctx context.Context, runnerScaleSetId int, sessionId *uuid.UUID) (*actions.RunnerScaleSetSession, error)
+}
+
+type Config struct {
+	Client     Client
+	ScaleSetID int
+	MinRunners int
+	MaxRunners int
+	Logger     logr.Logger
+	Metrics    metrics.Publisher
+}
+
+func (c *Config) Validate() error {
+	if c.Client == nil {
+		return errors.New("client is required")
+	}
+	if c.ScaleSetID == 0 {
+		return errors.New("scaleSetID is required")
+	}
+	if c.MinRunners < 0 {
+		return errors.New("minRunners must be greater than or equal to 0")
+	}
+	if c.MaxRunners < 0 {
+		return errors.New("maxRunners must be greater than or equal to 0")
+	}
+	if c.MaxRunners > 0 && c.MinRunners > c.MaxRunners {
+		return errors.New("minRunners must be less than or equal to maxRunners")
+	}
+	return nil
+}
+
+// The Listener's role is to manage all interactions with the actions service.
+// It receives messages and processes them using the given handler.
+type Listener struct {
+	// configured fields
+	scaleSetID int               // The ID of the scale set associated with the listener.
+	client     Client            // The client used to interact with the scale set.
+	metrics    metrics.Publisher // The publisher used to publish metrics.
+
+	// internal fields
+	logger   logr.Logger // The logger used for logging.
+	hostname string      // The hostname of the listener.
+
+	// updated fields
+	lastMessageID int64                          // The ID of the last processed message.
+	session       *actions.RunnerScaleSetSession // The session for managing the runner scale set.
+}
+
+func New(config Config) (*Listener, error) {
+	if err := config.Validate(); err != nil {
+		return nil, fmt.Errorf("invalid config: %w", err)
+	}
+
+	listener := &Listener{
+		scaleSetID: config.ScaleSetID,
+		client:     config.Client,
+		logger:     config.Logger,
+		metrics:    metrics.Discard,
+	}
+
+	if config.Metrics != nil {
+		listener.metrics = config.Metrics
+	}
+
+	listener.metrics.PublishStatic(config.MinRunners, config.MaxRunners)
+
+	hostname, err := os.Hostname()
+	if err != nil {
+		hostname = uuid.NewString()
+		listener.logger.Info("Failed to get hostname, fallback to uuid", "uuid", hostname, "error", err)
+	}
+	listener.hostname = hostname
+
+	return listener, nil
+}
+
+//go:generate mockery --name Handler --output ./mocks --outpkg mocks --case underscore
+type Handler interface {
+	HandleJobStarted(ctx context.Context, jobInfo *actions.JobStarted) error
+	HandleDesiredRunnerCount(ctx context.Context, desiredRunnerCount int) error
+}
+
+// Listen listens for incoming messages and handles them using the provided handler.
+// It continuously listens for messages until the context is cancelled.
+// The initial message contains the current statistics and acquirable jobs, if any.
+// The handler is responsible for handling the initial message and subsequent messages.
+// If an error occurs during any step, Listen returns an error.
+func (l *Listener) Listen(ctx context.Context, handler Handler) error {
+	if err := l.createSession(ctx); err != nil {
+		return fmt.Errorf("createSession failed: %w", err)
+	}
+
+	initialMessage := &actions.RunnerScaleSetMessage{
+		MessageId:   0,
+		MessageType: "RunnerScaleSetJobMessages",
+		Statistics:  l.session.Statistics,
+		Body:        "",
+	}
+
+	if l.session.Statistics.TotalAvailableJobs > 0 || l.session.Statistics.TotalAssignedJobs > 0 {
+		acquirableJobs, err := l.client.GetAcquirableJobs(ctx, l.scaleSetID)
+		if err != nil {
+			return fmt.Errorf("failed to call GetAcquirableJobs: %w", err)
+		}
+
+		acquirableJobsJson, err := json.Marshal(acquirableJobs)
+		if err != nil {
+			return fmt.Errorf("failed to marshal acquirable jobs: %w", err)
+		}
+
+		initialMessage.Body = string(acquirableJobsJson)
+	}
+
+	if err := handler.HandleDesiredRunnerCount(ctx, initialMessage.Statistics.TotalAssignedJobs); err != nil {
+		return fmt.Errorf("handling initial message failed: %w", err)
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("context cancelled: %w", ctx.Err())
+		default:
+		}
+
+		msg, err := l.getMessage(ctx)
+		if err != nil {
+			return fmt.Errorf("failed to get message: %w", err)
+		}
+
+		if msg == nil {
+			continue
+		}
+
+		statistics, jobsStarted, err := l.parseMessage(ctx, msg)
+		if err != nil {
+			return fmt.Errorf("failed to parse message: %w", err)
+		}
+
+		l.lastMessageID = msg.MessageId
+
+		if err := l.deleteLastMessage(ctx); err != nil {
+			return fmt.Errorf("failed to delete message: %w", err)
+		}
+
+		for _, jobStarted := range jobsStarted {
+			if err := handler.HandleJobStarted(ctx, jobStarted); err != nil {
+				return fmt.Errorf("failed to handle job started: %w", err)
+			}
+		}
+
+		if err := handler.HandleDesiredRunnerCount(ctx, statistics.TotalAssignedJobs); err != nil {
+			return fmt.Errorf("failed to handle desired runner count: %w", err)
+		}
+	}
+}
+
+func (l *Listener) createSession(ctx context.Context) error {
+	var session *actions.RunnerScaleSetSession
+	var retries int
+
+	for {
+		var err error
+		session, err = l.client.CreateMessageSession(ctx, l.scaleSetID, l.hostname)
+		if err == nil {
+			break
+		}
+
+		clientErr := &actions.HttpClientSideError{}
+		if !errors.As(err, &clientErr) {
+			return fmt.Errorf("failed to create session: %w", err)
+		}
+
+		if clientErr.Code != http.StatusConflict {
+			return fmt.Errorf("failed to create session: %w", err)
+		}
+
+		retries++
+		if retries >= sessionCreationMaxRetries {
+			return fmt.Errorf("failed to create session after %d retries: %w", retries, err)
+		}
+
+		l.logger.Info("Unable to create message session. Will try again in 30 seconds", "error", err.Error())
+
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("context cancelled: %w", ctx.Err())
+		case <-time.After(30 * time.Second):
+		}
+	}
+
+	statistics, err := json.Marshal(session.Statistics)
+	if err != nil {
+		return fmt.Errorf("failed to marshal statistics: %w", err)
+	}
+	l.logger.Info("Current runner scale set statistics.", "statistics", string(statistics))
+
+	l.session = session
+
+	return nil
+}
+
+func (l *Listener) getMessage(ctx context.Context) (*actions.RunnerScaleSetMessage, error) {
+	l.logger.Info("Getting next message", "lastMessageID", l.lastMessageID)
+	msg, err := l.client.GetMessage(ctx, l.session.MessageQueueUrl, l.session.MessageQueueAccessToken, l.lastMessageID)
+	if err == nil { // if NO error
+		return msg, nil
+	}
+
+	expiredError := &actions.MessageQueueTokenExpiredError{}
+	if !errors.As(err, &expiredError) {
+		return nil, fmt.Errorf("failed to get next message: %w", err)
+	}
+
+	if err := l.refreshSession(ctx); err != nil {
+		return nil, err
+	}
+
+	l.logger.Info("Getting next message", "lastMessageID", l.lastMessageID)
+
+	msg, err = l.client.GetMessage(ctx, l.session.MessageQueueUrl, l.session.MessageQueueAccessToken, l.lastMessageID)
+	if err != nil { // if NO error
+		return nil, fmt.Errorf("failed to get next message after message session refresh: %w", err)
+	}
+
+	return msg, nil
+
+}
+
+func (l *Listener) deleteLastMessage(ctx context.Context) error {
+	l.logger.Info("Deleting last message", "lastMessageID", l.lastMessageID)
+	if err := l.client.DeleteMessage(ctx, l.session.MessageQueueUrl, l.session.MessageQueueAccessToken, l.lastMessageID); err != nil {
+		return fmt.Errorf("failed to delete message: %w", err)
+	}
+
+	return nil
+}
+
+func (l *Listener) parseMessage(ctx context.Context, msg *actions.RunnerScaleSetMessage) (*actions.RunnerScaleSetStatistic, []*actions.JobStarted, error) {
+	l.logger.Info("Processing message", "messageId", msg.MessageId, "messageType", msg.MessageType)
+	if msg.Statistics == nil {
+		return nil, nil, fmt.Errorf("invalid message: statistics is nil")
+	}
+
+	l.logger.Info("New runner scale set statistics.", "statistics", msg.Statistics)
+
+	if msg.MessageType != "RunnerScaleSetJobMessages" {
+		l.logger.Info("Skipping message", "messageType", msg.MessageType)
+		return nil, nil, fmt.Errorf("invalid message type: %s", msg.MessageType)
+	}
+
+	var batchedMessages []json.RawMessage
+	if len(msg.Body) > 0 {
+		if err := json.Unmarshal([]byte(msg.Body), &batchedMessages); err != nil {
+			return nil, nil, fmt.Errorf("failed to unmarshal batched messages: %w", err)
+		}
+	}
+
+	var availableJobs []int64
+	var startedJobs []*actions.JobStarted
+	for _, msg := range batchedMessages {
+		var messageType actions.JobMessageType
+		if err := json.Unmarshal(msg, &messageType); err != nil {
+			return nil, nil, fmt.Errorf("failed to decode job message type: %w", err)
+		}
+
+		switch messageType.MessageType {
+		case messageTypeJobAvailable:
+			var jobAvailable actions.JobAvailable
+			if err := json.Unmarshal(msg, &jobAvailable); err != nil {
+				return nil, nil, fmt.Errorf("failed to decode job available: %w", err)
+			}
+
+			l.logger.Info("Job available message received", "jobId", jobAvailable.RunnerRequestId)
+			availableJobs = append(availableJobs, jobAvailable.RunnerRequestId)
+
+		case messageTypeJobAssigned:
+			var jobAssigned actions.JobAssigned
+			if err := json.Unmarshal(msg, &jobAssigned); err != nil {
+				return nil, nil, fmt.Errorf("failed to decode job assigned: %w", err)
+			}
+
+			l.logger.Info("Job assigned message received", "jobId", jobAssigned.RunnerRequestId)
+
+		case messageTypeJobStarted:
+			var jobStarted actions.JobStarted
+			if err := json.Unmarshal(msg, &jobStarted); err != nil {
+				return nil, nil, fmt.Errorf("could not decode job started message. %w", err)
+			}
+			l.logger.Info("Job started message received.", "RequestId", jobStarted.RunnerRequestId, "RunnerId", jobStarted.RunnerId)
+			startedJobs = append(startedJobs, &jobStarted)
+
+		case messageTypeJobCompleted:
+			var jobCompleted actions.JobCompleted
+			if err := json.Unmarshal(msg, &jobCompleted); err != nil {
+				return nil, nil, fmt.Errorf("failed to decode job completed: %w", err)
+			}
+
+			l.logger.Info("Job completed message received.", "RequestId", jobCompleted.RunnerRequestId, "Result", jobCompleted.Result, "RunnerId", jobCompleted.RunnerId, "RunnerName", jobCompleted.RunnerName)
+
+		default:
+			l.logger.Info("unknown job message type.", "messageType", messageType.MessageType)
+		}
+	}
+
+	l.logger.Info("Available jobs.", "count", len(availableJobs), "requestIds", fmt.Sprint(availableJobs))
+	if len(availableJobs) > 0 {
+		acquired, err := l.acquireAvailableJobs(ctx, availableJobs)
+		if err != nil {
+			return nil, nil, err
+		}
+
+		l.logger.Info("Jobs are acquired", "count", len(acquired), "requestIds", fmt.Sprint(acquired))
+	}
+
+	return msg.Statistics, startedJobs, nil
+}
+
+func (l *Listener) acquireAvailableJobs(ctx context.Context, availableJobs []int64) ([]int64, error) {
+	l.logger.Info("Acquiring jobs")
+
+	ids, err := l.client.AcquireJobs(ctx, l.scaleSetID, l.session.MessageQueueAccessToken, availableJobs)
+	if err == nil { // if NO errors
+		return ids, nil
+	}
+
+	expiredError := &actions.MessageQueueTokenExpiredError{}
+	if !errors.As(err, &expiredError) {
+		return nil, fmt.Errorf("failed to acquire jobs: %w", err)
+	}
+
+	if err := l.refreshSession(ctx); err != nil {
+		return nil, err
+	}
+
+	ids, err = l.client.AcquireJobs(ctx, l.scaleSetID, l.session.MessageQueueAccessToken, availableJobs)
+	if err != nil {
+		return nil, fmt.Errorf("failed to acquire jobs after session refresh: %w", err)
+	}
+
+	return ids, nil
+}
+
+func (l *Listener) refreshSession(ctx context.Context) error {
+	l.logger.Info("Message queue token is expired during GetNextMessage, refreshing...")
+	session, err := l.client.RefreshMessageSession(ctx, l.session.RunnerScaleSet.Id, l.session.SessionId)
+	if err != nil {
+		return fmt.Errorf("refresh message session failed. %w", err)
+	}
+
+	l.session = session
+	return nil
+}

cmd/ghalistener/listener/listener_test.go

@@ -0,0 +1,613 @@
+package listener
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"testing"
+	"time"
+
+	listenermocks "github.com/actions/actions-runner-controller/cmd/ghalistener/listener/mocks"
+	"github.com/actions/actions-runner-controller/cmd/ghalistener/metrics"
+	metricsmocks "github.com/actions/actions-runner-controller/cmd/ghalistener/metrics/mocks"
+	"github.com/actions/actions-runner-controller/github/actions"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNew(t *testing.T) {
+	t.Parallel()
+	t.Run("InvalidConfig", func(t *testing.T) {
+		t.Parallel()
+		var config Config
+		_, err := New(config)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("ValidConfig", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			Client:     listenermocks.NewClient(t),
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+		l, err := New(config)
+		assert.Nil(t, err)
+		assert.NotNil(t, l)
+	})
+
+	t.Run("SetStaticMetrics", func(t *testing.T) {
+		t.Parallel()
+
+		metrics := metricsmocks.NewPublisher(t)
+
+		metrics.On("PublishStatic", mock.Anything, mock.Anything).Once()
+
+		config := Config{
+			Client:     listenermocks.NewClient(t),
+			ScaleSetID: 1,
+			Metrics:    metrics,
+		}
+		l, err := New(config)
+		assert.Nil(t, err)
+		assert.NotNil(t, l)
+	})
+}
+
+func TestListener_createSession(t *testing.T) {
+	t.Parallel()
+	t.Run("FailOnce", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+		client.On("CreateMessageSession", ctx, mock.Anything, mock.Anything).Return(nil, assert.AnError).Once()
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		err = l.createSession(ctx)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("FailContext", func(t *testing.T) {
+		t.Parallel()
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+		client.On("CreateMessageSession", ctx, mock.Anything, mock.Anything).Return(nil,
+			&actions.HttpClientSideError{Code: http.StatusConflict}).Once()
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		err = l.createSession(ctx)
+		assert.True(t, errors.Is(err, context.DeadlineExceeded))
+	})
+
+	t.Run("SetsSession", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		uuid := uuid.New()
+		session := &actions.RunnerScaleSetSession{
+			SessionId:               &uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On("CreateMessageSession", mock.Anything, mock.Anything, mock.Anything).Return(session, nil).Once()
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		err = l.createSession(context.Background())
+		assert.Nil(t, err)
+		assert.Equal(t, session, l.session)
+	})
+}
+
+func TestListener_getMessage(t *testing.T) {
+	t.Parallel()
+
+	t.Run("ReceivesMessage", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+		want := &actions.RunnerScaleSetMessage{
+			MessageId: 1,
+		}
+		client.On("GetMessage", ctx, mock.Anything, mock.Anything, mock.Anything).Return(want, nil).Once()
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+		l.session = &actions.RunnerScaleSetSession{}
+
+		got, err := l.getMessage(ctx)
+		assert.Nil(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("NotExpiredError", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+		client.On("GetMessage", ctx, mock.Anything, mock.Anything, mock.Anything).Return(nil, &actions.HttpClientSideError{Code: http.StatusNotFound}).Once()
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		l.session = &actions.RunnerScaleSetSession{}
+
+		_, err = l.getMessage(ctx)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("RefreshAndSucceeds", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		uuid := uuid.New()
+		session := &actions.RunnerScaleSetSession{
+			SessionId:               &uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On("RefreshMessageSession", ctx, mock.Anything, mock.Anything).Return(session, nil).Once()
+
+		client.On("GetMessage", ctx, mock.Anything, mock.Anything, mock.Anything).Return(nil, &actions.MessageQueueTokenExpiredError{}).Once()
+
+		want := &actions.RunnerScaleSetMessage{
+			MessageId: 1,
+		}
+		client.On("GetMessage", ctx, mock.Anything, mock.Anything, mock.Anything).Return(want, nil).Once()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		l.session = &actions.RunnerScaleSetSession{
+			SessionId:      &uuid,
+			RunnerScaleSet: &actions.RunnerScaleSet{},
+		}
+
+		got, err := l.getMessage(ctx)
+		assert.Nil(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("RefreshAndFails", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		uuid := uuid.New()
+		session := &actions.RunnerScaleSetSession{
+			SessionId:               &uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On("RefreshMessageSession", ctx, mock.Anything, mock.Anything).Return(session, nil).Once()
+
+		client.On("GetMessage", ctx, mock.Anything, mock.Anything, mock.Anything).Return(nil, &actions.MessageQueueTokenExpiredError{}).Twice()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		l.session = &actions.RunnerScaleSetSession{
+			SessionId:      &uuid,
+			RunnerScaleSet: &actions.RunnerScaleSet{},
+		}
+
+		got, err := l.getMessage(ctx)
+		assert.NotNil(t, err)
+		assert.Nil(t, got)
+	})
+}
+
+func TestListener_refreshSession(t *testing.T) {
+	t.Parallel()
+
+	t.Run("SuccessfullyRefreshes", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		newUUID := uuid.New()
+		session := &actions.RunnerScaleSetSession{
+			SessionId:               &newUUID,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On("RefreshMessageSession", ctx, mock.Anything, mock.Anything).Return(session, nil).Once()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		oldUUID := uuid.New()
+		l.session = &actions.RunnerScaleSetSession{
+			SessionId:      &oldUUID,
+			RunnerScaleSet: &actions.RunnerScaleSet{},
+		}
+
+		err = l.refreshSession(ctx)
+		assert.Nil(t, err)
+		assert.Equal(t, session, l.session)
+	})
+
+	t.Run("FailsToRefresh", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		client.On("RefreshMessageSession", ctx, mock.Anything, mock.Anything).Return(nil, errors.New("error")).Once()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		oldUUID := uuid.New()
+		oldSession := &actions.RunnerScaleSetSession{
+			SessionId:      &oldUUID,
+			RunnerScaleSet: &actions.RunnerScaleSet{},
+		}
+		l.session = oldSession
+
+		err = l.refreshSession(ctx)
+		assert.NotNil(t, err)
+		assert.Equal(t, oldSession, l.session)
+	})
+}
+
+func TestListener_deleteLastMessage(t *testing.T) {
+	t.Parallel()
+
+	t.Run("SuccessfullyDeletes", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		client.On("DeleteMessage", ctx, mock.Anything, mock.Anything, mock.MatchedBy(func(lastMessageID any) bool {
+			return lastMessageID.(int64) == int64(5)
+		})).Return(nil).Once()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		l.session = &actions.RunnerScaleSetSession{}
+		l.lastMessageID = 5
+
+		err = l.deleteLastMessage(ctx)
+		assert.Nil(t, err)
+	})
+
+	t.Run("FailsToDelete", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		client.On("DeleteMessage", ctx, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")).Once()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		l.session = &actions.RunnerScaleSetSession{}
+		l.lastMessageID = 5
+
+		err = l.deleteLastMessage(ctx)
+		assert.NotNil(t, err)
+	})
+}
+
+func TestListener_Listen(t *testing.T) {
+	t.Parallel()
+
+	t.Run("CreateSessionFails", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+		client.On("CreateMessageSession", ctx, mock.Anything, mock.Anything).Return(nil, assert.AnError).Once()
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		err = l.Listen(ctx, nil)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("CallHandleRegardlessOfInitialMessage", func(t *testing.T) {
+		t.Parallel()
+		ctx, cancel := context.WithCancel(context.Background())
+
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		uuid := uuid.New()
+		session := &actions.RunnerScaleSetSession{
+			SessionId:               &uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              &actions.RunnerScaleSetStatistic{},
+		}
+		client.On("CreateMessageSession", ctx, mock.Anything, mock.Anything).Return(session, nil).Once()
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		var called bool
+		handler := listenermocks.NewHandler(t)
+		handler.On("HandleDesiredRunnerCount", mock.Anything, mock.Anything).
+			Return(nil).
+			Run(
+				func(mock.Arguments) {
+					called = true
+					cancel()
+				},
+			).
+			Once()
+
+		err = l.Listen(ctx, handler)
+		assert.True(t, errors.Is(err, context.Canceled))
+		assert.True(t, called)
+	})
+}
+
+func TestListener_acquireAvailableJobs(t *testing.T) {
+	t.Parallel()
+
+	t.Run("FailingToAcquireJobs", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		client.On("AcquireJobs", ctx, mock.Anything, mock.Anything, mock.Anything).Return(nil, assert.AnError).Once()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		uuid := uuid.New()
+		l.session = &actions.RunnerScaleSetSession{
+			SessionId:               &uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              &actions.RunnerScaleSetStatistic{},
+		}
+
+		_, err = l.acquireAvailableJobs(ctx, []int64{1, 2, 3})
+		assert.Error(t, err)
+	})
+
+	t.Run("SuccessfullyAcquiresJobsOnFirstRun", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		jobIDs := []int64{1, 2, 3}
+
+		client.On("AcquireJobs", ctx, mock.Anything, mock.Anything, mock.Anything).Return(jobIDs, nil).Once()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		uuid := uuid.New()
+		l.session = &actions.RunnerScaleSetSession{
+			SessionId:               &uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              &actions.RunnerScaleSetStatistic{},
+		}
+
+		acquiredJobIDs, err := l.acquireAvailableJobs(ctx, []int64{1, 2, 3})
+		assert.NoError(t, err)
+		assert.Equal(t, jobIDs, acquiredJobIDs)
+	})
+
+	t.Run("RefreshAndSucceeds", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		uuid := uuid.New()
+		session := &actions.RunnerScaleSetSession{
+			SessionId:               &uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On("RefreshMessageSession", ctx, mock.Anything, mock.Anything).Return(session, nil).Once()
+
+		// First call to AcquireJobs will fail with a token expired error
+		client.On("AcquireJobs", ctx, mock.Anything, mock.Anything, mock.Anything).Return(nil, &actions.MessageQueueTokenExpiredError{}).Once()
+
+		// Second call to AcquireJobs will succeed
+		want := []int64{1, 2, 3}
+		client.On("AcquireJobs", ctx, mock.Anything, mock.Anything, mock.Anything).Return(want, nil).Once()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		l.session = &actions.RunnerScaleSetSession{
+			SessionId:      &uuid,
+			RunnerScaleSet: &actions.RunnerScaleSet{},
+		}
+
+		got, err := l.acquireAvailableJobs(ctx, want)
+		assert.Nil(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("RefreshAndFails", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			Metrics:    metrics.Discard,
+		}
+
+		client := listenermocks.NewClient(t)
+
+		uuid := uuid.New()
+		session := &actions.RunnerScaleSetSession{
+			SessionId:               &uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &actions.RunnerScaleSet{},
+			MessageQueueUrl:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On("RefreshMessageSession", ctx, mock.Anything, mock.Anything).Return(session, nil).Once()
+
+		client.On("AcquireJobs", ctx, mock.Anything, mock.Anything, mock.Anything).Return(nil, &actions.MessageQueueTokenExpiredError{}).Twice()
+
+		config.Client = client
+
+		l, err := New(config)
+		require.Nil(t, err)
+
+		l.session = &actions.RunnerScaleSetSession{
+			SessionId:      &uuid,
+			RunnerScaleSet: &actions.RunnerScaleSet{},
+		}
+
+		got, err := l.acquireAvailableJobs(ctx, []int64{1, 2, 3})
+		assert.NotNil(t, err)
+		assert.Nil(t, got)
+	})
+}

cmd/ghalistener/listener/mocks/client.go

@@ -0,0 +1,176 @@
+// Code generated by mockery v2.36.1. DO NOT EDIT.
+
+package mocks
+
+import (
+	context "context"
+
+	actions "github.com/actions/actions-runner-controller/github/actions"
+
+	mock "github.com/stretchr/testify/mock"
+
+	uuid "github.com/google/uuid"
+)
+
+// Client is an autogenerated mock type for the Client type
+type Client struct {
+	mock.Mock
+}
+
+// AcquireJobs provides a mock function with given fields: ctx, runnerScaleSetId, messageQueueAccessToken, requestIds
+func (_m *Client) AcquireJobs(ctx context.Context, runnerScaleSetId int, messageQueueAccessToken string, requestIds []int64) ([]int64, error) {
+	ret := _m.Called(ctx, runnerScaleSetId, messageQueueAccessToken, requestIds)
+
+	var r0 []int64
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int, string, []int64) ([]int64, error)); ok {
+		return rf(ctx, runnerScaleSetId, messageQueueAccessToken, requestIds)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int, string, []int64) []int64); ok {
+		r0 = rf(ctx, runnerScaleSetId, messageQueueAccessToken, requestIds)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]int64)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int, string, []int64) error); ok {
+		r1 = rf(ctx, runnerScaleSetId, messageQueueAccessToken, requestIds)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// CreateMessageSession provides a mock function with given fields: ctx, runnerScaleSetId, owner
+func (_m *Client) CreateMessageSession(ctx context.Context, runnerScaleSetId int, owner string) (*actions.RunnerScaleSetSession, error) {
+	ret := _m.Called(ctx, runnerScaleSetId, owner)
+
+	var r0 *actions.RunnerScaleSetSession
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int, string) (*actions.RunnerScaleSetSession, error)); ok {
+		return rf(ctx, runnerScaleSetId, owner)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int, string) *actions.RunnerScaleSetSession); ok {
+		r0 = rf(ctx, runnerScaleSetId, owner)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*actions.RunnerScaleSetSession)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int, string) error); ok {
+		r1 = rf(ctx, runnerScaleSetId, owner)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// DeleteMessage provides a mock function with given fields: ctx, messageQueueUrl, messageQueueAccessToken, messageId
+func (_m *Client) DeleteMessage(ctx context.Context, messageQueueUrl string, messageQueueAccessToken string, messageId int64) error {
+	ret := _m.Called(ctx, messageQueueUrl, messageQueueAccessToken, messageId)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) error); ok {
+		r0 = rf(ctx, messageQueueUrl, messageQueueAccessToken, messageId)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// GetAcquirableJobs provides a mock function with given fields: ctx, runnerScaleSetId
+func (_m *Client) GetAcquirableJobs(ctx context.Context, runnerScaleSetId int) (*actions.AcquirableJobList, error) {
+	ret := _m.Called(ctx, runnerScaleSetId)
+
+	var r0 *actions.AcquirableJobList
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int) (*actions.AcquirableJobList, error)); ok {
+		return rf(ctx, runnerScaleSetId)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int) *actions.AcquirableJobList); ok {
+		r0 = rf(ctx, runnerScaleSetId)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*actions.AcquirableJobList)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int) error); ok {
+		r1 = rf(ctx, runnerScaleSetId)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetMessage provides a mock function with given fields: ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId
+func (_m *Client) GetMessage(ctx context.Context, messageQueueUrl string, messageQueueAccessToken string, lastMessageId int64) (*actions.RunnerScaleSetMessage, error) {
+	ret := _m.Called(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId)
+
+	var r0 *actions.RunnerScaleSetMessage
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) (*actions.RunnerScaleSetMessage, error)); ok {
+		return rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) *actions.RunnerScaleSetMessage); ok {
+		r0 = rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*actions.RunnerScaleSetMessage)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, string, string, int64) error); ok {
+		r1 = rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// RefreshMessageSession provides a mock function with given fields: ctx, runnerScaleSetId, sessionId
+func (_m *Client) RefreshMessageSession(ctx context.Context, runnerScaleSetId int, sessionId *uuid.UUID) (*actions.RunnerScaleSetSession, error) {
+	ret := _m.Called(ctx, runnerScaleSetId, sessionId)
+
+	var r0 *actions.RunnerScaleSetSession
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int, *uuid.UUID) (*actions.RunnerScaleSetSession, error)); ok {
+		return rf(ctx, runnerScaleSetId, sessionId)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int, *uuid.UUID) *actions.RunnerScaleSetSession); ok {
+		r0 = rf(ctx, runnerScaleSetId, sessionId)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*actions.RunnerScaleSetSession)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int, *uuid.UUID) error); ok {
+		r1 = rf(ctx, runnerScaleSetId, sessionId)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// NewClient creates a new instance of Client. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Client {
+	mock := &Client{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

cmd/ghalistener/listener/mocks/handler.go

@@ -0,0 +1,58 @@
+// Code generated by mockery v2.36.1. DO NOT EDIT.
+
+package mocks
+
+import (
+	context "context"
+
+	actions "github.com/actions/actions-runner-controller/github/actions"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// Handler is an autogenerated mock type for the Handler type
+type Handler struct {
+	mock.Mock
+}
+
+// HandleDesiredRunnerCount provides a mock function with given fields: ctx, desiredRunnerCount
+func (_m *Handler) HandleDesiredRunnerCount(ctx context.Context, desiredRunnerCount int) error {
+	ret := _m.Called(ctx, desiredRunnerCount)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, int) error); ok {
+		r0 = rf(ctx, desiredRunnerCount)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// HandleJobStarted provides a mock function with given fields: ctx, jobInfo
+func (_m *Handler) HandleJobStarted(ctx context.Context, jobInfo *actions.JobStarted) error {
+	ret := _m.Called(ctx, jobInfo)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *actions.JobStarted) error); ok {
+		r0 = rf(ctx, jobInfo)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// NewHandler creates a new instance of Handler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewHandler(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Handler {
+	mock := &Handler{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}