Add Task Root package info to UsageEvent

Bug: 113094946
Test: manual (use "adb shell dumpsys usagestats apptimelimit" to verify
apps at the root of tasks are considered active)
Test: atest cts/tests/tests/app.usage/src/android/app/usage/cts/UsageStatsTest.java#testTaskRootEventField
Test: atest cts/tests/tests/app.usage/src/android/app/usage/cts/UsageStatsTest.java#testUsageSourceAttribution
Test: atest UsageStatsDatabaseTest

Change-Id: I40f86743d33c13892de0e59ae02c9ebddb606ee7
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index 089847d..2f77ed6 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -2864,16 +2864,18 @@
      * @param userId
      * @param event
      * @param appToken ActivityRecord's appToken.
+     * @param taskRoot TaskRecord's root
      */
     public void updateActivityUsageStats(ComponentName activity, int userId, int event,
-            IBinder appToken) {
+            IBinder appToken, ComponentName taskRoot) {
         if (DEBUG_SWITCH) {
             Slog.d(TAG_SWITCH, "updateActivityUsageStats: comp="
                     + activity + " hash=" + appToken.hashCode() + " event=" + event);
         }
         synchronized (this) {
             if (mUsageStatsService != null) {
-                mUsageStatsService.reportEvent(activity, userId, event, appToken.hashCode());
+                mUsageStatsService.reportEvent(activity, userId, event, appToken.hashCode(),
+                        taskRoot);
             }
         }
     }
@@ -2911,7 +2913,7 @@
             if (mUsageStatsService != null) {
                 mUsageStatsService.reportEvent(service, userId,
                         started ? UsageEvents.Event.FOREGROUND_SERVICE_START
-                                : UsageEvents.Event.FOREGROUND_SERVICE_STOP, 0);
+                                : UsageEvents.Event.FOREGROUND_SERVICE_STOP, 0, null);
             }
         }
     }
