Enable testing updating estimated and transferred bytes.

Add shell commands to make it possible to test updating estimated
network bytes and transferred network bytes.

Bug: 255371817
Test: atest CtsJobSchedulerTestCases:DataTransferTest
Change-Id: I73a5a231ec2a4349eb00fa57c55f0a8850f43c73

4 files changed, 292 insertions, 14 deletions

diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java b/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
index 30986dde6b91..7897449c1e02 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
@@ -1688,6 +1688,40 @@ class JobConcurrencyManager {
         return foundSome;
     }
 
+    /**
+     * Returns the estimated network bytes if the job is running. Returns {@code null} if the job
+     * isn't running.
+     */
+    @Nullable
+    @GuardedBy("mLock")
+    Pair<Long, Long> getEstimatedNetworkBytesLocked(String pkgName, int uid, int jobId) {
+        for (int i = 0; i < mActiveServices.size(); i++) {
+            final JobServiceContext jc = mActiveServices.get(i);
+            final JobStatus js = jc.getRunningJobLocked();
+            if (js != null && js.matches(uid, jobId) && js.getSourcePackageName().equals(pkgName)) {
+                return jc.getEstimatedNetworkBytes();
+            }
+        }
+        return null;
+    }
+
+    /**
+     * Returns the transferred network bytes if the job is running. Returns {@code null} if the job
+     * isn't running.
+     */
+    @Nullable
+    @GuardedBy("mLock")
+    Pair<Long, Long> getTransferredNetworkBytesLocked(String pkgName, int uid, int jobId) {
+        for (int i = 0; i < mActiveServices.size(); i++) {
+            final JobServiceContext jc = mActiveServices.get(i);
+            final JobStatus js = jc.getRunningJobLocked();
+            if (js != null && js.matches(uid, jobId) && js.getSourcePackageName().equals(pkgName)) {
+                return jc.getTransferredNetworkBytes();
+            }
+        }
+        return null;
+    }
+
     @NonNull
     private JobServiceContext createNewJobServiceContext() {
         return mInjector.createJobServiceContext(mService, this, mNotificationCoordinator,
diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
index 6f58c7ce626c..4a961208265c 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
@@ -86,6 +86,7 @@ import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.IndentingPrintWriter;
 import android.util.Log;
+import android.util.Pair;
 import android.util.Slog;
 import android.util.SparseArray;
 import android.util.SparseBooleanArray;
@@ -4159,6 +4160,100 @@ public class JobSchedulerService extends com.android.server.SystemService
         }
     }
 
