Skip to content
116 changes: 112 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,99 @@ 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)
}

// TODO: determine reasonable behavior for internal errors (e.g., signing or marshal failures).
// 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.

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.
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("internal error: cannot find subsystem change %q", msg.ApplyChangeID)
}
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