@@ -17521,10 +17523,10 @@
 
         @Override
         public void updateActivityUsageStats(ComponentName activity, int userId, int event,
-                IBinder appToken) {
+                IBinder appToken, ComponentName taskRoot) {
             synchronized (ActivityManagerService.this) {
                 ActivityManagerService.this.updateActivityUsageStats(activity, userId, event,
-                        appToken);
+                        appToken, taskRoot);
             }
         }
 
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 61c4863..2affa97 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -5313,9 +5313,18 @@
     }
 
     void updateActivityUsageStats(ActivityRecord activity, int event) {
+        ComponentName taskRoot = null;
+        final TaskRecord task = activity.getTaskRecord();
+        if (task != null) {
+            final ActivityRecord rootActivity = task.getRootActivity();
+            if (rootActivity != null) {
+                taskRoot = rootActivity.mActivityComponent;
+            }
+        }
+
         final Message m = PooledLambda.obtainMessage(
                 ActivityManagerInternal::updateActivityUsageStats, mAmInternal,
-                activity.mActivityComponent, activity.mUserId, event, activity.appToken);
+                activity.mActivityComponent, activity.mUserId, event, activity.appToken, taskRoot);
         mH.sendMessage(m);
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/usage/UsageStatsDatabaseTest.java b/services/tests/servicestests/src/com/android/server/usage/UsageStatsDatabaseTest.java
index 860656b..8d9b3cf 100644
--- a/services/tests/servicestests/src/com/android/server/usage/UsageStatsDatabaseTest.java
+++ b/services/tests/servicestests/src/com/android/server/usage/UsageStatsDatabaseTest.java
@@ -45,6 +45,8 @@
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class UsageStatsDatabaseTest {
+
+    private static final int MAX_TESTED_VERSION = 4;
     protected Context mContext;
     private UsageStatsDatabase mUsageStatsDatabase;
     private File mTestDir;
@@ -131,8 +133,8 @@
 
         for (int i = 0; i < numberOfEvents; i++) {
             Event event = new Event();
-            final int packageInt = ((i / 3) % 7);
-            event.mPackage = "fake.package.name" + packageInt; //clusters of 3 events from 7 "apps"
+            final int packageInt = ((i / 3) % 7); //clusters of 3 events from 7 "apps"
+            event.mPackage = "fake.package.name" + packageInt;
             if (packageInt == 3) {
                 // Third app is an instant app
                 event.mFlags |= Event.FLAG_IS_PACKAGE_INSTANT_APP;
@@ -144,6 +146,13 @@
             event.mEventType = i % (MAX_EVENT_TYPE + 1); //"random" event type
             event.mInstanceId = instanceId;
 
+
+            final int rootPackageInt = (i % 5); // 5 "apps" start each task
+            event.mTaskRootPackage = "fake.package.name" + rootPackageInt;
+
+            final int rootClassInt = i % 6;
+            event.mTaskRootClass = ".fake.class.name" + rootClassInt;
+
             switch (event.mEventType) {
                 case Event.CONFIGURATION_CHANGE:
                     //empty config,
@@ -163,7 +172,7 @@
                     break;
             }
 
-            mIntervalStats.events.insert(event);
+            mIntervalStats.addEvent(event);
             mIntervalStats.update(event.mPackage, event.mClass, event.mTimeStamp, event.mEventType,
                     event.mInstanceId);
 
@@ -234,31 +243,40 @@
         assertEquals(us1.mChooserCounts, us2.mChooserCounts);
     }
 
-    void compareUsageEvent(Event e1, Event e2, int debugId) {
-        assertEquals(e1.mPackage, e2.mPackage, "Usage event " + debugId);
-        assertEquals(e1.mClass, e2.mClass, "Usage event " + debugId);
-        assertEquals(e1.mTimeStamp, e2.mTimeStamp, "Usage event " + debugId);
-        assertEquals(e1.mEventType, e2.mEventType, "Usage event " + debugId);
-        switch (e1.mEventType) {
-            case Event.CONFIGURATION_CHANGE:
-                assertEquals(e1.mConfiguration, e2.mConfiguration,
-                        "Usage event " + debugId + e2.mConfiguration.toString());
-                break;
-            case Event.SHORTCUT_INVOCATION:
-                assertEquals(e1.mShortcutId, e2.mShortcutId, "Usage event " + debugId);
-                break;
-            case Event.STANDBY_BUCKET_CHANGED:
-                assertEquals(e1.mBucketAndReason, e2.mBucketAndReason, "Usage event " + debugId);
-                break;
-            case Event.NOTIFICATION_INTERRUPTION:
-                assertEquals(e1.mNotificationChannelId, e2.mNotificationChannelId,
-                        "Usage event " + debugId);
-                break;
+    void compareUsageEvent(Event e1, Event e2, int debugId, int minVersion) {
+        switch (minVersion) {
+            case 4: // test fields added in version 4
+                assertEquals(e1.mInstanceId, e2.mInstanceId, "Usage event " + debugId);
+                assertEquals(e1.mTaskRootPackage, e2.mTaskRootPackage, "Usage event " + debugId);
+                assertEquals(e1.mTaskRootClass, e2.mTaskRootClass, "Usage event " + debugId);
+                // fallthrough
+            default:
+                assertEquals(e1.mPackage, e2.mPackage, "Usage event " + debugId);
+                assertEquals(e1.mClass, e2.mClass, "Usage event " + debugId);
+                assertEquals(e1.mTimeStamp, e2.mTimeStamp, "Usage event " + debugId);
+                assertEquals(e1.mEventType, e2.mEventType, "Usage event " + debugId);
+                switch (e1.mEventType) {
+                    case Event.CONFIGURATION_CHANGE:
+                        assertEquals(e1.mConfiguration, e2.mConfiguration,
+                                "Usage event " + debugId + e2.mConfiguration.toString());
+                        break;
+                    case Event.SHORTCUT_INVOCATION:
+                        assertEquals(e1.mShortcutId, e2.mShortcutId, "Usage event " + debugId);
+                        break;
+                    case Event.STANDBY_BUCKET_CHANGED:
+                        assertEquals(e1.mBucketAndReason, e2.mBucketAndReason,
+                                "Usage event " + debugId);
+                        break;
+                    case Event.NOTIFICATION_INTERRUPTION:
+                        assertEquals(e1.mNotificationChannelId, e2.mNotificationChannelId,
+                                "Usage event " + debugId);
+                        break;
+                }
+                assertEquals(e1.mFlags, e2.mFlags);
         }
-        assertEquals(e1.mFlags, e2.mFlags);
     }
 
-    void compareIntervalStats(IntervalStats stats1, IntervalStats stats2) {
+    void compareIntervalStats(IntervalStats stats1, IntervalStats stats2, int minVersion) {
         assertEquals(stats1.majorVersion, stats2.majorVersion);
         assertEquals(stats1.minorVersion, stats2.minorVersion);
         assertEquals(stats1.beginTime, stats2.beginTime);
@@ -311,7 +329,7 @@
         } else {
             assertEquals(stats1.events.size(), stats2.events.size());
             for (int i = 0; i < stats1.events.size(); i++) {
-                compareUsageEvent(stats1.events.get(i), stats2.events.get(i), i);
+                compareUsageEvent(stats1.events.get(i), stats2.events.get(i), i, minVersion);
             }
         }
     }
@@ -326,7 +344,7 @@
                 mIntervalStatsVerifier);
 
         assertEquals(1, stats.size());
-        compareIntervalStats(mIntervalStats, stats.get(0));
+        compareIntervalStats(mIntervalStats, stats.get(0), MAX_TESTED_VERSION);
     }
 
     /**
@@ -359,8 +377,10 @@
                 mIntervalStatsVerifier);
 
         assertEquals(1, stats.size());
+
+        final int minVersion = oldVersion < newVersion ? oldVersion : newVersion;
         // The written and read IntervalStats should match
-        compareIntervalStats(mIntervalStats, stats.get(0));
+        compareIntervalStats(mIntervalStats, stats.get(0), minVersion);
     }
 
     /**
@@ -401,7 +421,7 @@
         if (mIntervalStats.events != null) mIntervalStats.events.clear();
 
         // The written and read IntervalStats should match
-        compareIntervalStats(mIntervalStats, stats.get(0));
+        compareIntervalStats(mIntervalStats, stats.get(0), version);
     }
 
     /**
diff --git a/services/usage/java/com/android/server/usage/IntervalStats.java b/services/usage/java/com/android/server/usage/IntervalStats.java
index 9a5bd13..f1ddfe4 100644
--- a/services/usage/java/com/android/server/usage/IntervalStats.java
+++ b/services/usage/java/com/android/server/usage/IntervalStats.java
@@ -218,6 +218,14 @@
                 case (int) IntervalStatsProto.Event.INSTANCE_ID:
                     event.mInstanceId = parser.readInt(IntervalStatsProto.Event.INSTANCE_ID);
                     break;
+                case (int) IntervalStatsProto.Event.TASK_ROOT_PACKAGE_INDEX:
+                    event.mTaskRootPackage = getCachedStringRef(stringPool.get(
+                            parser.readInt(IntervalStatsProto.Event.TASK_ROOT_PACKAGE_INDEX) - 1));
+                    break;
+                case (int) IntervalStatsProto.Event.TASK_ROOT_CLASS_INDEX:
+                    event.mTaskRootClass = getCachedStringRef(stringPool.get(
+                            parser.readInt(IntervalStatsProto.Event.TASK_ROOT_CLASS_INDEX) - 1));
+                    break;
                 case ProtoInputStream.NO_MORE_FIELDS:
                     // Handle default values for certain events types
                     switch (event.mEventType) {
@@ -332,6 +340,12 @@
         if (event.mClass != null) {
             event.mClass = getCachedStringRef(event.mClass);
         }
+        if (event.mTaskRootPackage != null) {
+            event.mTaskRootPackage = getCachedStringRef(event.mTaskRootPackage);
+        }
+        if (event.mTaskRootClass != null) {
+            event.mTaskRootClass = getCachedStringRef(event.mTaskRootClass);
+        }
         if (event.mEventType == NOTIFICATION_INTERRUPTION) {
             event.mNotificationChannelId = getCachedStringRef(event.mNotificationChannelId);
         }
diff --git a/services/usage/java/com/android/server/usage/UsageStatsProto.java b/services/usage/java/com/android/server/usage/UsageStatsProto.java
index d706537..11d49eb 100644
--- a/services/usage/java/com/android/server/usage/UsageStatsProto.java
+++ b/services/usage/java/com/android/server/usage/UsageStatsProto.java
@@ -442,6 +442,28 @@
         proto.write(IntervalStatsProto.Event.FLAGS, event.mFlags);
         proto.write(IntervalStatsProto.Event.TYPE, event.mEventType);
         proto.write(IntervalStatsProto.Event.INSTANCE_ID, event.mInstanceId);
+        if (event.mTaskRootPackage != null) {
+            final int taskRootPackageIndex = stats.mStringCache.indexOf(event.mTaskRootPackage);
+            if (taskRootPackageIndex >= 0) {
+                proto.write(IntervalStatsProto.Event.TASK_ROOT_PACKAGE_INDEX,
+                        taskRootPackageIndex + 1);
+            } else {
+                // Task root package not in Stringpool for some reason.
+                Slog.w(TAG, "Usage event task root package name (" + event.mTaskRootPackage
+                        + ") not found in IntervalStats string cache");
+            }
+        }
+        if (event.mTaskRootClass != null) {
+            final int taskRootClassIndex = stats.mStringCache.indexOf(event.mTaskRootClass);
+            if (taskRootClassIndex >= 0) {
+                proto.write(IntervalStatsProto.Event.TASK_ROOT_CLASS_INDEX,
+                        taskRootClassIndex + 1);
+            } else {
+                // Task root class not in Stringpool for some reason.
+                Slog.w(TAG, "Usage event task root class name (" + event.mTaskRootClass
+                        + ") not found in IntervalStats string cache");
+            }
+        }
         switch (event.mEventType) {
             case UsageEvents.Event.CONFIGURATION_CHANGE:
                 if (event.mConfiguration != null) {
diff --git a/services/usage/java/com/android/server/usage/UsageStatsService.java b/services/usage/java/com/android/server/usage/UsageStatsService.java
index 76a3aa8..6ad698b 100644
--- a/services/usage/java/com/android/server/usage/UsageStatsService.java
+++ b/services/usage/java/com/android/server/usage/UsageStatsService.java
@@ -22,6 +22,8 @@
 import static android.app.usage.UsageEvents.Event.FLUSH_TO_DISK;
 import static android.app.usage.UsageEvents.Event.NOTIFICATION_INTERRUPTION;
 import static android.app.usage.UsageEvents.Event.SHORTCUT_INVOCATION;
+import static android.app.usage.UsageStatsManager.USAGE_SOURCE_CURRENT_ACTIVITY;
+import static android.app.usage.UsageStatsManager.USAGE_SOURCE_TASK_ROOT_ACTIVITY;
 
 import android.Manifest;
 import android.app.ActivityManager;
@@ -39,6 +41,7 @@
 import android.app.usage.UsageStats;
 import android.app.usage.UsageStatsManager;
 import android.app.usage.UsageStatsManager.StandbyBuckets;
+import android.app.usage.UsageStatsManager.UsageSource;
 import android.app.usage.UsageStatsManagerInternal;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
@@ -65,6 +68,7 @@
 import android.os.SystemProperties;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.provider.Settings;
 import android.util.ArraySet;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -132,6 +136,7 @@
     private File mUsageStatsDir;
     long mRealTimeSnapshot;
     long mSystemTimeSnapshot;
+    int mUsageSource;
 
     /** Manages the standby state of apps. */
     AppStandbyController mAppStandby;
@@ -258,6 +263,7 @@
             } else {
                 Slog.w(TAG, "Missing procfs interface: " + KERNEL_COUNTER_FILE);
             }
+            readUsageSourceSetting();
         }
     }
 
