Handle properly stopping and starting the ping

olm/olm.go

@@ -383,9 +383,9 @@ func (o *Olm) StartTunnel(config TunnelConfig) {
 		o.apiServer.SetConnectionStatus(true)
 
 		if o.connected {
-			logger.Debug("Already connected, skipping registration")
-			// Restart ping monitor on reconnect since the old one would have exited
 			o.websocket.StartPingMonitor()
+			
+			logger.Debug("Already connected, skipping registration")
 			return nil
 		}
 
@@ -686,6 +686,14 @@ func (o *Olm) SetPowerMode(mode string) error {
 
 		logger.Info("Switching to low power mode")
 
+		// Mark as disconnected so we re-register on reconnect
+		o.connected = false
+
+		// Update API server connection status
+		if o.apiServer != nil {
+			o.apiServer.SetConnectionStatus(false)
+		}
+
 		if o.websocket != nil {
 			logger.Info("Disconnecting websocket for low power mode")
 			if err := o.websocket.Disconnect(); err != nil {

websocket/client.go

@@ -102,6 +102,7 @@ type Client struct {
 	getPingData       func() map[string]any  // Callback to get additional ping data
 	pingStarted       bool                   // Flag to track if ping monitor has been started
 	pingStartedMux    sync.Mutex             // Protects pingStarted
+	pingDone          chan struct{}          // Channel to stop the ping monitor independently
 }
 
 type ClientOption func(*Client)
@@ -176,6 +177,7 @@ func NewClient(ID, secret, userToken, orgId, endpoint string, pingInterval time.
 		pingInterval:      pingInterval,
 		pingTimeout:       pingTimeout,
 		clientType:        "olm",
+		pingDone:          make(chan struct{}),
 	}
 
 	// Apply options before loading config
@@ -235,6 +237,9 @@ func (c *Client) Disconnect() error {
 	c.isDisconnected = true
 	c.setConnected(false)
 
+	// Stop the ping monitor
+	c.stopPingMonitor()
+
 	// Wait for any message currently being processed to complete
 	c.processingWg.Wait()
 
@@ -577,11 +582,6 @@ func (c *Client) establishConnection() error {
 	c.conn = conn
 	c.setConnected(true)
 
-	// Reset ping started flag on new connection
-	c.pingStartedMux.Lock()
-	c.pingStarted = false
-	c.pingStartedMux.Unlock()
-
 	// Note: ping monitor is NOT started here - it will be started when
 	// StartPingMonitor() is called after registration completes
 
@@ -722,6 +722,8 @@ func (c *Client) pingMonitor() {
 		select {
 		case <-c.done:
 			return
+		case <-c.pingDone:
+			return
 		case <-ticker.C:
 			c.sendPing()
 		}
@@ -740,6 +742,9 @@ func (c *Client) StartPingMonitor() {
 	}
 	c.pingStarted = true
 	
+	// Create a new pingDone channel for this ping monitor instance
+	c.pingDone = make(chan struct{})
+	
 	// Send an initial ping immediately
 	go func() {
 		c.sendPing()
@@ -747,6 +752,20 @@ func (c *Client) StartPingMonitor() {
 	}()
 }
 
+// stopPingMonitor stops the ping monitor goroutine if it's running.
+func (c *Client) stopPingMonitor() {
+	c.pingStartedMux.Lock()
+	defer c.pingStartedMux.Unlock()
+	
+	if !c.pingStarted {
+		return
+	}
+	
+	// Close the pingDone channel to stop the monitor
+	close(c.pingDone)
+	c.pingStarted = false
+}
+
 // GetConfigVersion returns the current config version
 func (c *Client) GetConfigVersion() int {
 	c.configVersionMux.RLock()