Wait until the remote side has loaded everything before start copying.

Also added 2 test cases for it.

Test: Manual tests and automatic tests pass.

Bug: 34277120
Change-Id: Ia587909f1490104f2823d67c6bf1bf5d86d86d9d
(cherry picked from commit 99d7f709997570b922181500068a04ad2c972952)
diff --git a/src/com/android/documentsui/services/CopyJob.java b/src/com/android/documentsui/services/CopyJob.java
index 7123033..f9218fe 100644
--- a/src/com/android/documentsui/services/CopyJob.java
+++ b/src/com/android/documentsui/services/CopyJob.java
@@ -38,9 +38,12 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.res.AssetFileDescriptor;
+import android.database.ContentObserver;
 import android.database.Cursor;
 import android.net.Uri;
 import android.os.CancellationSignal;
+import android.os.Handler;
+import android.os.Looper;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
 import android.provider.DocumentsContract;
@@ -75,6 +78,8 @@
 
     private static final String TAG = "CopyJob";
 
+    private static final long LOADING_TIMEOUT = 60000; // 1 min
+
     final ArrayList<DocumentInfo> convertedFiles = new ArrayList<>();
 
     private long mStartTime = -1;
@@ -460,10 +465,9 @@
         Cursor cursor = null;
         boolean success = true;
         // Iterate over srcs in the directory; copy to the destination directory.
