Delete message session when listener.Listen returns (#3240)

cmd/ghalistener/listener/listener.go

@@ -35,6 +35,7 @@ type Client interface {
 	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)
+	DeleteMessageSession(ctx context.Context, runnerScaleSetId int, sessionId *uuid.UUID) error
 }
 
 type Config struct {
@@ -126,6 +127,12 @@ func (l *Listener) Listen(ctx context.Context, handler Handler) error {
 		return fmt.Errorf("createSession failed: %w", err)
 	}
 
+	defer func() {
+		if err := l.deleteMessageSession(); err != nil {
+			l.logger.Error(err, "failed to delete message session")
+		}
+	}()
+
 	initialMessage := &actions.RunnerScaleSetMessage{
 		MessageId:   0,
 		MessageType: "RunnerScaleSetJobMessages",
@@ -409,3 +416,16 @@ func (l *Listener) refreshSession(ctx context.Context) error {
 	l.session = session
 	return nil
 }
+
+func (l *Listener) deleteMessageSession() error {
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	l.logger.Info("Deleting message session")
+
+	if err := l.client.DeleteMessageSession(ctx, l.session.RunnerScaleSet.Id, l.session.SessionId); err != nil {
+		return fmt.Errorf("failed to delete message session: %w", err)
+	}
+
+	return nil
+}

cmd/ghalistener/listener/listener_test.go

@@ -419,6 +419,7 @@ func TestListener_Listen(t *testing.T) {
 			Statistics:              &actions.RunnerScaleSetStatistic{},
 		}
 		client.On("CreateMessageSession", ctx, mock.Anything, mock.Anything).Return(session, nil).Once()
+		client.On("DeleteMessageSession", mock.Anything, session.RunnerScaleSet.Id, session.SessionId).Return(nil).Once()
 		config.Client = client
 
 		l, err := New(config)

cmd/ghalistener/listener/mocks/client.go

@@ -83,6 +83,20 @@ func (_m *Client) DeleteMessage(ctx context.Context, messageQueueUrl string, mes
 	return r0
 }
 
+// DeleteMessageSession provides a mock function with given fields: ctx, runnerScaleSetId, sessionId
+func (_m *Client) DeleteMessageSession(ctx context.Context, runnerScaleSetId int, sessionId *uuid.UUID) error {
+	ret := _m.Called(ctx, runnerScaleSetId, sessionId)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, int, *uuid.UUID) error); ok {
+		r0 = rf(ctx, runnerScaleSetId, sessionId)
+	} 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)

cmd/ghalistener/worker/worker.go

@@ -11,6 +11,7 @@ import (
 	"github.com/actions/actions-runner-controller/logging"
 	jsonpatch "github.com/evanphx/json-patch"
 	"github.com/go-logr/logr"
+	kerrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
@@ -141,6 +142,10 @@ func (w *Worker) HandleJobStarted(ctx context.Context, jobInfo *actions.JobStart
 		Do(ctx).
 		Into(patchedStatus)
 	if err != nil {
+		if kerrors.IsNotFound(err) {
+			w.logger.Info("Ephemeral runner not found, skipping patching of ephemeral runner status", "runnerName", jobInfo.RunnerName)
+			return nil
+		}
 		return fmt.Errorf("could not patch ephemeral runner status, patch JSON: %s, error: %w", string(mergePatch), err)
 	}