Ensure public methods are running on handler thread

Store the handler that should be running all the public methods and
use it by comparing its thread with the current thread.

Also refactor the tests to use and run on a test handler thread.

Bug: 273451360
Test: atest FrameworksNetTests
(cherry picked from https://android-review.googlesource.com/q/commit:1b71e3867345d6e1a3bc7c38a582ba9e9f9aee3d)
Merged-In: I053f1043a3b518728c4b30b94eec1d33fbfe39c0
Change-Id: I053f1043a3b518728c4b30b94eec1d33fbfe39c0
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index 7afd93f..fefedd4 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -179,7 +179,7 @@
     private static final int MAX_EVENTS_LOGS = 40;
     private final LocalLog mEventLog = new LocalLog(MAX_EVENTS_LOGS);
 
-    private final KeepaliveStatsTracker mKeepaliveStatsTracker = new KeepaliveStatsTracker();
+    private final KeepaliveStatsTracker mKeepaliveStatsTracker;
     /**
      * Information about a managed keepalive.
      *
@@ -307,6 +307,7 @@
                 mContext, mConnectivityServiceHandler);
 
         mAlarmManager = mDependencies.getAlarmManager(context);
+        mKeepaliveStatsTracker = new KeepaliveStatsTracker(handler);
     }
 
     private void startTcpPollingAlarm(@NonNull AutomaticOnOffKeepalive ki) {
diff --git a/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java b/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
index 290d201..07140c4 100644
--- a/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
@@ -16,6 +16,8 @@
 
 package com.android.server.connectivity;
 
+import android.annotation.NonNull;
+import android.os.Handler;
 import android.os.SystemClock;
 import android.util.Log;
 
@@ -26,6 +28,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 // TODO(b/273451360): Also track KeepaliveLifetimeForCarrier and DailykeepaliveInfoReported
 /**
@@ -38,7 +41,9 @@
 public class KeepaliveStatsTracker {
     private static final String TAG = KeepaliveStatsTracker.class.getSimpleName();
 
-    private final Dependencies mDependencies;
+    @NonNull private final Handler mConnectivityServiceHandler;
+    @NonNull private final Dependencies mDependencies;
+
     // List of duration stats metric where the index is the number of concurrent keepalives.
     // Each DurationForNumOfKeepalive message stores a registered duration and an active duration.
     // Registered duration is the total time spent with mNumRegisteredKeepalive == index.
@@ -50,7 +55,7 @@
     private int mNumActiveKeepalive = 0;
 
     // A timestamp of the most recent time the duration metrics was updated.
-    private long mTimestampSinceLastUpdateDurations;
+    private long mLastUpdateDurationsTimestamp;
 
     /** Dependency class */
     @VisibleForTesting
@@ -62,14 +67,16 @@
         }
     }
 
-    public KeepaliveStatsTracker() {
-        this(new Dependencies());
+    public KeepaliveStatsTracker(@NonNull Handler handler) {
+        this(handler, new Dependencies());
     }
 
     @VisibleForTesting
-    public KeepaliveStatsTracker(Dependencies dependencies) {
-        mDependencies = dependencies;
-        mTimestampSinceLastUpdateDurations = mDependencies.getUptimeMillis();
+    public KeepaliveStatsTracker(@NonNull Handler handler, @NonNull Dependencies dependencies) {
+        mDependencies = Objects.requireNonNull(dependencies);
+        mConnectivityServiceHandler = Objects.requireNonNull(handler);
+
+        mLastUpdateDurationsTimestamp = mDependencies.getUptimeMillis();
     }
 
     /** Ensures the list of duration metrics is large enough for number of registered keepalives. */
@@ -107,7 +114,7 @@
         }
         ensureDurationPerNumOfKeepaliveSize();
 
-        final int durationIncrease = (int) (timeNow - mTimestampSinceLastUpdateDurations);
+        final int durationIncrease = (int) (timeNow - mLastUpdateDurationsTimestamp);
         final DurationForNumOfKeepalive.Builder durationForNumOfRegisteredKeepalive =
                 mDurationPerNumOfKeepalive.get(mNumRegisteredKeepalive);
 
@@ -122,11 +129,13 @@
                 durationForNumOfActiveKeepalive.getKeepaliveActiveDurationsMsec()
                         + durationIncrease);
 
