Skip to content
112 changes: 108 additions & 4 deletions overlord/devicemgmtstate/devicemgmtmgr.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
package devicemgmtstate

import (
"encoding/json"
"errors"
"fmt"
"sort"
Expand Down Expand Up @@ -58,6 +59,8 @@ var (
maxSequences = 256
maxBlockedMessagesPerSequence = 8

awaitSubsystemRetryInterval = 30 * time.Second

deviceMgmtExchangeChangeKind = swfeats.RegisterChangeKind("device-management-exchange")
)

Expand Down Expand Up @@ -180,6 +183,23 @@ func (ms *deviceMgmtState) getRequestMessage(id string) (*RequestMessage, error)
return nil, fmt.Errorf("cannot find message %q", id)
}

// removeRequestMessage removes a processed request message from its sequence,
// leaving the sequence entry in place so its Applied progress is preserved for
// later messages in the same sequence.
func (ms *deviceMgmtState) removeRequestMessage(msg *RequestMessage) {
seq := ms.Sequences[msg.BaseID]
if seq == nil {
return
}

for i, m := range seq.Messages {
if m.SeqNum == msg.SeqNum {
seq.Messages = append(seq.Messages[:i], seq.Messages[i+1:]...)
return
}
}
}

// enqueueRequestMessages queues incoming request messages for processing
// and updates polling state accordingly.
func (ms *deviceMgmtState) enqueueRequestMessages(pollResp *store.MessageExchangeResponse) {
Expand Down Expand Up @@ -479,8 +499,6 @@ func (m *DeviceMgmtManager) dispatchSequence(dispatchTask *state.Task, seq *sequ
// the final task so callers can chain subsequent messages after it.
func (m *DeviceMgmtManager) dispatchMessage(prevTask *state.Task, msg *RequestMessage) *state.Task {
chg := prevTask.Change()
// TODO: add tests verifying that a failure in one message's task chain does not
// affect other messages (lanes provide this isolation, but it needs test coverage).
lane := m.state.NewLane()

addTask := func(kind, summary string) {
Expand Down Expand Up @@ -586,9 +604,95 @@ func (m *DeviceMgmtManager) doApplyMessage(t *state.Task, _ *tomb.Tomb) error {
}

// doQueueResponse builds a response, signs it, and queues it for transmission on the next exchange.
// Retries until subsystem change completes.
// Retries until the subsystem change (if any) completes.
func (m *DeviceMgmtManager) doQueueResponse(t *state.Task, _ *tomb.Tomb) error {
// TODO: implement this task, no-op for now.
m.state.Lock()
defer m.state.Unlock()

ms, err := m.getState()
if err != nil {
return err
}

var msgID string
err = t.Get("message-id", &msgID)
if err != nil {
return err
}

msg, err := ms.getRequestMessage(msgID)
if err != nil {
// Message already processed on a prior run.
return nil
}
Comment thread
st3v3nmw marked this conversation as resolved.

err = m.setMessageResponseFromChange(msg)
if err != nil {
return err
}

bodyBytes, err := json.Marshal(msg.ResponseBody)
if err != nil {
return fmt.Errorf("cannot marshal response body: %w", err)
}

resAs, err := m.signer.SignResponseMessage(msg.AccountID, msg.ID(), msg.ResponseStatus, bodyBytes)
if err != nil {
return fmt.Errorf("cannot sign response message: %w", err)
}
Comment thread
miguelpires marked this conversation as resolved.

ms.ReadyResponses[msg.ID()] = store.Message{
Format: "assertion",
Data: string(asserts.Encode(resAs)),
}

// TODO: rejecting sequences currently happens in 2 ways:
// 1. doDispatchMessage can evict the sequence immediately if it's rejected early.
// 2. If it errors elsewhere (in validate, apply, or queue-response), we end
// up not advancing Applied, which means we accumulate messages until we
// hit the sequence cap.
// Refactor sequence rejection to always evict immediately and determine the
// right behavior for internal errors.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we definitely need to see what to do for internal errors?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also the comment makes it sound like it applies only to sequences, but reasonable behavior for internal error is also important for not-sequenced messages

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we definitely need to see what to do for internal errors?

Yes. Since tasks are idempotent, a failed message will not be re-dispatched or re-applied on the next change, but the request message remains in state until doQueueResponse completes, so such failures leave it hanging indefinitely.

also the comment makes it sound like it applies only to sequences, but reasonable behavior for internal error is also important for not-sequenced messages

I've fixed it to be clearer.

if msg.SeqNum > 0 && msg.ResponseStatus == asserts.MessageStatusSuccess {
ms.Sequences[msg.BaseID].Applied = msg.SeqNum
}
ms.removeRequestMessage(msg)

m.setState(ms)

return nil
}

// setMessageResponseFromChange populates msg's response fields from the completed apply change.
func (m *DeviceMgmtManager) setMessageResponseFromChange(msg *RequestMessage) error {
if msg.ResponseStatus != "" {
return nil
}

handler, ok := m.handlers[msg.Kind]
if !ok {
msg.ResponseStatus = asserts.MessageStatusError
msg.ResponseBody = map[string]any{"message": fmt.Sprintf("cannot find handler for message kind %q", msg.Kind)}
return nil
}

change := m.state.Change(msg.ApplyChangeID)
if change == nil {
return fmt.Errorf("cannot find subsystem change %q", msg.ApplyChangeID)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this seems also an internal error, in which scenario would this happen?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've marked it as internal error: .... The branch is a defensive guard in the rare case the change is not in state anymore. For instance, if a ready subsystem change is pruned before this task fires. I'm looking at the pruning code to see if this is even possible.. say if the computer is shut down mid-device-management-change and then booted after a week.

}
if !change.Status().Ready() {
return &state.Retry{After: awaitSubsystemRetryInterval}
}

body, err := handler.ResultFromChange(change)
if err != nil {
msg.ResponseStatus = asserts.MessageStatusError
msg.ResponseBody = map[string]any{"message": fmt.Sprintf("cannot process message: %v", err)}
} else {
msg.ResponseStatus = asserts.MessageStatusSuccess
msg.ResponseBody = body
}

return nil
}

Expand Down
Loading
Loading