fix: reconnect agent after phoenix reload

This commit is contained in:
Graham McIntire 2026-03-24 12:35:44 -05:00
parent f8ea46d075
commit ee86b480d0
No known key found for this signature in database
2 changed files with 58 additions and 16 deletions

View file

@ -22,6 +22,7 @@ import (
var doSelfUpdate = selfUpdate var doSelfUpdate = selfUpdate
var errRestartRequested = fmt.Errorf("restart requested") var errRestartRequested = fmt.Errorf("restart requested")
var errChannelReloaded = fmt.Errorf("channel reloaded")
var joinTimeout = 10 * time.Second var joinTimeout = 10 * time.Second
var heartbeatInterval = 60 * time.Second var heartbeatInterval = 60 * time.Second
var channelHeartbeatInterval = 25 * time.Second var channelHeartbeatInterval = 25 * time.Second
@ -320,9 +321,10 @@ func runSession(ctx context.Context, baseURL, token string) error {
slog.Warn("invalid message", "error", err) slog.Warn("invalid message", "error", err)
continue continue
} }
if handleMessage(sessionCtx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh, checkResultCh, lldpTopologyResultCh) { shouldEnd, endErr := handleMessage(sessionCtx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh, checkResultCh, lldpTopologyResultCh)
if shouldEnd {
flushSnmpBatch() flushSnmpBatch()
return errRestartRequested return endErr
} }
case result := <-snmpResultCh: case result := <-snmpResultCh:
@ -382,7 +384,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
} }
// handleMessage dispatches incoming channel messages. // handleMessage dispatches incoming channel messages.
// Returns true if the session should end (e.g. restart requested). // Returns whether the session should end and the reason for reconnecting.
func handleMessage( func handleMessage(
ctx context.Context, ctx context.Context,
msg channelMsg, msg channelMsg,
@ -393,32 +395,38 @@ func handleMessage(
monitoringCheckCh chan<- *pb.MonitoringCheck, monitoringCheckCh chan<- *pb.MonitoringCheck,
checkResultCh chan<- *pb.CheckResult, checkResultCh chan<- *pb.CheckResult,
lldpTopologyResultCh chan<- *pb.LldpTopologyResult, lldpTopologyResultCh chan<- *pb.LldpTopologyResult,
) bool { ) (bool, error) {
switch msg.Event { switch msg.Event {
case "phx_reply": case "phx_reply":
slog.Debug("channel reply", "topic", msg.Topic) slog.Debug("channel reply", "topic", msg.Topic)
case "phx_error", "phx_close":
slog.Warn("phoenix channel ended, reconnecting",
"event", msg.Event,
"topic", msg.Topic)
return true, errChannelReloaded
case "jobs", "discovery_job", "backup_job": case "jobs", "discovery_job", "backup_job":
var payload struct { var payload struct {
Binary string `json:"binary"` Binary string `json:"binary"`
} }
if err := json.Unmarshal(msg.Payload, &payload); err != nil { if err := json.Unmarshal(msg.Payload, &payload); err != nil {
slog.Error("decode job payload", "error", err) slog.Error("decode job payload", "error", err)
return false return false, nil
} }
if len(payload.Binary) > maxJobPayloadBytes { if len(payload.Binary) > maxJobPayloadBytes {
slog.Error("job payload too large", "size", len(payload.Binary), "max", maxJobPayloadBytes) slog.Error("job payload too large", "size", len(payload.Binary), "max", maxJobPayloadBytes)
return false return false, nil
} }
bin, err := base64.StdEncoding.DecodeString(payload.Binary) bin, err := base64.StdEncoding.DecodeString(payload.Binary)
if err != nil { if err != nil {
slog.Error("decode base64", "error", err) slog.Error("decode base64", "error", err)
return false return false, nil
} }
var jobList pb.AgentJobList var jobList pb.AgentJobList
if err := proto.Unmarshal(bin, &jobList); err != nil { if err := proto.Unmarshal(bin, &jobList); err != nil {
slog.Error("unmarshal job list", "error", err) slog.Error("unmarshal job list", "error", err)
return false return false, nil
} }
slog.Info("received jobs", "count", len(jobList.Jobs)) slog.Info("received jobs", "count", len(jobList.Jobs))
for _, job := range jobList.Jobs { for _, job := range jobList.Jobs {
@ -431,21 +439,21 @@ func handleMessage(
} }
if err := json.Unmarshal(msg.Payload, &payload); err != nil { if err := json.Unmarshal(msg.Payload, &payload); err != nil {
slog.Error("decode check_jobs payload", "error", err) slog.Error("decode check_jobs payload", "error", err)
return false return false, nil
} }
if len(payload.Binary) > maxJobPayloadBytes { if len(payload.Binary) > maxJobPayloadBytes {
slog.Error("check_jobs payload too large", "size", len(payload.Binary), "max", maxJobPayloadBytes) slog.Error("check_jobs payload too large", "size", len(payload.Binary), "max", maxJobPayloadBytes)
return false return false, nil
} }
bin, err := base64.StdEncoding.DecodeString(payload.Binary) bin, err := base64.StdEncoding.DecodeString(payload.Binary)
if err != nil { if err != nil {
slog.Error("decode check_jobs base64", "error", err) slog.Error("decode check_jobs base64", "error", err)
return false return false, nil
} }
var checkList pb.CheckList var checkList pb.CheckList
if err := proto.Unmarshal(bin, &checkList); err != nil { if err := proto.Unmarshal(bin, &checkList); err != nil {
slog.Error("unmarshal check list", "error", err) slog.Error("unmarshal check list", "error", err)
return false return false, nil
} }
slog.Info("received checks", "count", len(checkList.Checks)) slog.Info("received checks", "count", len(checkList.Checks))
for _, check := range checkList.Checks { for _, check := range checkList.Checks {
@ -454,7 +462,7 @@ func handleMessage(
case "restart": case "restart":
slog.Info("restart requested by server") slog.Info("restart requested by server")
return true return true, errRestartRequested
case "update": case "update":
var payload struct { var payload struct {
@ -463,7 +471,7 @@ func handleMessage(
} }
if err := json.Unmarshal(msg.Payload, &payload); err != nil || payload.URL == "" || payload.Checksum == "" { if err := json.Unmarshal(msg.Payload, &payload); err != nil || payload.URL == "" || payload.Checksum == "" {
slog.Error("invalid update payload") slog.Error("invalid update payload")
return false return false, nil
} }
slog.Info("update requested", "url", sanitizeURL(payload.URL)) slog.Info("update requested", "url", sanitizeURL(payload.URL))
if err := doSelfUpdate(payload.URL, payload.Checksum); err != nil { if err := doSelfUpdate(payload.URL, payload.Checksum); err != nil {
@ -473,7 +481,7 @@ func handleMessage(
default: default:
slog.Debug("ignoring event", "event", msg.Event) slog.Debug("ignoring event", "event", msg.Event)
} }
return false return false, nil
} }
// jobPools holds the worker pools for each job type. // jobPools holds the worker pools for each job type.

View file

@ -182,11 +182,45 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1) credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1) monCh := make(chan *pb.MonitoringCheck, 1)
checkCh := make(chan *pb.CheckResult, 1) checkCh := make(chan *pb.CheckResult, 1)
shouldEnd := handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1)) shouldEnd, reason := handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
if !shouldEnd { if !shouldEnd {
t.Error("expected handleMessage to return true for restart") t.Error("expected handleMessage to return true for restart")
} }
if reason != errRestartRequested {
t.Fatalf("expected restart reason %v, got %v", errRestartRequested, reason)
}
})
t.Run("phoenix channel teardown reconnects", func(t *testing.T) {
events := []string{"phx_error", "phx_close"}
for _, event := range events {
snmpCh := make(chan *pb.SnmpResult, 1)
mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
checkCh := make(chan *pb.CheckResult, 1)
shouldEnd, reason := handleMessage(
context.Background(),
channelMsg{Topic: "agent:test", Event: event, Payload: json.RawMessage(`{}`)},
testPools(t),
snmpCh,
mtCh,
credCh,
monCh,
checkCh,
make(chan *pb.LldpTopologyResult, 1),
)
if !shouldEnd {
t.Fatalf("expected handleMessage to end session for %s", event)
}
if reason != errChannelReloaded {
t.Fatalf("expected reload reason %v for %s, got %v", errChannelReloaded, event, reason)
}
}
}) })
t.Run("update success", func(t *testing.T) { t.Run("update success", func(t *testing.T) {