-        mTimestampSinceLastUpdateDurations = timeNow;
+        mLastUpdateDurationsTimestamp = timeNow;
     }
 
     /** Inform the KeepaliveStatsTracker a keepalive has just started and is active. */
     public void onStartKeepalive() {
+        ensureRunningOnHandlerThread();
+
         final long timeNow = mDependencies.getUptimeMillis();
         updateDurationsPerNumOfKeepalive(timeNow);
 
@@ -136,6 +145,8 @@
 
     /** Inform the KeepaliveStatsTracker a keepalive has just been paused. */
     public void onPauseKeepalive() {
+        ensureRunningOnHandlerThread();
+
         final long timeNow = mDependencies.getUptimeMillis();
         updateDurationsPerNumOfKeepalive(timeNow);
 
@@ -144,6 +155,8 @@
 
     /** Inform the KeepaliveStatsTracker a keepalive has just been resumed. */
     public void onResumeKeepalive() {
+        ensureRunningOnHandlerThread();
+
         final long timeNow = mDependencies.getUptimeMillis();
         updateDurationsPerNumOfKeepalive(timeNow);
 
@@ -152,6 +165,8 @@
 
     /** Inform the KeepaliveStatsTracker a keepalive has just been stopped. */
     public void onStopKeepalive(boolean wasActive) {
+        ensureRunningOnHandlerThread();
+
         final long timeNow = mDependencies.getUptimeMillis();
         updateDurationsPerNumOfKeepalive(timeNow);
 
@@ -163,6 +178,8 @@
      * Builds and returns DailykeepaliveInfoReported proto.
      */
     public DailykeepaliveInfoReported buildKeepaliveMetrics() {
+        ensureRunningOnHandlerThread();
+
         final long timeNow = mDependencies.getUptimeMillis();
         updateDurationsPerNumOfKeepalive(timeNow);
 
@@ -185,7 +202,16 @@
 
     /** Resets the stored metrics but maintains the state of keepalives */
     public void resetMetrics() {
+        ensureRunningOnHandlerThread();
+
         mDurationPerNumOfKeepalive.clear();
         ensureDurationPerNumOfKeepaliveSize();
     }
+
+    private void ensureRunningOnHandlerThread() {
+        if (mConnectivityServiceHandler.getLooper().getThread() != Thread.currentThread()) {
+            throw new IllegalStateException(
+                    "Not running on handler thread: " + Thread.currentThread().getName());
+        }
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java b/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
index d262255..2e9bf26 100644
--- a/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
@@ -16,12 +16,17 @@
 
 package com.android.server.connectivity;
 
+import static com.android.testutils.HandlerUtils.visibleOnHandlerThread;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.doReturn;
 
 import android.os.Build;
+import android.os.Handler;
+import android.os.HandlerThread;
 
 import androidx.test.filters.SmallTest;
 
@@ -41,22 +46,84 @@
 @SmallTest
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.TIRAMISU)
 public class KeepaliveStatsTrackerTest {
-    private static final int TEST_UID = 1234;
+    private HandlerThread mHandlerThread;
+    private Handler mTestHandler;
 
     private KeepaliveStatsTracker mKeepaliveStatsTracker;
-    @Mock KeepaliveStatsTracker.Dependencies mDependencies;
+
+    @Mock private KeepaliveStatsTracker.Dependencies mDependencies;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+
+        mHandlerThread = new HandlerThread("KeepaliveStatsTrackerTest");
+        mHandlerThread.start();
+        mTestHandler = new Handler(mHandlerThread.getLooper());
+
         setUptimeMillis(0);
-        mKeepaliveStatsTracker = new KeepaliveStatsTracker(mDependencies);
+        mKeepaliveStatsTracker = new KeepaliveStatsTracker(mTestHandler, mDependencies);
     }
 
     private void setUptimeMillis(long time) {
         doReturn(time).when(mDependencies).getUptimeMillis();
     }
 
+    private DailykeepaliveInfoReported buildKeepaliveMetrics(long time) {
+        setUptimeMillis(time);
+
+        return visibleOnHandlerThread(
+                mTestHandler, () -> mKeepaliveStatsTracker.buildKeepaliveMetrics());
+    }
+
+    private DailykeepaliveInfoReported buildAndResetMetrics(long time) {
+        setUptimeMillis(time);
+
+        return visibleOnHandlerThread(
+                mTestHandler,
+                () -> {
+                    final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
+                            mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                    mKeepaliveStatsTracker.resetMetrics();
+                    return dailyKeepaliveInfoReported;
+                });
+    }
+
+    private void onStartKeepalive(long time) {
+        setUptimeMillis(time);
+        visibleOnHandlerThread(mTestHandler, () -> mKeepaliveStatsTracker.onStartKeepalive());
+    }
+
+    private void onPauseKeepalive(long time) {
+        setUptimeMillis(time);
+        visibleOnHandlerThread(mTestHandler, () -> mKeepaliveStatsTracker.onPauseKeepalive());
+    }
+
+    private void onResumeKeepalive(long time) {
+        setUptimeMillis(time);
+        visibleOnHandlerThread(mTestHandler, () -> mKeepaliveStatsTracker.onResumeKeepalive());
+    }
+
+    private void onStopKeepalive(long time, boolean wasActive) {
+        setUptimeMillis(time);
+        visibleOnHandlerThread(
+                mTestHandler, () -> mKeepaliveStatsTracker.onStopKeepalive(wasActive));
+    }
+
+    @Test
+    public void testEnsureRunningOnHandlerThread() {
+        // Not running on handler thread
+        assertThrows(IllegalStateException.class, () -> mKeepaliveStatsTracker.onStartKeepalive());
+        assertThrows(IllegalStateException.class, () -> mKeepaliveStatsTracker.onPauseKeepalive());
+        assertThrows(IllegalStateException.class, () -> mKeepaliveStatsTracker.onResumeKeepalive());
+        assertThrows(
+                IllegalStateException.class, () -> mKeepaliveStatsTracker.onStopKeepalive(true));
+        assertThrows(
+                IllegalStateException.class, () -> mKeepaliveStatsTracker.buildKeepaliveMetrics());
+        assertThrows(
+                IllegalStateException.class, () -> mKeepaliveStatsTracker.resetMetrics());
+    }
+
     /**
      * Asserts that a DurationPerNumOfKeepalive contains expected values
      *
@@ -111,9 +178,8 @@
     public void testNoKeepalive() {
         final int writeTime = 5000;
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
 
         // Expect that the durations are all in numOfKeepalive = 0.
         final int[] expectRegisteredDurations = new int[] {writeTime};
@@ -137,12 +203,10 @@
         final int startTime = 1000;
         final int writeTime = 5000;
 
-        setUptimeMillis(startTime);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime);
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
 
         // The keepalive is never stopped, expect the duration for numberOfKeepalive of 1 to range
         // from startTime to writeTime.
@@ -167,15 +231,12 @@
         final int pauseTime = 2030;
         final int writeTime = 5000;
 
-        setUptimeMillis(startTime);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime);
 
-        setUptimeMillis(pauseTime);
-        mKeepaliveStatsTracker.onPauseKeepalive();
+        onPauseKeepalive(pauseTime);
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
 
         // The keepalive is paused but not stopped, expect the registered duration for
         // numberOfKeepalive of 1 to still range from startTime to writeTime while the active
@@ -203,18 +264,14 @@
         final int resumeTime = 3450;
         final int writeTime = 5000;
 
-        setUptimeMillis(startTime);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime);
 
-        setUptimeMillis(pauseTime);
-        mKeepaliveStatsTracker.onPauseKeepalive();
+        onPauseKeepalive(pauseTime);
 
-        setUptimeMillis(resumeTime);
-        mKeepaliveStatsTracker.onResumeKeepalive();
+        onResumeKeepalive(resumeTime);
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
 
         // The keepalive is paused and resumed but not stopped, expect the registered duration for
         // numberOfKeepalive of 1 to still range from startTime to writeTime while the active
@@ -246,21 +303,16 @@
         final int stopTime = 4157;
         final int writeTime = 5000;
 
-        setUptimeMillis(startTime);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime);
 
-        setUptimeMillis(pauseTime);
-        mKeepaliveStatsTracker.onPauseKeepalive();
+        onPauseKeepalive(pauseTime);
 
-        setUptimeMillis(resumeTime);
-        mKeepaliveStatsTracker.onResumeKeepalive();
+        onResumeKeepalive(resumeTime);
 
-        setUptimeMillis(stopTime);
-        mKeepaliveStatsTracker.onStopKeepalive(/* wasActive= */ true);
+        onStopKeepalive(stopTime, /* wasActive= */ true);
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
 
         // The keepalive is now stopped, expect the registered duration for numberOfKeepalive of 1
         // to now range from startTime to stopTime while the active duration stops at pauseTime but
@@ -292,18 +344,14 @@
         final int stopTime = 4157;
         final int writeTime = 5000;
 
-        setUptimeMillis(startTime);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime);
 
-        setUptimeMillis(pauseTime);
-        mKeepaliveStatsTracker.onPauseKeepalive();
+        onPauseKeepalive(pauseTime);
 
-        setUptimeMillis(stopTime);
-        mKeepaliveStatsTracker.onStopKeepalive(/* wasActive= */ false);
+        onStopKeepalive(stopTime, /* wasActive= */ false);
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
 
         // The keepalive is stopped while paused, expect the registered duration for
         // numberOfKeepalive of 1 to range from startTime to stopTime while the active duration
@@ -333,24 +381,20 @@
         final int stopTime = 4000;
         final int writeTime = 5000;
 
-        setUptimeMillis(startTime);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime);
 
         for (int i = 0; i < pauseResumeTimes.length; i++) {
-            setUptimeMillis(pauseResumeTimes[i]);
             if (i % 2 == 0) {
-                mKeepaliveStatsTracker.onPauseKeepalive();
+                onPauseKeepalive(pauseResumeTimes[i]);
             } else {
-                mKeepaliveStatsTracker.onResumeKeepalive();
+                onResumeKeepalive(pauseResumeTimes[i]);
             }
         }
 
-        setUptimeMillis(stopTime);
-        mKeepaliveStatsTracker.onStopKeepalive(/* wasActive= */ true);
+        onStopKeepalive(stopTime, /* wasActive= */ true);
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
 
         final int[] expectRegisteredDurations =
                 new int[] {startTime + (writeTime - stopTime), stopTime - startTime};
@@ -387,30 +431,22 @@
         final int stopTime1 = 4157;
         final int writeTime = 5000;
 
-        setUptimeMillis(startTime1);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime1);
 
-        setUptimeMillis(pauseTime1);
-        mKeepaliveStatsTracker.onPauseKeepalive();
+        onPauseKeepalive(pauseTime1);
 
-        setUptimeMillis(startTime2);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime2);
 
