Add work source uid collection to LooperStats

Add ThreadLocalWorkSourceUid and Message.workSourceUid
which will store the UID which caused, directly or
indirectly, the Message to be enqueued.

This was needed because using Binder.getCallingUid() does
not give us enough data since quite a few Binder services
call Binder.clearCallingIdentity() when processing the
call.

Test: UT and manual
Change-Id: I35af3a9ca5193477070990b41ff5d9c52f420069
diff --git a/cmds/statsd/src/atoms.proto b/cmds/statsd/src/atoms.proto
index b8f19ee..c6c10ec 100644
--- a/cmds/statsd/src/atoms.proto
+++ b/cmds/statsd/src/atoms.proto
@@ -2339,7 +2339,7 @@
  * Next tag: 11
  */
 message LooperStats {
-    // Currently not collected and always set to 0.
+    // The uid that made a call to the System Server and caused the message to be enqueued.
     optional int32 uid = 1 [(is_uid) = true];
 
     // Fully qualified class name of the handler target class.
diff --git a/core/java/android/os/Binder.java b/core/java/android/os/Binder.java
index b9d9007..f947b5e 100644
--- a/core/java/android/os/Binder.java
+++ b/core/java/android/os/Binder.java
@@ -773,6 +773,7 @@
             if (tracingEnabled) {
                 Trace.traceBegin(Trace.TRACE_TAG_ALWAYS, getClass().getName() + ":" + code);
             }
+            ThreadLocalWorkSourceUid.set(Binder.getCallingUid());
             res = onTransact(code, data, reply, flags);
         } catch (RemoteException|RuntimeException e) {
             if (observer != null) {
@@ -793,6 +794,7 @@
             }
             res = true;
         } finally {
+            ThreadLocalWorkSourceUid.clear();
             if (tracingEnabled) {
                 Trace.traceEnd(Trace.TRACE_TAG_ALWAYS);
             }
diff --git a/core/java/android/os/Handler.java b/core/java/android/os/Handler.java
index e03af9d..f3a9a50 100644
--- a/core/java/android/os/Handler.java
+++ b/core/java/android/os/Handler.java
@@ -739,6 +739,8 @@
 
     private boolean enqueueMessage(MessageQueue queue, Message msg, long uptimeMillis) {
         msg.target = this;
+        msg.workSourceUid = ThreadLocalWorkSourceUid.get();
+
         if (mAsynchronous) {
             msg.setAsynchronous(true);
         }
diff --git a/core/java/android/os/Looper.java b/core/java/android/os/Looper.java
index ded3a19..5b8abab 100644
--- a/core/java/android/os/Looper.java
+++ b/core/java/android/os/Looper.java
@@ -205,6 +205,7 @@
                 token = observer.messageDispatchStarting();
             }
             try {
+                ThreadLocalWorkSourceUid.set(msg.workSourceUid);
                 msg.target.dispatchMessage(msg);
                 if (observer != null) {
                     observer.messageDispatched(token, msg);
@@ -216,6 +217,7 @@
                 }
                 throw exception;
             } finally {
+                ThreadLocalWorkSourceUid.clear();
                 if (traceTag != 0) {
                     Trace.traceEnd(traceTag);
                 }
diff --git a/core/java/android/os/Message.java b/core/java/android/os/Message.java
index 455d8ed..cd3f301 100644
--- a/core/java/android/os/Message.java
+++ b/core/java/android/os/Message.java
@@ -16,7 +16,6 @@
 
 package android.os;
 
-import android.os.MessageProto;
 import android.util.TimeUtils;
 import android.util.proto.ProtoOutputStream;
 
@@ -74,11 +73,25 @@
     public Messenger replyTo;
 
     /**
+     * Indicates that the uid is not set;
+     *
+     * @hide Only for use within the system server.
+     */
+    public static final int UID_NONE = -1;
+
+    /**
      * Optional field indicating the uid that sent the message.  This is
      * only valid for messages posted by a {@link Messenger}; otherwise,
      * it will be -1.
      */
-    public int sendingUid = -1;
+    public int sendingUid = UID_NONE;
+
+    /**
+     * Optional field indicating the uid that caused this message to be enqueued.
+     *
+     * @hide Only for use within the system server.
+     */
+    public int workSourceUid = UID_NONE;
 
     /** If set message is in use.
      * This flag is set when the message is enqueued and remains set while it
@@ -151,6 +164,7 @@
         m.obj = orig.obj;
         m.replyTo = orig.replyTo;
         m.sendingUid = orig.sendingUid;
+        m.workSourceUid = orig.workSourceUid;
         if (orig.data != null) {
             m.data = new Bundle(orig.data);
         }
@@ -301,7 +315,8 @@
         arg2 = 0;
         obj = null;
         replyTo = null;
-        sendingUid = -1;
+        sendingUid = UID_NONE;
+        workSourceUid = UID_NONE;
         when = 0;
         target = null;
         callback = null;
@@ -329,6 +344,7 @@
         this.obj = o.obj;
         this.replyTo = o.replyTo;
         this.sendingUid = o.sendingUid;
+        this.workSourceUid = o.workSourceUid;
 
         if (o.data != null) {
             this.data = (Bundle) o.data.clone();
@@ -612,6 +628,7 @@
         dest.writeBundle(data);
         Messenger.writeMessengerOrNullToParcel(replyTo, dest);
         dest.writeInt(sendingUid);
+        dest.writeInt(workSourceUid);
     }
 
     private void readFromParcel(Parcel source) {
@@ -625,5 +642,6 @@
         data = source.readBundle();
         replyTo = Messenger.readMessengerOrNullFromParcel(source);
         sendingUid = source.readInt();
+        workSourceUid = source.readInt();
     }
 }
diff --git a/core/java/android/os/ThreadLocalWorkSourceUid.java b/core/java/android/os/ThreadLocalWorkSourceUid.java
new file mode 100644
index 0000000..df1d275
--- /dev/null
+++ b/core/java/android/os/ThreadLocalWorkSourceUid.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2018 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 android.os;
+
+/**
+ * @hide Only for use within system server.
+ */
+public final class ThreadLocalWorkSourceUid {
+    public static final int UID_NONE = Message.UID_NONE;
+    private static final ThreadLocal<Integer> sWorkSourceUid =
+            ThreadLocal.withInitial(() -> UID_NONE);
+
+    /** Returns the original work source uid. */
+    public static int get() {
+        return sWorkSourceUid.get();
+    }
+
+    /** Sets the original work source uid. */
+    public static void set(int uid) {
+        sWorkSourceUid.set(uid);
+    }
+
+    /** Clears the stored work source uid. */
+    public static void clear() {
+        sWorkSourceUid.set(UID_NONE);
+    }
+
+    private ThreadLocalWorkSourceUid() {
+    }
+}
diff --git a/core/java/com/android/internal/os/LooperStats.java b/core/java/com/android/internal/os/LooperStats.java
index 0650d0af..b8e1806 100644
--- a/core/java/com/android/internal/os/LooperStats.java
+++ b/core/java/com/android/internal/os/LooperStats.java
@@ -176,7 +176,8 @@
             }
         }
 
