Merge "Use WakeupMesage for delayed events in VcnGatewayConnection." am: b244bfa4d6

Original change: https://android-review.googlesource.com/c/platform/frameworks/base/+/1574005

MUST ONLY BE SUBMITTED BY AUTOMERGER

Change-Id: I4ff4e8c5b7ec0cf8cf32cdb9f9728fe53c7a7f28

services/core/java/com/android/server/vcn/VcnGatewayConnection.java

@@ -61,6 +61,7 @@ import android.os.Message;
 import android.os.ParcelUuid;
 import android.os.PowerManager;
 import android.os.PowerManager.WakeLock;
+import android.os.SystemClock;
 import android.util.ArraySet;
 import android.util.Slog;
 
@@ -68,6 +69,7 @@ import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.annotations.VisibleForTesting.Visibility;
 import com.android.internal.util.State;
 import com.android.internal.util.StateMachine;
+import com.android.internal.util.WakeupMessage;
 import com.android.server.vcn.TelephonySubscriptionTracker.TelephonySubscriptionSnapshot;
 import com.android.server.vcn.UnderlyingNetworkTracker.UnderlyingNetworkRecord;
 import com.android.server.vcn.UnderlyingNetworkTracker.UnderlyingNetworkTrackerCallback;
@@ -129,8 +131,8 @@ import java.util.concurrent.TimeUnit;
  * lack of WakeLocks).
  *
  * <p>Any attempt to remove messages from the Handler should be done using {@link
- * #removeEqualMessages(int, Object)}. This is necessary to ensure that the WakeLock is correctly
+ * #removeEqualMessages}. This is necessary to ensure that the WakeLock is correctly released when
- * released when no messages remain in the Handler queue.
+ * no messages remain in the Handler queue.
  *
  * @hide
  */
@@ -140,6 +142,15 @@ public class VcnGatewayConnection extends StateMachine {
     @VisibleForTesting(visibility = Visibility.PRIVATE)
     static final InetAddress DUMMY_ADDR = InetAddresses.parseNumericAddress("192.0.2.0");
 
+    @VisibleForTesting(visibility = Visibility.PRIVATE)
+    static final String TEARDOWN_TIMEOUT_ALARM = TAG + "_TEARDOWN_TIMEOUT_ALARM";
+
+    @VisibleForTesting(visibility = Visibility.PRIVATE)
+    static final String DISCONNECT_REQUEST_ALARM = TAG + "_DISCONNECT_REQUEST_ALARM";
+
+    @VisibleForTesting(visibility = Visibility.PRIVATE)
+    static final String RETRY_TIMEOUT_ALARM = TAG + "_RETRY_TIMEOUT_ALARM";
+
     private static final int[] MERGED_CAPABILITIES =
             new int[] {NET_CAPABILITY_NOT_METERED, NET_CAPABILITY_NOT_ROAMING};
     private static final int ARG_NOT_PRESENT = Integer.MIN_VALUE;
@@ -150,7 +161,8 @@ public class VcnGatewayConnection extends StateMachine {
     private static final String DISCONNECT_REASON_TEARDOWN = "teardown() called on VcnTunnel";
     private static final int TOKEN_ALL = Integer.MIN_VALUE;
 
-    private static final int NETWORK_LOSS_DISCONNECT_TIMEOUT_SECONDS = 30;
+    @VisibleForTesting(visibility = Visibility.PRIVATE)
+    static final int NETWORK_LOSS_DISCONNECT_TIMEOUT_SECONDS = 30;
 
     @VisibleForTesting(visibility = Visibility.PRIVATE)
     static final int TEARDOWN_TIMEOUT_SECONDS = 5;
@@ -510,6 +522,10 @@ public class VcnGatewayConnection extends StateMachine {
      */
     private NetworkAgent mNetworkAgent;
 
+    @Nullable private WakeupMessage mTeardownTimeoutAlarm;
+    @Nullable private WakeupMessage mDisconnectRequestAlarm;
+    @Nullable private WakeupMessage mRetryTimeoutAlarm;
+
     public VcnGatewayConnection(
             @NonNull VcnContext vcnContext,
             @NonNull ParcelUuid subscriptionGroup,
@@ -592,6 +608,10 @@ public class VcnGatewayConnection extends StateMachine {
 
         releaseWakeLock();
 
+        cancelTeardownTimeoutAlarm();
+        cancelDisconnectRequestAlarm();
+        cancelRetryTimeoutAlarm();
+
         mUnderlyingNetworkTracker.teardown();
     }
 
@@ -623,17 +643,11 @@ public class VcnGatewayConnection extends StateMachine {
             // If underlying is null, all underlying networks have been lost. Disconnect VCN after a
             // timeout.
             if (underlying == null) {
-                sendMessageDelayed(
+                setDisconnectRequestAlarm();
-                        EVENT_DISCONNECT_REQUESTED,
-                        TOKEN_ALL,
-                        new EventDisconnectRequestedInfo(DISCONNECT_REASON_UNDERLYING_NETWORK_LOST),
-                        TimeUnit.SECONDS.toMillis(NETWORK_LOSS_DISCONNECT_TIMEOUT_SECONDS));
             } else {
-                // Cancel any existing disconnect due to previous loss of underlying network
+                // Received a new Network so any previous alarm is irrelevant - cancel + clear it,
-                removeEqualMessages(
+                // and cancel any queued EVENT_DISCONNECT_REQUEST messages
-                        EVENT_DISCONNECT_REQUESTED,
+                cancelDisconnectRequestAlarm();
-                        new EventDisconnectRequestedInfo(
-                                DISCONNECT_REASON_UNDERLYING_NETWORK_LOST));
             }
 
             sendMessageAndAcquireWakeLock(
@@ -724,6 +738,10 @@ public class VcnGatewayConnection extends StateMachine {
         super.sendMessage(msg);
     }
 
+    // TODO(b/180146061): also override and Log.wtf() other Message handling methods
+    // In mind are sendMessageDelayed(), sendMessageAtFrontOfQueue, removeMessages, and
+    // removeDeferredMessages
+
     /**
      * WakeLock-based alternative to {@link #sendMessage}. Use to guarantee that the device will not
      * go to sleep before processing the sent message.
@@ -751,14 +769,13 @@ public class VcnGatewayConnection extends StateMachine {
         super.sendMessage(what, token, arg2, data);
     }
 
-    // TODO: remove this method once WakupMessage is used instead
+    /**
-    private void sendMessageDelayed(int what, int token, EventInfo data, long timeout) {
+     * WakeLock-based alternative to {@link #sendMessage}. Use to guarantee that the device will not
-        super.sendMessageDelayed(what, token, ARG_NOT_PRESENT, data, timeout);
+     * go to sleep before processing the sent message.
-    }
+     */
-
+    private void sendMessageAndAcquireWakeLock(Message msg) {
-    // TODO: remove this method once WakupMessage is used instead
+        acquireWakeLock();
-    private void sendMessageDelayed(int what, int token, int arg2, EventInfo data, long timeout) {
+        super.sendMessage(msg);
-        super.sendMessageDelayed(what, token, arg2, data, timeout);
     }
 
     /**
@@ -788,6 +805,100 @@ public class VcnGatewayConnection extends StateMachine {
         maybeReleaseWakeLock();
     }
 
+    private WakeupMessage createScheduledAlarm(
+            @NonNull String cmdName, Message delayedMessage, long delay) {
+        // WakeupMessage uses Handler#dispatchMessage() to immediately handle the specified Runnable
+        // at the scheduled time. dispatchMessage() immediately executes and there may be queued
+        // events that resolve the scheduled alarm pending in the queue. So, use the Runnable to
+        // place the alarm event at the end of the queue with sendMessageAndAcquireWakeLock (which
+        // guarantees the device will stay awake).
+        final WakeupMessage alarm =
+                mDeps.newWakeupMessage(
+                        mVcnContext,
+                        getHandler(),
+                        cmdName,
+                        () -> sendMessageAndAcquireWakeLock(delayedMessage));
+        alarm.schedule(mDeps.getElapsedRealTime() + delay);
+        return alarm;
+    }
+
+    private void setTeardownTimeoutAlarm() {
+        // Safe to assign this alarm because it is either 1) already null, or 2) already fired. In
+        // either case, there is nothing to cancel.
+        if (mTeardownTimeoutAlarm != null) {
+            Slog.wtf(TAG, "mTeardownTimeoutAlarm should be null before being set");
+        }
+
+        final Message delayedMessage = obtainMessage(EVENT_TEARDOWN_TIMEOUT_EXPIRED, mCurrentToken);
+        mTeardownTimeoutAlarm =
+                createScheduledAlarm(
+                        TEARDOWN_TIMEOUT_ALARM,
+                        delayedMessage,
+                        TimeUnit.SECONDS.toMillis(TEARDOWN_TIMEOUT_SECONDS));
+    }
+
+    private void cancelTeardownTimeoutAlarm() {
+        if (mTeardownTimeoutAlarm != null) {
+            mTeardownTimeoutAlarm.cancel();
+            mTeardownTimeoutAlarm = null;
+        }
+
+        // Cancel any existing teardown timeouts
+        removeEqualMessages(EVENT_TEARDOWN_TIMEOUT_EXPIRED);
+    }
+
+    private void setDisconnectRequestAlarm() {
+        // Only schedule a NEW alarm if none is already set.
+        if (mDisconnectRequestAlarm != null) {
+            return;
+        }
+
+        final Message delayedMessage =
+                obtainMessage(
+                        EVENT_DISCONNECT_REQUESTED,
+                        TOKEN_ALL,
+                        0 /* arg2 */,
+                        new EventDisconnectRequestedInfo(
+                                DISCONNECT_REASON_UNDERLYING_NETWORK_LOST));
+        mDisconnectRequestAlarm =
+                createScheduledAlarm(
+                        DISCONNECT_REQUEST_ALARM,
+                        delayedMessage,
+                        TimeUnit.SECONDS.toMillis(NETWORK_LOSS_DISCONNECT_TIMEOUT_SECONDS));
+    }
+
+    private void cancelDisconnectRequestAlarm() {
+        if (mDisconnectRequestAlarm != null) {
+            mDisconnectRequestAlarm.cancel();
+            mDisconnectRequestAlarm = null;
+        }
+
+        // Cancel any existing disconnect due to previous loss of underlying network
+        removeEqualMessages(
+                EVENT_DISCONNECT_REQUESTED,
+                new EventDisconnectRequestedInfo(DISCONNECT_REASON_UNDERLYING_NETWORK_LOST));
+    }
+
+    private void setRetryTimeoutAlarm(long delay) {
+        // Safe to assign this alarm because it is either 1) already null, or 2) already fired. In
+        // either case, there is nothing to cancel.
+        if (mRetryTimeoutAlarm != null) {
+            Slog.wtf(TAG, "mRetryTimeoutAlarm should be null before being set");
+        }
+
+        final Message delayedMessage = obtainMessage(EVENT_RETRY_TIMEOUT_EXPIRED, mCurrentToken);
+        mRetryTimeoutAlarm = createScheduledAlarm(RETRY_TIMEOUT_ALARM, delayedMessage, delay);
+    }
+
+    private void cancelRetryTimeoutAlarm() {
+        if (mRetryTimeoutAlarm != null) {
+            mRetryTimeoutAlarm.cancel();
+            mRetryTimeoutAlarm = null;
+        }
+
+        removeEqualMessages(EVENT_RETRY_TIMEOUT_EXPIRED);
+    }
+
     private void sessionLost(int token, @Nullable Exception exception) {
         sendMessageAndAcquireWakeLock(
                 EVENT_SESSION_LOST, token, new EventSessionLostInfo(exception));
@@ -847,8 +958,8 @@ public class VcnGatewayConnection extends StateMachine {
          * builds.
          *
          * <p>Here be dragons: processMessage() is final to ensure that mWakeLock is released once
-         * the Handler queue is empty. Future changes to processMessage() (or overrides to
+         * the Handler queue is empty. Future changes (or overrides) to processMessage() to MUST
-         * processMessage) MUST ensure that mWakeLock is correctly released.
+         * ensure that mWakeLock is correctly released.
          */
         @Override
         public final boolean processMessage(Message msg) {
@@ -1030,10 +1141,9 @@ public class VcnGatewayConnection extends StateMachine {
             }
 
             mIkeSession.close();
-            sendMessageDelayed(
+
-                    EVENT_TEARDOWN_TIMEOUT_EXPIRED,
+            // Safe to blindly set up, as it is cancelled and cleared on exiting this state
-                    mCurrentToken,
+            setTeardownTimeoutAlarm();
-                    TimeUnit.SECONDS.toMillis(TEARDOWN_TIMEOUT_SECONDS));
         }
 
         @Override
@@ -1084,6 +1194,8 @@ public class VcnGatewayConnection extends StateMachine {
         @Override
         protected void exitState() throws Exception {
             mSkipRetryTimeout = false;
+
+            cancelTeardownTimeoutAlarm();
         }
     }
 
@@ -1385,8 +1497,8 @@ public class VcnGatewayConnection extends StateMachine {
                 Slog.wtf(TAG, "Underlying network was null in retry state");
                 transitionTo(mDisconnectedState);
             } else {
-                sendMessageDelayed(
+                // Safe to blindly set up, as it is cancelled and cleared on exiting this state
-                        EVENT_RETRY_TIMEOUT_EXPIRED, mCurrentToken, getNextRetryIntervalsMs());
+                setRetryTimeoutAlarm(getNextRetryIntervalsMs());
             }
         }
 
@@ -1399,8 +1511,6 @@ public class VcnGatewayConnection extends StateMachine {
 
                     // If new underlying is null, all networks were lost; go back to disconnected.
                     if (mUnderlying == null) {
-                        removeEqualMessages(EVENT_RETRY_TIMEOUT_EXPIRED);
-
                         transitionTo(mDisconnectedState);
                         return;
                     } else if (oldUnderlying != null
@@ -1411,8 +1521,6 @@ public class VcnGatewayConnection extends StateMachine {
 
                     // Fallthrough
                 case EVENT_RETRY_TIMEOUT_EXPIRED:
-                    removeEqualMessages(EVENT_RETRY_TIMEOUT_EXPIRED);
-
                     transitionTo(mConnectingState);
                     break;
                 case EVENT_DISCONNECT_REQUESTED:
@@ -1424,6 +1532,11 @@ public class VcnGatewayConnection extends StateMachine {
             }
         }
 
+        @Override
+        public void exitState() {
+            cancelRetryTimeoutAlarm();
+        }
+
         private long getNextRetryIntervalsMs() {
             final int retryDelayIndex = mFailedAttempts - 1;
             final long[] retryIntervalsMs = mConnectionConfig.getRetryIntervalsMs();
@@ -1702,6 +1815,20 @@ public class VcnGatewayConnection extends StateMachine {
                 @NonNull Context context, int wakeLockFlag, @NonNull String wakeLockTag) {
             return new VcnWakeLock(context, wakeLockFlag, wakeLockTag);
         }
+
+        /** Builds a new WakeupMessage. */
+        public WakeupMessage newWakeupMessage(
+                @NonNull VcnContext vcnContext,
+                @NonNull Handler handler,
+                @NonNull String tag,
+                @NonNull Runnable runnable) {
+            return new WakeupMessage(vcnContext.getContext(), handler, tag, runnable);
+        }
+
+        /** Gets the elapsed real time since boot, in millis. */
+        public long getElapsedRealTime() {
+            return SystemClock.elapsedRealtime();
+        }
     }
 
     /**

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionConnectedStateTest.java

@@ -81,6 +81,7 @@ public class VcnGatewayConnectionConnectedStateTest extends VcnGatewayConnection
 
         assertEquals(mGatewayConnection.mConnectedState, mGatewayConnection.getCurrentState());
         verify(mIkeSession, never()).close();
+        verifyDisconnectRequestAlarmAndGetCallback(false /* expectCanceled */);
     }
 
     @Test
@@ -170,6 +171,7 @@ public class VcnGatewayConnectionConnectedStateTest extends VcnGatewayConnection
         mTestLooper.dispatchAll();
 
         assertEquals(mGatewayConnection.mDisconnectingState, mGatewayConnection.getCurrentState());
+        verifyTeardownTimeoutAlarmAndGetCallback(false /* expectCanceled */);
     }
 
     @Test
@@ -179,5 +181,6 @@ public class VcnGatewayConnectionConnectedStateTest extends VcnGatewayConnection
 
         assertEquals(mGatewayConnection.mRetryTimeoutState, mGatewayConnection.getCurrentState());
         verify(mIkeSession).close();
+        verifyTeardownTimeoutAlarmAndGetCallback(true /* expectCanceled */);
     }
 }

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionConnectingStateTest.java

@@ -61,6 +61,7 @@ public class VcnGatewayConnectionConnectingStateTest extends VcnGatewayConnectio
 
         assertEquals(mGatewayConnection.mDisconnectingState, mGatewayConnection.getCurrentState());
         verify(mIkeSession).kill();
+        verifyDisconnectRequestAlarmAndGetCallback(false /* expectCanceled */);
     }
 
     @Test
@@ -73,6 +74,7 @@ public class VcnGatewayConnectionConnectingStateTest extends VcnGatewayConnectio
         assertEquals(mGatewayConnection.mDisconnectingState, mGatewayConnection.getCurrentState());
         verify(mIkeSession).close();
         verify(mIkeSession, never()).kill();
+        verifyTeardownTimeoutAlarmAndGetCallback(false /* expectCanceled */);
     }
 
     @Test
@@ -92,6 +94,7 @@ public class VcnGatewayConnectionConnectingStateTest extends VcnGatewayConnectio
 
         assertEquals(mGatewayConnection.mDisconnectingState, mGatewayConnection.getCurrentState());
         verify(mIkeSession).close();
+        verifyTeardownTimeoutAlarmAndGetCallback(false /* expectCanceled */);
     }
 
     @Test
@@ -101,5 +104,6 @@ public class VcnGatewayConnectionConnectingStateTest extends VcnGatewayConnectio
 
         assertEquals(mGatewayConnection.mRetryTimeoutState, mGatewayConnection.getCurrentState());
         verify(mIkeSession).close();
+        verifyTeardownTimeoutAlarmAndGetCallback(true /* expectCanceled */);
     }
 }

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionDisconnectedStateTest.java