-        final Uri queryUri = buildChildDocumentsUri(srcDir.authority, srcDir.documentId);
         try {
             try {
-                cursor = getClient(srcDir).query(queryUri, queryColumns, null, null, null);
+                cursor = queryChildren(srcDir, queryColumns);
             } catch (RemoteException | RuntimeException e) {
                 Metrics.logFileOperationFailure(
                         appContext, Metrics.SUBFILEOP_QUERY_CHILDREN, srcDir.derivedUri);
@@ -669,7 +673,6 @@
     long calculateFileSizesRecursively(
             ContentProviderClient client, Uri uri) throws ResourceException {
         final String authority = uri.getAuthority();
-        final Uri queryUri = buildChildDocumentsUri(authority, getDocumentId(uri));
         final String queryColumns[] = new String[] {
                 Document.COLUMN_DOCUMENT_ID,
                 Document.COLUMN_MIME_TYPE,
@@ -679,7 +682,7 @@
         long result = 0;
         Cursor cursor = null;
         try {
-            cursor = client.query(queryUri, queryColumns, null, null, null);
+            cursor = queryChildren(client, uri, queryColumns);
             while (cursor.moveToNext() && !isCanceled()) {
                 if (Document.MIME_TYPE_DIR.equals(
                         getCursorString(cursor, Document.COLUMN_MIME_TYPE))) {
@@ -704,6 +707,69 @@
     }
 
     /**
+     * Queries children documents.
+     *
+     * SAF allows {@link DocumentsContract#EXTRA_LOADING} in {@link Cursor#getExtras()} to indicate
+     * there are more data to be loaded. Wait until {@link DocumentsContract#EXTRA_LOADING} is
+     * false and then return the cursor.
+     *
+     * @param srcDir the directory whose children are being loading
+     * @param queryColumns columns of metadata to load
+     * @return cursor of all children documents
+     * @throws RemoteException when the remote throws or waiting for update times out
+     */
+    private Cursor queryChildren(DocumentInfo srcDir, String[] queryColumns)
+            throws RemoteException {
+        try (final ContentProviderClient client = getClient(srcDir)) {
+            return queryChildren(client, srcDir.derivedUri, queryColumns);
+        }
+    }
+
+    /**
+     * Queries children documents.
+     *
+     * SAF allows {@link DocumentsContract#EXTRA_LOADING} in {@link Cursor#getExtras()} to indicate
+     * there are more data to be loaded. Wait until {@link DocumentsContract#EXTRA_LOADING} is
+     * false and then return the cursor.
+     *
+     * @param client the {@link ContentProviderClient} to use to query children
+     * @param dirDocUri the document Uri of the directory whose children are being loaded
+     * @param queryColumns columns of metadata to load
+     * @return cursor of all children documents
+     * @throws RemoteException when the remote throws or waiting for update times out
+     */
+    private Cursor queryChildren(ContentProviderClient client, Uri dirDocUri, String[] queryColumns)
+            throws RemoteException {
+        // TODO (b/34459983): Optimize this performance by processing partial result first while provider is loading
+        // more data. Note we need to skip size calculation to achieve it.
+        final Uri queryUri = buildChildDocumentsUri(dirDocUri.getAuthority(), getDocumentId(dirDocUri));
+        Cursor cursor = client.query(
+                queryUri, queryColumns, (String) null, null, null);
+        while (cursor.getExtras().getBoolean(DocumentsContract.EXTRA_LOADING)) {
+            cursor.registerContentObserver(new DirectoryChildrenObserver(queryUri));
+            try {
+                long start = System.currentTimeMillis();
+                synchronized (queryUri) {
+                    queryUri.wait(LOADING_TIMEOUT);
+                }
+                if (System.currentTimeMillis() - start > LOADING_TIMEOUT) {
+                    // Timed out
+                    throw new RemoteException("Timed out waiting on update for " + queryUri);
+                }
+            } catch (InterruptedException e) {
+                // Should never happen
+                throw new RuntimeException(e);
+            }
+
+            // Make another query
+            cursor = client.query(
+                    queryUri, queryColumns, (String) null, null, null);
+        }
+
+        return cursor;
+    }
+
+    /**
      * Returns true if {@code doc} is a descendant of {@code parentDoc}.
      * @throws ResourceException
      */
@@ -733,4 +799,22 @@
                 .append("}")
                 .toString();
     }
+
+    private static class DirectoryChildrenObserver extends ContentObserver {
+
+        private final Object mNotifier;
+
+        private DirectoryChildrenObserver(Object notifier) {
+            super(new Handler(Looper.getMainLooper()));
+            assert(notifier != null);
+            mNotifier = notifier;
+        }
+
+        @Override
+        public void onChange(boolean selfChange, Uri uri) {
+            synchronized (mNotifier) {
+                mNotifier.notify();
+            }
+        }
+    }
 }
diff --git a/tests/common/com/android/documentsui/DocumentsProviderHelper.java b/tests/common/com/android/documentsui/DocumentsProviderHelper.java
index beb7e5c..cfed56a 100644
--- a/tests/common/com/android/documentsui/DocumentsProviderHelper.java
+++ b/tests/common/com/android/documentsui/DocumentsProviderHelper.java
@@ -304,4 +304,9 @@
         return DocumentsContract.buildDocumentUri(mAuthority, documentId);
     }
 
+    public void setLoadingDuration(long duration) throws RemoteException {
+        final Bundle extra = new Bundle();
+        extra.putLong(DocumentsContract.EXTRA_LOADING, duration);
+        mClient.call("setLoadingDuration", null, extra);
+    }
 }
diff --git a/tests/common/com/android/documentsui/StubProvider.java b/tests/common/com/android/documentsui/StubProvider.java
index dbd948c..62e6659 100644
--- a/tests/common/com/android/documentsui/StubProvider.java
+++ b/tests/common/com/android/documentsui/StubProvider.java
@@ -16,6 +16,7 @@
 
 package com.android.documentsui;
 
+import android.content.ContentResolver;
 import android.content.Context;
 import android.content.SharedPreferences;
 import android.content.pm.ProviderInfo;
@@ -25,10 +26,7 @@
 import android.database.MatrixCursor.RowBuilder;
 import android.graphics.Point;
 import android.net.Uri;
-import android.os.Bundle;
-import android.os.CancellationSignal;
-import android.os.FileUtils;
-import android.os.ParcelFileDescriptor;
+import android.os.*;
 import android.provider.DocumentsContract;
 import android.provider.DocumentsContract.Document;
 import android.provider.DocumentsContract.Root;
@@ -91,6 +89,7 @@
     private String mAuthority = DEFAULT_AUTHORITY;
     private SharedPreferences mPrefs;
     private Set<String> mSimulateReadErrorIds = new HashSet<>();
+    private long mLoadingDuration = 0;
 
     @Override
     public void attachInfo(Context context, ProviderInfo info) {
@@ -134,6 +133,8 @@
             mStorage.put(rootInfo.document.documentId, rootInfo.document);
             mRoots.put(rootId, rootInfo);
         }
+
+        mLoadingDuration = 0;
     }
 
     /**
@@ -225,22 +226,38 @@
     @Override
     public Cursor queryChildDocuments(String parentDocumentId, String[] projection, String sortOrder)
             throws FileNotFoundException {
-        final StubDocument parentDocument = mStorage.get(parentDocumentId);
-        if (parentDocument == null || parentDocument.file.isFile()) {
-            throw new FileNotFoundException();
-        }
-        final MatrixCursor result = new MatrixCursor(projection != null ? projection
-                : DEFAULT_DOCUMENT_PROJECTION);
-        result.setNotificationUri(getContext().getContentResolver(),
-                DocumentsContract.buildChildDocumentsUri(mAuthority, parentDocumentId));
-        StubDocument document;
-        for (File file : parentDocument.file.listFiles()) {
-            document = mStorage.get(getDocumentIdForFile(file));
-            if (document != null) {
-                includeDocument(result, document);
+        if (mLoadingDuration > 0) {
+            final Uri notifyUri = DocumentsContract.buildDocumentUri(mAuthority, parentDocumentId);
+            final ContentResolver resolver = getContext().getContentResolver();
+            new Handler(Looper.getMainLooper()).postDelayed(
+                    () -> resolver.notifyChange(notifyUri, null, false),
+                    mLoadingDuration);
+            mLoadingDuration = 0;
+
+            MatrixCursor cursor = new MatrixCursor(projection != null ? projection : DEFAULT_DOCUMENT_PROJECTION);
+            Bundle bundle = new Bundle();
+            bundle.putBoolean(DocumentsContract.EXTRA_LOADING, true);
+            cursor.setExtras(bundle);
+            cursor.setNotificationUri(resolver, notifyUri);
+            return cursor;
+        } else {
+            final StubDocument parentDocument = mStorage.get(parentDocumentId);
+            if (parentDocument == null || parentDocument.file.isFile()) {
+                throw new FileNotFoundException();
             }
+            final MatrixCursor result = new MatrixCursor(projection != null ? projection
+                    : DEFAULT_DOCUMENT_PROJECTION);
+            result.setNotificationUri(getContext().getContentResolver(),
+                    DocumentsContract.buildChildDocumentsUri(mAuthority, parentDocumentId));
+            StubDocument document;
+            for (File file : parentDocument.file.listFiles()) {
+                document = mStorage.get(getDocumentIdForFile(file));
+                if (document != null) {
+                    includeDocument(result, document);
+                }
+            }
+            return result;
         }
-        return result;
     }
 
     @Override
@@ -489,6 +506,9 @@
                 return null;
             case "createDocumentWithFlags":
                 return dispatchCreateDocumentWithFlags(extras);
+            case "setLoadingDuration":
+                mLoadingDuration = extras.getLong(DocumentsContract.EXTRA_LOADING);
+                return null;
         }
 
         return null;
diff --git a/tests/unit/com/android/documentsui/services/CopyJobTest.java b/tests/unit/com/android/documentsui/services/CopyJobTest.java
index a6b9562..ccf90f3 100644
--- a/tests/unit/com/android/documentsui/services/CopyJobTest.java
+++ b/tests/unit/com/android/documentsui/services/CopyJobTest.java
@@ -66,6 +66,11 @@
         runCopyDirRecursivelyTest();
     }
 
+    public void testCopyDirRecursively_loadingInFirstCursor() throws Exception {
+        mDocs.setLoadingDuration(500);
+        testCopyDirRecursively();
+    }
+
     public void testNoCopyDirToSelf() throws Exception {
         runNoCopyDirToSelfTest();
     }
diff --git a/tests/unit/com/android/documentsui/services/MoveJobTest.java b/tests/unit/com/android/documentsui/services/MoveJobTest.java
index 8607576..99f1be6 100644
--- a/tests/unit/com/android/documentsui/services/MoveJobTest.java
+++ b/tests/unit/com/android/documentsui/services/MoveJobTest.java
@@ -102,6 +102,11 @@
         mDocs.assertChildCount(mSrcRoot, 0);
     }
 
+    public void testMoveDirRecursively_loadingInFirstCursor() throws Exception {
+        mDocs.setLoadingDuration(500);
+        testMoveDirRecursively();
+    }
+
     public void testNoMoveDirToSelf() throws Exception {
         runNoCopyDirToSelfTest();