-        if (entry.handler.getClass() != msg.getTarget().getClass()
+        if (entry.workSourceUid != msg.workSourceUid
+                || entry.handler.getClass() != msg.getTarget().getClass()
                 || entry.handler.getLooper().getThread() != msg.getTarget().getLooper().getThread()
                 || entry.isInteractive != isInteractive) {
             // If a hash collision happened, track totals under a single entry.
@@ -210,6 +211,7 @@
     }
 
     private static class Entry {
+        public final int workSourceUid;
         public final Handler handler;
         public final String messageName;
         public final boolean isInteractive;
@@ -222,12 +224,14 @@
         public long maxCpuUsageMicro;
 
         Entry(Message msg, boolean isInteractive) {
+            this.workSourceUid = msg.workSourceUid;
             this.handler = msg.getTarget();
             this.messageName = handler.getMessageName(msg);
             this.isInteractive = isInteractive;
         }
 
         Entry(String specialEntryName) {
+            this.workSourceUid = Message.UID_NONE;
             this.messageName = specialEntryName;
             this.handler = null;
             this.isInteractive = false;
@@ -245,6 +249,7 @@
 
         static int idFor(Message msg, boolean isInteractive) {
             int result = 7;
+            result = 31 * result + msg.workSourceUid;
             result = 31 * result + msg.getTarget().getLooper().getThread().hashCode();
             result = 31 * result + msg.getTarget().getClass().hashCode();
             result = 31 * result + (isInteractive ? 1231 : 1237);
@@ -258,6 +263,7 @@
 
     /** Aggregated data of Looper message dispatching in the in the current process. */
     public static class ExportedEntry {
+        public final int workSourceUid;
         public final String handlerClassName;
         public final String threadName;
         public final String messageName;
@@ -271,6 +277,7 @@
         public final long maxCpuUsageMicros;
 
         ExportedEntry(Entry entry) {
+            this.workSourceUid = entry.workSourceUid;
             if (entry.handler != null) {
                 this.handlerClassName = entry.handler.getClass().getName();
                 this.threadName = entry.handler.getLooper().getThread().getName();
diff --git a/core/tests/coretests/src/com/android/internal/os/LooperStatsTest.java b/core/tests/coretests/src/com/android/internal/os/LooperStatsTest.java
index 565a3ec..1008a90 100644
--- a/core/tests/coretests/src/com/android/internal/os/LooperStatsTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/LooperStatsTest.java
@@ -21,6 +21,7 @@
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.Looper;
+import android.os.Message;
 import android.platform.test.annotations.Presubmit;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
@@ -74,14 +75,17 @@
     public void testSingleMessageDispatched() {
         TestableLooperStats looperStats = new TestableLooperStats(1, 100);
 
+        Message message = mHandlerFirst.obtainMessage(1000);
+        message.workSourceUid = 1000;
         Object token = looperStats.messageDispatchStarting();
         looperStats.tickRealtime(100);
         looperStats.tickThreadTime(10);
-        looperStats.messageDispatched(token, mHandlerFirst.obtainMessage(1000));
+        looperStats.messageDispatched(token, message);
 
         List<LooperStats.ExportedEntry> entries = looperStats.getEntries();
         assertThat(entries).hasSize(1);
         LooperStats.ExportedEntry entry = entries.get(0);
+        assertThat(entry.workSourceUid).isEqualTo(1000);
         assertThat(entry.threadName).isEqualTo("TestThread1");
         assertThat(entry.handlerClassName).isEqualTo(
                 "com.android.internal.os.LooperStatsTest$TestHandlerFirst");
@@ -100,15 +104,17 @@
     public void testThrewException() {
         TestableLooperStats looperStats = new TestableLooperStats(1, 100);
 
+        Message message = mHandlerFirst.obtainMessage(7);
+        message.workSourceUid = 123;
         Object token = looperStats.messageDispatchStarting();
         looperStats.tickRealtime(100);
         looperStats.tickThreadTime(10);
-        looperStats.dispatchingThrewException(token, mHandlerFirst.obtainMessage(7),
-                new ArithmeticException());
+        looperStats.dispatchingThrewException(token, message, new ArithmeticException());
 
         List<LooperStats.ExportedEntry> entries = looperStats.getEntries();
         assertThat(entries).hasSize(1);
         LooperStats.ExportedEntry entry = entries.get(0);
+        assertThat(entry.workSourceUid).isEqualTo(123);
         assertThat(entry.threadName).isEqualTo("TestThread1");
         assertThat(entry.handlerClassName).isEqualTo(
                 "com.android.internal.os.LooperStatsTest$TestHandlerFirst");
@@ -158,6 +164,7 @@
 
         // Captures data for token4 call.
         LooperStats.ExportedEntry entry1 = entries.get(0);
+        assertThat(entry1.workSourceUid).isEqualTo(-1);
         assertThat(entry1.threadName).isEqualTo("TestThread1");
         assertThat(entry1.handlerClassName).isEqualTo("com.android.internal.os.LooperStatsTest$1");
         assertThat(entry1.messageName).isEqualTo("0x1" /* 1 in hex */);
@@ -171,6 +178,7 @@
 
         // Captures data for token1 and token2 calls.
         LooperStats.ExportedEntry entry2 = entries.get(1);
+        assertThat(entry2.workSourceUid).isEqualTo(-1);
         assertThat(entry2.threadName).isEqualTo("TestThread1");
         assertThat(entry2.handlerClassName).isEqualTo(
                 "com.android.internal.os.LooperStatsTest$TestHandlerFirst");
@@ -185,6 +193,7 @@
 
         // Captures data for token3 call.
         LooperStats.ExportedEntry entry3 = entries.get(2);
+        assertThat(entry3.workSourceUid).isEqualTo(-1);
         assertThat(entry3.threadName).isEqualTo("TestThread2");
         assertThat(entry3.handlerClassName).isEqualTo(
                 "com.android.internal.os.LooperStatsTest$TestHandlerSecond");
diff --git a/services/core/java/com/android/server/LooperStatsService.java b/services/core/java/com/android/server/LooperStatsService.java
index ee01d86..4f0e170 100644
--- a/services/core/java/com/android/server/LooperStatsService.java
+++ b/services/core/java/com/android/server/LooperStatsService.java
@@ -93,10 +93,12 @@
         if (!DumpUtils.checkDumpPermission(mContext, TAG, pw)) return;
         List<LooperStats.ExportedEntry> entries = mStats.getEntries();
         entries.sort(Comparator
-                .comparing((LooperStats.ExportedEntry entry) -> entry.threadName)
+                .comparing((LooperStats.ExportedEntry entry) -> entry.workSourceUid)
+                .thenComparing(entry -> entry.threadName)
                 .thenComparing(entry -> entry.handlerClassName)
                 .thenComparing(entry -> entry.messageName));
         String header = String.join(",", Arrays.asList(
+                "work_source_uid",
                 "thread_name",
                 "handler_class",
                 "message_name",
@@ -110,11 +112,11 @@
                 "exception_count"));
         pw.println(header);
         for (LooperStats.ExportedEntry entry : entries) {
-            pw.printf("%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s\n", entry.threadName,
-                    entry.handlerClassName, entry.messageName, entry.isInteractive,
-                    entry.messageCount, entry.recordedMessageCount, entry.totalLatencyMicros,
-                    entry.maxLatencyMicros, entry.cpuUsageMicros, entry.maxCpuUsageMicros,
-                    entry.exceptionCount);
+            pw.printf("%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s\n", entry.workSourceUid,
+                    entry.threadName, entry.handlerClassName, entry.messageName,
+                    entry.isInteractive, entry.messageCount, entry.recordedMessageCount,
+                    entry.totalLatencyMicros, entry.maxLatencyMicros, entry.cpuUsageMicros,
+                    entry.maxCpuUsageMicros, entry.exceptionCount);
         }
     }
 
diff --git a/services/core/java/com/android/server/stats/StatsCompanionService.java b/services/core/java/com/android/server/stats/StatsCompanionService.java
index 5e3fe0a..bfa03ca 100644
--- a/services/core/java/com/android/server/stats/StatsCompanionService.java
+++ b/services/core/java/com/android/server/stats/StatsCompanionService.java
@@ -1031,7 +1031,7 @@
         looperStats.reset();
         for (LooperStats.ExportedEntry entry : entries) {
             StatsLogEventWrapper e = new StatsLogEventWrapper(tagId, elapsedNanos, wallClockNanos);
-            e.writeInt(1000); // uid collection not implemented yet
+            e.writeInt(entry.workSourceUid);
             e.writeString(entry.handlerClassName);
             e.writeString(entry.threadName);
             e.writeString(entry.messageName);