From a293f760af653afb0afe7635680c0ff5a7527128 Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Tue, 2 Dec 2025 16:30:15 +0100 Subject: [PATCH] [client] Add conditional peer removal logic during shutdown (#4897) --- client/internal/engine.go | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/client/internal/engine.go b/client/internal/engine.go index 0ff1006cd..84bb8beea 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -292,6 +292,12 @@ func (e *Engine) Stop() error { } log.Info("Network monitor: stopped") + if os.Getenv("NB_REMOVE_BEFORE_DNS") == "true" && os.Getenv("NB_REMOVE_BEFORE_ROUTES") != "true" { + log.Info("removing peers before dns") + if err := e.removeAllPeers(); err != nil { + return fmt.Errorf("failed to remove all peers: %s", err) + } + } if err := e.stopSSHServer(); err != nil { log.Warnf("failed to stop SSH server: %v", err) } @@ -310,6 +316,13 @@ func (e *Engine) Stop() error { e.stopDNSForwarder() + if os.Getenv("NB_REMOVE_BEFORE_ROUTES") == "true" && os.Getenv("NB_REMOVE_BEFORE_DNS") != "true" { + log.Info("removing peers before routes") + if err := e.removeAllPeers(); err != nil { + return fmt.Errorf("failed to remove all peers: %s", err) + } + } + if e.routeManager != nil { e.routeManager.Stop(e.stateManager) } @@ -317,13 +330,16 @@ func (e *Engine) Stop() error { if e.srWatcher != nil { e.srWatcher.Close() } - + log.Info("cleaning up status recorder states") e.statusRecorder.ReplaceOfflinePeers([]peer.State{}) e.statusRecorder.UpdateDNSStates([]peer.NSGroupState{}) e.statusRecorder.UpdateRelayStates([]relay.ProbeResult{}) - if err := e.removeAllPeers(); err != nil { - return fmt.Errorf("failed to remove all peers: %s", err) + if os.Getenv("NB_REMOVE_BEFORE_DNS") != "true" && os.Getenv("NB_REMOVE_BEFORE_ROUTES") != "true" { + log.Info("removing peers after dns and routes") + if err := e.removeAllPeers(); err != nil { + return fmt.Errorf("failed to remove all peers: %s", err) + } } if e.cancel != nil {