@@ -268,6 +274,13 @@
         return mDpmInternal;
     }
 
+    private void readUsageSourceSetting() {
+        synchronized (mLock) {
+            mUsageSource = Settings.Global.getInt(getContext().getContentResolver(),
+                    Settings.Global.APP_TIME_LIMIT_USAGE_SOURCE, USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+        }
+    }
+
     private class UserActionsReceiver extends BroadcastReceiver {
         @Override
         public void onReceive(Context context, Intent intent) {
@@ -459,12 +472,28 @@
             service.reportEvent(event);
 
             mAppStandby.reportEvent(event, elapsedRealtime, userId);
+
+            String packageName;
+
+            switch(mUsageSource) {
+                case USAGE_SOURCE_CURRENT_ACTIVITY:
+                    packageName = event.getPackageName();
+                    break;
+                case USAGE_SOURCE_TASK_ROOT_ACTIVITY:
+                default:
+                    packageName = event.getTaskRootPackageName();
+                    if (packageName == null) {
+                        packageName = event.getPackageName();
+                    }
+                    break;
+            }
+
             switch (event.mEventType) {
                 case Event.ACTIVITY_RESUMED:
                     synchronized (mVisibleActivities) {
                         mVisibleActivities.put(event.mInstanceId, event.getClassName());
                         try {
-                            mAppTimeLimit.noteUsageStart(event.getPackageName(), userId);
+                            mAppTimeLimit.noteUsageStart(packageName, userId);
                         } catch (IllegalArgumentException iae) {
                             Slog.e(TAG, "Failed to note usage start", iae);
                         }
@@ -496,7 +525,7 @@
                     synchronized (mVisibleActivities) {
                         if (mVisibleActivities.removeReturnOld(event.mInstanceId) != null) {
                             try {
-                                mAppTimeLimit.noteUsageStop(event.getPackageName(), userId);
+                                mAppTimeLimit.noteUsageStop(packageName, userId);
                             } catch (IllegalArgumentException iae) {
                                 Slog.w(TAG, "Failed to note usage stop", iae);
                             }
@@ -638,7 +667,7 @@
      * Called by the Binder stub.
      */
     UsageEvents queryEventsForPackage(int userId, long beginTime, long endTime,
-            String packageName) {
+            String packageName, boolean includeTaskRoot) {
         synchronized (mLock) {
             final long timeNow = checkAndGetTimeLocked();
             if (!validRange(timeNow, beginTime, endTime)) {
@@ -647,7 +676,7 @@
 
             final UserUsageStatsService service =
                     getUserDataAndInitializeIfNeededLocked(userId, timeNow);
-            return service.queryEventsForPackage(beginTime, endTime, packageName);
+            return service.queryEventsForPackage(beginTime, endTime, packageName, includeTaskRoot);
         }
     }
 
@@ -738,6 +767,10 @@
                 mAppStandby.dumpState(args, pw);
             }
 
+            idpw.println();
+            idpw.printPair("Usage Source", UsageStatsManager.usageSourceToString(mUsageSource));
+            idpw.println();
+
             mAppTimeLimit.dump(null, pw);
         }
     }
@@ -808,7 +841,7 @@
             return mode == AppOpsManager.MODE_ALLOWED;
         }
 
-        private boolean hasObserverPermission(String callingPackage) {
+        private boolean hasObserverPermission() {
             final int callingUid = Binder.getCallingUid();
             DevicePolicyManagerInternal dpmInternal = getDpmInternal();
             if (callingUid == Process.SYSTEM_UID
@@ -939,10 +972,12 @@
             final int callingUserId = UserHandle.getUserId(callingUid);
 
             checkCallerIsSameApp(callingPackage);
+            final boolean includeTaskRoot = hasPermission(callingPackage);
+
             final long token = Binder.clearCallingIdentity();
             try {
                 return UsageStatsService.this.queryEventsForPackage(callingUserId, beginTime,
-                        endTime, callingPackage);
+                        endTime, callingPackage, includeTaskRoot);
             } finally {
                 Binder.restoreCallingIdentity(token);
             }
@@ -989,7 +1024,7 @@
             final long token = Binder.clearCallingIdentity();
             try {
                 return UsageStatsService.this.queryEventsForPackage(userId, beginTime,
-                        endTime, pkg);
+                        endTime, pkg, true);
             } finally {
                 Binder.restoreCallingIdentity(token);
             }
@@ -1229,7 +1264,7 @@
         public void registerAppUsageObserver(int observerId,
                 String[] packages, long timeLimitMs, PendingIntent
                 callbackIntent, String callingPackage) {
-            if (!hasObserverPermission(callingPackage)) {
+            if (!hasObserverPermission()) {
                 throw new SecurityException("Caller doesn't have OBSERVE_APP_USAGE permission");
             }
 
@@ -1252,7 +1287,7 @@
 
         @Override
         public void unregisterAppUsageObserver(int observerId, String callingPackage) {
-            if (!hasObserverPermission(callingPackage)) {
+            if (!hasObserverPermission()) {
                 throw new SecurityException("Caller doesn't have OBSERVE_APP_USAGE permission");
             }
 
@@ -1271,7 +1306,7 @@
                 long timeLimitMs, long sessionThresholdTimeMs,
                 PendingIntent limitReachedCallbackIntent, PendingIntent sessionEndCallbackIntent,
                 String callingPackage) {
-            if (!hasObserverPermission(callingPackage)) {
+            if (!hasObserverPermission()) {
                 throw new SecurityException("Caller doesn't have OBSERVE_APP_USAGE permission");
             }
 
@@ -1295,7 +1330,7 @@
 
         @Override
         public void unregisterUsageSessionObserver(int sessionObserverId, String callingPackage) {
-            if (!hasObserverPermission(callingPackage)) {
+            if (!hasObserverPermission()) {
                 throw new SecurityException("Caller doesn't have OBSERVE_APP_USAGE permission");
             }
 
@@ -1373,6 +1408,21 @@
                 Binder.restoreCallingIdentity(binderToken);
             }
         }
+
+        @Override
+        public @UsageSource int getUsageSource() {
+            if (!hasObserverPermission()) {
+                throw new SecurityException("Caller doesn't have OBSERVE_APP_USAGE permission");
+            }
+            synchronized (mLock) {
+                return mUsageSource;
+            }
+        }
+
+        @Override
+        public void forceUsageSourceSettingRead() {
+            readUsageSourceSetting();
+        }
     }
 
     void registerAppUsageObserver(int callingUid, int observerId, String[] packages,
@@ -1406,7 +1456,7 @@
 
         @Override
         public void reportEvent(ComponentName component, int userId, int eventType,
-                int instanceId) {
+                int instanceId, ComponentName taskRoot) {
             if (component == null) {
                 Slog.w(TAG, "Event reported without a component name");
                 return;
@@ -1416,6 +1466,13 @@
             event.mPackage = component.getPackageName();
             event.mClass = component.getClassName();
             event.mInstanceId = instanceId;
+            if (taskRoot == null) {
+                event.mTaskRootPackage = null;
+                event.mTaskRootClass = null;
+            } else {
+                event.mTaskRootPackage = taskRoot.getPackageName();
+                event.mTaskRootClass = taskRoot.getClassName();
+            }
             mHandler.obtainMessage(MSG_REPORT_EVENT, userId, 0, event).sendToTarget();
         }
 
diff --git a/services/usage/java/com/android/server/usage/UserUsageStatsService.java b/services/usage/java/com/android/server/usage/UserUsageStatsService.java
index 2d1098c7..d52d32f 100644
--- a/services/usage/java/com/android/server/usage/UserUsageStatsService.java
+++ b/services/usage/java/com/android/server/usage/UserUsageStatsService.java
@@ -401,7 +401,7 @@
     }
 
     UsageEvents queryEvents(final long beginTime, final long endTime,
-            boolean obfuscateInstantApps) {
+                            boolean obfuscateInstantApps) {
         final ArraySet<String> names = new ArraySet<>();
         List<Event> results = queryStats(INTERVAL_DAILY,
                 beginTime, endTime, new StatCombiner<Event>() {
@@ -425,6 +425,12 @@
                             if (event.mClass != null) {
                                 names.add(event.mClass);
                             }
+                            if (event.mTaskRootPackage != null) {
+                                names.add(event.mTaskRootPackage);
+                            }
+                            if (event.mTaskRootClass != null) {
+                                names.add(event.mTaskRootClass);
+                            }
                             accumulatedResult.add(event);
                         }
                     }
@@ -436,11 +442,11 @@
 
         String[] table = names.toArray(new String[names.size()]);
         Arrays.sort(table);
-        return new UsageEvents(results, table);
+        return new UsageEvents(results, table, true);
     }
 
     UsageEvents queryEventsForPackage(final long beginTime, final long endTime,
-            final String packageName) {
+            final String packageName, boolean includeTaskRoot) {
         final ArraySet<String> names = new ArraySet<>();
         names.add(packageName);
         final List<Event> results = queryStats(INTERVAL_DAILY,
@@ -459,6 +465,12 @@
                         if (event.mClass != null) {
                             names.add(event.mClass);
                         }
+                        if (includeTaskRoot && event.mTaskRootPackage != null) {
+                            names.add(event.mTaskRootPackage);
+                        }
+                        if (includeTaskRoot && event.mTaskRootClass != null) {
+                            names.add(event.mTaskRootClass);
+                        }
                         accumulatedResult.add(event);
                     }
                 });
@@ -469,7 +481,7 @@
 
         final String[] table = names.toArray(new String[names.size()]);
         Arrays.sort(table);
-        return new UsageEvents(results, table);
+        return new UsageEvents(results, table, includeTaskRoot);
     }
 
     void persistActiveStats() {
@@ -684,6 +696,14 @@
             pw.printPair("instanceId", event.getInstanceId());
         }
 
+        if (event.getTaskRootPackageName() != null) {
+            pw.printPair("taskRootPackage", event.getTaskRootPackageName());
+        }
+
+        if (event.getTaskRootClassName() != null) {
+            pw.printPair("taskRootClass", event.getTaskRootClassName());
+        }
+
         if (event.mNotificationChannelId != null) {
             pw.printPair("channelId", event.mNotificationChannelId);
         }