+    int getEstimatedNetworkBytes(PrintWriter pw, String pkgName, int userId, int jobId,
+            int byteOption) {
+        try {
+            final int uid = AppGlobals.getPackageManager().getPackageUid(pkgName, 0,
+                    userId != UserHandle.USER_ALL ? userId : UserHandle.USER_SYSTEM);
+            if (uid < 0) {
+                pw.print("unknown(");
+                pw.print(pkgName);
+                pw.println(")");
+                return JobSchedulerShellCommand.CMD_ERR_NO_PACKAGE;
+            }
+
+            synchronized (mLock) {
+                final JobStatus js = mJobs.getJobByUidAndJobId(uid, jobId);
+                if (DEBUG) {
+                    Slog.d(TAG, "get-estimated-network-bytes " + uid + "/" + jobId + ": " + js);
+                }
+                if (js == null) {
+                    pw.print("unknown("); UserHandle.formatUid(pw, uid);
+                    pw.print("/jid"); pw.print(jobId); pw.println(")");
+                    return JobSchedulerShellCommand.CMD_ERR_NO_JOB;
+                }
+
+                final long downloadBytes;
+                final long uploadBytes;
+                final Pair<Long, Long> bytes =
+                        mConcurrencyManager.getEstimatedNetworkBytesLocked(pkgName, uid, jobId);
+                if (bytes == null) {
+                    downloadBytes = js.getEstimatedNetworkDownloadBytes();
+                    uploadBytes = js.getEstimatedNetworkUploadBytes();
+                } else {
+                    downloadBytes = bytes.first;
+                    uploadBytes = bytes.second;
+                }
+                if (byteOption == JobSchedulerShellCommand.BYTE_OPTION_DOWNLOAD) {
+                    pw.println(downloadBytes);
+                } else {
+                    pw.println(uploadBytes);
+                }
+                pw.println();
+            }
+        } catch (RemoteException e) {
+            // can't happen
+        }
+        return 0;
+    }
+
+    int getTransferredNetworkBytes(PrintWriter pw, String pkgName, int userId, int jobId,
+            int byteOption) {
+        try {
+            final int uid = AppGlobals.getPackageManager().getPackageUid(pkgName, 0,
+                    userId != UserHandle.USER_ALL ? userId : UserHandle.USER_SYSTEM);
+            if (uid < 0) {
+                pw.print("unknown(");
+                pw.print(pkgName);
+                pw.println(")");
+                return JobSchedulerShellCommand.CMD_ERR_NO_PACKAGE;
+            }
+
+            synchronized (mLock) {
+                final JobStatus js = mJobs.getJobByUidAndJobId(uid, jobId);
+                if (DEBUG) {
+                    Slog.d(TAG, "get-transferred-network-bytes " + uid + "/" + jobId + ": " + js);
+                }
+                if (js == null) {
+                    pw.print("unknown("); UserHandle.formatUid(pw, uid);
+                    pw.print("/jid"); pw.print(jobId); pw.println(")");
+                    return JobSchedulerShellCommand.CMD_ERR_NO_JOB;
+                }
+
+                final long downloadBytes;
+                final long uploadBytes;
+                final Pair<Long, Long> bytes =
+                        mConcurrencyManager.getTransferredNetworkBytesLocked(pkgName, uid, jobId);
+                if (bytes == null) {
+                    downloadBytes = 0;
+                    uploadBytes = 0;
+                } else {
+                    downloadBytes = bytes.first;
+                    uploadBytes = bytes.second;
+                }
+                if (byteOption == JobSchedulerShellCommand.BYTE_OPTION_DOWNLOAD) {
+                    pw.println(downloadBytes);
+                } else {
+                    pw.println(uploadBytes);
+                }
+                pw.println();
+            }
+        } catch (RemoteException e) {
+            // can't happen
+        }
+        return 0;
+    }
+
     private boolean checkRunLongJobsPermission(int packageUid, String packageName) {
         // Returns true if both the appop and permission are granted.
         return PermissionChecker.checkPermissionForPreflight(getTestableContext(),
diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerShellCommand.java b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerShellCommand.java
index 27268d267001..36ba8dd10bd5 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerShellCommand.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerShellCommand.java
@@ -32,6 +32,9 @@ public final class JobSchedulerShellCommand extends BasicShellCommandHandler {
     public static final int CMD_ERR_NO_JOB = -1001;
     public static final int CMD_ERR_CONSTRAINTS = -1002;
 
+    static final int BYTE_OPTION_DOWNLOAD = 0;
+    static final int BYTE_OPTION_UPLOAD = 1;
+
     JobSchedulerService mInternal;
     IPackageManager mPM;
 
@@ -59,10 +62,18 @@ public final class JobSchedulerShellCommand extends BasicShellCommandHandler {
                     return getBatteryCharging(pw);
                 case "get-battery-not-low":
                     return getBatteryNotLow(pw);
+                case "get-estimated-download-bytes":
+                    return getEstimatedNetworkBytes(pw, BYTE_OPTION_DOWNLOAD);
+                case "get-estimated-upload-bytes":
+                    return getEstimatedNetworkBytes(pw, BYTE_OPTION_UPLOAD);
                 case "get-storage-seq":
                     return getStorageSeq(pw);
                 case "get-storage-not-low":
                     return getStorageNotLow(pw);
+                case "get-transferred-download-bytes":
+                    return getTransferredNetworkBytes(pw, BYTE_OPTION_DOWNLOAD);
+                case "get-transferred-upload-bytes":
+                    return getTransferredNetworkBytes(pw, BYTE_OPTION_UPLOAD);
                 case "get-job-state":
                     return getJobState(pw);
                 case "heartbeat":
@@ -304,6 +315,43 @@ public final class JobSchedulerShellCommand extends BasicShellCommandHandler {
         return 0;
     }
 
+    private int getEstimatedNetworkBytes(PrintWriter pw, int byteOption) throws Exception {
+        checkPermission("get estimated bytes");
+
+        int userId = UserHandle.USER_SYSTEM;
+
+        String opt;
+        while ((opt = getNextOption()) != null) {
+            switch (opt) {
+                case "-u":
+                case "--user":
+                    userId = UserHandle.parseUserArg(getNextArgRequired());
+                    break;
+
+                default:
+                    pw.println("Error: unknown option '" + opt + "'");
+                    return -1;
+            }
+        }
+
+        if (userId == UserHandle.USER_CURRENT) {
+            userId = ActivityManager.getCurrentUser();
+        }
+
+        final String pkgName = getNextArgRequired();
+        final String jobIdStr = getNextArgRequired();
+        final int jobId = Integer.parseInt(jobIdStr);
+
+        final long ident = Binder.clearCallingIdentity();
+        try {
+            int ret = mInternal.getEstimatedNetworkBytes(pw, pkgName, userId, jobId, byteOption);
+            printError(ret, pkgName, userId, jobId);
+            return ret;
+        } finally {
+            Binder.restoreCallingIdentity(ident);
+        }
+    }
+
     private int getStorageSeq(PrintWriter pw) {
         int seq = mInternal.getStorageSeq();
         pw.println(seq);
@@ -316,8 +364,45 @@ public final class JobSchedulerShellCommand extends BasicShellCommandHandler {
         return 0;
     }
 
+    private int getTransferredNetworkBytes(PrintWriter pw, int byteOption) throws Exception {
+        checkPermission("get transferred bytes");
+
+        int userId = UserHandle.USER_SYSTEM;
+
+        String opt;
+        while ((opt = getNextOption()) != null) {
+            switch (opt) {
+                case "-u":
+                case "--user":
+                    userId = UserHandle.parseUserArg(getNextArgRequired());
+                    break;
+
+                default:
+                    pw.println("Error: unknown option '" + opt + "'");
+                    return -1;
+            }
+        }
+
+        if (userId == UserHandle.USER_CURRENT) {
+            userId = ActivityManager.getCurrentUser();
+        }
+
+        final String pkgName = getNextArgRequired();
+        final String jobIdStr = getNextArgRequired();
+        final int jobId = Integer.parseInt(jobIdStr);
+
+        final long ident = Binder.clearCallingIdentity();
+        try {
+            int ret = mInternal.getTransferredNetworkBytes(pw, pkgName, userId, jobId, byteOption);
+            printError(ret, pkgName, userId, jobId);
+            return ret;
+        } finally {
+            Binder.restoreCallingIdentity(ident);
+        }
+    }
+
     private int getJobState(PrintWriter pw) throws Exception {
-        checkPermission("force timeout jobs");
+        checkPermission("get job state");
 
         int userId = UserHandle.USER_SYSTEM;
 
@@ -473,10 +558,30 @@ public final class JobSchedulerShellCommand extends BasicShellCommandHandler {
         pw.println("    Return whether the battery is currently considered to be charging.");
         pw.println("  get-battery-not-low");
         pw.println("    Return whether the battery is currently considered to not be low.");
+        pw.println("  get-estimated-download-bytes [-u | --user USER_ID] PACKAGE JOB_ID");
+        pw.println("    Return the most recent estimated download bytes for the job.");
+        pw.println("    Options:");
+        pw.println("      -u or --user: specify which user's job is to be run; the default is");
+        pw.println("         the primary or system user");
+        pw.println("  get-estimated-upload-bytes [-u | --user USER_ID] PACKAGE JOB_ID");
+        pw.println("    Return the most recent estimated upload bytes for the job.");
+        pw.println("    Options:");
+        pw.println("      -u or --user: specify which user's job is to be run; the default is");
+        pw.println("         the primary or system user");
         pw.println("  get-storage-seq");
         pw.println("    Return the last storage update sequence number that was received.");
         pw.println("  get-storage-not-low");
         pw.println("    Return whether storage is currently considered to not be low.");
+        pw.println("  get-transferred-download-bytes [-u | --user USER_ID] PACKAGE JOB_ID");
+        pw.println("    Return the most recent transferred download bytes for the job.");
+        pw.println("    Options:");
+        pw.println("      -u or --user: specify which user's job is to be run; the default is");
+        pw.println("         the primary or system user");
+        pw.println("  get-transferred-upload-bytes [-u | --user USER_ID] PACKAGE JOB_ID");
+        pw.println("    Return the most recent transferred upload bytes for the job.");
+        pw.println("    Options:");
+        pw.println("      -u or --user: specify which user's job is to be run; the default is");
+        pw.println("         the primary or system user");
         pw.println("  get-job-state [-u | --user USER_ID] PACKAGE JOB_ID");
         pw.println("    Return the current state of a job, may be any combination of:");
         pw.println("      pending: currently on the pending list, waiting to be active");
@@ -493,5 +598,4 @@ public final class JobSchedulerShellCommand extends BasicShellCommandHandler {
         pw.println("    Trigger wireless charging dock state.  Active by default.");
         pw.println();
     }
-
 }
diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobServiceContext.java b/apex/jobscheduler/service/java/com/android/server/job/JobServiceContext.java
index df47f1787fdc..0dcfd24b263c 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobServiceContext.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobServiceContext.java
@@ -49,6 +49,7 @@ import android.os.Trace;
 import android.os.UserHandle;
 import android.util.EventLog;
 import android.util.IndentingPrintWriter;
+import android.util.Pair;
 import android.util.Slog;
 import android.util.TimeUtils;
 
@@ -171,6 +172,11 @@ public final class JobServiceContext implements ServiceConnection {
     /** The absolute maximum amount of time the job can run */
     private long mMaxExecutionTimeMillis;
 
+    private long mEstimatedDownloadBytes;
+    private long mEstimatedUploadBytes;
+    private long mTransferredDownloadBytes;
+    private long mTransferredUploadBytes;
+
     /**
      * The stop reason for a pending cancel. If there's not pending cancel, then the value should be
      * {@link JobParameters#STOP_REASON_UNDEFINED}.
@@ -306,6 +312,9 @@ public final class JobServiceContext implements ServiceConnection {
             mMinExecutionGuaranteeMillis = mService.getMinJobExecutionGuaranteeMs(job);
             mMaxExecutionTimeMillis =
                     Math.max(mService.getMaxJobExecutionTimeMs(job), mMinExecutionGuaranteeMillis);
+            mEstimatedDownloadBytes = job.getEstimatedNetworkDownloadBytes();
+            mEstimatedUploadBytes = job.getEstimatedNetworkUploadBytes();
+            mTransferredDownloadBytes = mTransferredUploadBytes = 0;
 
             final long whenDeferred = job.getWhenStandbyDeferred();
             if (whenDeferred > 0) {
@@ -524,6 +533,16 @@ public final class JobServiceContext implements ServiceConnection {
         return false;
     }
 
+    @GuardedBy("mLock")
+    Pair<Long, Long> getEstimatedNetworkBytes() {
+        return Pair.create(mEstimatedDownloadBytes, mEstimatedUploadBytes);
+    }
+
+    @GuardedBy("mLock")
+    Pair<Long, Long> getTransferredNetworkBytes() {
+        return Pair.create(mTransferredDownloadBytes, mTransferredUploadBytes);
+    }
+
     void doJobFinished(JobCallback cb, int jobId, boolean reschedule) {
         final long ident = Binder.clearCallingIdentity();
         try {
@@ -541,14 +560,26 @@ public final class JobServiceContext implements ServiceConnection {
         }
     }
 
-    private void doAcknowledgeGetTransferredDownloadBytesMessage(JobCallback jobCallback, int jobId,
+    private void doAcknowledgeGetTransferredDownloadBytesMessage(JobCallback cb, int jobId,
             int workId, @BytesLong long transferredBytes) {
         // TODO(255393346): Make sure apps call this appropriately and monitor for abuse
+        synchronized (mLock) {
+            if (!verifyCallerLocked(cb)) {
+                return;
+            }
+            mTransferredDownloadBytes = transferredBytes;
+        }
     }
 
-    private void doAcknowledgeGetTransferredUploadBytesMessage(JobCallback jobCallback, int jobId,
+    private void doAcknowledgeGetTransferredUploadBytesMessage(JobCallback cb, int jobId,
             int workId, @BytesLong long transferredBytes) {
         // TODO(255393346): Make sure apps call this appropriately and monitor for abuse
+        synchronized (mLock) {
+            if (!verifyCallerLocked(cb)) {
+                return;
+            }
+            mTransferredUploadBytes = transferredBytes;
+        }
     }
 
     void doAcknowledgeStopMessage(JobCallback cb, int jobId, boolean reschedule) {
@@ -603,6 +634,30 @@ public final class JobServiceContext implements ServiceConnection {
         }
     }
 
+    private void doUpdateEstimatedNetworkBytes(JobCallback cb, int jobId,
+            @Nullable JobWorkItem item, long downloadBytes, long uploadBytes) {
+        // TODO(255393346): Make sure apps call this appropriately and monitor for abuse
+        synchronized (mLock) {
+            if (!verifyCallerLocked(cb)) {
+                return;
+            }
+            mEstimatedDownloadBytes = downloadBytes;
+            mEstimatedUploadBytes = uploadBytes;
+        }
+    }
+
+    private void doUpdateTransferredNetworkBytes(JobCallback cb, int jobId,
+            @Nullable JobWorkItem item, long downloadBytes, long uploadBytes) {
+        // TODO(255393346): Make sure apps call this appropriately and monitor for abuse
+        synchronized (mLock) {
+            if (!verifyCallerLocked(cb)) {
+                return;
+            }
+            mTransferredDownloadBytes = downloadBytes;
+            mTransferredUploadBytes = uploadBytes;
+        }
+    }
+
     private void doSetNotification(JobCallback cb, int jodId, int notificationId,
             Notification notification, int jobEndNotificationPolicy) {
         final int callingPid = Binder.getCallingPid();
@@ -627,16 +682,6 @@ public final class JobServiceContext implements ServiceConnection {
         }
     }
 
-    private void doUpdateTransferredNetworkBytes(JobCallback jobCallback, int jobId,
-            @Nullable JobWorkItem item, long downloadBytes, long uploadBytes) {
-        // TODO(255393346): Make sure apps call this appropriately and monitor for abuse
-    }
-
-    private void doUpdateEstimatedNetworkBytes(JobCallback jobCallback, int jobId,
-            @Nullable JobWorkItem item, long downloadBytes, long uploadBytes) {
-        // TODO(255393346): Make sure apps call this appropriately and monitor for abuse
-    }
-
     /**
      * We acquire/release a wakelock on onServiceConnected/unbindService. This mirrors the work
      * we intend to send to the client - we stop sending work when the service is unbound so until