Use direct-callback alarms in AppTimeLimitController.

Update SessionUsageGroups to use direct-callback alarms.

Bug: 181983817
Test: atest AppTimeLimitControllerTests
Change-Id: I73a755cb9a893952cef33c0d7cc1b2c0bf2a1318

3 files changed, 80 insertions, 27 deletions

diff --git a/services/tests/servicestests/src/com/android/server/usage/AppTimeLimitControllerTests.java b/services/tests/servicestests/src/com/android/server/usage/AppTimeLimitControllerTests.java
index 7d8696139245..b2e5ea07e626 100644
--- a/services/tests/servicestests/src/com/android/server/usage/AppTimeLimitControllerTests.java
+++ b/services/tests/servicestests/src/com/android/server/usage/AppTimeLimitControllerTests.java
@@ -21,13 +21,21 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
-
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.app.AlarmManager;
 import android.app.PendingIntent;
 import android.app.usage.UsageStatsManagerInternal;
 import android.os.HandlerThread;
 import android.os.Looper;
 import android.os.UserHandle;
 
+import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -35,6 +43,9 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 import java.io.PrintWriter;
 import java.io.StringWriter;
@@ -89,6 +100,8 @@ public class AppTimeLimitControllerTests {
 
     private AppTimeLimitController mController;
 
+    @Mock private AlarmManager mMockAlarmManager;
+
     private HandlerThread mThread;
 
     private long mElapsedTime;
@@ -112,7 +125,12 @@ public class AppTimeLimitControllerTests {
     class MyAppTimeLimitController extends AppTimeLimitController {
         MyAppTimeLimitController(AppTimeLimitController.TimeLimitCallbackListener listener,
                 Looper looper) {
-            super(listener, looper);
+            super(InstrumentationRegistry.getContext(), listener, looper);
+        }
+
+        @Override
+        protected AlarmManager getAlarmManager() {
+            return mMockAlarmManager;
         }
 
         @Override
@@ -146,6 +164,8 @@ public class AppTimeLimitControllerTests {
         mThread = new HandlerThread("Test");
         mThread.start();
         mController = new MyAppTimeLimitController(mListener, mThread.getLooper());
+
+        MockitoAnnotations.initMocks(this);
     }
 
     @After
@@ -486,9 +506,14 @@ public class AppTimeLimitControllerTests {
         setTime(6_000L);
         assertTrue(mLimitReachedLatch.await(6_000L, TimeUnit.MILLISECONDS));
         stopUsage(PKG_SOC1);
+        final ArgumentCaptor<AlarmManager.OnAlarmListener> onAlarmListenerArgumentCaptor =
+                ArgumentCaptor.forClass(AlarmManager.OnAlarmListener.class);
+        verify(mMockAlarmManager).setExact(eq(AlarmManager.ELAPSED_REALTIME), anyLong(),
+                anyString(), onAlarmListenerArgumentCaptor.capture(), any());
         // Usage has stopped, Session should end in a second. Verify session end occurs in a second
         // (+/- 100ms, which is hopefully not too slim a margin)
         assertFalse(mSessionEndLatch.await(900L, TimeUnit.MILLISECONDS));
+        onAlarmListenerArgumentCaptor.getValue().onAlarm();
         assertTrue(mSessionEndLatch.await(200L, TimeUnit.MILLISECONDS));
         // Verify that the observer was not removed
         assertTrue(hasUsageSessionObserver(UID, OBS_ID1));
@@ -597,9 +622,14 @@ public class AppTimeLimitControllerTests {
         // Should call back by 11 seconds (6 earlier + 5 now)
         assertTrue(mLimitReachedLatch.await(5_000L, TimeUnit.MILLISECONDS));
         stopUsage(PKG_SOC1);
+        final ArgumentCaptor<AlarmManager.OnAlarmListener> onAlarmListenerArgumentCaptor =
+                ArgumentCaptor.forClass(AlarmManager.OnAlarmListener.class);
+        verify(mMockAlarmManager).setExact(eq(AlarmManager.ELAPSED_REALTIME), anyLong(),
+                anyString(), onAlarmListenerArgumentCaptor.capture(), any());
         // Usage has stopped, Session should end in a second. Verify session end occurs in a second
         // (+/- 100ms, which is hopefully not too slim a margin)
         assertFalse(mSessionEndLatch.await(900L, TimeUnit.MILLISECONDS));
+        onAlarmListenerArgumentCaptor.getValue().onAlarm();
         assertTrue(mSessionEndLatch.await(200L, TimeUnit.MILLISECONDS));
         // Verify that the observer was removed
         assertTrue(hasUsageSessionObserver(UID, OBS_ID1));
@@ -849,9 +879,14 @@ public class AppTimeLimitControllerTests {
         setTime(12_000L);
         assertTrue(mLimitReachedLatch.await(1_000L, TimeUnit.MILLISECONDS));
         stopUsage(PKG_SOC1);
+        final ArgumentCaptor<AlarmManager.OnAlarmListener> onAlarmListenerArgumentCaptor =
+                ArgumentCaptor.forClass(AlarmManager.OnAlarmListener.class);
+        verify(mMockAlarmManager).setExact(eq(AlarmManager.ELAPSED_REALTIME), anyLong(),
+                anyString(), onAlarmListenerArgumentCaptor.capture(), any());
         // Usage has stopped, Session should end in 2 seconds. Verify session end occurs
         // (+/- 100ms, which is hopefully not too slim a margin)
         assertFalse(mSessionEndLatch.await(1_900L, TimeUnit.MILLISECONDS));
+        onAlarmListenerArgumentCaptor.getValue().onAlarm();
         assertTrue(mSessionEndLatch.await(200L, TimeUnit.MILLISECONDS));
         // Verify that the observer was not removed
         assertTrue(hasUsageSessionObserver(UID, OBS_ID1));
@@ -882,9 +917,14 @@ public class AppTimeLimitControllerTests {
         setTime(18_000L);
         assertTrue(mLimitReachedLatch.await(2000L, TimeUnit.MILLISECONDS));
         stopUsage(PKG_SOC1);
+        final ArgumentCaptor<AlarmManager.OnAlarmListener> onAlarmListenerArgumentCaptor =
+                ArgumentCaptor.forClass(AlarmManager.OnAlarmListener.class);
+        verify(mMockAlarmManager).setExact(eq(AlarmManager.ELAPSED_REALTIME), anyLong(),
+                anyString(), onAlarmListenerArgumentCaptor.capture(), any());
         // Usage has stopped, Session should end in 2 seconds. Verify session end occurs
         // (+/- 100ms, which is hopefully not too slim a margin)
         assertFalse(mSessionEndLatch.await(900L, TimeUnit.MILLISECONDS));
+        onAlarmListenerArgumentCaptor.getValue().onAlarm();
         assertTrue(mSessionEndLatch.await(200L, TimeUnit.MILLISECONDS));
         // Verify that the observer was not removed
         assertTrue(hasUsageSessionObserver(UID, OBS_ID1));
@@ -903,9 +943,14 @@ public class AppTimeLimitControllerTests {
         setTime(11_000L);
         assertTrue(mLimitReachedLatch.await(2_000L, TimeUnit.MILLISECONDS));
         stopUsage(PKG_SOC1);
+        final ArgumentCaptor<AlarmManager.OnAlarmListener> onAlarmListenerArgumentCaptor =
+                ArgumentCaptor.forClass(AlarmManager.OnAlarmListener.class);
+        verify(mMockAlarmManager).setExact(eq(AlarmManager.ELAPSED_REALTIME), anyLong(),
+                anyString(), onAlarmListenerArgumentCaptor.capture(), any());
         // Usage has stopped, Session should end in 1 seconds. Verify session end occurs
         // (+/- 100ms, which is hopefully not too slim a margin)
         assertFalse(mSessionEndLatch.await(900L, TimeUnit.MILLISECONDS));
+        onAlarmListenerArgumentCaptor.getValue().onAlarm();
         assertTrue(mSessionEndLatch.await(200L, TimeUnit.MILLISECONDS));
 
         // Rearm the countdown latches
@@ -921,7 +966,10 @@ public class AppTimeLimitControllerTests {
         setTime(31_000L);
         assertTrue(mLimitReachedLatch.await(2_000L, TimeUnit.MILLISECONDS));
         stopUsage(PKG_SOC1);
+        verify(mMockAlarmManager, times(2)).setExact(eq(AlarmManager.ELAPSED_REALTIME), anyLong(),
+                anyString(), onAlarmListenerArgumentCaptor.capture(), any());
         assertFalse(mSessionEndLatch.await(900L, TimeUnit.MILLISECONDS));
+        onAlarmListenerArgumentCaptor.getValue().onAlarm();
         assertTrue(mSessionEndLatch.await(200L, TimeUnit.MILLISECONDS));
         assertTrue(hasUsageSessionObserver(UID, OBS_ID1));
     }
diff --git a/services/usage/java/com/android/server/usage/AppTimeLimitController.java b/services/usage/java/com/android/server/usage/AppTimeLimitController.java
index 1dc1e7776b01..f169926fb3ac 100644
--- a/services/usage/java/com/android/server/usage/AppTimeLimitController.java
+++ b/services/usage/java/com/android/server/usage/AppTimeLimitController.java
@@ -17,8 +17,10 @@
 package com.android.server.usage;
 
 import android.annotation.UserIdInt;
+import android.app.AlarmManager;
 import android.app.PendingIntent;
 import android.app.usage.UsageStatsManagerInternal;
+import android.content.Context;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
@@ -58,6 +60,10 @@ public class AppTimeLimitController {
 
     private final MyHandler mHandler;
 
+    private final Context mContext;
+
+    private AlarmManager mAlarmManager;
+
     private TimeLimitCallbackListener mListener;
 
     private static final long MAX_OBSERVER_PER_UID = 1000;
@@ -434,7 +440,7 @@ public class AppTimeLimitController {
         }
     }
 
-    class SessionUsageGroup extends UsageGroup {
+    class SessionUsageGroup extends UsageGroup implements AlarmManager.OnAlarmListener {
         private long mNewSessionThresholdMs;
         private PendingIntent mSessionEndCallback;
 
@@ -466,7 +472,7 @@ public class AppTimeLimitController {
                     // New session has started, clear usage time.
                     mUsageTimeMs = 0;
                 }
-                AppTimeLimitController.this.cancelInformSessionEndListener(this);
+                getAlarmManager().cancel(this);
             }
             super.noteUsageStart(startTimeMs, currentTimeMs);
         }
@@ -479,10 +485,9 @@ public class AppTimeLimitController {
                 if (mUsageTimeMs >= mTimeLimitMs) {
                     // Usage has ended. Schedule the session end callback to be triggered once
                     // the new session threshold has been reached
-                    AppTimeLimitController.this.postInformSessionEndListenerLocked(this,
-                            mNewSessionThresholdMs);
+                    getAlarmManager().setExact(AlarmManager.ELAPSED_REALTIME,
+                            getElapsedRealtime() + mNewSessionThresholdMs, TAG, this, mHandler);
                 }
-
             }
         }
 
@@ -499,6 +504,13 @@ public class AppTimeLimitController {
         }
 
         @Override
+        public void onAlarm() {
+            synchronized (mLock) {
+                onSessionEnd();
+            }
+        }
+
+        @Override
         @GuardedBy("mLock")
         void dump(PrintWriter pw) {
             super.dump(pw);
@@ -546,7 +558,6 @@ public class AppTimeLimitController {
     private class MyHandler extends Handler {
         static final int MSG_CHECK_TIMEOUT = 1;
         static final int MSG_INFORM_LIMIT_REACHED_LISTENER = 2;
-        static final int MSG_INFORM_SESSION_END = 3;
 
         MyHandler(Looper looper) {
             super(looper);
@@ -565,11 +576,6 @@ public class AppTimeLimitController {
                         ((UsageGroup) msg.obj).onLimitReached();
                     }
                     break;
-                case MSG_INFORM_SESSION_END:
-                    synchronized (mLock) {
-                        ((SessionUsageGroup) msg.obj).onSessionEnd();
-                    }
-                    break;
                 default:
                     super.handleMessage(msg);
                     break;
@@ -577,13 +583,24 @@ public class AppTimeLimitController {
         }
     }
 
-    public AppTimeLimitController(TimeLimitCallbackListener listener, Looper looper) {
+    public AppTimeLimitController(Context context, TimeLimitCallbackListener listener,
+            Looper looper) {
+        mContext = context;
         mHandler = new MyHandler(looper);
         mListener = listener;
     }
 
     /** Overrideable by a test */
     @VisibleForTesting
+    protected AlarmManager getAlarmManager() {
+        if (mAlarmManager == null) {
+            mAlarmManager = mContext.getSystemService(AlarmManager.class);
+        }
+        return mAlarmManager;
+    }
+
+    /** Overrideable by a test */
+    @VisibleForTesting
     protected long getElapsedRealtime() {
         return SystemClock.elapsedRealtime();
     }
@@ -985,18 +1002,6 @@ public class AppTimeLimitController {
     }
 
     @GuardedBy("mLock")
-    private void postInformSessionEndListenerLocked(SessionUsageGroup group, long timeout) {
-        mHandler.sendMessageDelayed(
-                mHandler.obtainMessage(MyHandler.MSG_INFORM_SESSION_END, group),
-                timeout);
-    }
-
-    @GuardedBy("mLock")
-    private void cancelInformSessionEndListener(SessionUsageGroup group) {
-        mHandler.removeMessages(MyHandler.MSG_INFORM_SESSION_END, group);
-    }
-
-    @GuardedBy("mLock")
     private void postCheckTimeoutLocked(UsageGroup group, long timeout) {
         mHandler.sendMessageDelayed(mHandler.obtainMessage(MyHandler.MSG_CHECK_TIMEOUT, group),
                 timeout);
diff --git a/services/usage/java/com/android/server/usage/UsageStatsService.java b/services/usage/java/com/android/server/usage/UsageStatsService.java
index 309673d72dd4..f7e63757bd77 100644
--- a/services/usage/java/com/android/server/usage/UsageStatsService.java
+++ b/services/usage/java/com/android/server/usage/UsageStatsService.java
@@ -252,7 +252,7 @@ public class UsageStatsService extends SystemService implements
 
         mAppStandby = mInjector.getAppStandbyController(getContext());
 
-        mAppTimeLimit = new AppTimeLimitController(
+        mAppTimeLimit = new AppTimeLimitController(getContext(),
                 new AppTimeLimitController.TimeLimitCallbackListener() {
                     @Override
                     public void onLimitReached(int observerId, int userId, long timeLimit,