diff --git a/core/jni/android_util_Binder.cpp b/core/jni/android_util_Binder.cpp
index 4f8bbc1396c8a..3329e2047085f 100644
--- a/core/jni/android_util_Binder.cpp
+++ b/core/jni/android_util_Binder.cpp
@@ -919,7 +919,7 @@ static jlong android_os_Binder_clearCallingWorkSource()
     return IPCThreadState::self()->clearCallingWorkSource();
 }
 
-static void android_os_Binder_restoreCallingWorkSource(long token)
+static void android_os_Binder_restoreCallingWorkSource(jlong token)
 {
     IPCThreadState::self()->restoreCallingWorkSource(token);
 }
diff --git a/core/tests/coretests/src/android/os/BinderWorkSourceTest.java b/core/tests/coretests/src/android/os/BinderWorkSourceTest.java
index d1dbd3ccba33d..ef14b00f97754 100644
--- a/core/tests/coretests/src/android/os/BinderWorkSourceTest.java
+++ b/core/tests/coretests/src/android/os/BinderWorkSourceTest.java
@@ -125,8 +125,10 @@ public class BinderWorkSourceTest {
         Binder.setCallingWorkSourceUid(UID);
         long token = Binder.clearCallingWorkSource();
         Binder.restoreCallingWorkSource(token);
+        assertEquals(UID, Binder.getCallingWorkSourceUid());
 
         assertEquals(UID, mService.getIncomingWorkSourceUid());
+        // Still the same after the binder transaction.
         assertEquals(UID, Binder.getCallingWorkSourceUid());
     }