-        setUptimeMillis(resumeTime1);
-        mKeepaliveStatsTracker.onResumeKeepalive();
+        onResumeKeepalive(resumeTime1);
 
-        setUptimeMillis(pauseTime2);
-        mKeepaliveStatsTracker.onPauseKeepalive();
+        onPauseKeepalive(pauseTime2);
 
-        setUptimeMillis(resumeTime2);
-        mKeepaliveStatsTracker.onResumeKeepalive();
+        onResumeKeepalive(resumeTime2);
 
-        setUptimeMillis(stopTime1);
-        mKeepaliveStatsTracker.onStopKeepalive(/* wasActive= */ true);
+        onStopKeepalive(stopTime1, /* wasActive= */ true);
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
 
         // With two keepalives, the number of concurrent keepalives can vary from 0-2 depending on
         // both keepalive states.
@@ -458,12 +494,10 @@
         final int stopTime = 7000;
         final int writeTime2 = 10000;
 
-        setUptimeMillis(startTime);
-        mKeepaliveStatsTracker.onStartKeepalive();
+        onStartKeepalive(startTime);
 
-        setUptimeMillis(writeTime);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildAndResetMetrics(writeTime);
 
         // Same expect as testOneKeepalive_startOnly
         final int[] expectRegisteredDurations = new int[] {startTime, writeTime - startTime};
@@ -473,11 +507,9 @@
                 expectRegisteredDurations,
                 expectActiveDurations);
 
-        // Reset metrics
-        mKeepaliveStatsTracker.resetMetrics();
-
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported2 =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime);
+
         // Expect the stored durations to be 0 but still contain the number of keepalive = 1.
         assertDailyKeepaliveInfoReported(
                 dailyKeepaliveInfoReported2,
@@ -485,12 +517,10 @@
                 /* expectActiveDurations= */ new int[] {0, 0});
 
         // Expect that the keepalive is still registered after resetting so it can be stopped.
-        setUptimeMillis(stopTime);
-        mKeepaliveStatsTracker.onStopKeepalive(/* wasActive= */ true);
+        onStopKeepalive(stopTime, /* wasActive= */ true);
 
-        setUptimeMillis(writeTime2);
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported3 =
-                mKeepaliveStatsTracker.buildKeepaliveMetrics();
+                buildKeepaliveMetrics(writeTime2);
 
         final int[] expectRegisteredDurations2 =
                 new int[] {writeTime2 - stopTime, stopTime - writeTime};