@@ -88,6 +88,7 @@ public class VcnGatewayConnectionDisconnectedStateTest extends VcnGatewayConnect
         mTestLooper.dispatchAll();
 
         assertEquals(mGatewayConnection.mDisconnectedState, mGatewayConnection.getCurrentState());
+        verifyDisconnectRequestAlarmAndGetCallback(false /* expectCanceled */);
     }
 
     @Test

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionDisconnectingStateTest.java

@@ -16,9 +16,9 @@
 
 package com.android.server.vcn;
 
-import static com.android.server.vcn.VcnGatewayConnection.TEARDOWN_TIMEOUT_SECONDS;
 
 import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 import androidx.test.filters.SmallTest;
@@ -28,8 +28,6 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import java.util.concurrent.TimeUnit;
-
 /** Tests for VcnGatewayConnection.DisconnectedState */
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -40,6 +38,9 @@ public class VcnGatewayConnectionDisconnectingStateTest extends VcnGatewayConnec
 
         mGatewayConnection.setIkeSession(mGatewayConnection.buildIkeSession());
 
+        // ensure that mGatewayConnection has an underlying Network before entering
+        // DisconnectingState
+        mGatewayConnection.setUnderlyingNetwork(TEST_UNDERLYING_NETWORK_RECORD_2);
         mGatewayConnection.transitionTo(mGatewayConnection.mDisconnectingState);
         mTestLooper.dispatchAll();
     }
