diff --git a/client.go b/client.go index dd31472..8ef78e4 100644 --- a/client.go +++ b/client.go @@ -631,14 +631,75 @@ func (c *Client) GetMessage(ctx context.Context, messageQueueURL, messageQueueAc } } - var message *RunnerScaleSetMessage - if err := json.NewDecoder(resp.Body).Decode(&message); err != nil { + message, err := c.parseRunnerScaleSetMessageResponse(resp.Body) + if err != nil { return nil, &ActionsError{ StatusCode: resp.StatusCode, ActivityID: resp.Header.Get(headerActionsActivityID), Err: err, } } + + return message, nil +} + +func (c *Client) parseRunnerScaleSetMessageResponse(respBody io.Reader) (*RunnerScaleSetMessage, error) { + var messageResponse runnerScaleSetMessageResponse + if err := json.NewDecoder(respBody).Decode(&messageResponse); err != nil { + return nil, fmt.Errorf("failed to decode runner scale set message response: %w", err) + } + + if messageResponse.MessageType != "RunnerScaleSetJobMessages" { + return nil, fmt.Errorf("unsupported message type: %s", messageResponse.MessageType) + } + + message := &RunnerScaleSetMessage{ + MessageID: messageResponse.MessageID, + Statistics: messageResponse.Statistics, + } + + var batchedMessages []json.RawMessage + if len(messageResponse.Body) > 0 { + if err := json.Unmarshal([]byte(messageResponse.Body), &batchedMessages); err != nil { + return nil, fmt.Errorf("failed to unmarshal batched messages: %w", err) + } + } + + for _, msg := range batchedMessages { + var messageType JobMessageType + if err := json.Unmarshal(msg, &messageType); err != nil { + return nil, fmt.Errorf("failed to decode job message type: %w", err) + } + + switch messageType.MessageType { + case MessageTypeJobAssigned: + var jobAssigned JobAssigned + if err := json.Unmarshal(msg, &jobAssigned); err != nil { + return nil, fmt.Errorf("failed to decode job assigned: %w", err) + } + + message.JobAssignedMessages = append(message.JobAssignedMessages, &jobAssigned) + + case MessageTypeJobStarted: + var jobStarted JobStarted + if err := json.Unmarshal(msg, &jobStarted); err != nil { + return nil, fmt.Errorf("could not decode job started message. %w", err) + } + + message.JobStartedMessages = append(message.JobStartedMessages, &jobStarted) + + case MessageTypeJobCompleted: + var jobCompleted JobCompleted + if err := json.Unmarshal(msg, &jobCompleted); err != nil { + return nil, fmt.Errorf("failed to decode job completed: %w", err) + } + + message.JobCompletedMessages = append(message.JobCompletedMessages, &jobCompleted) + + default: + } + } + return message, nil } diff --git a/client_test.go b/client_test.go index 5e99922..b1db86b 100644 --- a/client_test.go +++ b/client_test.go @@ -1085,13 +1085,12 @@ func TestGetMessage(t *testing.T) { token := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwiaWF0IjoxNTE2MjM5MDIyLCJleHAiOjI1MTYyMzkwMjJ9.tlrHslTmDkoqnc4Kk9ISoKoUNDfHo-kjlH-ByISBqzE" runnerScaleSetMessage := &RunnerScaleSetMessage{ - MessageID: 1, - MessageType: "rssType", + MessageID: 1, } t.Run("Get Runner Scale Set Message", func(t *testing.T) { want := runnerScaleSetMessage - response := []byte(`{"messageId":1,"messageType":"rssType"}`) + response := []byte(`{"messageId":1,"messageType":"RunnerScaleSetJobMessages"}`) s := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Write(response) })) @@ -1106,7 +1105,7 @@ func TestGetMessage(t *testing.T) { t.Run("GetMessage sets the last message id if not 0", func(t *testing.T) { want := runnerScaleSetMessage - response := []byte(`{"messageId":1,"messageType":"rssType"}`) + response := []byte(`{"messageId":1,"messageType":"RunnerScaleSetJobMessages"}`) s := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { q := r.URL.Query() assert.Equal(t, "1", q.Get("lastMessageId")) @@ -1220,8 +1219,7 @@ func TestDeleteMessage(t *testing.T) { token := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwiaWF0IjoxNTE2MjM5MDIyLCJleHAiOjI1MTYyMzkwMjJ9.tlrHslTmDkoqnc4Kk9ISoKoUNDfHo-kjlH-ByISBqzE" runnerScaleSetMessage := &RunnerScaleSetMessage{ - MessageID: 1, - MessageType: "rssType", + MessageID: 1, } t.Run("Delete existing message", func(t *testing.T) { diff --git a/error.go b/error.go index 062fd0b..c2c4b17 100644 --- a/error.go +++ b/error.go @@ -36,7 +36,7 @@ type ActionsError struct { } func (e *ActionsError) Error() string { - return fmt.Sprintf("actions error: StatusCode %d, AcivityId %q: %v", e.StatusCode, e.ActivityID, e.Err) + return fmt.Sprintf("actions error: StatusCode %d, ActivityId %q: %v", e.StatusCode, e.ActivityID, e.Err) } func (e *ActionsError) Unwrap() error { @@ -111,7 +111,7 @@ type MessageQueueTokenExpiredError struct { } func (e *MessageQueueTokenExpiredError) Error() string { - return fmt.Sprintf("MessageQueueTokenExpiredError: AcivityId %q, StatusCode %d: %s", e.activityID, e.statusCode, e.msg) + return fmt.Sprintf("MessageQueueTokenExpiredError: ActivityId %q, StatusCode %d: %s", e.activityID, e.statusCode, e.msg) } type HttpClientSideError struct { diff --git a/errors_test.go b/errors_test.go index eb31a69..07c0a69 100644 --- a/errors_test.go +++ b/errors_test.go @@ -21,7 +21,7 @@ func TestActionsError(t *testing.T) { s := err.Error() assert.Contains(t, s, "StatusCode 404") - assert.Contains(t, s, "AcivityId \"activity-id\"") + assert.Contains(t, s, "ActivityId \"activity-id\"") assert.Contains(t, s, "example error description") }) diff --git a/listener/listener.go b/listener/listener.go index 972465a..e06618f 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -160,29 +160,23 @@ func (l *Listener) Run(ctx context.Context, scaler Scaler) error { } func (l *Listener) handleMessage(ctx context.Context, handler Scaler, msg *scaleset.RunnerScaleSetMessage) error { - parsedMsg, 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 parsedMsg.jobsStarted { + for _, jobStarted := range msg.JobStartedMessages { if err := handler.HandleJobStarted(ctx, jobStarted); err != nil { return fmt.Errorf("failed to handle job started: %w", err) } } - for _, jobCompleted := range parsedMsg.jobsCompleted { + for _, jobCompleted := range msg.JobCompletedMessages { if err := handler.HandleJobCompleted(ctx, jobCompleted); err != nil { return fmt.Errorf("failed to handle job completed: %w", err) } } - if _, err := handler.HandleDesiredRunnerCount(ctx, parsedMsg.statistics.TotalAssignedJobs); err != nil { + if _, err := handler.HandleDesiredRunnerCount(ctx, msg.Statistics.TotalAssignedJobs); err != nil { return fmt.Errorf("failed to handle desired runner count: %w", err) } @@ -312,76 +306,6 @@ type parsedMessage struct { jobsCompleted []*scaleset.JobCompleted } -func (l *Listener) parseMessage(ctx context.Context, msg *scaleset.RunnerScaleSetMessage) (*parsedMessage, error) { - if msg.MessageType != "RunnerScaleSetJobMessages" { - l.logger.Info("Skipping message", "messageType", msg.MessageType) - return nil, fmt.Errorf("invalid message type: %s", msg.MessageType) - } - - l.logger.Info("Processing message", "messageId", msg.MessageID, "messageType", msg.MessageType) - if msg.Statistics == nil { - return nil, fmt.Errorf("invalid message: statistics is nil") - } - - l.logger.Info("New runner scale set statistics.", "statistics", msg.Statistics) - - var batchedMessages []json.RawMessage - if len(msg.Body) > 0 { - if err := json.Unmarshal([]byte(msg.Body), &batchedMessages); err != nil { - return nil, fmt.Errorf("failed to unmarshal batched messages: %w", err) - } - } - - parsedMsg := &parsedMessage{ - statistics: msg.Statistics, - } - - for _, msg := range batchedMessages { - var messageType scaleset.JobMessageType - if err := json.Unmarshal(msg, &messageType); err != nil { - return nil, fmt.Errorf("failed to decode job message type: %w", err) - } - - switch messageType.MessageType { - case scaleset.MessageTypeJobAssigned: - var jobAssigned scaleset.JobAssigned - if err := json.Unmarshal(msg, &jobAssigned); err != nil { - return nil, fmt.Errorf("failed to decode job assigned: %w", err) - } - - l.logger.Info("Job assigned message received", "jobId", jobAssigned.JobID) - - case scaleset.MessageTypeJobStarted: - var jobStarted scaleset.JobStarted - if err := json.Unmarshal(msg, &jobStarted); err != nil { - return nil, fmt.Errorf("could not decode job started message. %w", err) - } - l.logger.Info("Job started message received.", "JobID", jobStarted.JobID, "RunnerId", jobStarted.RunnerID) - parsedMsg.jobsStarted = append(parsedMsg.jobsStarted, &jobStarted) - - case scaleset.MessageTypeJobCompleted: - var jobCompleted scaleset.JobCompleted - if err := json.Unmarshal(msg, &jobCompleted); err != nil { - return nil, fmt.Errorf("failed to decode job completed: %w", err) - } - - l.logger.Info( - "Job completed message received.", - "JobID", jobCompleted.JobID, - "Result", jobCompleted.Result, - "RunnerId", jobCompleted.RunnerID, - "RunnerName", jobCompleted.RunnerName, - ) - parsedMsg.jobsCompleted = append(parsedMsg.jobsCompleted, &jobCompleted) - - default: - l.logger.Info("unknown job message type.", "messageType", messageType.MessageType) - } - } - - return parsedMsg, ctx.Err() -} - func (l *Listener) refreshSession(ctx context.Context) error { l.logger.Info("Message queue token is expired during GetNextMessage, refreshing...") session, err := l.client.RefreshMessageSession( diff --git a/types.go b/types.go index fec367c..bfc3e77 100644 --- a/types.go +++ b/types.go @@ -99,13 +99,21 @@ type RunnerScaleSetJitRunnerSetting struct { WorkFolder string `json:"workFolder"` } -type RunnerScaleSetMessage struct { +type runnerScaleSetMessageResponse struct { MessageID uint64 `json:"messageId"` MessageType string `json:"messageType"` Body string `json:"body"` Statistics *RunnerScaleSetStatistic `json:"statistics"` } +type RunnerScaleSetMessage struct { + MessageID uint64 + Statistics *RunnerScaleSetStatistic + JobAssignedMessages []*JobAssigned + JobStartedMessages []*JobStarted + JobCompletedMessages []*JobCompleted +} + type runnerScaleSetsResponse struct { Count uint64 `json:"count"` RunnerScaleSets []RunnerScaleSet `json:"value"`