Fix UI stuck on “Disconnected” during network-change engine restart

tool/src/main/java/io/netbird/client/tool/EngineRestarter.java

@@ -5,7 +5,11 @@
 import android.os.Looper;
 import android.util.Log;
 
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+
 import io.netbird.client.tool.networks.NetworkToggleListener;
+import io.netbird.gomobile.android.ConnectionListener;
 
 /**
  * <p>EngineRestarter restarts the Go engine.</p>
@@ -22,10 +26,25 @@ class EngineRestarter implements NetworkToggleListener {
     private ServiceStateListener currentListener;
 
     private volatile boolean isRestartInProgress = false;
+    private volatile boolean restartScheduled = false;
+    private final Runnable connectedObserver = this::onEngineReconnected;
+
     public EngineRestarter(EngineRunner engineRunner) {
         this.engineRunner = engineRunner;
         this.handler = new Handler(Looper.getMainLooper());
         this.restartRunnable = this::restartEngine;
+        this.engineRunner.addOnConnectedObserver(connectedObserver);
+    }
+
+    private void onEngineReconnected() {
+        // The Go core reconnected on its own; the pending restart is no
+        // longer needed. Cancel the debounced restart so we do not tear
+        // down a working connection.
+        if (restartScheduled) {
+            Log.d(LOGTAG, "engine reconnected on its own, cancelling pending restart");
+            restartScheduled = false;
+            handler.removeCallbacks(restartRunnable);
+        }
     }
 
     /**
@@ -35,6 +54,8 @@ public EngineRestarter(EngineRunner engineRunner) {
      * <p>If the engine isn't running, this method does nothing.</p>
      */
     private void restartEngine() {
+        restartScheduled = false;
+
         // Prevent concurrent restarts
         if (isRestartInProgress) {
             Log.d(LOGTAG, "restart already in progress, ignoring duplicate request");
@@ -48,10 +69,29 @@ private void restartEngine() {
 
         isRestartInProgress = true;
 
+        // Snapshot the current listener and wrap it so disconnect events from
+        // the old engine teardown — and the synthetic Disconnected the new
+        // engine emits before its first ClientStart() — do not reach the UI.
+        ConnectionListener savedListener = engineRunner.getConnectionListener();
+        FilteringConnectionListener filteringListener =
+                savedListener != null ? new FilteringConnectionListener(savedListener) : null;
+        if (filteringListener != null) {
+            engineRunner.setConnectionListener(filteringListener);
+        }
+
+        // Hold a reference to suppressed external listeners so we can
+        // unsuppress them on completion, error, or timeout.
+        AtomicReference<List<ServiceStateListener>> suppressedHolder = new AtomicReference<>();
+
         timeoutCallback = () -> {
             if (isRestartInProgress) {
                 Log.e(LOGTAG, "engine restart timeout - forcing flag reset");
                 isRestartInProgress = false;
+                if (filteringListener != null) {
+                    filteringListener.allowAll();
+                }
+                unsuppressAll(suppressedHolder.get());
+                notifyDisconnected(savedListener);
             }
         };
 
@@ -67,6 +107,14 @@ public void onStarted() {
                 isRestartInProgress = false;  // Reset flag on success
                 handler.removeCallbacks(timeoutCallback);  // Cancel timeout
                 engineRunner.removeServiceStateListener(this);
+                // The Go ClientStart() will fire OnConnecting shortly; from
+                // that point onward we want the listener to see real state
+                // again. The filter stays in place until the first Connecting
+                // / Connected event passes through.
+                if (filteringListener != null) {
+                    filteringListener.allowAfterFirstConnectingOrConnected();
+                }
+                unsuppressAll(suppressedHolder.get());
             }
 
             @Override
@@ -81,6 +129,11 @@ public void onError(String msg) {
                 isRestartInProgress = false; // Resetting flag on error as well
                 handler.removeCallbacks(timeoutCallback);  // Cancel timeout
                 engineRunner.removeServiceStateListener(this);
+                if (filteringListener != null) {
+                    filteringListener.allowAll();
+                }
+                unsuppressAll(suppressedHolder.get());
+                notifyDisconnected(savedListener);
             }
         };
         currentListener = serviceStateListener;
@@ -90,28 +143,183 @@ public void onError(String msg) {
             Log.d(LOGTAG, "engine stopped before restart could begin - aborting");
             handler.removeCallbacks(timeoutCallback);
             isRestartInProgress = false;
+            if (filteringListener != null) {
+                engineRunner.setConnectionListener(savedListener);
+            }
             return;
         }
 
+        // Suppress external service-state listeners so the old engine's
+        // onStopped (and the new engine's onStarted) do not reach the UI;
+        // we drive UI state through ConnectionListener exclusively during
+        // the restart window.
+        List<ServiceStateListener> suppressed =
+                engineRunner.snapshotExternalListeners(serviceStateListener);
+        for (ServiceStateListener s : suppressed) {
+            engineRunner.suppressServiceStateListener(s);
+        }
+        suppressedHolder.set(suppressed);
+
         Log.d(LOGTAG, "engine is running, stopping due to network change");
+        notifyConnecting(savedListener);
         engineRunner.stop();
     }
 
+    private void unsuppressAll(List<ServiceStateListener> suppressed) {
+        if (suppressed == null) return;
+        for (ServiceStateListener s : suppressed) {
+            engineRunner.unsuppressServiceStateListener(s);
+        }
+    }
+
+    private void notifyConnecting(ConnectionListener listener) {
+        if (listener == null) {
+            return;
+        }
+        try {
+            listener.onConnecting();
+        } catch (Exception e) {
+            Log.w(LOGTAG, "onConnecting notification failed: " + e.getMessage());
+        }
+    }
+
+    private void notifyDisconnected(ConnectionListener listener) {
+        if (listener == null) {
+            return;
+        }
+        try {
+            listener.onDisconnected();
+        } catch (Exception e) {
+            Log.w(LOGTAG, "onDisconnected notification failed: " + e.getMessage());
+        }
+    }
+
+    /**
+     * Wraps a ConnectionListener and drops Disconnecting/Disconnected events
+     * during a restart. Disconnects from the old engine's teardown — and the
+     * default-state replay the Go notifier sends to a listener attached
+     * before the new engine's ClientStart() — would otherwise flash the UI
+     * to Disconnected. Filtering ends as soon as the first real Connecting
+     * or Connected event arrives from the new engine.
+     */
+    private static final class FilteringConnectionListener implements ConnectionListener {
+        private final ConnectionListener delegate;
+        private volatile boolean dropDisconnects = true;
+        private volatile boolean releaseAfterFirstActive = false;
+
+        FilteringConnectionListener(ConnectionListener delegate) {
+            this.delegate = delegate;
+        }
+
+        void allowAll() {
+            dropDisconnects = false;
+            releaseAfterFirstActive = false;
+        }
+
+        void allowAfterFirstConnectingOrConnected() {
+            releaseAfterFirstActive = true;
+        }
+
+        @Override
+        public void onConnecting() {
+            if (releaseAfterFirstActive) {
+                dropDisconnects = false;
+            }
+            try {
+                delegate.onConnecting();
+            } catch (Exception e) {
+                Log.w(LOGTAG, "delegate onConnecting failed: " + e.getMessage());
+            }
+        }
+
+        @Override
+        public void onConnected() {
+            if (releaseAfterFirstActive) {
+                dropDisconnects = false;
+            }
+            try {
+                delegate.onConnected();
+            } catch (Exception e) {
+                Log.w(LOGTAG, "delegate onConnected failed: " + e.getMessage());
+            }
+        }
+
+        @Override
+        public void onDisconnecting() {
+            if (dropDisconnects) {
+                Log.d(LOGTAG, "filtered onDisconnecting during restart");
+                return;
+            }
+            try {
+                delegate.onDisconnecting();
+            } catch (Exception e) {
+                Log.w(LOGTAG, "delegate onDisconnecting failed: " + e.getMessage());
+            }
+        }
+
+        @Override
+        public void onDisconnected() {
+            if (dropDisconnects) {
+                Log.d(LOGTAG, "filtered onDisconnected during restart");
+                return;
+            }
+            try {
+                delegate.onDisconnected();
+            } catch (Exception e) {
+                Log.w(LOGTAG, "delegate onDisconnected failed: " + e.getMessage());
+            }
+        }
+
+        @Override
+        public void onAddressChanged(String fqdn, String ip) {
+            try {
+                delegate.onAddressChanged(fqdn, ip);
+            } catch (Exception e) {
+                Log.w(LOGTAG, "delegate onAddressChanged failed: " + e.getMessage());
+            }
+        }
+
+        @Override
+        public void onPeersListChanged(long numberOfPeers) {
+            try {
+                delegate.onPeersListChanged(numberOfPeers);
+            } catch (Exception e) {
+                Log.w(LOGTAG, "delegate onPeersListChanged failed: " + e.getMessage());
+            }
+        }
+    }
+
     @Override
     public void onNetworkTypeChanged() {
         Log.d(LOGTAG, "network type changed, scheduling restart with "
                 + DEBOUNCE_DELAY_MS + "ms debounce.");
 
+        restartScheduled = true;
         handler.removeCallbacks(restartRunnable);
         handler.postDelayed(restartRunnable, DEBOUNCE_DELAY_MS);
     }
 
+    /**
+     * Cancels any pending debounced restart. Called whenever an external
+     * actor (typically a user-driven Connect/Disconnect) takes over the
+     * engine lifecycle, so the network-change-driven restart does not
+     * interfere with that explicit action.
+     */
+    public void cancelPendingRestart() {
+        if (restartScheduled) {
+            Log.d(LOGTAG, "external action took over engine lifecycle; cancelling pending restart");
+            handler.removeCallbacks(restartRunnable);
+            restartScheduled = false;
+        }
+    }
+
     /**
      * <p>Cleans up resources, like the restart runnable and timeout callback.</p>
      * <p>Call this when the EngineRestarter is no longer needed to prevent memory leaks.</p>
      */
     public void cleanup() {
         handler.removeCallbacks(restartRunnable);
+        restartScheduled = false;
 
         if (timeoutCallback != null) {
             handler.removeCallbacks(timeoutCallback);
@@ -122,6 +330,8 @@ public void cleanup() {
             currentListener = null;
         }
 
+        engineRunner.removeOnConnectedObserver(connectedObserver);
+
         isRestartInProgress = false;
     }
 }

tool/src/main/java/io/netbird/client/tool/EngineRunner.java

@@ -28,7 +28,10 @@ class EngineRunner {
     private final ProfileManagerWrapper profileManager;
     private boolean engineIsRunning = false;
     Set<ServiceStateListener> serviceStateListeners = ConcurrentHashMap.newKeySet();
+    private final Set<ServiceStateListener> suppressedServiceStateListeners = ConcurrentHashMap.newKeySet();
+    private final Set<Runnable> connectedObservers = ConcurrentHashMap.newKeySet();
     private final Client goClient;
+    private ConnectionListener connectionListener;
 
     public EngineRunner(Context context, NetworkChangeListener networkChangeListener, TunAdapter tunAdapter,
                         IFaceDiscover iFaceDiscover, String versionName, boolean isTraceLogEnabled, boolean isDebuggable,
@@ -124,13 +127,45 @@ public synchronized boolean isRunning() {
     }
 
     public synchronized void setConnectionListener(ConnectionListener listener) {
-        goClient.setConnectionListener(listener);
+        ConnectionListener wrapped = listener == null ? null : new ConnectionListener() {
+            @Override public void onConnecting() { listener.onConnecting(); }
+            @Override public void onConnected() {
+                listener.onConnected();
+                for (Runnable obs : connectedObservers) {
+                    try { obs.run(); } catch (Exception e) { Log.w(LOGTAG, "connected observer failed", e); }
+                }
+            }
+            @Override public void onDisconnecting() { listener.onDisconnecting(); }
+            @Override public void onDisconnected() { listener.onDisconnected(); }
+            @Override public void onAddressChanged(String f, String i) { listener.onAddressChanged(f, i); }
+            @Override public void onPeersListChanged(long n) { listener.onPeersListChanged(n); }
+        };
+        this.connectionListener = wrapped;
+        goClient.setConnectionListener(wrapped);
     }
 
     public synchronized void removeStatusListener() {
+        this.connectionListener = null;
         goClient.removeConnectionListener();
     }
 
+    synchronized ConnectionListener getConnectionListener() {
+        return connectionListener;
+    }
+
+    /**
+     * Registers a callback that fires every time the engine reports
+     * OnConnected. EngineRestarter uses this to cancel a pending restart
+     * when the Go core has already reconnected on its own.
+     */
+    public void addOnConnectedObserver(Runnable observer) {
+        connectedObservers.add(observer);
+    }
+
+    public void removeOnConnectedObserver(Runnable observer) {
+        connectedObservers.remove(observer);
+    }
+
     public synchronized void addServiceStateListener(ServiceStateListener serviceStateListener) {
         if (engineIsRunning) {
             serviceStateListener.onStarted();
@@ -157,6 +192,29 @@ public synchronized boolean addServiceStateListenerForRestart(ServiceStateListen
 
     public synchronized void removeServiceStateListener(ServiceStateListener serviceStateListener) {
         serviceStateListeners.remove(serviceStateListener);
+        suppressedServiceStateListeners.remove(serviceStateListener);
+    }
+
+    /**
+     * Marks a listener as suppressed: it will not receive onStarted / onStopped
+     * notifications until {@link #unsuppressServiceStateListener} is called.
+     * Used by EngineRestarter to hide the engine teardown from external UI
+     * listeners during a restart.
+     */
+    public synchronized void suppressServiceStateListener(ServiceStateListener listener) {
+        suppressedServiceStateListeners.add(listener);
+    }
+
+    public synchronized void unsuppressServiceStateListener(ServiceStateListener listener) {
+        suppressedServiceStateListeners.remove(listener);
+    }
+
+    public synchronized java.util.List<ServiceStateListener> snapshotExternalListeners(ServiceStateListener exclude) {
+        java.util.List<ServiceStateListener> out = new java.util.ArrayList<>();
+        for (ServiceStateListener s : serviceStateListeners) {
+            if (s != exclude) out.add(s);
+        }
+        return out;
     }
 
     public synchronized void stop() {
@@ -184,6 +242,9 @@ private synchronized void notifyError(Exception e) {
 
     private synchronized void notifyServiceStateListeners(boolean engineIsRunning) {
         for (ServiceStateListener s : serviceStateListeners) {
+            if (suppressedServiceStateListeners.contains(s)) {
+                continue;
+            }
             if (engineIsRunning) {
                 s.onStarted();
             } else {