@@ -49,12 +50,22 @@ public class VcnGatewayConnectionDisconnectingStateTest extends VcnGatewayConnec
         getIkeSessionCallback().onClosed();
         mTestLooper.dispatchAll();
 
-        assertEquals(mGatewayConnection.mDisconnectedState, mGatewayConnection.getCurrentState());
+        assertEquals(mGatewayConnection.mRetryTimeoutState, mGatewayConnection.getCurrentState());
+        verify(mMockIkeSession).close();
+        verify(mMockIkeSession, never()).kill();
+        verifyTeardownTimeoutAlarmAndGetCallback(true /* expectCanceled */);
     }
 
     @Test
     public void testTimeoutExpired() throws Exception {
-        mTestLooper.moveTimeForward(TimeUnit.SECONDS.toMillis(TEARDOWN_TIMEOUT_SECONDS));
+        Runnable delayedEvent =
+                verifyTeardownTimeoutAlarmAndGetCallback(false /* expectCanceled */);
+
+        // Can't use mTestLooper to advance the time since VcnGatewayConnection uses WakeupMessages
+        // (which are mocked here). Directly invoke the runnable instead. This is still sufficient,
+        // since verifyTeardownTimeoutAlarmAndGetCallback() verifies the WakeupMessage was scheduled
+        // with the correct delay.
+        delayedEvent.run();
         mTestLooper.dispatchAll();
 
         verify(mMockIkeSession).kill();
@@ -67,5 +78,6 @@ public class VcnGatewayConnectionDisconnectingStateTest extends VcnGatewayConnec
 
         // Should do nothing; already tearing down.
         assertEquals(mGatewayConnection.mDisconnectingState, mGatewayConnection.getCurrentState());
+        verifyTeardownTimeoutAlarmAndGetCallback(false /* expectCanceled */);
     }
 }

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionRetryTimeoutStateTest.java

