Merge "Fix potential OOM caused by DataProcessManager" into main
diff --git a/src/com/android/settings/fuelgauge/batteryusage/BatteryUsageDataLoader.java b/src/com/android/settings/fuelgauge/batteryusage/BatteryUsageDataLoader.java
index df84aba..b917d1f 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/BatteryUsageDataLoader.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/BatteryUsageDataLoader.java
@@ -118,6 +118,7 @@
         final BatteryLevelData batteryLevelData =
                 DataProcessManager.getBatteryLevelData(
                         context,
+                        null,
                         userIdsSeries,
                         /* isFromPeriodJob= */ true,
                         batteryDiffDataMap -> {
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java b/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
index fd548ab..9992313 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
@@ -18,12 +18,12 @@
 
 import android.app.usage.UsageEvents;
 import android.content.Context;
-import android.os.AsyncTask;
 import android.util.ArrayMap;
 import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.lifecycle.Lifecycle;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.settings.fuelgauge.PowerUsageFeatureProvider;
@@ -80,6 +80,7 @@
     private final long mLastFullChargeTimestamp;
     private final boolean mIsFromPeriodJob;
     private final Context mContext;
+    private final @Nullable Lifecycle mLifecycle;
     private final UserIdsSeries mUserIdsSeries;
     private final OnBatteryDiffDataMapLoadedListener mCallbackFunction;
     private final List<AppUsageEvent> mAppUsageEventList = new ArrayList<>();
@@ -120,6 +121,7 @@
     /** Constructor when there exists battery level data. */
     DataProcessManager(
             Context context,
+            @Nullable Lifecycle lifecycle,
             final UserIdsSeries userIdsSeries,
             final boolean isFromPeriodJob,
             final long rawStartTimestamp,
@@ -128,6 +130,7 @@
             @NonNull final List<BatteryLevelData.PeriodBatteryLevelData> hourlyBatteryLevelsPerDay,
             @NonNull final Map<Long, Map<String, BatteryHistEntry>> batteryHistoryMap) {
         mContext = context.getApplicationContext();
+        mLifecycle = lifecycle;
         mUserIdsSeries = userIdsSeries;
         mIsFromPeriodJob = isFromPeriodJob;
         mRawStartTimestamp = rawStartTimestamp;
@@ -140,9 +143,11 @@
     /** Constructor when there is no battery level data. */
     DataProcessManager(
             Context context,
+            @Nullable Lifecycle lifecycle,
             final UserIdsSeries userIdsSeries,
             @NonNull final OnBatteryDiffDataMapLoadedListener callbackFunction) {
         mContext = context.getApplicationContext();
+        mLifecycle = lifecycle;
         mUserIdsSeries = userIdsSeries;
         mCallbackFunction = callbackFunction;
         mIsFromPeriodJob = false;
@@ -223,7 +228,7 @@
     }
 
     private void loadCurrentBatteryHistoryMap() {
-        new AsyncTask<Void, Void, Map<String, BatteryHistEntry>>() {
+        new LifecycleAwareAsyncTask<Map<String, BatteryHistEntry>>(mLifecycle) {
             @Override
             protected Map<String, BatteryHistEntry> doInBackground(Void... voids) {
                 final long startTime = System.currentTimeMillis();
@@ -242,6 +247,7 @@
             @Override
             protected void onPostExecute(
                     final Map<String, BatteryHistEntry> currentBatteryHistoryMap) {
+                super.onPostExecute(currentBatteryHistoryMap);
                 if (mBatteryHistoryMap != null) {
                     // Replaces the placeholder in mBatteryHistoryMap.
                     for (Map.Entry<Long, Map<String, BatteryHistEntry>> mapEntry :
@@ -256,11 +262,11 @@
                 mIsCurrentBatteryHistoryLoaded = true;
                 tryToGenerateFinalDataAndApplyCallback();
             }
-        }.execute();
+        }.start();
     }
 
     private void loadCurrentAppUsageList() {
-        new AsyncTask<Void, Void, List<AppUsageEvent>>() {
+        new LifecycleAwareAsyncTask<List<AppUsageEvent>>(mLifecycle) {
             @Override
             @Nullable
             protected List<AppUsageEvent> doInBackground(Void... voids) {
@@ -299,6 +305,7 @@
 
             @Override
             protected void onPostExecute(final List<AppUsageEvent> currentAppUsageList) {
+                super.onPostExecute(currentAppUsageList);
                 if (currentAppUsageList == null || currentAppUsageList.isEmpty()) {
                     Log.d(TAG, "currentAppUsageList is null or empty");
                 } else {
@@ -307,11 +314,11 @@
                 mIsCurrentAppUsageLoaded = true;
                 tryToProcessAppUsageData();
             }
-        }.execute();
+        }.start();
     }
 
     private void loadDatabaseAppUsageList() {
-        new AsyncTask<Void, Void, List<AppUsageEvent>>() {
+        new LifecycleAwareAsyncTask<List<AppUsageEvent>>(mLifecycle) {
             @Override
             protected List<AppUsageEvent> doInBackground(Void... voids) {
                 if (!shouldLoadAppUsageData()) {
@@ -337,6 +344,7 @@
 
             @Override
             protected void onPostExecute(final List<AppUsageEvent> databaseAppUsageList) {
+                super.onPostExecute(databaseAppUsageList);
                 if (databaseAppUsageList == null || databaseAppUsageList.isEmpty()) {
                     Log.d(TAG, "databaseAppUsageList is null or empty");
                 } else {
@@ -345,11 +353,11 @@
                 mIsDatabaseAppUsageLoaded = true;
                 tryToProcessAppUsageData();
             }
-        }.execute();
+        }.start();
     }
 
     private void loadPowerConnectionBatteryEventList() {
-        new AsyncTask<Void, Void, List<BatteryEvent>>() {
+        new LifecycleAwareAsyncTask<List<BatteryEvent>>(mLifecycle) {
             @Override
             protected List<BatteryEvent> doInBackground(Void... voids) {
                 final long startTime = System.currentTimeMillis();
@@ -370,6 +378,7 @@
 
             @Override
             protected void onPostExecute(final List<BatteryEvent> batteryEventList) {
+                super.onPostExecute(batteryEventList);
                 if (batteryEventList == null || batteryEventList.isEmpty()) {
                     Log.d(TAG, "batteryEventList is null or empty");
                 } else {
@@ -379,11 +388,11 @@
                 mIsBatteryEventLoaded = true;
                 tryToProcessAppUsageData();
             }
-        }.execute();
+        }.start();
     }
 
     private void loadBatteryUsageSlotList() {
-        new AsyncTask<Void, Void, List<BatteryUsageSlot>>() {
+        new LifecycleAwareAsyncTask<List<BatteryUsageSlot>>(mLifecycle) {
             @Override
             protected List<BatteryUsageSlot> doInBackground(Void... voids) {
                 final long startTime = System.currentTimeMillis();
@@ -402,6 +411,7 @@
 
             @Override
             protected void onPostExecute(final List<BatteryUsageSlot> batteryUsageSlotList) {
+                super.onPostExecute(batteryUsageSlotList);
                 if (batteryUsageSlotList == null || batteryUsageSlotList.isEmpty()) {
                     Log.d(TAG, "batteryUsageSlotList is null or empty");
                 } else {
@@ -411,11 +421,11 @@
                 mIsBatteryUsageSlotLoaded = true;
                 tryToGenerateFinalDataAndApplyCallback();
             }
-        }.execute();
+        }.start();
     }
 
     private void loadAndApplyBatteryMapFromServiceOnly() {
-        new AsyncTask<Void, Void, Map<Long, BatteryDiffData>>() {
+        new LifecycleAwareAsyncTask<Map<Long, BatteryDiffData>>(mLifecycle) {
             @Override
             protected Map<Long, BatteryDiffData> doInBackground(Void... voids) {
                 final long startTime = System.currentTimeMillis();
@@ -437,11 +447,12 @@
 
             @Override
             protected void onPostExecute(final Map<Long, BatteryDiffData> batteryDiffDataMap) {
+                super.onPostExecute(batteryDiffDataMap);
                 if (mCallbackFunction != null) {
                     mCallbackFunction.onBatteryDiffDataMapLoaded(batteryDiffDataMap);
                 }
             }
-        }.execute();
+        }.start();
     }
 
     private void tryToProcessAppUsageData() {
@@ -481,7 +492,7 @@
     }
 
     private synchronized void generateFinalDataAndApplyCallback() {
-        new AsyncTask<Void, Void, Map<Long, BatteryDiffData>>() {
+        new LifecycleAwareAsyncTask<Map<Long, BatteryDiffData>>(mLifecycle) {
             @Override
             protected Map<Long, BatteryDiffData> doInBackground(Void... voids) {
                 final long startTime = System.currentTimeMillis();
@@ -523,11 +534,12 @@
 
             @Override
             protected void onPostExecute(final Map<Long, BatteryDiffData> batteryDiffDataMap) {
+                super.onPostExecute(batteryDiffDataMap);
                 if (mCallbackFunction != null) {
                     mCallbackFunction.onBatteryDiffDataMapLoaded(batteryDiffDataMap);
                 }
             }
-        }.execute();
+        }.start();
     }
 
     // Whether we should load app usage data from service or database.
@@ -566,6 +578,7 @@
     @Nullable
     public static BatteryLevelData getBatteryLevelData(
             Context context,
+            @Nullable Lifecycle lifecycle,
             final UserIdsSeries userIdsSeries,
             final boolean isFromPeriodJob,
             final OnBatteryDiffDataMapLoadedListener onBatteryUsageMapLoadedListener) {
@@ -585,6 +598,7 @@
         final BatteryLevelData batteryLevelData =
                 getPeriodBatteryLevelData(
                         context,
+                        lifecycle,
                         userIdsSeries,
                         startTimestamp,
                         lastFullChargeTime,
@@ -604,6 +618,7 @@
 
     private static BatteryLevelData getPeriodBatteryLevelData(
             Context context,
+            @Nullable Lifecycle lifecycle,
             final UserIdsSeries userIdsSeries,
             final long startTimestamp,
             final long lastFullChargeTime,
@@ -631,7 +646,8 @@
                                 lastFullChargeTime);
         if (batteryHistoryMap == null || batteryHistoryMap.isEmpty()) {
             Log.d(TAG, "batteryHistoryMap is null in getPeriodBatteryLevelData()");
-            new DataProcessManager(context, userIdsSeries, onBatteryDiffDataMapLoadedListener)
+            new DataProcessManager(context, lifecycle, userIdsSeries,
+                    onBatteryDiffDataMapLoadedListener)
                     .start();
             return null;
         }
@@ -660,7 +676,8 @@
                 DataProcessor.getLevelDataThroughProcessedHistoryMap(
                         context, processedBatteryHistoryMap);
         if (batteryLevelData == null) {
-            new DataProcessManager(context, userIdsSeries, onBatteryDiffDataMapLoadedListener)
+            new DataProcessManager(context, lifecycle, userIdsSeries,
+                    onBatteryDiffDataMapLoadedListener)
                     .start();
             Log.d(TAG, "getBatteryLevelData() returns null");
             return null;
@@ -669,6 +686,7 @@
         // Start the async task to compute diff usage data and load labels and icons.
         new DataProcessManager(
                         context,
+                        lifecycle,
                         userIdsSeries,
                         isFromPeriodJob,
                         startTimestamp,
diff --git a/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTask.kt b/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTask.kt
new file mode 100644
index 0000000..a715cb7
--- /dev/null
+++ b/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTask.kt
@@ -0,0 +1,61 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.fuelgauge.batteryusage
+
+import android.os.AsyncTask
+import androidx.annotation.CallSuper
+import androidx.lifecycle.DefaultLifecycleObserver
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import com.android.settingslib.datastore.HandlerExecutor.Companion.main as mainExecutor
+
+/**
+ * Lifecycle aware [AsyncTask] to cancel task automatically when [lifecycle] is stopped.
+ *
+ * Must call [start] instead of [execute] to run the task.
+ */
+abstract class LifecycleAwareAsyncTask<Result>(private val lifecycle: Lifecycle?) :
+    AsyncTask<Void, Void, Result>(), DefaultLifecycleObserver {
+
+    @CallSuper
+    override fun onPostExecute(result: Result) {
+        lifecycle?.removeObserver(this)
+    }
+
+    override fun onStop(owner: LifecycleOwner) {
+        cancel(false)
+        lifecycle?.removeObserver(this)
+    }
+
+    /**
+     * Starts the task, which invokes [execute] (cannot override [execute] as it is final).
+     *
+     * This method is expected to be invoked from main thread but current usage might call from
+     * background thread.
+     */
+    fun start() {
+        execute() // expects main thread
+        val lifecycle = lifecycle ?: return
+        mainExecutor.execute {
+            // Status is updated to FINISHED if onPoseExecute happened before. And task is cancelled
+            // if lifecycle is stopped.
+            if (status == Status.RUNNING && !isCancelled) {
+                lifecycle.addObserver(this) // requires main thread
+            }
+        }
+    }
+}
diff --git a/src/com/android/settings/fuelgauge/batteryusage/PowerUsageAdvanced.java b/src/com/android/settings/fuelgauge/batteryusage/PowerUsageAdvanced.java
index 1ed6a74..9943d0b 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/PowerUsageAdvanced.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/PowerUsageAdvanced.java
@@ -503,9 +503,11 @@
 
                 @Override
                 public BatteryLevelData loadInBackground() {
+                    Context context = getContext();
                     return DataProcessManager.getBatteryLevelData(
-                            getContext(),
-                            new UserIdsSeries(getContext(), /* isNonUIRequest= */ false),
+                            context,
+                            getLifecycle(),
+                            new UserIdsSeries(context, /* isNonUIRequest= */ false),
                             /* isFromPeriodJob= */ false,
                             PowerUsageAdvanced.this::onBatteryDiffDataMapUpdate);
                 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessManagerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessManagerTest.java
index 1fed13f..e71d6a3 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessManagerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessManagerTest.java
@@ -110,6 +110,7 @@
         mDataProcessManager =
                 new DataProcessManager(
                         mContext,
+                        null,
                         mUserIdsSeries,
                         /* isFromPeriodJob= */ false,
                         /* rawStartTimestamp= */ 0L,
@@ -130,6 +131,7 @@
         final DataProcessManager dataProcessManager =
                 new DataProcessManager(
                         mContext,
+                        null,
                         mUserIdsSeries,
                         /* callbackFunction= */ null);
         assertThat(dataProcessManager.getShowScreenOnTime()).isFalse();
@@ -255,6 +257,7 @@
         final DataProcessManager dataProcessManager =
                 new DataProcessManager(
                         mContext,
+                        null,
                         mUserIdsSeries,
                         /* isFromPeriodJob= */ false,
                         /* rawStartTimestamp= */ 2L,
@@ -346,6 +349,7 @@
         assertThat(
                         DataProcessManager.getBatteryLevelData(
                                 mContext,
+                                null,
                                 mUserIdsSeries,
                                 /* isFromPeriodJob= */ false,
                                 /* asyncResponseDelegate= */ null))
@@ -353,6 +357,7 @@
         assertThat(
                         DataProcessManager.getBatteryLevelData(
                                 mContext,
+                                null,
                                 mUserIdsSeries,
                                 /* isFromPeriodJob= */ true,
                                 /* asyncResponseDelegate= */ null))
@@ -374,6 +379,7 @@
         final BatteryLevelData resultData =
                 DataProcessManager.getBatteryLevelData(
                         mContext,
+                        null,
                         mUserIdsSeries,
                         /* isFromPeriodJob= */ false,
                         /* asyncResponseDelegate= */ null);
@@ -402,6 +408,7 @@
         final BatteryLevelData resultData =
                 DataProcessManager.getBatteryLevelData(
                         mContext,
+                        null,
                         mUserIdsSeries,
                         /* isFromPeriodJob= */ false,
                         /* asyncResponseDelegate= */ null);
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTaskTest.kt b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTaskTest.kt
new file mode 100644
index 0000000..fde45b7
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTaskTest.kt
@@ -0,0 +1,150 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.fuelgauge.batteryusage
+
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleObserver
+import androidx.lifecycle.LifecycleOwner
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+import java.util.concurrent.CountDownLatch
+
+@RunWith(AndroidJUnit4::class)
+class LifecycleAwareAsyncTaskTest {
+    private val instrumentation = InstrumentationRegistry.getInstrumentation()
+    private val lifecycle = mock<Lifecycle>()
+    private val lifecycleOwner = mock<LifecycleOwner>()
+
+    @Test
+    fun addObserver_onPostExecute_onStop() {
+        val taskBeginCountDownLatch = CountDownLatch(1)
+        val taskEndCountDownLatch = CountDownLatch(1)
+        val asyncTask =
+            object : LifecycleAwareAsyncTask<Void?>(lifecycle) {
+                override fun doInBackground(vararg params: Void): Void? {
+                    taskBeginCountDownLatch.await()
+                    taskEndCountDownLatch.countDown()
+                    return null
+                }
+            }
+
+        asyncTask.start()
+        taskBeginCountDownLatch.countDown()
+        verify(lifecycle).addObserver(asyncTask)
+
+        taskEndCountDownLatch.await()
+        asyncTask.onStop(lifecycleOwner)
+        assertThat(asyncTask.isCancelled).isTrue()
+        verify(lifecycle).removeObserver(asyncTask)
+    }
+
+    @Test
+    fun addObserver_onStop() {
+        val executorBlocker = CountDownLatch(1)
+        object : LifecycleAwareAsyncTask<Void?>(null) {
+                override fun doInBackground(vararg params: Void?): Void? {
+                    executorBlocker.await()
+                    return null
+                }
+            }
+            .start()
+
+        val asyncTask =
+            object : LifecycleAwareAsyncTask<Void?>(lifecycle) {
+                override fun doInBackground(vararg params: Void) = null
+            }
+
+        asyncTask.start()
+        verify(lifecycle).addObserver(asyncTask)
+
+        asyncTask.onStop(lifecycleOwner)
+        executorBlocker.countDown()
+        assertThat(asyncTask.isCancelled).isTrue()
+        verify(lifecycle).removeObserver(asyncTask)
+    }
+
+    @Test
+    fun onPostExecute_addObserver() {
+        val observers = mutableListOf<LifecycleObserver>()
+        val lifecycle =
+            object : Lifecycle() {
+                override val currentState: State
+                    get() = State.RESUMED
+
+                override fun addObserver(observer: LifecycleObserver) {
+                    observers.add(observer)
+                }
+
+                override fun removeObserver(observer: LifecycleObserver) {
+                    observers.remove(observer)
+                }
+            }
+        val asyncTask =
+            object : LifecycleAwareAsyncTask<Void?>(lifecycle) {
+                override fun doInBackground(vararg params: Void) = null
+            }
+
+        Thread { asyncTask.start() }.start()
+        idleAsyncTaskExecutor()
+        instrumentation.waitForIdleSync()
+
+        assertThat(observers).isEmpty()
+    }
+
+    private fun idleAsyncTaskExecutor() {
+        val taskCountDownLatch = CountDownLatch(1)
+        object : LifecycleAwareAsyncTask<Void?>(null) {
+                override fun doInBackground(vararg params: Void): Void? {
+                    taskCountDownLatch.countDown()
+                    return null
+                }
+            }
+            .start()
+        taskCountDownLatch.await()
+    }
+
+    @Test
+    fun onStop_addObserver() {
+        val executorBlocker = CountDownLatch(1)
+        object : LifecycleAwareAsyncTask<Void?>(null) {
+                override fun doInBackground(vararg params: Void?): Void? {
+                    executorBlocker.await()
+                    return null
+                }
+            }
+            .start()
+
+        val asyncTask =
+            object : LifecycleAwareAsyncTask<Void?>(lifecycle) {
+                override fun doInBackground(vararg params: Void) = null
+            }
+
+        asyncTask.onStop(lifecycleOwner)
+        assertThat(asyncTask.isCancelled).isTrue()
+        verify(lifecycle).removeObserver(asyncTask)
+
+        asyncTask.start()
+        verify(lifecycle, never()).addObserver(asyncTask)
+        executorBlocker.countDown()
+    }
+}