Run framework migration on a background thread.

Bug: 183177268
Test: AppSearchSchemaMigrationCtsTest
Change-Id: Ie0566ebf1e669cacc0e1a37f5131d49b0256dde6

1 files changed, 129 insertions, 125 deletions

diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
index 0f6468a62794..547aa16ee6e4 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
@@ -32,13 +32,11 @@ import com.android.internal.infra.AndroidFuture;
 import com.android.internal.util.Preconditions;
 
 import java.io.Closeable;
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
@@ -158,68 +156,16 @@ public final class AppSearchSession implements Closeable {
                     schemasPackageAccessibleBundles,
                     callbackExecutor,
                     callback);
-            return;
-        }
-
-        try {
-            // Migration process
-            // 1. Generate the current and the final version map.
-            // TODO(b/182855402) Release binder thread and move the heavy work into worker thread.
-            AndroidFuture<AppSearchResult<GetSchemaResponse>> future = new AndroidFuture<>();
-            getSchema(callbackExecutor, future::complete);
-            AppSearchResult<GetSchemaResponse> getSchemaResult = future.get();
-            if (!getSchemaResult.isSuccess()) {
-                callback.accept(AppSearchResult.newFailedResult(getSchemaResult));
-                return;
-            }
-            GetSchemaResponse getSchemaResponse = getSchemaResult.getResultValue();
-            Set<AppSearchSchema> currentSchemas = getSchemaResponse.getSchemas();
-            Map<String, Integer> currentVersionMap =
-                    SchemaMigrationUtil.buildVersionMap(currentSchemas,
-                            getSchemaResponse.getVersion());
-            Map<String, Integer> finalVersionMap =
-                    SchemaMigrationUtil.buildVersionMap(request.getSchemas(), request.getVersion());
-
-            // 2. SetSchema with forceOverride=false, to retrieve the list of incompatible/deleted
-            // types.
-            mService.setSchema(
-                    mPackageName,
-                    mDatabaseName,
+        } else {
+            setSchemaWithMigrations(
+                    request,
                     schemaBundles,
-                    new ArrayList<>(request.getSchemasNotDisplayedBySystem()),
                     schemasPackageAccessibleBundles,
-                    /*forceOverride=*/ false,
-                    mUserId,
-                    request.getVersion(),
-                    new IAppSearchResultCallback.Stub() {
-                        public void onResult(AppSearchResult result) {
-                            callbackExecutor.execute(() -> {
-                                if (result.isSuccess()) {
-                                    // TODO(b/183177268): once migration is implemented, run
-                                    //  it on workExecutor.
-                                    try {
-                                        Bundle bundle = (Bundle) result.getResultValue();
-                                        SetSchemaResponse setSchemaResponse =
-                                                new SetSchemaResponse(bundle);
-                                        setSchemaMigration(
-                                                request, setSchemaResponse, schemaBundles,
-                                                schemasPackageAccessibleBundles, currentVersionMap,
-                                                finalVersionMap, callback);
-                                    } catch (Throwable t) {
-                                        callback.accept(AppSearchResult.throwableToFailedResult(t));
-                                    }
-                                } else {
-                                    callback.accept(result);
-                                }
-                            });
-                        }
-                    });
-            mIsMutated = true;
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }  catch (Throwable t) {
-            callback.accept(AppSearchResult.throwableToFailedResult(t));
+                    workExecutor,
+                    callbackExecutor,
+                    callback);
         }
+        mIsMutated = true;
     }
 
     /**
@@ -667,7 +613,8 @@ public final class AppSearchSession implements Closeable {
      * <p>We only need one time {@link #setSchema} call for no-migration scenario by using the
      * forceoverride in the request.
      */
-    private void setSchemaNoMigrations(@NonNull SetSchemaRequest request,
+    private void setSchemaNoMigrations(
+            @NonNull SetSchemaRequest request,
             @NonNull List<Bundle> schemaBundles,
             @NonNull Map<String, List<Bundle>> schemasPackageAccessibleBundles,
             @NonNull @CallbackExecutor Executor executor,
@@ -709,7 +656,6 @@ public final class AppSearchSession implements Closeable {
                             });
                         }
                     });