@@ -29,10 +29,14 @@ import org.junit.runner.RunWith;
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnectionTestBase {
+    private long mFirstRetryInterval;
+
     @Before
     public void setUp() throws Exception {
         super.setUp();
 
+        mFirstRetryInterval = mConfig.getRetryInterval()[0];
+
         mGatewayConnection.setUnderlyingNetwork(TEST_UNDERLYING_NETWORK_RECORD_1);
         mGatewayConnection.transitionTo(mGatewayConnection.mRetryTimeoutState);
         mTestLooper.dispatchAll();
@@ -46,6 +50,7 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
         mTestLooper.dispatchAll();
 
         assertEquals(mGatewayConnection.mConnectingState, mGatewayConnection.getCurrentState());
+        verifyRetryTimeoutAlarmAndGetCallback(mFirstRetryInterval, true /* expectCanceled */);
     }
 
     @Test
@@ -56,6 +61,7 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
         mTestLooper.dispatchAll();
 
         assertEquals(mGatewayConnection.mRetryTimeoutState, mGatewayConnection.getCurrentState());
+        verifyRetryTimeoutAlarmAndGetCallback(mFirstRetryInterval, false /* expectCanceled */);
     }
 
     @Test
@@ -66,13 +72,23 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
         mTestLooper.dispatchAll();
 
         assertEquals(mGatewayConnection.mDisconnectedState, mGatewayConnection.getCurrentState());
