Rename WorkSource methods on Binder and IPCThreadState.

This change only renames methods, there is no behavior changes except
using the new restore methods instead of clear.

Test: unit tests
Change-Id: I35ae966461657e2e2a67e916d752b9ee53381c83

4 files changed, 62 insertions, 28 deletions

diff --git a/core/java/android/os/Binder.java b/core/java/android/os/Binder.java
index da4b82308313..c7184c0743ce 100644
--- a/core/java/android/os/Binder.java
+++ b/core/java/android/os/Binder.java
@@ -382,7 +382,9 @@ public class Binder implements IBinder {
     /**
      * Sets the work source for this thread.
      *
-     * <p>All the following binder calls on this thread will use the provided work source.
+     * <p>All the following binder calls on this thread will use the provided work source. If this
+     * is called during an on-going binder transaction, all the following binder calls will use the
+     * work source until the end of the transaction.
      *
      * <p>The concept of worksource is similar to {@link WorkSource}. However, for performance
      * reasons, we only support one UID. This UID represents the original user responsible for the
@@ -390,20 +392,20 @@ public class Binder implements IBinder {
      *
      * <p>A typical use case would be
      * <pre>
-     * Binder.setThreadWorkSource(uid);
+     * long token = Binder.setCallingWorkSourceUid(uid);
      * try {
      *   // Call an API.
      * } finally {
-     *   Binder.clearThreadWorkSource();
+     *   Binder.restoreCallingWorkSource(token);
      * }
      * </pre>
      *
      * @param workSource The original UID responsible for the binder call.
-     * @return The previously set work source.
+     * @return token to restore original work source.
      * @hide
      **/
     @CriticalNative
-    public static final native int setThreadWorkSource(int workSource);
+    public static final native long setCallingWorkSourceUid(int workSource);
 
     /**
      * Returns the work source set by the caller.
@@ -416,16 +418,34 @@ public class Binder implements IBinder {
      * @hide
      */
     @CriticalNative
-    public static final native int getThreadWorkSource();
+    public static final native int getCallingWorkSourceUid();
 
     /**
      * Clears the work source on this thread.
      *
-     * @return The previously set work source.
+     * @return token to restore original work source.
      * @hide
      **/
     @CriticalNative
-    public static final native int clearThreadWorkSource();
+    public static final native long clearCallingWorkSource();
+
+    /**
+     * Restores the work source on this thread using a token returned by
+     * {@link #setCallingWorkSourceUid(int) or {@link clearCallingWorkSource()}.
+     *
+     * <p>A typical use case would be
+     * <pre>
+     * long token = Binder.setCallingWorkSourceUid(uid);
+     * try {
+     *   // Call an API.
+     * } finally {
+     *   Binder.restoreCallingWorkSource(token);
+     * }
+     * </pre>
+     * @hide
+     **/
+    @CriticalNative
+    public static final native void restoreCallingWorkSource(long token);
 
     /**
      * Flush any Binder commands pending in the current thread to the kernel
@@ -586,7 +606,7 @@ public class Binder implements IBinder {
      *
      * <li>By default, this listener will propagate the worksource if the outgoing call happens on
      * the same thread as the incoming binder call.
-     * <li>Custom attribution can be done by calling {@link ThreadLocalWorkSourceUid#set(int)}.
+     * <li>Custom attribution can be done by calling {@link ThreadLocalWorkSource#setUid(int)}.
      * @hide
      */
     public static class PropagateWorkSourceTransactListener implements ProxyTransactListener {
@@ -595,12 +615,11 @@ public class Binder implements IBinder {
            // Note that {@link Binder#getCallingUid()} is already set to the UID of the current
            // process when this method is called.
            //
-           // We use ThreadLocalWorkSourceUid instead. It also allows feature owners to set
-           // {@link ThreadLocalWorkSourceUid#set(int) manually to attribute resources to a UID.
-            int uid = ThreadLocalWorkSourceUid.get();
-            if (uid >= 0) {
-                int originalUid = Binder.setThreadWorkSource(uid);
-                return Integer.valueOf(originalUid);
+           // We use ThreadLocalWorkSource instead. It also allows feature owners to set
+           // {@link ThreadLocalWorkSource#set(int) manually to attribute resources to a UID.
+            int uid = ThreadLocalWorkSource.getUid();
+            if (uid != ThreadLocalWorkSource.UID_NONE) {
+                return Binder.setCallingWorkSourceUid(uid);
             }
             return null;
         }
@@ -608,8 +627,8 @@ public class Binder implements IBinder {
         @Override
         public void onTransactEnded(Object session) {
             if (session != null) {
-                int uid = (int) session;
-                Binder.setThreadWorkSource(uid);
+                long token = (long) session;
+                Binder.restoreCallingWorkSource(token);
             }
         }
     }
@@ -897,11 +916,11 @@ public class Binder implements IBinder {
         // Log any exceptions as warnings, don't silently suppress them.
         // If the call was FLAG_ONEWAY then these exceptions disappear into the ether.
         final boolean tracingEnabled = Binder.isTracingEnabled();
+        final long origWorkSource = ThreadLocalWorkSource.setUid(Binder.getCallingUid());
         try {
             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) {
@@ -922,7 +941,7 @@ public class Binder implements IBinder {
             }
             res = true;
         } finally {
-            ThreadLocalWorkSourceUid.clear();
+            ThreadLocalWorkSource.restore(origWorkSource);
             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 f3a9a504a765..e8704afd0139 100644
--- a/core/java/android/os/Handler.java
+++ b/core/java/android/os/Handler.java
@@ -739,7 +739,7 @@ public class Handler {
 
     private boolean enqueueMessage(MessageQueue queue, Message msg, long uptimeMillis) {
         msg.target = this;
-        msg.workSourceUid = ThreadLocalWorkSourceUid.get();
+        msg.workSourceUid = ThreadLocalWorkSource.getUid();
 
         if (mAsynchronous) {
             msg.setAsynchronous(true);
diff --git a/core/java/android/os/Looper.java b/core/java/android/os/Looper.java
index 5b8ababb8ca4..a8d1215f48e6 100644
--- a/core/java/android/os/Looper.java
+++ b/core/java/android/os/Looper.java
@@ -204,8 +204,8 @@ public final class Looper {
             if (observer != null) {
                 token = observer.messageDispatchStarting();
             }
+            long origWorkSource = ThreadLocalWorkSource.setUid(msg.workSourceUid);
             try {
-                ThreadLocalWorkSourceUid.set(msg.workSourceUid);
                 msg.target.dispatchMessage(msg);
                 if (observer != null) {
                     observer.messageDispatched(token, msg);
@@ -217,7 +217,7 @@ public final class Looper {
                 }
                 throw exception;
             } finally {
-                ThreadLocalWorkSourceUid.clear();
+                ThreadLocalWorkSource.restore(origWorkSource);
                 if (traceTag != 0) {
                     Trace.traceEnd(traceTag);
                 }
diff --git a/core/java/android/os/ThreadLocalWorkSourceUid.java b/core/java/android/os/ThreadLocalWorkSource.java
index df1d275f358d..53dd460d62c4 100644
--- a/core/java/android/os/ThreadLocalWorkSourceUid.java
+++ b/core/java/android/os/ThreadLocalWorkSource.java
@@ -19,26 +19,41 @@ package android.os;
 /**
  * @hide Only for use within system server.
  */
-public final class ThreadLocalWorkSourceUid {
+public final class ThreadLocalWorkSource {
     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() {
+    public static int getUid() {
         return sWorkSourceUid.get();
     }
 
     /** Sets the original work source uid. */
-    public static void set(int uid) {
+    public static long setUid(int uid) {
+        final long token = getToken();
         sWorkSourceUid.set(uid);
+        return token;
+    }
+
+    /** Restores the state using the provided token. */
+    public static void restore(long token) {
+        sWorkSourceUid.set(parseUidFromToken(token));
     }
 
     /** Clears the stored work source uid. */
-    public static void clear() {
-        sWorkSourceUid.set(UID_NONE);
+    public static long clear() {
+        return setUid(UID_NONE);
+    }
+
+    private static int parseUidFromToken(long token) {
+        return (int) token;
+    }
+
+    private static long getToken() {
+        return sWorkSourceUid.get();
     }
 
-    private ThreadLocalWorkSourceUid() {
+    private ThreadLocalWorkSource() {
     }
 }