Restore job acquisition flow (#90)

Author: rentziass

client.go

@@ -510,6 +510,14 @@ func parseRunnerScaleSetMessageResponse(respBody io.Reader) (*RunnerScaleSetMess
 		}
 
 		switch messageType.MessageType {
+		case MessageTypeJobAvailable:
+			var jobAvailable JobAvailable
+			if err := json.Unmarshal(msg, &jobAvailable); err != nil {
+				return nil, fmt.Errorf("failed to decode job available: %w", err)
+			}
+
+			message.JobAvailableMessages = append(message.JobAvailableMessages, &jobAvailable)
+
 		case MessageTypeJobAssigned:
 			var jobAssigned JobAssigned
 			if err := json.Unmarshal(msg, &jobAssigned); err != nil {

listener/listener.go

@@ -49,6 +49,7 @@ func (c *Config) Validate() error {
 type Client interface {
 	GetMessage(ctx context.Context, lastMessageID, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error)
 	DeleteMessage(ctx context.Context, messageID int) error
+	AcquireJobs(ctx context.Context, requestIDs []int64) ([]int64, error)
 	Session() scaleset.RunnerScaleSetSession
 }
 
@@ -210,6 +211,12 @@ func (l *Listener) handleMessage(ctx context.Context, handler Scaler, msg *scale
 		return fmt.Errorf("failed to delete message: %w", err)
 	}
 
+	if len(msg.JobAvailableMessages) > 0 {
+		if err := l.acquireAvailableJobs(ctx, msg.JobAvailableMessages); err != nil {
+			return fmt.Errorf("failed to acquire available jobs: %w", err)
+		}
+	}
+
 	for _, jobStarted := range msg.JobStartedMessages {
 		l.metricsRecorder.RecordJobStarted(jobStarted)
 		if err := handler.HandleJobStarted(ctx, jobStarted); err != nil {
@@ -232,6 +239,23 @@ func (l *Listener) handleMessage(ctx context.Context, handler Scaler, msg *scale
 	return nil
 }
 
+func (l *Listener) acquireAvailableJobs(ctx context.Context, jobsAvailable []*scaleset.JobAvailable) error {
+	ids := make([]int64, 0, len(jobsAvailable))
+	for _, job := range jobsAvailable {
+		ids = append(ids, job.RunnerRequestID)
+	}
+
+	l.logger.Info("Acquiring jobs", slog.Int("count", len(ids)))
+
+	acquired, err := l.client.AcquireJobs(ctx, ids)
+	if err != nil {
+		return fmt.Errorf("acquiring jobs: %w", err)
+	}
+
+	l.logger.Info("Jobs acquired", slog.Int("count", len(acquired)))
+	return nil
+}
+
 func (l *Listener) handleStatistics(ctx context.Context, msg *scaleset.RunnerScaleSetStatistic) {
 	l.latestStatistics = msg
 	l.metricsRecorder.RecordStatistics(msg)

listener/mocks_test.go

@@ -234,6 +234,67 @@ func (c *MessageSessionClient) Session() RunnerScaleSetSession {
 	return *c.session
 }
 
+// AcquireJobs acquires the given job request IDs from the runner scale set.
+// If the current session token is expired, it refreshes the session and tries one more time.
+func (c *MessageSessionClient) AcquireJobs(ctx context.Context, requestIDs []int64) ([]int64, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	ids, err := c.acquireJobs(ctx, requestIDs)
+	if err == nil {
+		return ids, nil
+	}
+
+	if !errors.Is(err, MessageQueueTokenExpiredError) {
+		return nil, fmt.Errorf("failed to acquire jobs: %w", err)
+	}
+
+	if err := c.refreshMessageSession(ctx); err != nil {
+		return nil, fmt.Errorf("failed to refresh message session: %w", err)
+	}
+
+	return c.acquireJobs(ctx, requestIDs)
+}
+
+func (c *MessageSessionClient) acquireJobs(ctx context.Context, requestIDs []int64) ([]int64, error) {
+	body, err := json.Marshal(requestIDs)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal request ids: %w", err)
+	}
+
+	path := fmt.Sprintf("/%s/%d/acquirejobs", scaleSetEndpoint, c.scaleSetID)
+
+	c.innerClient.mu.Lock()
+	req, err := c.innerClient.newActionsServiceRequest(ctx, http.MethodPost, path, bytes.NewBuffer(body))
+	c.innerClient.mu.Unlock()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create acquire jobs request: %w", err)
+	}
+
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.session.MessageQueueAccessToken))
+
+	resp, err := c.commonClient.do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to issue acquire jobs request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusUnauthorized {
+		return nil, newRequestResponseError(req, resp, MessageQueueTokenExpiredError)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, newRequestResponseError(req, resp, fmt.Errorf("unexpected status code %s", resp.Status))
+	}
+
+	var result acquireJobsResponse
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, newRequestResponseError(req, resp, fmt.Errorf("failed to decode acquire jobs response: %w", err))
+	}
+
+	return result.Value, nil
+}
+
 func (c *MessageSessionClient) doSessionRequest(ctx context.Context, method, path string, requestData io.Reader, expectedResponseStatusCode int, responseUnmarshalTarget any) error {
 	c.innerClient.mu.Lock()
 	defer c.innerClient.mu.Unlock()