+        verifyRetryTimeoutAlarmAndGetCallback(mFirstRetryInterval, true /* expectCanceled */);
     }
 
     @Test
     public void testTimeoutElapsingTriggersRetry() throws Exception {
-        mTestLooper.moveTimeForward(mConfig.getRetryIntervalsMs()[0]);
+        final Runnable delayedEvent =
+                verifyRetryTimeoutAlarmAndGetCallback(
+                        mFirstRetryInterval, false /* expectCanceled */);
+
+        // Can't use mTestLooper to advance the time since VcnGatewayConnection uses WakeupMessages
+        // (which are mocked here). Directly invoke the runnable instead. This is still sufficient,
+        // since verifyRetryTimeoutAlarmAndGetCallback() verifies the WakeupMessage was scheduled
+        // with the correct delay.
+        delayedEvent.run();
         mTestLooper.dispatchAll();
 
         assertEquals(mGatewayConnection.mConnectingState, mGatewayConnection.getCurrentState());
+        verifyRetryTimeoutAlarmAndGetCallback(mFirstRetryInterval, true /* expectCanceled */);
     }
 }

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java

@@ -145,4 +145,19 @@ public class VcnGatewayConnectionTest extends VcnGatewayConnectionTestBase {
 
         verifyWakeLockReleased();
     }
