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

core/java/android/os/Binder.java

@@ -773,6 +773,7 @@ public class Binder implements IBinder {
             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 @@ public class Binder implements IBinder {
             }
             res = true;
         } finally {
+            ThreadLocalWorkSourceUid.clear();
             if (tracingEnabled) {
                 Trace.traceEnd(Trace.TRACE_TAG_ALWAYS);
             }

core/java/android/os/Handler.java

@@ -739,6 +739,8 @@ public class Handler {
 
     private boolean enqueueMessage(MessageQueue queue, Message msg, long uptimeMillis) {
         msg.target = this;
+        msg.workSourceUid = ThreadLocalWorkSourceUid.get();
+
         if (mAsynchronous) {
             msg.setAsynchronous(true);
         }

core/java/android/os/Looper.java

@@ -205,6 +205,7 @@ public final class Looper {
                 token = observer.messageDispatchStarting();
             }
             try {
+                ThreadLocalWorkSourceUid.set(msg.workSourceUid);
                 msg.target.dispatchMessage(msg);
                 if (observer != null) {
                     observer.messageDispatched(token, msg);
@@ -216,6 +217,7 @@ public final class Looper {
                 }
                 throw exception;
             } finally {
+                ThreadLocalWorkSourceUid.clear();
                 if (traceTag != 0) {
                     Trace.traceEnd(traceTag);
                 }

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;
 
@@ -73,12 +72,26 @@ public final class Message implements Parcelable {
      */
     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 @@ public final class Message implements Parcelable {
         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 @@ public final class Message implements Parcelable {
         arg2 = 0;
         obj = null;
         replyTo = null;
-        sendingUid = -1;
+        sendingUid = UID_NONE;
+        workSourceUid = UID_NONE;
         when = 0;
         target = null;
         callback = null;
@@ -329,6 +344,7 @@ public final class Message implements Parcelable {
         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 @@ public final class Message implements Parcelable {
         dest.writeBundle(data);
         Messenger.writeMessengerOrNullToParcel(replyTo, dest);
         dest.writeInt(sendingUid);
+        dest.writeInt(workSourceUid);
     }
 
     private void readFromParcel(Parcel source) {
@@ -625,5 +642,6 @@ public final class Message implements Parcelable {
         data = source.readBundle();
         replyTo = Messenger.readMessengerOrNullFromParcel(source);
         sendingUid = source.readInt();
+        workSourceUid = source.readInt();
     }
 }

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() {
+    }
+}