-            mIsMutated = true;
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
@@ -721,86 +667,144 @@ public final class AppSearchSession implements Closeable {
      * <p>First time {@link #setSchema} call with forceOverride is false gives us all incompatible
      * changes. After trigger migrations, the second time call {@link #setSchema} will actually
      * apply the changes.
-     *
-     * @param setSchemaResponse the result of the first setSchema call with forceOverride=false.
      */
-    private void setSchemaMigration(@NonNull SetSchemaRequest request,
-            @NonNull SetSchemaResponse setSchemaResponse,
+    private void setSchemaWithMigrations(
+            @NonNull SetSchemaRequest request,
             @NonNull List<Bundle> schemaBundles,
             @NonNull Map<String, List<Bundle>> schemasPackageAccessibleBundles,
-            @NonNull Map<String, Integer> currentVersionMap, Map<String, Integer> finalVersionMap,
-            @NonNull Consumer<AppSearchResult<SetSchemaResponse>> callback)
-            throws AppSearchException, IOException, RemoteException, ExecutionException,
-            InterruptedException {
-        // 1. If forceOverride is false, check that all incompatible types will be migrated.
-        // If some aren't we must throw an error, rather than proceeding and deleting those
-        // types.
-        if (!request.isForceOverride()) {
-            Set<String> unmigratedTypes = SchemaMigrationUtil.getUnmigratedIncompatibleTypes(
-                    setSchemaResponse.getIncompatibleTypes(),
-                    request.getMigrators(),
-                    currentVersionMap,
-                    finalVersionMap);
-            // check if there are any unmigrated types or deleted types. If there are, we will throw
-            // an exception.
-            // Since the force override is false, the schema will not have been set if there are any
-            // incompatible or deleted types.
-            checkDeletedAndIncompatible(setSchemaResponse.getDeletedTypes(),
-                    unmigratedTypes);
-        }
-
-        try (AppSearchMigrationHelper migrationHelper =
-                     new AppSearchMigrationHelper(mService, mUserId, currentVersionMap,
-                             finalVersionMap, mPackageName, mDatabaseName)) {
-            Map<String, Migrator> migratorMap = request.getMigrators();
-
-            // 2. Trigger migration for all migrators.
-            // TODO(b/177266929) trigger migration for all types together rather than separately.
-            Set<String> migratedTypes = new ArraySet<>();
-            for (Map.Entry<String, Migrator> entry : migratorMap.entrySet()) {
-                String schemaType = entry.getKey();
-                Migrator migrator = entry.getValue();
-                if (SchemaMigrationUtil.shouldTriggerMigration(
-                        schemaType, migrator, currentVersionMap, finalVersionMap)) {
-                    migrationHelper.queryAndTransform(schemaType, migrator);
-                    migratedTypes.add(schemaType);
+            @NonNull Executor workExecutor,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull Consumer<AppSearchResult<SetSchemaResponse>> callback) {
+        workExecutor.execute(() -> {
+            try {
+                // Migration process
+                // 1. Generate the current and the final version map.
+                AndroidFuture<AppSearchResult<GetSchemaResponse>> getSchemaFuture =
+                        new AndroidFuture<>();
+                getSchema(callbackExecutor, getSchemaFuture::complete);
+                AppSearchResult<GetSchemaResponse> getSchemaResult = getSchemaFuture.get();
+                if (!getSchemaResult.isSuccess()) {
+                    callbackExecutor.execute(() ->
+                            callback.accept(AppSearchResult.newFailedResult(getSchemaResult)));
+                    return;
                 }
-            }
-
-            // 3. SetSchema a second time with forceOverride=true if the first attempted failed.
-            if (!setSchemaResponse.getIncompatibleTypes().isEmpty()
-                    || !setSchemaResponse.getDeletedTypes().isEmpty()) {
-                AndroidFuture<AppSearchResult<SetSchemaResponse>> future = new AndroidFuture<>();
-                // only trigger second setSchema() call if the first one is fail.
+                GetSchemaResponse getSchemaResponse = getSchemaResult.getResultValue();
+                Set<AppSearchSchema> currentSchemas = getSchemaResponse.getSchemas();
+                Map<String, Integer> currentVersionMap = SchemaMigrationUtil.buildVersionMap(
+                        currentSchemas, getSchemaResponse.getVersion());
+                Map<String, Integer> finalVersionMap = SchemaMigrationUtil.buildVersionMap(
+                        request.getSchemas(), request.getVersion());
+
+                // 2. SetSchema with forceOverride=false, to retrieve the list of
+                // incompatible/deleted types.
+                AndroidFuture<AppSearchResult<Bundle>> setSchemaFuture = new AndroidFuture<>();
                 mService.setSchema(
                         mPackageName,
                         mDatabaseName,
                         schemaBundles,
                         new ArrayList<>(request.getSchemasNotDisplayedBySystem()),
                         schemasPackageAccessibleBundles,
-                        /*forceOverride=*/ true,
+                        /*forceOverride=*/ false,
                         mUserId,
                         request.getVersion(),
                         new IAppSearchResultCallback.Stub() {
-                            @Override
-                            public void onResult(AppSearchResult result) throws RemoteException {
-                                future.complete(result);
+                            public void onResult(AppSearchResult result) {
+                                setSchemaFuture.complete(result);
                             }
                         });
-                AppSearchResult<SetSchemaResponse> secondSetSchemaResult = future.get();
-                if (!secondSetSchemaResult.isSuccess()) {
-                    // we failed to set the schema in second time with force override = true, which
-                    // is an impossible case. Since we only swallow the incompatible error in the
-                    // first setSchema call, all other errors will be thrown at the first time.
-                    callback.accept(secondSetSchemaResult);
+                AppSearchResult<Bundle> setSchemaResult = setSchemaFuture.get();
+                if (!setSchemaResult.isSuccess()) {
+                    callbackExecutor.execute(() ->
+                            callback.accept(AppSearchResult.newFailedResult(setSchemaResult)));
                     return;
                 }
-            }
+                SetSchemaResponse setSchemaResponse =
+                        new SetSchemaResponse(setSchemaResult.getResultValue());
+
+                // 1. If forceOverride is false, check that all incompatible types will be migrated.
+                // If some aren't we must throw an error, rather than proceeding and deleting those
+                // types.
+                if (!request.isForceOverride()) {
+                    Set<String> unmigratedTypes =
+                            SchemaMigrationUtil.getUnmigratedIncompatibleTypes(
+                                    setSchemaResponse.getIncompatibleTypes(),
+                                    request.getMigrators(),
+                                    currentVersionMap,
+                                    finalVersionMap);
+
+                    // check if there are any unmigrated types or deleted types. If there are, we
+                    // will throw an exception.
+                    // Since the force override is false, the schema will not have been set if there
+                    // are any incompatible or deleted types.
+                    checkDeletedAndIncompatible(
+                            setSchemaResponse.getDeletedTypes(), unmigratedTypes);
+                }
 
-            SetSchemaResponse.Builder responseBuilder = setSchemaResponse.toBuilder()
-                    .addMigratedTypes(migratedTypes);
-            callback.accept(migrationHelper.putMigratedDocuments(responseBuilder));
-        }
+                try (AppSearchMigrationHelper migrationHelper =
+                             new AppSearchMigrationHelper(
+                                     mService, mUserId, currentVersionMap, finalVersionMap,
+                                     mPackageName, mDatabaseName)) {
+                    Map<String, Migrator> migratorMap = request.getMigrators();
+
+                    // 2. Trigger migration for all migrators.
+                    // TODO(b/177266929) trigger migration for all types together rather than
+                    //  separately.
+                    Set<String> migratedTypes = new ArraySet<>();
+                    for (Map.Entry<String, Migrator> entry : migratorMap.entrySet()) {
+                        String schemaType = entry.getKey();
+                        Migrator migrator = entry.getValue();
+                        if (SchemaMigrationUtil.shouldTriggerMigration(
+                                schemaType, migrator, currentVersionMap, finalVersionMap)) {
+                            migrationHelper.queryAndTransform(schemaType, migrator);
+                            migratedTypes.add(schemaType);
+                        }
+                    }
+
+                    // 3. SetSchema a second time with forceOverride=true if the first attempted
+                    // failed.
+                    if (!setSchemaResponse.getIncompatibleTypes().isEmpty()
+                            || !setSchemaResponse.getDeletedTypes().isEmpty()) {
+                        AndroidFuture<AppSearchResult<Bundle>> setSchema2Future =
+                                new AndroidFuture<>();
+                        // only trigger second setSchema() call if the first one is fail.
+                        mService.setSchema(
+                                mPackageName,
+                                mDatabaseName,
+                                schemaBundles,
+                                new ArrayList<>(request.getSchemasNotDisplayedBySystem()),
+                                schemasPackageAccessibleBundles,
+                                /*forceOverride=*/ true,
+                                request.getVersion(),
+                                mUserId,
+                                new IAppSearchResultCallback.Stub() {
+                                    @Override
+                                    public void onResult(AppSearchResult result) {
+                                        setSchema2Future.complete(result);
+                                    }
+                                });
+                        AppSearchResult<Bundle> setSchema2Result = setSchema2Future.get();
+                        if (!setSchema2Result.isSuccess()) {
+                            // we failed to set the schema in second time with forceOverride = true,
+                            // which is an impossible case. Since we only swallow the incompatible
+                            // error in the first setSchema call, all other errors will be thrown at
+                            // the first time.
+                            callbackExecutor.execute(() -> callback.accept(
+                                    AppSearchResult.newFailedResult(setSchemaResult)));
+                            return;
+                        }
+                    }
+
+                    SetSchemaResponse.Builder responseBuilder = setSchemaResponse.toBuilder()
+                            .addMigratedTypes(migratedTypes);
+                    AppSearchResult<SetSchemaResponse> putResult =
+                            migrationHelper.putMigratedDocuments(responseBuilder);
+                    callbackExecutor.execute(() -> callback.accept(putResult));
+                }
+            } catch (Throwable t) {
+                callbackExecutor.execute(() -> callback.accept(
+                        AppSearchResult.throwableToFailedResult(t)));
+            }
+        });
     }
 
     /**  Checks the setSchema() call won't delete any types or has incompatible types. */