+
+    @Test
+    public void testNonNullUnderlyingNetworkRecordUpdateCancelsAlarm() {
+        mGatewayConnection
+                .getUnderlyingNetworkTrackerCallback()
+                .onSelectedUnderlyingNetworkChanged(null);
+
+        verifyDisconnectRequestAlarmAndGetCallback(false /* expectCanceled */);
+
+        mGatewayConnection
+                .getUnderlyingNetworkTrackerCallback()
+                .onSelectedUnderlyingNetworkChanged(TEST_UNDERLYING_NETWORK_RECORD_1);
+
+        verify(mDisconnectRequestAlarm).cancel();
+    }
 }

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java

@@ -22,8 +22,11 @@ import static com.android.server.vcn.VcnTestUtils.setupIpSecManager;
 
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
@@ -47,6 +50,7 @@ import android.os.ParcelUuid;
 import android.os.PowerManager;
 import android.os.test.TestLooper;
 
+import com.android.internal.util.WakeupMessage;
 import com.android.server.IpSecService;
 import com.android.server.vcn.TelephonySubscriptionTracker.TelephonySubscriptionSnapshot;
 import com.android.server.vcn.Vcn.VcnGatewayStatusCallback;
@@ -59,6 +63,7 @@ import org.mockito.ArgumentCaptor;
 import java.net.InetAddress;
 import java.util.Collections;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 
 public class VcnGatewayConnectionTestBase {
     protected static final ParcelUuid TEST_SUB_GRP = new ParcelUuid(UUID.randomUUID());
@@ -72,6 +77,7 @@ public class VcnGatewayConnectionTestBase {
     protected static final int TEST_IPSEC_TRANSFORM_RESOURCE_ID = 2;
     protected static final int TEST_IPSEC_TUNNEL_RESOURCE_ID = 3;
     protected static final int TEST_SUB_ID = 5;
+    protected static final long ELAPSED_REAL_TIME = 123456789L;
     protected static final String TEST_IPSEC_TUNNEL_IFACE = "IPSEC_IFACE";
     protected static final UnderlyingNetworkRecord TEST_UNDERLYING_NETWORK_RECORD_1 =
             new UnderlyingNetworkRecord(
@@ -99,6 +105,9 @@ public class VcnGatewayConnectionTestBase {
     @NonNull protected final VcnGatewayConnection.Dependencies mDeps;
     @NonNull protected final UnderlyingNetworkTracker mUnderlyingNetworkTracker;
     @NonNull protected final VcnWakeLock mWakeLock;
+    @NonNull protected final WakeupMessage mTeardownTimeoutAlarm;
+    @NonNull protected final WakeupMessage mDisconnectRequestAlarm;
+    @NonNull protected final WakeupMessage mRetryTimeoutAlarm;
 
     @NonNull protected final IpSecService mIpSecSvc;
     @NonNull protected final ConnectivityManager mConnMgr;
@@ -116,6 +125,9 @@ public class VcnGatewayConnectionTestBase {
         mDeps = mock(VcnGatewayConnection.Dependencies.class);
         mUnderlyingNetworkTracker = mock(UnderlyingNetworkTracker.class);
         mWakeLock = mock(VcnWakeLock.class);
+        mTeardownTimeoutAlarm = mock(WakeupMessage.class);
+        mDisconnectRequestAlarm = mock(WakeupMessage.class);
+        mRetryTimeoutAlarm = mock(WakeupMessage.class);
 
         mIpSecSvc = mock(IpSecService.class);
         setupIpSecManager(mContext, mIpSecSvc);
@@ -134,6 +146,16 @@ public class VcnGatewayConnectionTestBase {
         doReturn(mWakeLock)
                 .when(mDeps)
                 .newWakeLock(eq(mContext), eq(PowerManager.PARTIAL_WAKE_LOCK), any());
+
+        setUpWakeupMessage(mTeardownTimeoutAlarm, VcnGatewayConnection.TEARDOWN_TIMEOUT_ALARM);
+        setUpWakeupMessage(mDisconnectRequestAlarm, VcnGatewayConnection.DISCONNECT_REQUEST_ALARM);
+        setUpWakeupMessage(mRetryTimeoutAlarm, VcnGatewayConnection.RETRY_TIMEOUT_ALARM);
+
+        doReturn(ELAPSED_REAL_TIME).when(mDeps).getElapsedRealTime();
+    }
+
+    private void setUpWakeupMessage(@NonNull WakeupMessage msg, @NonNull String cmdName) {
+        doReturn(msg).when(mDeps).newWakeupMessage(eq(mVcnContext), any(), eq(cmdName), any());
     }
 
     @Before
@@ -190,4 +212,45 @@ public class VcnGatewayConnectionTestBase {
         verify(mWakeLock).release();
         verifyNoMoreInteractions(mWakeLock);
     }
+
+    private Runnable verifyWakeupMessageSetUpAndGetCallback(
+            @NonNull String tag,
+            @NonNull WakeupMessage msg,
+            long delayInMillis,
+            boolean expectCanceled) {
+        ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mDeps).newWakeupMessage(eq(mVcnContext), any(), eq(tag), runnableCaptor.capture());
+
+        verify(mDeps, atLeastOnce()).getElapsedRealTime();
+        verify(msg).schedule(ELAPSED_REAL_TIME + delayInMillis);
+        verify(msg, expectCanceled ? times(1) : never()).cancel();
+
+        return runnableCaptor.getValue();
+    }
+
+    protected Runnable verifyTeardownTimeoutAlarmAndGetCallback(boolean expectCanceled) {
+        return verifyWakeupMessageSetUpAndGetCallback(
+                VcnGatewayConnection.TEARDOWN_TIMEOUT_ALARM,
+                mTeardownTimeoutAlarm,
+                TimeUnit.SECONDS.toMillis(VcnGatewayConnection.TEARDOWN_TIMEOUT_SECONDS),
+                expectCanceled);
+    }
+
+    protected Runnable verifyDisconnectRequestAlarmAndGetCallback(boolean expectCanceled) {
+        return verifyWakeupMessageSetUpAndGetCallback(
+                VcnGatewayConnection.DISCONNECT_REQUEST_ALARM,
+                mDisconnectRequestAlarm,
+                TimeUnit.SECONDS.toMillis(
+                        VcnGatewayConnection.NETWORK_LOSS_DISCONNECT_TIMEOUT_SECONDS),
+                expectCanceled);
+    }
+
+    protected Runnable verifyRetryTimeoutAlarmAndGetCallback(
+            long delayInMillis, boolean expectCanceled) {
+        return verifyWakeupMessageSetUpAndGetCallback(
+                VcnGatewayConnection.RETRY_TIMEOUT_ALARM,
+                mRetryTimeoutAlarm,
+                delayInMillis,
+                expectCanceled);
+    }
 }