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 errRestartRequested = fmt.Errorf("restart requested")
var errChannelReloaded = fmt.Errorf("channel reloaded")
var joinTimeout = 10 * time.Second
var heartbeatInterval = 60 * 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)
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()
return errRestartRequested
return endErr
}
case result := <-snmpResultCh:
@ -382,7 +384,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
}
// 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(
ctx context.Context,
msg channelMsg,
@ -393,32 +395,38 @@ func handleMessage(
monitoringCheckCh chan<- *pb.MonitoringCheck,
checkResultCh chan<- *pb.CheckResult,
lldpTopologyResultCh chan<- *pb.LldpTopologyResult,
) bool {
) (bool, error) {
switch msg.Event {
case "phx_reply":
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":
var payload struct {
Binary string `json:"binary"`
}
if err := json.Unmarshal(msg.Payload, &payload); err != nil {
slog.Error("decode job payload", "error", err)
return false
return false, nil
}
if len(payload.Binary) > 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)
if err != nil {
slog.Error("decode base64", "error", err)
return false
return false, nil
}
var jobList pb.AgentJobList
if err := proto.Unmarshal(bin, &jobList); err != nil {
slog.Error("unmarshal job list", "error", err)
return false
return false, nil
}
slog.Info("received jobs", "count", len(jobList.Jobs))
for _, job := range jobList.Jobs {
@ -431,21 +439,21 @@ func handleMessage(
}
if err := json.Unmarshal(msg.Payload, &payload); err != nil {
slog.Error("decode check_jobs payload", "error", err)
return false
return false, nil
}
if len(payload.Binary) > 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)
if err != nil {
slog.Error("decode check_jobs base64", "error", err)
return false
return false, nil
}
var checkList pb.CheckList
if err := proto.Unmarshal(bin, &checkList); err != nil {
slog.Error("unmarshal check list", "error", err)
return false
return false, nil
}
slog.Info("received checks", "count", len(checkList.Checks))
for _, check := range checkList.Checks {
@ -454,7 +462,7 @@ func handleMessage(
case "restart":
slog.Info("restart requested by server")
return true
return true, errRestartRequested
case "update":
var payload struct {
@ -463,7 +471,7 @@ func handleMessage(
}
if err := json.Unmarshal(msg.Payload, &payload); err != nil || payload.URL == "" || payload.Checksum == "" {
slog.Error("invalid update payload")
return false
return false, nil
}
slog.Info("update requested", "url", sanitizeURL(payload.URL))
if err := doSelfUpdate(payload.URL, payload.Checksum); err != nil {
@ -473,7 +481,7 @@ func handleMessage(
default:
slog.Debug("ignoring event", "event", msg.Event)
}
return false
return false, nil
}
// 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)
monCh := make(chan *pb.MonitoringCheck, 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 {
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) {