Merge "Update Framework from Jetpack." into sc-dev
diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchBatchResult.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchBatchResult.java
index 97cfe36..cd75b14 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchBatchResult.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchBatchResult.java
@@ -37,24 +37,36 @@
 public final class AppSearchBatchResult<KeyType, ValueType> implements Parcelable {
     @NonNull private final Map<KeyType, ValueType> mSuccesses;
     @NonNull private final Map<KeyType, AppSearchResult<ValueType>> mFailures;
+    @NonNull private final Map<KeyType, AppSearchResult<ValueType>> mAll;
 
     AppSearchBatchResult(
             @NonNull Map<KeyType, ValueType> successes,
-            @NonNull Map<KeyType, AppSearchResult<ValueType>> failures) {
+            @NonNull Map<KeyType, AppSearchResult<ValueType>> failures,
+            @NonNull Map<KeyType, AppSearchResult<ValueType>> all) {
         mSuccesses = successes;
         mFailures = failures;
+        mAll = all;
     }
 
     private AppSearchBatchResult(@NonNull Parcel in) {
-        mSuccesses = Collections.unmodifiableMap(in.readHashMap(/*loader=*/ null));
-        mFailures = Collections.unmodifiableMap(in.readHashMap(/*loader=*/ null));
+        mAll = Collections.unmodifiableMap(in.readHashMap(/*loader=*/ null));
+        Map<KeyType, ValueType> successes = new ArrayMap<>();
+        Map<KeyType, AppSearchResult<ValueType>> failures = new ArrayMap<>();
+        for (Map.Entry<KeyType, AppSearchResult<ValueType>> entry : mAll.entrySet()) {
+            if (entry.getValue().isSuccess()) {
+                successes.put(entry.getKey(), entry.getValue().getResultValue());
+            } else {
+                failures.put(entry.getKey(), entry.getValue());
+            }
+        }
+        mSuccesses = Collections.unmodifiableMap(successes);
+        mFailures = Collections.unmodifiableMap(failures);
     }
 
     /** @hide */
     @Override
     public void writeToParcel(@NonNull Parcel dest, int flags) {
-        dest.writeMap(mSuccesses);
-        dest.writeMap(mFailures);
+        dest.writeMap(mAll);
     }
 
     /** Returns {@code true} if this {@link AppSearchBatchResult} has no failures. */
@@ -63,8 +75,8 @@
     }
 
     /**
-     * Returns a {@link Map} of all successful keys mapped to the successful
-     * {@link AppSearchResult}s they produced.
+     * Returns a {@link Map} of all successful keys mapped to the successful {@link
+     * AppSearchResult}s they produced.
      *
      * <p>The values of the {@link Map} will not be {@code null}.
      */
@@ -85,7 +97,19 @@
     }
 
     /**
+     * Returns a {@link Map} of all keys mapped to the {@link AppSearchResult}s they produced.
+     *
+     * <p>The values of the {@link Map} will not be {@code null}.
+     * @hide
+     */
+    @NonNull
+    public Map<KeyType, AppSearchResult<ValueType>> getAll() {
+        return mAll;
+    }
+
+    /**
      * Asserts that this {@link AppSearchBatchResult} has no failures.
+     *
      * @hide
      */
     public void checkSuccess() {
@@ -133,6 +157,7 @@
     public static final class Builder<KeyType, ValueType> {
         private final Map<KeyType, ValueType> mSuccesses = new ArrayMap<>();
         private final Map<KeyType, AppSearchResult<ValueType>> mFailures = new ArrayMap<>();
+        private final Map<KeyType, AppSearchResult<ValueType>> mAll = new ArrayMap<>();
         private boolean mBuilt = false;
 
         /**
@@ -181,6 +206,7 @@
                 mFailures.put(key, result);
                 mSuccesses.remove(key);
             }
+            mAll.put(key, result);
             return this;
         }
 
@@ -189,7 +215,7 @@
         public AppSearchBatchResult<KeyType, ValueType> build() {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
             mBuilt = true;
-            return new AppSearchBatchResult<>(mSuccesses, mFailures);
+            return new AppSearchBatchResult<>(mSuccesses, mFailures, mAll);
         }
     }
 }
diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchResult.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchResult.java
index 76225e4..440f633 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchResult.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchResult.java
@@ -22,6 +22,7 @@
 import android.app.appsearch.exceptions.AppSearchException;
 import android.os.Parcel;
 import android.os.Parcelable;
+import android.util.Log;
 
 import java.io.IOException;
 import java.lang.annotation.Retention;
@@ -35,19 +36,21 @@
  */
 public final class AppSearchResult<ValueType> implements Parcelable {
     /**
-     * Result codes from {@link AppSearchManager} methods.
+     * Result codes from {@link AppSearchSession} methods.
+     *
      * @hide
      */
-    @IntDef(value = {
-            RESULT_OK,
-            RESULT_UNKNOWN_ERROR,
-            RESULT_INTERNAL_ERROR,
-            RESULT_INVALID_ARGUMENT,
-            RESULT_IO_ERROR,
-            RESULT_OUT_OF_SPACE,
-            RESULT_NOT_FOUND,
-            RESULT_INVALID_SCHEMA,
-    })
+    @IntDef(
+            value = {
+                RESULT_OK,
+                RESULT_UNKNOWN_ERROR,
+                RESULT_INTERNAL_ERROR,
+                RESULT_INVALID_ARGUMENT,
+                RESULT_IO_ERROR,
+                RESULT_OUT_OF_SPACE,
+                RESULT_NOT_FOUND,
+                RESULT_INVALID_SCHEMA,
+            })
     @Retention(RetentionPolicy.SOURCE)
     public @interface ResultCode {}
 
@@ -60,21 +63,21 @@
     /**
      * An internal error occurred within AppSearch, which the caller cannot address.
      *
-     * This error may be considered similar to {@link IllegalStateException}
+     * <p>This error may be considered similar to {@link IllegalStateException}
      */
     public static final int RESULT_INTERNAL_ERROR = 2;
 
     /**
      * The caller supplied invalid arguments to the call.
      *
-     * This error may be considered similar to {@link IllegalArgumentException}.
+     * <p>This error may be considered similar to {@link IllegalArgumentException}.
      */
     public static final int RESULT_INVALID_ARGUMENT = 3;
 
     /**
      * An issue occurred reading or writing to storage. The call might succeed if repeated.
      *
-     * This error may be considered similar to {@link java.io.IOException}.
+     * <p>This error may be considered similar to {@link java.io.IOException}.
      */
     public static final int RESULT_IO_ERROR = 4;
 
@@ -127,7 +130,7 @@
     /**
      * Returns the result value associated with this result, if it was successful.
      *
-     * <p>See the documentation of the particular {@link AppSearchManager} call producing this
+     * <p>See the documentation of the particular {@link AppSearchSession} call producing this
      * {@link AppSearchResult} for what is placed in the result value by that call.
      *
      * @throws IllegalStateException if this {@link AppSearchResult} is not successful.
@@ -145,8 +148,8 @@
      *
      * <p>If {@link #isSuccess} is {@code true}, the error message is always {@code null}. The error
      * message may be {@code null} even if {@link #isSuccess} is {@code false}. See the
-     * documentation of the particular {@link AppSearchManager} call producing this
-     * {@link AppSearchResult} for what is returned by {@link #getErrorMessage}.
+     * documentation of the particular {@link AppSearchSession} call producing this {@link
+     * AppSearchResult} for what is returned by {@link #getErrorMessage}.
      */
     @Nullable
     public String getErrorMessage() {
@@ -205,6 +208,7 @@
 
     /**
      * Creates a new successful {@link AppSearchResult}.
+     *
      * @hide
      */
     @NonNull
@@ -215,6 +219,7 @@
 
     /**
      * Creates a new failed {@link AppSearchResult}.
+     *
      * @hide
      */
     @NonNull
@@ -227,6 +232,8 @@
     @NonNull
     public static <ValueType> AppSearchResult<ValueType> throwableToFailedResult(
             @NonNull Throwable t) {
+        Log.d("AppSearchResult", "Converting throwable to failed result.", t);
+
         if (t instanceof AppSearchException) {
             return ((AppSearchException) t).toAppSearchResult();
         }
@@ -241,6 +248,6 @@
         } else {
             resultCode = AppSearchResult.RESULT_UNKNOWN_ERROR;
         }
-        return AppSearchResult.newFailedResult(resultCode, t.toString());
+        return AppSearchResult.newFailedResult(resultCode, t.getMessage());
     }
 }
diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
index 8fcd2f9..73ca0cc 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
@@ -108,54 +108,85 @@
      * to {@link #setSchema}, if any, to determine how to treat existing documents. The following
      * types of schema modifications are always safe and are made without deleting any existing
      * documents:
+     *
      * <ul>
-     *     <li>Addition of new types
-     *     <li>Addition of new
-     *         {@link AppSearchSchema.PropertyConfig#CARDINALITY_OPTIONAL OPTIONAL} or
-     *         {@link AppSearchSchema.PropertyConfig#CARDINALITY_REPEATED REPEATED} properties to a
-     *         type
-     *     <li>Changing the cardinality of a data type to be less restrictive (e.g. changing an
-     *         {@link AppSearchSchema.PropertyConfig#CARDINALITY_OPTIONAL OPTIONAL} property into a
-     *         {@link AppSearchSchema.PropertyConfig#CARDINALITY_REPEATED REPEATED} property.
+     *   <li>Addition of new types
+     *   <li>Addition of new {@link AppSearchSchema.PropertyConfig#CARDINALITY_OPTIONAL OPTIONAL} or
+     *       {@link AppSearchSchema.PropertyConfig#CARDINALITY_REPEATED REPEATED} properties to a
+     *       type
+     *   <li>Changing the cardinality of a data type to be less restrictive (e.g. changing an {@link
+     *       AppSearchSchema.PropertyConfig#CARDINALITY_OPTIONAL OPTIONAL} property into a {@link
+     *       AppSearchSchema.PropertyConfig#CARDINALITY_REPEATED REPEATED} property.
      * </ul>
      *
      * <p>The following types of schema changes are not backwards-compatible:
-     * <ul>
-     *     <li>Removal of an existing type
-     *     <li>Removal of a property from a type
-     *     <li>Changing the data type ({@code boolean}, {@code long}, etc.) of an existing property
-     *     <li>For properties of {@code Document} type, changing the schema type of
-     *         {@code Document}s of that property
-     *     <li>Changing the cardinality of a data type to be more restrictive (e.g. changing an
-     *         {@link AppSearchSchema.PropertyConfig#CARDINALITY_OPTIONAL OPTIONAL} property into a
-     *         {@link AppSearchSchema.PropertyConfig#CARDINALITY_REQUIRED REQUIRED} property).
-     *     <li>Adding a
-     *         {@link AppSearchSchema.PropertyConfig#CARDINALITY_REQUIRED REQUIRED} property.
-     * </ul>
-     * <p>Supplying a schema with such changes will, by default, result in this call returning an
-     * {@link AppSearchResult} with a code of {@link AppSearchResult#RESULT_INVALID_SCHEMA} and an
-     * error message describing the incompatibility. In this case the previously set schema will
-     * remain active.
      *
-     * <p>If you need to make non-backwards-compatible changes as described above, you can set the
-     * {@link SetSchemaRequest.Builder#setForceOverride} method to {@code true}. In this case,
-     * instead of returning an {@link AppSearchResult} with the
-     * {@link AppSearchResult#RESULT_INVALID_SCHEMA} error code, all documents which are not
-     * compatible with the new schema will be deleted and the incompatible schema will be applied.
+     * <ul>
+     *   <li>Removal of an existing type
+     *   <li>Removal of a property from a type
+     *   <li>Changing the data type ({@code boolean}, {@code long}, etc.) of an existing property
+     *   <li>For properties of {@code Document} type, changing the schema type of {@code Document}s
+     *       of that property
+     *   <li>Changing the cardinality of a data type to be more restrictive (e.g. changing an {@link
+     *       AppSearchSchema.PropertyConfig#CARDINALITY_OPTIONAL OPTIONAL} property into a {@link
+     *       AppSearchSchema.PropertyConfig#CARDINALITY_REQUIRED REQUIRED} property).
+     *   <li>Adding a {@link AppSearchSchema.PropertyConfig#CARDINALITY_REQUIRED REQUIRED} property.
+     * </ul>
+     *
+     * <p>Supplying a schema with such changes will, by default, result in this call completing its
+     * future with an {@link android.app.appsearch.exceptions.AppSearchException} with a code of
+     * {@link AppSearchResult#RESULT_INVALID_SCHEMA} and a message describing the incompatibility.
+     * In this case the previously set schema will remain active.
+     *
+     * <p>If you need to make non-backwards-compatible changes as described above, you can either:
+     *
+     * <ul>
+     *   <li>Set the {@link SetSchemaRequest.Builder#setForceOverride} method to {@code true}. In
+     *       this case, instead of completing its future with an {@link
+     *       android.app.appsearch.exceptions.AppSearchException} with the {@link
+     *       AppSearchResult#RESULT_INVALID_SCHEMA} error code, all documents which are not
+     *       compatible with the new schema will be deleted and the incompatible schema will be
+     *       applied. Incompatible types and deleted types will be set into {@link
+     *       SetSchemaResponse#getIncompatibleTypes()} and {@link
+     *       SetSchemaResponse#getDeletedTypes()}, respectively.
+     *   <li>Add a {@link android.app.appsearch.AppSearchSchema.Migrator} for each incompatible type
+     *       and make no deletion. The migrator will migrate documents from it's old schema version
+     *       to the new version. Migrated types will be set into both {@link
+     *       SetSchemaResponse#getIncompatibleTypes()} and {@link
+     *       SetSchemaResponse#getMigratedTypes()}. See the migration section below.
+     * </ul>
      *
      * <p>It is a no-op to set the same schema as has been previously set; this is handled
      * efficiently.
      *
-     * <p>By default, documents are visible on platform surfaces. To opt out, call
-     * {@link SetSchemaRequest.Builder#setSchemaTypeVisibilityForSystemUi} with {@code visible} as
-     * false. Any visibility settings apply only to the schemas that are included in the
-     * {@code request}. Visibility settings for a schema type do not persist across
-     * {@link #setSchema} calls.
+     * <p>By default, documents are visible on platform surfaces. To opt out, call {@code
+     * SetSchemaRequest.Builder#setPlatformSurfaceable} with {@code surfaceable} as false. Any
+     * visibility settings apply only to the schemas that are included in the {@code request}.
+     * Visibility settings for a schema type do not apply or persist across {@link
+     * SetSchemaRequest}s.
      *
-     * @param request  The schema update request.
+     * <p>Migration: make non-backwards-compatible changes will delete all stored documents in old
+     * schema. You can save your documents by setting {@link
+     * android.app.appsearch.AppSearchSchema.Migrator} via the {@link
+     * SetSchemaRequest.Builder#setMigrator} for each type you want to save.
+     *
+     * <p>{@link android.app.appsearch.AppSearchSchema.Migrator#onDowngrade} or {@link
+     * android.app.appsearch.AppSearchSchema.Migrator#onUpgrade} will be triggered if the version
+     * number of the schema stored in AppSearch is different with the version in the request.
+     *
+     * <p>If any error or Exception occurred in the {@link
+     * android.app.appsearch.AppSearchSchema.Migrator#onDowngrade}, {@link
+     * android.app.appsearch.AppSearchSchema.Migrator#onUpgrade} or {@link
+     * android.app.appsearch.AppSearchMigrationHelper.Transformer#transform}, the migration will be
+     * terminated, the setSchema request will be rejected unless the schema changes are
+     * backwards-compatible, and stored documents won't have any observable changes.
+     *
+     * @param request The schema update request.
      * @param executor Executor on which to invoke the callback.
      * @param callback Callback to receive errors resulting from setting the schema. If the
      *                 operation succeeds, the callback will be invoked with {@code null}.
+     * @see android.app.appsearch.AppSearchSchema.Migrator
+     * @see android.app.appsearch.AppSearchMigrationHelper.Transformer
      */
     public void setSchema(
             @NonNull SetSchemaRequest request,
@@ -193,11 +224,9 @@
                             executor.execute(() -> {
                                 if (result.isSuccess()) {
                                     callback.accept(
-                                            // TODO(b/151178558) implement Migration in platform.
+                                            // TODO(b/177266929) implement Migration in platform.
                                             AppSearchResult.newSuccessfulResult(
-                                                    new SetSchemaResponse.Builder().setResultCode(
-                                                            result.getResultCode())
-                                                            .build()));
+                                                    new SetSchemaResponse.Builder().build()));
                                 } else {
                                     callback.accept(result);
                                 }
@@ -256,7 +285,7 @@
      * <p>Each {@link GenericDocument}'s {@code schemaType} field must be set to the name of a
      * schema type previously registered via the {@link #setSchema} method.
      *
-     * @param request  {@link PutDocumentsRequest} containing documents to be indexed
+     * @param request {@link PutDocumentsRequest} containing documents to be indexed
      * @param executor Executor on which to invoke the callback.
      * @param callback Callback to receive pending result of performing this operation. The keys
      *                 of the returned {@link AppSearchBatchResult} are the URIs of the input
@@ -297,9 +326,10 @@
     }
 
     /**
-     * Retrieves {@link GenericDocument}s by URI.
+     * Gets {@link GenericDocument} objects by URIs and namespace from the {@link AppSearchSession}
+     * database.
      *
-     * @param request  {@link GetByUriRequest} containing URIs to be retrieved.
+     * @param request a request containing URIs and namespace to get documents for.
      * @param executor Executor on which to invoke the callback.
      * @param callback Callback to receive the pending result of performing this operation. The keys
      *                 of the returned {@link AppSearchBatchResult} are the input URIs. The values
@@ -377,48 +407,65 @@
     }
 
     /**
-     * Searches a document based on a given query string.
+     * Retrieves documents from the open {@link AppSearchSession} that match a given query string
+     * and type of search provided.
      *
-     * <p>Currently we support following features in the raw query format:
+     * <p>Query strings can be empty, contain one term with no operators, or contain multiple terms
+     * and operators.
+     *
+     * <p>For query strings that are empty, all documents that match the {@link SearchSpec} will be
+     * returned.
+     *
+     * <p>For query strings with a single term and no operators, documents that match the provided
+     * query string and {@link SearchSpec} will be returned.
+     *
+     * <p>The following operators are supported:
+     *
      * <ul>
-     *     <li>AND
-     *     <p>AND joins (e.g. “match documents that have both the terms ‘dog’ and
-     *     ‘cat’”).
-     *     Example: hello world matches documents that have both ‘hello’ and ‘world’
-     *     <li>OR
-     *     <p>OR joins (e.g. “match documents that have either the term ‘dog’ or
-     *     ‘cat’”).
-     *     Example: dog OR puppy
-     *     <li>Exclusion
-     *     <p>Exclude a term (e.g. “match documents that do
-     *     not have the term ‘dog’”).
-     *     Example: -dog excludes the term ‘dog’
-     *     <li>Grouping terms
-     *     <p>Allow for conceptual grouping of subqueries to enable hierarchical structures (e.g.
-     *     “match documents that have either ‘dog’ or ‘puppy’, and either ‘cat’ or ‘kitten’”).
-     *     Example: (dog puppy) (cat kitten) two one group containing two terms.
-     *     <li>Property restricts
-     *     <p> Specifies which properties of a document to specifically match terms in (e.g.
-     *     “match documents where the ‘subject’ property contains ‘important’”).
-     *     Example: subject:important matches documents with the term ‘important’ in the
-     *     ‘subject’ property
-     *     <li>Schema type restricts
-     *     <p>This is similar to property restricts, but allows for restricts on top-level document
-     *     fields, such as schema_type. Clients should be able to limit their query to documents of
-     *     a certain schema_type (e.g. “match documents that are of the ‘Email’ schema_type”).
-     *     Example: { schema_type_filters: “Email”, “Video”,query: “dog” } will match documents
-     *     that contain the query term ‘dog’ and are of either the ‘Email’ schema type or the
-     *     ‘Video’ schema type.
+     *   <li>AND (implicit)
+     *       <p>AND is an operator that matches documents that contain <i>all</i> provided terms.
+     *       <p><b>NOTE:</b> A space between terms is treated as an "AND" operator. Explicitly
+     *       including "AND" in a query string will treat "AND" as a term, returning documents that
+     *       also contain "AND".
+     *       <p>Example: "apple AND banana" matches documents that contain the terms "apple", "and",
+     *       "banana".
+     *       <p>Example: "apple banana" matches documents that contain both "apple" and "banana".
+     *       <p>Example: "apple banana cherry" matches documents that contain "apple", "banana", and
+     *       "cherry".
+     *   <li>OR
+     *       <p>OR is an operator that matches documents that contain <i>any</i> provided term.
+     *       <p>Example: "apple OR banana" matches documents that contain either "apple" or
+     *       "banana".
+     *       <p>Example: "apple OR banana OR cherry" matches documents that contain any of "apple",
+     *       "banana", or "cherry".
+     *   <li>Exclusion (-)
+     *       <p>Exclusion (-) is an operator that matches documents that <i>do not</i> contain the
+     *       provided term.
+     *       <p>Example: "-apple" matches documents that do not contain "apple".
+     *   <li>Grouped Terms
+     *       <p>For queries that require multiple operators and terms, terms can be grouped into
+     *       subqueries. Subqueries are contained within an open "(" and close ")" parenthesis.
+     *       <p>Example: "(donut OR bagel) (coffee OR tea)" matches documents that contain either
+     *       "donut" or "bagel" and either "coffee" or "tea".
+     *   <li>Property Restricts
+     *       <p>For queries that require a term to match a specific {@link AppSearchSchema} property
+     *       of a document, a ":" must be included between the property name and the term.
+     *       <p>Example: "subject:important" matches documents that contain the term "important" in
+     *       the "subject" property.
      * </ul>
      *
-     * <p> This method is lightweight. The heavy work will be done in
-     * {@link SearchResults#getNextPage}.
+     * <p>Additional search specifications, such as filtering by {@link AppSearchSchema} type or
+     * adding projection, can be set by calling the corresponding {@link SearchSpec.Builder} setter.
      *
-     * @param queryExpression Query String to search.
-     * @param searchSpec      Spec for setting filters, raw query etc.
+     * <p>This method is lightweight. The heavy work will be done in {@link
+     * SearchResults#getNextPage}.
+     *
+     * @param queryExpression query string to search.
+     * @param searchSpec spec for setting document filters, adding projection, setting term match
+     *     type, etc.
      * @param executor        Executor on which to invoke the callback of the following request
      *                        {@link SearchResults#getNextPage}.
-     * @return The search result of performing this operation.
+     * @return a {@link SearchResults} object for retrieved matched documents.
      */
     @NonNull
     public SearchResults search(
@@ -440,8 +487,8 @@
      *
      * <p>For each call to {@link #reportUsage}, AppSearch updates usage count and usage recency
      * metrics for that particular document. These metrics are used for ordering {@link #search}
-     * results by the {@link SearchSpec#RANKING_STRATEGY_USAGE_COUNT} and
-     * {@link SearchSpec#RANKING_STRATEGY_USAGE_LAST_USED_TIMESTAMP} ranking strategies.
+     * results by the {@link SearchSpec#RANKING_STRATEGY_USAGE_COUNT} and {@link
+     * SearchSpec#RANKING_STRATEGY_USAGE_LAST_USED_TIMESTAMP} ranking strategies.
      *
      * <p>Reporting usage of a document is optional.
      *
@@ -479,9 +526,17 @@
     }
 
     /**
-     * Removes {@link GenericDocument}s from the index by URI.
+     * Removes {@link GenericDocument} objects by URIs and namespace from the {@link
+     * AppSearchSession} database.
      *
-     * @param request  Request containing URIs to be removed.
+     * <p>Removed documents will no longer be surfaced by {@link #search} or {@link #getByUri}
+     * calls.
+     *
+     * <p><b>NOTE:</b>By default, documents are removed via a soft delete operation. Once the
+     * document crosses the count threshold or byte usage threshold, the documents will be removed
+     * from disk.
+     *
+     * @param request {@link RemoveByUriRequest} with URIs and namespace to remove from the index.
      * @param executor Executor on which to invoke the callback.
      * @param callback Callback to receive the pending result of performing this operation. The keys
      *                 of the returned {@link AppSearchBatchResult} are the input URIs. The values
@@ -520,19 +575,18 @@
 
     /**
      * Removes {@link GenericDocument}s from the index by Query. Documents will be removed if they
-     * match the {@code queryExpression} in given namespaces and schemaTypes which is set via
-     * {@link SearchSpec.Builder#addFilterNamespaces} and
-     * {@link SearchSpec.Builder#addFilterSchemas}.
+     * match the {@code queryExpression} in given namespaces and schemaTypes which is set via {@link
+     * SearchSpec.Builder#addFilterNamespaces} and {@link SearchSpec.Builder#addFilterSchemas}.
      *
-     * <p> An empty {@code queryExpression} matches all documents.
+     * <p>An empty {@code queryExpression} matches all documents.
      *
-     * <p> An empty set of namespaces or schemaTypes matches all namespaces or schemaTypes in
-     * the current database.
+     * <p>An empty set of namespaces or schemaTypes matches all namespaces or schemaTypes in the
+     * current database.
      *
      * @param queryExpression Query String to search.
-     * @param searchSpec      Spec containing schemaTypes, namespaces and query expression indicates
-     *                        how document will be removed. All specific about how to scoring,
-     *                        ordering, snippeting and resulting will be ignored.
+     * @param searchSpec Spec containing schemaTypes, namespaces and query expression indicates how
+     *     document will be removed. All specific about how to scoring, ordering, snippeting and
+     *     resulting will be ignored.
      * @param executor        Executor on which to invoke the callback.
      * @param callback        Callback to receive errors resulting from removing the documents. If
      *                        the operation succeeds, the callback will be invoked with
diff --git a/apex/appsearch/framework/java/android/app/appsearch/GlobalSearchSession.java b/apex/appsearch/framework/java/android/app/appsearch/GlobalSearchSession.java
index 8651834..09bca4f 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/GlobalSearchSession.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/GlobalSearchSession.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * Copyright 2020 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -32,7 +32,7 @@
 /**
  * This class provides global access to the centralized AppSearch index maintained by the system.
  *
- * <p>Apps can retrieve indexed documents through the query API.
+ * <p>Apps can retrieve indexed documents through the {@link #search} API.
  */
 public class GlobalSearchSession implements Closeable {
 
@@ -90,48 +90,26 @@
     }
 
     /**
-     * Searches across all documents in the storage based on a given query string.
+     * Retrieves documents from all AppSearch databases that the querying application has access to.
      *
-     * <p>Currently we support following features in the raw query format:
-     * <ul>
-     *     <li>AND
-     *     <p>AND joins (e.g. “match documents that have both the terms ‘dog’ and
-     *     ‘cat’”).
-     *     Example: hello world matches documents that have both ‘hello’ and ‘world’
-     *     <li>OR
-     *     <p>OR joins (e.g. “match documents that have either the term ‘dog’ or
-     *     ‘cat’”).
-     *     Example: dog OR puppy
-     *     <li>Exclusion
-     *     <p>Exclude a term (e.g. “match documents that do
-     *     not have the term ‘dog’”).
-     *     Example: -dog excludes the term ‘dog’
-     *     <li>Grouping terms
-     *     <p>Allow for conceptual grouping of subqueries to enable hierarchical structures (e.g.
-     *     “match documents that have either ‘dog’ or ‘puppy’, and either ‘cat’ or ‘kitten’”).
-     *     Example: (dog puppy) (cat kitten) two one group containing two terms.
-     *     <li>Property restricts
-     *     <p> Specifies which properties of a document to specifically match terms in (e.g.
-     *     “match documents where the ‘subject’ property contains ‘important’”).
-     *     Example: subject:important matches documents with the term ‘important’ in the
-     *     ‘subject’ property
-     *     <li>Schema type restricts
-     *     <p>This is similar to property restricts, but allows for restricts on top-level document
-     *     fields, such as schema_type. Clients should be able to limit their query to documents of
-     *     a certain schema_type (e.g. “match documents that are of the ‘Email’ schema_type”).
-     *     Example: { schema_type_filters: “Email”, “Video”,query: “dog” } will match documents
-     *     that contain the query term ‘dog’ and are of either the ‘Email’ schema type or the
-     *     ‘Video’ schema type.
-     * </ul>
+     * <p>Applications can be granted access to documents by specifying {@link
+     * SetSchemaRequest.Builder#setSchemaTypeVisibilityForPackage} when building a schema.
      *
-     * <p> This method is lightweight. The heavy work will be done in
-     * {@link SearchResults#getNextPage}.
+     * <p>Document access can also be granted to system UIs by specifying {@link
+     * SetSchemaRequest.Builder#setSchemaTypeVisibilityForSystemUi} when building a schema.
      *
-     * @param queryExpression Query String to search.
-     * @param searchSpec      Spec for setting filters, raw query etc.
+     * <p>See {@link AppSearchSession#search} for a detailed explanation on
+     * forming a query string.
+     *
+     * <p>This method is lightweight. The heavy work will be done in {@link
+     * SearchResults#getNextPage}.
+     *
+     * @param queryExpression query string to search.
+     * @param searchSpec spec for setting document filters, adding projection, setting term match
+     *     type, etc.
      * @param executor        Executor on which to invoke the callback of the following request
      *                        {@link SearchResults#getNextPage}.
-     * @return The search result of performing this operation.
+     * @return a {@link SearchResults} object for retrieved matched documents.
      */
     @NonNull
     public SearchResults search(
diff --git a/apex/appsearch/framework/java/android/app/appsearch/SearchResults.java b/apex/appsearch/framework/java/android/app/appsearch/SearchResults.java
index 704509b..a63e015 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/SearchResults.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/SearchResults.java
@@ -35,8 +35,8 @@
 /**
  * SearchResults are a returned object from a query API.
  *
- * <p>Each {@link SearchResult} contains a document and may contain other fields like snippets
- * based on request.
+ * <p>Each {@link SearchResult} contains a document and may contain other fields like snippets based
+ * on request.
  *
  * <p>Should close this object after finish fetching results.
  *
@@ -89,8 +89,8 @@
     /**
      * Gets a whole page of {@link SearchResult}s.
      *
-     * <p>Re-call this method to get next page of {@link SearchResult}, until it returns an
-     * empty list.
+     * <p>Re-call this method to get next page of {@link SearchResult}, until it returns an empty
+     * list.
      *
      * <p>The page size is set by {@link SearchSpec.Builder#setResultCountPerPage}.
      *
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchSchema.java b/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchSchema.java
index e94b3b2..8bf438d 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchSchema.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchSchema.java
@@ -90,6 +90,7 @@
      * <p>This method creates a new list when called.
      */
     @NonNull
+    @SuppressWarnings("MixedMutabilityReturnType")
     public List<PropertyConfig> getProperties() {
         ArrayList<Bundle> propertyBundles =
                 mBundle.getParcelableArrayList(AppSearchSchema.PROPERTIES_FIELD);
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/GenericDocument.java b/apex/appsearch/framework/java/external/android/app/appsearch/GenericDocument.java
index 4c11514..138eb23 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/GenericDocument.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/GenericDocument.java
@@ -36,7 +36,8 @@
 /**
  * Represents a document unit.
  *
- * <p>Documents are constructed via {@link GenericDocument.Builder}.
+ * <p>Documents contain structured data conforming to their {@link AppSearchSchema} type. Each
+ * document is uniquely identified by a URI and namespace.
  *
  * @see AppSearchSession#put
  * @see AppSearchSession#getByUri
@@ -48,16 +49,10 @@
     /** The default empty namespace. */
     public static final String DEFAULT_NAMESPACE = "";
 
-    /**
-     * The maximum number of elements in a repeatable field. Will reject the request if exceed this
-     * limit.
-     */
+    /** The maximum number of elements in a repeatable field. */
     private static final int MAX_REPEATED_PROPERTY_LENGTH = 100;
 
-    /**
-     * The maximum {@link String#length} of a {@link String} field. Will reject the request if
-     * {@link String}s longer than this.
-     */
+    /** The maximum {@link String#length} of a {@link String} field. */
     private static final int MAX_STRING_LENGTH = 20_000;
 
     /** The maximum number of indexed properties a document can have. */
@@ -149,7 +144,7 @@
         return mBundle.getString(NAMESPACE_FIELD, DEFAULT_NAMESPACE);
     }
 
-    /** Returns the schema type of the {@link GenericDocument}. */
+    /** Returns the {@link AppSearchSchema} type of the {@link GenericDocument}. */
     @NonNull
     public String getSchemaType() {
         return mSchemaType;
@@ -165,14 +160,14 @@
     }
 
     /**
-     * Returns the TTL (Time To Live) of the {@link GenericDocument}, in milliseconds.
+     * Returns the TTL (time-to-live) of the {@link GenericDocument}, in milliseconds.
      *
      * <p>The TTL is measured against {@link #getCreationTimestampMillis}. At the timestamp of
      * {@code creationTimestampMillis + ttlMillis}, measured in the {@link System#currentTimeMillis}
      * time base, the document will be auto-deleted.
      *
      * <p>The default value is 0, which means the document is permanent and won't be auto-deleted
-     * until the app is uninstalled.
+     * until the app is uninstalled or {@link AppSearchSession#remove} is called.
      */
     public long getTtlMillis() {
         return mBundle.getLong(TTL_MILLIS_FIELD, DEFAULT_TTL_MILLIS);
@@ -182,12 +177,12 @@
      * Returns the score of the {@link GenericDocument}.
      *
      * <p>The score is a query-independent measure of the document's quality, relative to other
-     * {@link GenericDocument}s of the same type.
+     * {@link GenericDocument} objects of the same {@link AppSearchSchema} type.
      *
      * <p>Results may be sorted by score using {@link SearchSpec.Builder#setRankingStrategy}.
      * Documents with higher scores are considered better than documents with lower scores.
      *
-     * <p>Any nonnegative integer can be used a score.
+     * <p>Any non-negative integer can be used a score.
      */
     public int getScore() {
         return mBundle.getInt(SCORE_FIELD, DEFAULT_SCORE);
@@ -355,7 +350,7 @@
     }
 
     /**
-     * Retrieves a repeated {@link String} property by key.
+     * Retrieves a repeated {@code long[]} property by key.
      *
      * @param key The key to look for.
      * @return The {@code long[]} associated with the given key, or {@code null} if no value is set
@@ -580,14 +575,17 @@
         private boolean mBuilt = false;
 
         /**
-         * Create a new {@link GenericDocument.Builder}.
+         * Creates a new {@link GenericDocument.Builder}.
          *
-         * @param uri The uri of {@link GenericDocument}.
-         * @param schemaType The schema type of the {@link GenericDocument}. The passed-in {@code
-         *     schemaType} must be defined using {@link AppSearchSession#setSchema} prior to
-         *     inserting a document of this {@code schemaType} into the AppSearch index using {@link
-         *     AppSearchSession#put}. Otherwise, the document will be rejected by {@link
-         *     AppSearchSession#put}.
+         * <p>Once {@link #build} is called, the instance can no longer be used.
+         *
+         * @param uri the URI to set for the {@link GenericDocument}.
+         * @param schemaType the {@link AppSearchSchema} type of the {@link GenericDocument}. The
+         *     provided {@code schemaType} must be defined using {@link AppSearchSession#setSchema}
+         *     prior to inserting a document of this {@code schemaType} into the AppSearch index
+         *     using {@link AppSearchSession#put}. Otherwise, the document will
+         *     be rejected by {@link AppSearchSession#put} with result code
+         *     {@link AppSearchResult#RESULT_NOT_FOUND}.
          */
         @SuppressWarnings("unchecked")
         public Builder(@NonNull String uri, @NonNull String schemaType) {
@@ -606,15 +604,18 @@
         }
 
         /**
-         * Sets the app-defined namespace this Document resides in. No special values are reserved
+         * Sets the app-defined namespace this document resides in. No special values are reserved
          * or understood by the infrastructure.
          *
          * <p>URIs are unique within a namespace.
          *
          * <p>The number of namespaces per app should be kept small for efficiency reasons.
+         *
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setNamespace(@NonNull String namespace) {
+            Preconditions.checkState(!mBuilt, "Builder has already been used");
             mBundle.putString(GenericDocument.NAMESPACE_FIELD, namespace);
             return mBuilderTypeInstance;
         }
@@ -623,14 +624,15 @@
          * Sets the score of the {@link GenericDocument}.
          *
          * <p>The score is a query-independent measure of the document's quality, relative to other
-         * {@link GenericDocument}s of the same type.
+         * {@link GenericDocument} objects of the same {@link AppSearchSchema} type.
          *
          * <p>Results may be sorted by score using {@link SearchSpec.Builder#setRankingStrategy}.
          * Documents with higher scores are considered better than documents with lower scores.
          *
-         * <p>Any nonnegative integer can be used a score.
+         * <p>Any non-negative integer can be used a score. By default, scores are set to 0.
          *
-         * @throws IllegalArgumentException If the provided value is negative.
+         * @param score any non-negative {@code int} representing the document's score.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setScore(@IntRange(from = 0, to = Integer.MAX_VALUE) int score) {
@@ -645,8 +647,11 @@
         /**
          * Sets the creation timestamp of the {@link GenericDocument}, in milliseconds.
          *
-         * <p>Should be set using a value obtained from the {@link System#currentTimeMillis} time
-         * base.
+         * <p>This should be set using a value obtained from the {@link System#currentTimeMillis}
+         * time base.
+         *
+         * @param creationTimestampMillis a creation timestamp in milliseconds.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setCreationTimestampMillis(long creationTimestampMillis) {
@@ -657,17 +662,17 @@
         }
 
         /**
-         * Sets the TTL (Time To Live) of the {@link GenericDocument}, in milliseconds.
+         * Sets the TTL (time-to-live) of the {@link GenericDocument}, in milliseconds.
          *
          * <p>The TTL is measured against {@link #getCreationTimestampMillis}. At the timestamp of
          * {@code creationTimestampMillis + ttlMillis}, measured in the {@link
          * System#currentTimeMillis} time base, the document will be auto-deleted.
          *
          * <p>The default value is 0, which means the document is permanent and won't be
-         * auto-deleted until the app is uninstalled.
+         * auto-deleted until the app is uninstalled or {@link AppSearchSession#remove} is called.
          *
-         * @param ttlMillis A non-negative duration in milliseconds.
-         * @throws IllegalArgumentException If the provided value is negative.
+         * @param ttlMillis a non-negative duration in milliseconds.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setTtlMillis(long ttlMillis) {
@@ -682,8 +687,11 @@
         /**
          * Sets one or multiple {@code String} values for a property, replacing its previous values.
          *
-         * @param key The key associated with the {@code values}.
-         * @param values The {@code String} values of the property.
+         * @param key the key associated with the {@code values}.
+         * @param values the {@code String} values of the property.
+         * @throws IllegalArgumentException if no values are provided, if provided values exceed
+         *     maximum repeated property length, or if a passed in {@code String} is {@code null}.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setPropertyString(@NonNull String key, @NonNull String... values) {
@@ -698,8 +706,11 @@
          * Sets one or multiple {@code boolean} values for a property, replacing its previous
          * values.
          *
-         * @param key The key associated with the {@code values}.
-         * @param values The {@code boolean} values of the property.
+         * @param key the key associated with the {@code values}.
+         * @param values the {@code boolean} values of the property.
+         * @throws IllegalArgumentException if no values are provided or if values exceed maximum
+         *     repeated property length.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setPropertyBoolean(@NonNull String key, @NonNull boolean... values) {
@@ -713,8 +724,11 @@
         /**
          * Sets one or multiple {@code long} values for a property, replacing its previous values.
          *
-         * @param key The key associated with the {@code values}.
-         * @param values The {@code long} values of the property.
+         * @param key the key associated with the {@code values}.
+         * @param values the {@code long} values of the property.
+         * @throws IllegalArgumentException if no values are provided or if values exceed maximum
+         *     repeated property length.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setPropertyLong(@NonNull String key, @NonNull long... values) {
@@ -728,8 +742,11 @@
         /**
          * Sets one or multiple {@code double} values for a property, replacing its previous values.
          *
-         * @param key The key associated with the {@code values}.
-         * @param values The {@code double} values of the property.
+         * @param key the key associated with the {@code values}.
+         * @param values the {@code double} values of the property.
+         * @throws IllegalArgumentException if no values are provided or if values exceed maximum
+         *     repeated property length.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setPropertyDouble(@NonNull String key, @NonNull double... values) {
@@ -743,8 +760,11 @@
         /**
          * Sets one or multiple {@code byte[]} for a property, replacing its previous values.
          *
-         * @param key The key associated with the {@code values}.
-         * @param values The {@code byte[]} of the property.
+         * @param key the key associated with the {@code values}.
+         * @param values the {@code byte[]} of the property.
+         * @throws IllegalArgumentException if no values are provided, if provided values exceed
+         *     maximum repeated property length, or if a passed in {@code byte[]} is {@code null}.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setPropertyBytes(@NonNull String key, @NonNull byte[]... values) {
@@ -759,8 +779,12 @@
          * Sets one or multiple {@link GenericDocument} values for a property, replacing its
          * previous values.
          *
-         * @param key The key associated with the {@code values}.
-         * @param values The {@link GenericDocument} values of the property.
+         * @param key the key associated with the {@code values}.
+         * @param values the {@link GenericDocument} values of the property.
+         * @throws IllegalArgumentException if no values are provided, if provided values exceed if
+         *     provided values exceed maximum repeated property length, or if a passed in {@link
+         *     GenericDocument} is {@code null}.
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public BuilderType setPropertyDocument(
@@ -853,7 +877,11 @@
             }
         }
 
-        /** Builds the {@link GenericDocument} object. */
+        /**
+         * Builds the {@link GenericDocument} object.
+         *
+         * @throws IllegalStateException if the builder has already been used.
+         */
         @NonNull
         public GenericDocument build() {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/GetByUriRequest.java b/apex/appsearch/framework/java/external/android/app/appsearch/GetByUriRequest.java
index 656608d..c927e34 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/GetByUriRequest.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/GetByUriRequest.java
@@ -31,7 +31,8 @@
 import java.util.Set;
 
 /**
- * Encapsulates a request to retrieve documents by namespace and URI.
+ * Encapsulates a request to retrieve documents by namespace and URIs from the {@link
+ * AppSearchSession} database.
  *
  * @see AppSearchSession#getByUri
  */
@@ -56,13 +57,13 @@
         mTypePropertyPathsMap = Preconditions.checkNotNull(typePropertyPathsMap);
     }
 
-    /** Returns the namespace to get documents from. */
+    /** Returns the namespace attached to the request. */
     @NonNull
     public String getNamespace() {
         return mNamespace;
     }
 
-    /** Returns the URIs to get from the namespace. */
+    /** Returns the set of URIs attached to the request. */
     @NonNull
     public Set<String> getUris() {
         return Collections.unmodifiableSet(mUris);
@@ -100,7 +101,11 @@
         return mTypePropertyPathsMap;
     }
 
-    /** Builder for {@link GetByUriRequest} objects. */
+    /**
+     * Builder for {@link GetByUriRequest} objects.
+     *
+     * <p>Once {@link #build} is called, the instance can no longer be used.
+     */
     public static final class Builder {
         private String mNamespace = GenericDocument.DEFAULT_NAMESPACE;
         private final Set<String> mUris = new ArraySet<>();
@@ -108,9 +113,12 @@
         private boolean mBuilt = false;
 
         /**
-         * Sets which namespace these documents will be retrieved from.
+         * Sets the namespace to retrieve documents for.
          *
-         * <p>If this is not set, it defaults to {@link GenericDocument#DEFAULT_NAMESPACE}.
+         * <p>If this is not called, the namespace defaults to {@link
+         * GenericDocument#DEFAULT_NAMESPACE}.
+         *
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public Builder setNamespace(@NonNull String namespace) {
@@ -120,14 +128,22 @@
             return this;
         }
 
-        /** Adds one or more URIs to the request. */
+        /**
+         * Adds one or more URIs to the request.
+         *
+         * @throws IllegalStateException if the builder has already been used.
+         */
         @NonNull
         public Builder addUris(@NonNull String... uris) {
             Preconditions.checkNotNull(uris);
             return addUris(Arrays.asList(uris));
         }
 
-        /** Adds one or more URIs to the request. */
+        /**
+         * Adds a collection of URIs to the request.
+         *
+         * @throws IllegalStateException if the builder has already been used.
+         */
         @NonNull
         public Builder addUris(@NonNull Collection<String> uris) {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
@@ -149,7 +165,8 @@
          * GetByUriRequest#PROJECTION_SCHEMA_TYPE_WILDCARD}, then those property paths will apply to
          * all results, excepting any types that have their own, specific property paths set.
          *
-         * <p>{@see SearchSpec.Builder#addProjection(String, String...)}
+         * @throws IllegalStateException if the builder has already been used.
+         *     <p>{@see SearchSpec.Builder#addProjection(String, String...)}
          */
         @NonNull
         public Builder addProjection(@NonNull String schemaType, @NonNull String... propertyPaths) {
@@ -170,7 +187,8 @@
          * GetByUriRequest#PROJECTION_SCHEMA_TYPE_WILDCARD}, then those property paths will apply to
          * all results, excepting any types that have their own, specific property paths set.
          *
-         * <p>{@see SearchSpec.Builder#addProjection(String, String...)}
+         * @throws IllegalStateException if the builder has already been used.
+         *     <p>{@see SearchSpec.Builder#addProjection(String, String...)}
          */
         @NonNull
         public Builder addProjection(
@@ -187,7 +205,11 @@
             return this;
         }
 
-        /** Builds a new {@link GetByUriRequest}. */
+        /**
+         * Builds a new {@link GetByUriRequest}.
+         *
+         * @throws IllegalStateException if the builder has already been used.
+         */
         @NonNull
         public GetByUriRequest build() {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/RemoveByUriRequest.java b/apex/appsearch/framework/java/external/android/app/appsearch/RemoveByUriRequest.java
index 198eee8..39b53b6 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/RemoveByUriRequest.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/RemoveByUriRequest.java
@@ -27,7 +27,8 @@
 import java.util.Set;
 
 /**
- * Encapsulates a request to remove documents by namespace and URI.
+ * Encapsulates a request to remove documents by namespace and URIs from the {@link
+ * AppSearchSession} database.
  *
  * @see AppSearchSession#remove
  */
@@ -46,22 +47,28 @@
         return mNamespace;
     }
 
-    /** Returns the URIs of documents to remove from the namespace. */
+    /** Returns the set of URIs attached to the request. */
     @NonNull
     public Set<String> getUris() {
         return Collections.unmodifiableSet(mUris);
     }
 
-    /** Builder for {@link RemoveByUriRequest} objects. */
+    /**
+     * Builder for {@link RemoveByUriRequest} objects.
+     *
+     * <p>Once {@link #build} is called, the instance can no longer be used.
+     */
     public static final class Builder {
         private String mNamespace = GenericDocument.DEFAULT_NAMESPACE;
         private final Set<String> mUris = new ArraySet<>();
         private boolean mBuilt = false;
 
         /**
-         * Sets which namespace these documents will be removed from.
+         * Sets the namespace to remove documents for.
          *
          * <p>If this is not set, it defaults to {@link GenericDocument#DEFAULT_NAMESPACE}.
+         *
+         * @throws IllegalStateException if the builder has already been used.
          */
         @NonNull
         public Builder setNamespace(@NonNull String namespace) {
@@ -71,14 +78,22 @@
             return this;
         }
 
-        /** Adds one or more URIs to the request. */
+        /**
+         * Adds one or more URIs to the request.
+         *
+         * @throws IllegalStateException if the builder has already been used.
+         */
         @NonNull
         public Builder addUris(@NonNull String... uris) {
             Preconditions.checkNotNull(uris);
             return addUris(Arrays.asList(uris));
         }
 
-        /** Adds one or more URIs to the request. */
+        /**
+         * Adds a collection of URIs to the request.
+         *
+         * @throws IllegalStateException if the builder has already been used.
+         */
         @NonNull
         public Builder addUris(@NonNull Collection<String> uris) {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
@@ -87,7 +102,11 @@
             return this;
         }
 
-        /** Builds a new {@link RemoveByUriRequest}. */
+        /**
+         * Builds a new {@link RemoveByUriRequest}.
+         *
+         * @throws IllegalStateException if the builder has already been used.
+         */
         @NonNull
         public RemoveByUriRequest build() {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResponse.java b/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResponse.java
index 4869aa3..bc99d4f 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResponse.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResponse.java
@@ -16,10 +16,9 @@
 
 package android.app.appsearch;
 
-import static android.app.appsearch.AppSearchResult.RESULT_OK;
-
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.os.Bundle;
 import android.util.ArraySet;
 
 import com.android.internal.util.Preconditions;
@@ -32,23 +31,58 @@
 
 /** The response class of {@link AppSearchSession#setSchema} */
 public class SetSchemaResponse {
-    private final List<MigrationFailure> mMigrationFailures;
-    private final Set<String> mDeletedTypes;
-    private final Set<String> mMigratedTypes;
-    private final Set<String> mIncompatibleTypes;
-    private final @AppSearchResult.ResultCode int mResultCode;
 
-    SetSchemaResponse(
-            @NonNull List<MigrationFailure> migrationFailures,
-            @NonNull Set<String> deletedTypes,
-            @NonNull Set<String> migratedTypes,
-            @NonNull Set<String> incompatibleTypes,
-            @AppSearchResult.ResultCode int resultCode) {
+    private static final String DELETED_TYPES_FIELD = "deletedTypes";
+    private static final String INCOMPATIBLE_TYPES_FIELD = "incompatibleTypes";
+    private static final String MIGRATED_TYPES_FIELD = "migratedTypes";
+
+    private final Bundle mBundle;
+    /**
+     * The migrationFailures won't be saved in the bundle. Since:
+     *
+     * <ul>
+     *   <li>{@link MigrationFailure} is generated in {@link AppSearchSession} which will be the SDK
+     *       side in platform. We don't need to pass it from service side via binder.
+     *   <li>Translate multiple {@link MigrationFailure}s to bundles in {@link Builder} and then
+     *       back in constructor will be a huge waste.
+     * </ul>
+     */
+    private final List<MigrationFailure> mMigrationFailures;
+
+    /** Cache of the inflated deleted schema types. Comes from inflating mBundles at first use. */
+    @Nullable private Set<String> mDeletedTypes;
+
+    /** Cache of the inflated migrated schema types. Comes from inflating mBundles at first use. */
+    @Nullable private Set<String> mMigratedTypes;
+
+    /**
+     * Cache of the inflated incompatible schema types. Comes from inflating mBundles at first use.
+     */
+    @Nullable private Set<String> mIncompatibleTypes;
+
+    SetSchemaResponse(@NonNull Bundle bundle, @NonNull List<MigrationFailure> migrationFailures) {
+        mBundle = Preconditions.checkNotNull(bundle);
         mMigrationFailures = Preconditions.checkNotNull(migrationFailures);
-        mDeletedTypes = Preconditions.checkNotNull(deletedTypes);
-        mMigratedTypes = Preconditions.checkNotNull(migratedTypes);
-        mIncompatibleTypes = Preconditions.checkNotNull(incompatibleTypes);
-        mResultCode = resultCode;
+    }
+
+    SetSchemaResponse(@NonNull Bundle bundle) {
+        this(bundle, /*migrationFailures=*/ Collections.emptyList());
+    }
+
+    /**
+     * Returns the {@link Bundle} populated by this builder.
+     *
+     * @hide
+     */
+    @NonNull
+    public Bundle getBundle() {
+        return mBundle;
+    }
+
+    /** TODO(b/177266929): Remove this deprecated method */
+    //@Deprecated
+    public boolean isSuccess() {
+        return true;
     }
 
     /**
@@ -72,6 +106,12 @@
      */
     @NonNull
     public Set<String> getDeletedTypes() {
+        if (mDeletedTypes == null) {
+            mDeletedTypes =
+                    new ArraySet<>(
+                            Preconditions.checkNotNull(
+                                    mBundle.getStringArrayList(DELETED_TYPES_FIELD)));
+        }
         return Collections.unmodifiableSet(mDeletedTypes);
     }
 
@@ -81,6 +121,12 @@
      */
     @NonNull
     public Set<String> getMigratedTypes() {
+        if (mMigratedTypes == null) {
+            mMigratedTypes =
+                    new ArraySet<>(
+                            Preconditions.checkNotNull(
+                                    mBundle.getStringArrayList(MIGRATED_TYPES_FIELD)));
+        }
         return Collections.unmodifiableSet(mMigratedTypes);
     }
 
@@ -96,22 +142,28 @@
      */
     @NonNull
     public Set<String> getIncompatibleTypes() {
+        if (mIncompatibleTypes == null) {
+            mIncompatibleTypes =
+                    new ArraySet<>(
+                            Preconditions.checkNotNull(
+                                    mBundle.getStringArrayList(INCOMPATIBLE_TYPES_FIELD)));
+        }
         return Collections.unmodifiableSet(mIncompatibleTypes);
     }
 
-    /** Returns {@code true} if all {@link AppSearchSchema}s are successful set to the system. */
-    public boolean isSuccess() {
-        return mResultCode == RESULT_OK;
-    }
-
-    @Override
+    /**
+     * Translates the {@link SetSchemaResponse}'s bundle to {@link Builder}.
+     *
+     * @hide
+     */
     @NonNull
-    public String toString() {
-        return "{\n  Does setSchema success? : "
-                + isSuccess()
-                + "\n  failures: "
-                + mMigrationFailures
-                + "\n}";
+    // TODO(b/179302942) change to Builder(mBundle) powered by mBundle.deepCopy
+    public Builder toBuilder() {
+        return new Builder()
+                .addDeletedTypes(getDeletedTypes())
+                .addIncompatibleTypes(getIncompatibleTypes())
+                .addMigratedTypes(getMigratedTypes())
+                .addMigrationFailures(mMigrationFailures);
     }
 
     /**
@@ -120,44 +172,26 @@
      * @hide
      */
     public static class Builder {
-        private final List<MigrationFailure> mMigrationFailures = new ArrayList<>();
-        private final Set<String> mDeletedTypes = new ArraySet<>();
-        private final Set<String> mMigratedTypes = new ArraySet<>();
-        private final Set<String> mIncompatibleTypes = new ArraySet<>();
-        private @AppSearchResult.ResultCode int mResultCode = RESULT_OK;
+        private final ArrayList<MigrationFailure> mMigrationFailures = new ArrayList<>();
+        private final ArrayList<String> mDeletedTypes = new ArrayList<>();
+        private final ArrayList<String> mMigratedTypes = new ArrayList<>();
+        private final ArrayList<String> mIncompatibleTypes = new ArrayList<>();
         private boolean mBuilt = false;
 
-        /** Adds a {@link MigrationFailure}. */
+        /** Adds {@link MigrationFailure}s to the list of migration failures. */
         @NonNull
-        public Builder setFailure(
-                @NonNull String schemaType,
-                @NonNull String namespace,
-                @NonNull String uri,
-                @NonNull AppSearchResult<Void> failureResult) {
+        public Builder addMigrationFailures(
+                @NonNull Collection<MigrationFailure> migrationFailures) {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
-            Preconditions.checkNotNull(schemaType);
-            Preconditions.checkNotNull(namespace);
-            Preconditions.checkNotNull(uri);
-            Preconditions.checkNotNull(failureResult);
-            Preconditions.checkState(!failureResult.isSuccess());
-            mMigrationFailures.add(new MigrationFailure(schemaType, namespace, uri, failureResult));
+            mMigrationFailures.addAll(Preconditions.checkNotNull(migrationFailures));
             return this;
         }
 
-        /** Adds a {@link MigrationFailure}. */
+        /** Adds a {@link MigrationFailure} to the list of migration failures. */
         @NonNull
-        public Builder setFailure(
-                @NonNull String schemaType,
-                @NonNull String namespace,
-                @NonNull String uri,
-                @AppSearchResult.ResultCode int resultCode,
-                @Nullable String errorMessage) {
-            mMigrationFailures.add(
-                    new MigrationFailure(
-                            schemaType,
-                            namespace,
-                            uri,
-                            AppSearchResult.newFailedResult(resultCode, errorMessage)));
+        public Builder addMigrationFailure(@NonNull MigrationFailure migrationFailure) {
+            Preconditions.checkState(!mBuilt, "Builder has already been used");
+            mMigrationFailures.add(Preconditions.checkNotNull(migrationFailure));
             return this;
         }
 
@@ -169,6 +203,14 @@
             return this;
         }
 
+        /** Adds one deletedType to the list of deleted schema types. */
+        @NonNull
+        public Builder addDeletedType(@NonNull String deletedType) {
+            Preconditions.checkState(!mBuilt, "Builder has already been used");
+            mDeletedTypes.add(Preconditions.checkNotNull(deletedType));
+            return this;
+        }
+
         /** Adds incompatibleTypes to the list of incompatible schema types. */
         @NonNull
         public Builder addIncompatibleTypes(@NonNull Collection<String> incompatibleTypes) {
@@ -177,6 +219,14 @@
             return this;
         }
 
+        /** Adds one incompatibleType to the list of incompatible schema types. */
+        @NonNull
+        public Builder addIncompatibleType(@NonNull String incompatibleType) {
+            Preconditions.checkState(!mBuilt, "Builder has already been used");
+            mIncompatibleTypes.add(Preconditions.checkNotNull(incompatibleType));
+            return this;
+        }
+
         /** Adds migratedTypes to the list of migrated schema types. */
         @NonNull
         public Builder addMigratedTypes(@NonNull Collection<String> migratedTypes) {
@@ -185,11 +235,11 @@
             return this;
         }
 
-        /** Sets the {@link AppSearchResult.ResultCode} of the response. */
+        /** Adds one migratedType to the list of migrated schema types. */
         @NonNull
-        public Builder setResultCode(@AppSearchResult.ResultCode int resultCode) {
+        public Builder addMigratedType(@NonNull String migratedType) {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
-            mResultCode = resultCode;
+            mMigratedTypes.add(Preconditions.checkNotNull(migratedType));
             return this;
         }
 
@@ -197,13 +247,15 @@
         @NonNull
         public SetSchemaResponse build() {
             Preconditions.checkState(!mBuilt, "Builder has already been used");
+            Bundle bundle = new Bundle();
+            bundle.putStringArrayList(INCOMPATIBLE_TYPES_FIELD, mIncompatibleTypes);
+            bundle.putStringArrayList(DELETED_TYPES_FIELD, mDeletedTypes);
+            bundle.putStringArrayList(MIGRATED_TYPES_FIELD, mMigratedTypes);
             mBuilt = true;
-            return new SetSchemaResponse(
-                    mMigrationFailures,
-                    mDeletedTypes,
-                    mMigratedTypes,
-                    mIncompatibleTypes,
-                    mResultCode);
+            // Avoid converting the potential thousands of MigrationFailures to Pracelable and
+            // back just for put in bundle. In platform, we should set MigrationFailures in
+            // AppSearchSession after we pass SetSchemaResponse via binder.
+            return new SetSchemaResponse(bundle, mMigrationFailures);
         }
     }
 
@@ -212,38 +264,44 @@
      * {@link AppSearchSession#setSchema}.
      */
     public static class MigrationFailure {
-        private final String mSchemaType;
-        private final String mNamespace;
-        private final String mUri;
-        AppSearchResult<Void> mFailureResult;
+        private static final String SCHEMA_TYPE_FIELD = "schemaType";
+        private static final String NAMESPACE_FIELD = "namespace";
+        private static final String URI_FIELD = "uri";
+        private static final String ERROR_MESSAGE_FIELD = "errorMessage";
+        private static final String RESULT_CODE_FIELD = "resultCode";
 
-        MigrationFailure(
-                @NonNull String schemaType,
-                @NonNull String namespace,
-                @NonNull String uri,
-                @NonNull AppSearchResult<Void> result) {
-            mSchemaType = schemaType;
-            mNamespace = namespace;
-            mUri = uri;
-            mFailureResult = result;
+        private final Bundle mBundle;
+
+        MigrationFailure(@NonNull Bundle bundle) {
+            mBundle = bundle;
+        }
+
+        /**
+         * Returns the Bundle of the {@link MigrationFailure}.
+         *
+         * @hide
+         */
+        @NonNull
+        public Bundle getBundle() {
+            return mBundle;
         }
 
         /** Returns the schema type of the {@link GenericDocument} that fails to be migrated. */
         @NonNull
         public String getSchemaType() {
-            return mSchemaType;
+            return mBundle.getString(SCHEMA_TYPE_FIELD, /*defaultValue=*/ "");
         }
 
         /** Returns the namespace of the {@link GenericDocument} that fails to be migrated. */
         @NonNull
         public String getNamespace() {
-            return mNamespace;
+            return mBundle.getString(NAMESPACE_FIELD, /*defaultValue=*/ "");
         }
 
         /** Returns the uri of the {@link GenericDocument} that fails to be migrated. */
         @NonNull
         public String getUri() {
-            return mUri;
+            return mBundle.getString(URI_FIELD, /*defaultValue=*/ "");
         }
 
         /**
@@ -252,7 +310,69 @@
          */
         @NonNull
         public AppSearchResult<Void> getAppSearchResult() {
-            return mFailureResult;
+            return AppSearchResult.newFailedResult(
+                    mBundle.getInt(RESULT_CODE_FIELD),
+                    mBundle.getString(ERROR_MESSAGE_FIELD, /*defaultValue=*/ ""));
+        }
+
+        /**
+         * Builder for {@link MigrationFailure} objects.
+         *
+         * @hide
+         */
+        public static class Builder {
+            private String mSchemaType;
+            private String mNamespace;
+            private String mUri;
+            private final Bundle mBundle = new Bundle();
+            private AppSearchResult<Void> mFailureResult;
+            private boolean mBuilt = false;
+
+            /** Sets the schema type for the {@link MigrationFailure}. */
+            @NonNull
+            public Builder setSchemaType(@NonNull String schemaType) {
+                Preconditions.checkState(!mBuilt, "Builder has already been used");
+                mSchemaType = Preconditions.checkNotNull(schemaType);
+                return this;
+            }
+
+            /** Sets the namespace for the {@link MigrationFailure}. */
+            @NonNull
+            public Builder setNamespace(@NonNull String namespace) {
+                Preconditions.checkState(!mBuilt, "Builder has already been used");
+                mNamespace = Preconditions.checkNotNull(namespace);
+                return this;
+            }
+
+            /** Sets the uri for the {@link MigrationFailure}. */
+            @NonNull
+            public Builder setUri(@NonNull String uri) {
+                Preconditions.checkState(!mBuilt, "Builder has already been used");
+                mUri = Preconditions.checkNotNull(uri);
+                return this;
+            }
+
+            /** Sets the failure {@link AppSearchResult} for the {@link MigrationFailure}. */
+            @NonNull
+            public Builder setAppSearchResult(@NonNull AppSearchResult<Void> appSearchResult) {
+                Preconditions.checkState(!mBuilt, "Builder has already been used");
+                Preconditions.checkState(!appSearchResult.isSuccess(), "Input a success result");
+                mFailureResult = Preconditions.checkNotNull(appSearchResult);
+                return this;
+            }
+
+            /** Builds a {@link MigrationFailure} object. */
+            @NonNull
+            public MigrationFailure build() {
+                Preconditions.checkState(!mBuilt, "Builder has already been used");
+                mBundle.putString(SCHEMA_TYPE_FIELD, mSchemaType);
+                mBundle.putString(NAMESPACE_FIELD, mNamespace);
+                mBundle.putString(URI_FIELD, mUri);
+                mBundle.putString(ERROR_MESSAGE_FIELD, mFailureResult.getErrorMessage());
+                mBundle.putInt(RESULT_CODE_FIELD, mFailureResult.getResultCode());
+                mBuilt = true;
+                return new MigrationFailure(mBundle);
+            }
         }
     }
 }
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResult.java b/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResult.java
deleted file mode 100644
index f04ace6..0000000
--- a/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResult.java
+++ /dev/null
@@ -1,125 +0,0 @@
-/*
- * Copyright 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.app.appsearch;
-
-import android.annotation.NonNull;
-import android.os.Bundle;
-
-import com.android.internal.util.Preconditions;
-
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-
-/**
- * This class represents the results of setSchema().
- *
- * @hide
- */
-public class SetSchemaResult {
-
-    public static final String DELETED_SCHEMA_TYPES_FIELD = "deletedSchemaTypes";
-    public static final String INCOMPATIBLE_SCHEMA_TYPES_FIELD = "incompatibleSchemaTypes";
-    public static final String RESULT_CODE_FIELD = "resultCode";
-    private final List<String> mDeletedSchemaTypes;
-    private final List<String> mIncompatibleSchemaTypes;
-    private final Bundle mBundle;
-
-    SetSchemaResult(@NonNull Bundle bundle) {
-        mBundle = Preconditions.checkNotNull(bundle);
-        mDeletedSchemaTypes =
-                Preconditions.checkNotNull(mBundle.getStringArrayList(DELETED_SCHEMA_TYPES_FIELD));
-        mIncompatibleSchemaTypes =
-                Preconditions.checkNotNull(
-                        mBundle.getStringArrayList(INCOMPATIBLE_SCHEMA_TYPES_FIELD));
-    }
-
-    /** Returns the {@link Bundle} of this class. */
-    @NonNull
-    public Bundle getBundle() {
-        return mBundle;
-    }
-
-    /** returns all deleted schema types in this setSchema call. */
-    @NonNull
-    public List<String> getDeletedSchemaTypes() {
-        return Collections.unmodifiableList(mDeletedSchemaTypes);
-    }
-
-    /** returns all incompatible schema types in this setSchema call. */
-    @NonNull
-    public List<String> getIncompatibleSchemaTypes() {
-        return Collections.unmodifiableList(mIncompatibleSchemaTypes);
-    }
-
-    /**
-     * returns the {@link android.app.appsearch.AppSearchResult.ResultCode} of the {@link
-     * AppSearchSession#setSchema} call.
-     */
-    public int getResultCode() {
-        return mBundle.getInt(RESULT_CODE_FIELD);
-    }
-
-    /** Builder for {@link SetSchemaResult} objects. */
-    public static final class Builder {
-        private final ArrayList<String> mDeletedSchemaTypes = new ArrayList<>();
-        private final ArrayList<String> mIncompatibleSchemaTypes = new ArrayList<>();
-        @AppSearchResult.ResultCode private int mResultCode;
-        private boolean mBuilt = false;
-
-        /** Adds a deletedSchemaTypes to the {@link SetSchemaResult}. */
-        @NonNull
-        public Builder addDeletedSchemaType(@NonNull String deletedSchemaType) {
-            Preconditions.checkState(!mBuilt, "Builder has already been used");
-            mDeletedSchemaTypes.add(Preconditions.checkNotNull(deletedSchemaType));
-            return this;
-        }
-
-        /** Adds a incompatible SchemaTypes to the {@link SetSchemaResult}. */
-        @NonNull
-        public Builder addIncompatibleSchemaType(@NonNull String incompatibleSchemaTypes) {
-            Preconditions.checkState(!mBuilt, "Builder has already been used");
-            mIncompatibleSchemaTypes.add(Preconditions.checkNotNull(incompatibleSchemaTypes));
-            return this;
-        }
-
-        /**
-         * Sets the {@link android.app.appsearch.AppSearchResult.ResultCode} of the {@link
-         * AppSearchSession#setSchema} call to the {@link SetSchemaResult}
-         */
-        @NonNull
-        public Builder setResultCode(@AppSearchResult.ResultCode int resultCode) {
-            Preconditions.checkState(!mBuilt, "Builder has already been used");
-            mResultCode = resultCode;
-            return this;
-        }
-
-        /** Builds a {@link SetSchemaResult}. */
-        @NonNull
-        public SetSchemaResult build() {
-            Preconditions.checkState(!mBuilt, "Builder has already been used");
-            Bundle bundle = new Bundle();
-            bundle.putStringArrayList(
-                    SetSchemaResult.DELETED_SCHEMA_TYPES_FIELD, mDeletedSchemaTypes);
-            bundle.putStringArrayList(
-                    SetSchemaResult.INCOMPATIBLE_SCHEMA_TYPES_FIELD, mIncompatibleSchemaTypes);
-            bundle.putInt(RESULT_CODE_FIELD, mResultCode);
-            mBuilt = true;
-            return new SetSchemaResult(bundle);
-        }
-    }
-}
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
index 2f1817e..6c2e30e 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
@@ -25,7 +25,7 @@
 import android.app.appsearch.PackageIdentifier;
 import android.app.appsearch.SearchResultPage;
 import android.app.appsearch.SearchSpec;
-import android.app.appsearch.SetSchemaResult;
+import android.app.appsearch.SetSchemaResponse;
 import android.app.appsearch.exceptions.AppSearchException;
 import android.content.Context;
 import android.os.Bundle;
@@ -41,7 +41,7 @@
 import com.android.server.appsearch.external.localstorage.converter.SchemaToProtoConverter;
 import com.android.server.appsearch.external.localstorage.converter.SearchResultToProtoConverter;
 import com.android.server.appsearch.external.localstorage.converter.SearchSpecToProtoConverter;
-import com.android.server.appsearch.external.localstorage.converter.SetSchemaResultToProtoConverter;
+import com.android.server.appsearch.external.localstorage.converter.SetSchemaResponseToProtoConverter;
 import com.android.server.appsearch.external.localstorage.converter.TypePropertyPathToProtoConverter;
 
 import com.google.android.icing.IcingSearchEngine;
@@ -73,6 +73,7 @@
 import com.google.android.icing.proto.TypePropertyMask;
 import com.google.android.icing.proto.UsageReport;
 
+import java.io.Closeable;
 import java.io.File;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -119,7 +120,7 @@
  * @hide
  */
 @WorkerThread
-public final class AppSearchImpl {
+public final class AppSearchImpl implements Closeable {
     private static final String TAG = "AppSearchImpl";
 
     @VisibleForTesting static final char DATABASE_DELIMITER = '/';
@@ -151,12 +152,16 @@
     private final Map<String, Set<String>> mNamespaceMapLocked = new HashMap<>();
 
     /**
-     * The counter to check when to call {@link #checkForOptimizeLocked(boolean)}. The interval is
-     * {@link #CHECK_OPTIMIZE_INTERVAL}.
+     * The counter to check when to call {@link #checkForOptimize}. The interval is {@link
+     * #CHECK_OPTIMIZE_INTERVAL}.
      */
     @GuardedBy("mReadWriteLock")
     private int mOptimizeIntervalCountLocked = 0;
 
+    /** Whether this instance has been closed, and therefore unusable. */
+    @GuardedBy("mReadWriteLock")
+    private boolean mClosedLocked = false;
+
     /**
      * Creates and initializes an instance of {@link AppSearchImpl} which writes data to the given
      * folder.
@@ -208,7 +213,7 @@
             } catch (AppSearchException e) {
                 Log.w(TAG, "Error initializing, resetting IcingSearchEngine.", e);
                 // Some error. Reset and see if it fixes it.
-                reset();
+                resetLocked();
                 return;
             }
 
@@ -222,11 +227,6 @@
             for (String prefixedNamespace : getAllNamespacesResultProto.getNamespacesList()) {
                 addToMap(mNamespaceMapLocked, getPrefix(prefixedNamespace), prefixedNamespace);
             }
-
-            // TODO(b/155939114): It's possible to optimize after init, which would reduce the time
-            //   to when we're able to serve queries. Consider moving this optimize call out.
-            checkForOptimizeLocked(/* force= */ true);
-
         } finally {
             mReadWriteLock.writeLock().unlock();
         }
@@ -240,12 +240,45 @@
     void initializeVisibilityStore() throws AppSearchException {
         mReadWriteLock.writeLock().lock();
         try {
+            throwIfClosedLocked();
+
             mVisibilityStoreLocked.initialize();
         } finally {
             mReadWriteLock.writeLock().unlock();
         }
     }
 
+    @GuardedBy("mReadWriteLock")
+    private void throwIfClosedLocked() {
+        if (mClosedLocked) {
+            throw new IllegalStateException("Trying to use a closed AppSearchImpl instance.");
+        }
+    }
+
+    /**
+     * Persists data to disk and closes the instance.
+     *
+     * <p>This instance is no longer usable after it's been closed. Call {@link #create} to create a
+     * new, usable instance.
+     */
+    @Override
+    public void close() {
+        mReadWriteLock.writeLock().lock();
+        try {
+            if (mClosedLocked) {
+                return;
+            }
+
+            persistToDisk();
+            mIcingSearchEngineLocked.close();
+            mClosedLocked = true;
+        } catch (AppSearchException e) {
+            Log.w(TAG, "Error when closing AppSearchImpl.", e);
+        } finally {
+            mReadWriteLock.writeLock().unlock();
+        }
+    }
+
     /**
      * Updates the AppSearch schema for this app.
      *
@@ -259,10 +292,14 @@
      * @param schemasPackageAccessible Schema types that are visible to the specified packages.
      * @param forceOverride Whether to force-apply the schema even if it is incompatible. Documents
      *     which do not comply with the new schema will be deleted.
-     * @throws AppSearchException on IcingSearchEngine error.
+     * @throws AppSearchException On IcingSearchEngine error. If the status code is
+     *     FAILED_PRECONDITION for the incompatible change, the exception will be converted to the
+     *     SetSchemaResponse.
+     * @return The response contains deleted schema types and incompatible schema types of this
+     *     call.
      */
     @NonNull
-    public SetSchemaResult setSchema(
+    public SetSchemaResponse setSchema(
             @NonNull String packageName,
             @NonNull String databaseName,
             @NonNull List<AppSearchSchema> schemas,
@@ -272,6 +309,8 @@
             throws AppSearchException {
         mReadWriteLock.writeLock().lock();
         try {
+            throwIfClosedLocked();
+
             SchemaProto.Builder existingSchemaBuilder = getSchemaProtoLocked().toBuilder();
 
             SchemaProto.Builder newSchemaBuilder = SchemaProto.newBuilder();
@@ -297,9 +336,16 @@
             try {
                 checkSuccess(setSchemaResultProto.getStatus());
             } catch (AppSearchException e) {
-                if (setSchemaResultProto.getDeletedSchemaTypesCount() > 0
-                        || setSchemaResultProto.getIncompatibleSchemaTypesCount() > 0) {
-                    return SetSchemaResultToProtoConverter.toSetSchemaResult(
+                // Swallow the exception for the incompatible change case. We will propagate
+                // those deleted schemas and incompatible types to the SetSchemaResponse.
+                boolean isFailedPrecondition =
+                        setSchemaResultProto.getStatus().getCode()
+                                == StatusProto.Code.FAILED_PRECONDITION;
+                boolean isIncompatible =
+                        setSchemaResultProto.getDeletedSchemaTypesCount() > 0
+                                || setSchemaResultProto.getIncompatibleSchemaTypesCount() > 0;
+                if (isFailedPrecondition && isIncompatible) {
+                    return SetSchemaResponseToProtoConverter.toSetSchemaResponse(
                             setSchemaResultProto, prefix);
                 } else {
                     throw e;
@@ -328,16 +374,8 @@
                     prefixedSchemasNotPlatformSurfaceable,
                     prefixedSchemasPackageAccessible);
 
-            // Determine whether to schedule an immediate optimize.
-            if (setSchemaResultProto.getDeletedSchemaTypesCount() > 0
-                    || (setSchemaResultProto.getIncompatibleSchemaTypesCount() > 0
-                            && forceOverride)) {
-                // Any existing schemas which is not in 'schemas' will be deleted, and all
-                // documents of these types were also deleted. And so well if we force override
-                // incompatible schemas.
-                checkForOptimizeLocked(/* force= */ true);
-            }
-            return SetSchemaResultToProtoConverter.toSetSchemaResult(setSchemaResultProto, prefix);
+            return SetSchemaResponseToProtoConverter.toSetSchemaResponse(
+                    setSchemaResultProto, prefix);
         } finally {
             mReadWriteLock.writeLock().unlock();
         }
@@ -355,44 +393,47 @@
     @NonNull
     public List<AppSearchSchema> getSchema(
             @NonNull String packageName, @NonNull String databaseName) throws AppSearchException {
-        SchemaProto fullSchema;
         mReadWriteLock.readLock().lock();
         try {
-            fullSchema = getSchemaProtoLocked();
+            throwIfClosedLocked();
+
+            SchemaProto fullSchema = getSchemaProtoLocked();
+
+            String prefix = createPrefix(packageName, databaseName);
+            List<AppSearchSchema> result = new ArrayList<>();
+            for (int i = 0; i < fullSchema.getTypesCount(); i++) {
+                String typePrefix = getPrefix(fullSchema.getTypes(i).getSchemaType());
+                if (!prefix.equals(typePrefix)) {
+                    continue;
+                }
+                // Rewrite SchemaProto.types.schema_type
+                SchemaTypeConfigProto.Builder typeConfigBuilder =
+                        fullSchema.getTypes(i).toBuilder();
+                String newSchemaType = typeConfigBuilder.getSchemaType().substring(prefix.length());
+                typeConfigBuilder.setSchemaType(newSchemaType);
+
+                // Rewrite SchemaProto.types.properties.schema_type
+                for (int propertyIdx = 0;
+                        propertyIdx < typeConfigBuilder.getPropertiesCount();
+                        propertyIdx++) {
+                    PropertyConfigProto.Builder propertyConfigBuilder =
+                            typeConfigBuilder.getProperties(propertyIdx).toBuilder();
+                    if (!propertyConfigBuilder.getSchemaType().isEmpty()) {
+                        String newPropertySchemaType =
+                                propertyConfigBuilder.getSchemaType().substring(prefix.length());
+                        propertyConfigBuilder.setSchemaType(newPropertySchemaType);
+                        typeConfigBuilder.setProperties(propertyIdx, propertyConfigBuilder);
+                    }
+                }
+
+                AppSearchSchema schema =
+                        SchemaToProtoConverter.toAppSearchSchema(typeConfigBuilder);
+                result.add(schema);
+            }
+            return result;
         } finally {
             mReadWriteLock.readLock().unlock();
         }
-
-        String prefix = createPrefix(packageName, databaseName);
-        List<AppSearchSchema> result = new ArrayList<>();
-        for (int i = 0; i < fullSchema.getTypesCount(); i++) {
-            String typePrefix = getPrefix(fullSchema.getTypes(i).getSchemaType());
-            if (!prefix.equals(typePrefix)) {
-                continue;
-            }
-            // Rewrite SchemaProto.types.schema_type
-            SchemaTypeConfigProto.Builder typeConfigBuilder = fullSchema.getTypes(i).toBuilder();
-            String newSchemaType = typeConfigBuilder.getSchemaType().substring(prefix.length());
-            typeConfigBuilder.setSchemaType(newSchemaType);
-
-            // Rewrite SchemaProto.types.properties.schema_type
-            for (int propertyIdx = 0;
-                    propertyIdx < typeConfigBuilder.getPropertiesCount();
-                    propertyIdx++) {
-                PropertyConfigProto.Builder propertyConfigBuilder =
-                        typeConfigBuilder.getProperties(propertyIdx).toBuilder();
-                if (!propertyConfigBuilder.getSchemaType().isEmpty()) {
-                    String newPropertySchemaType =
-                            propertyConfigBuilder.getSchemaType().substring(prefix.length());
-                    propertyConfigBuilder.setSchemaType(newPropertySchemaType);
-                    typeConfigBuilder.setProperties(propertyIdx, propertyConfigBuilder);
-                }
-            }
-
-            AppSearchSchema schema = SchemaToProtoConverter.toAppSearchSchema(typeConfigBuilder);
-            result.add(schema);
-        }
-        return result;
     }
 
     /**
@@ -410,23 +451,22 @@
             @NonNull String databaseName,
             @NonNull GenericDocument document)
             throws AppSearchException {
-        DocumentProto.Builder documentBuilder =
-                GenericDocumentToProtoConverter.toDocumentProto(document).toBuilder();
-        String prefix = createPrefix(packageName, databaseName);
-        addPrefixToDocument(documentBuilder, prefix);
-
-        PutResultProto putResultProto;
         mReadWriteLock.writeLock().lock();
         try {
-            putResultProto = mIcingSearchEngineLocked.put(documentBuilder.build());
+            throwIfClosedLocked();
+
+            DocumentProto.Builder documentBuilder =
+                    GenericDocumentToProtoConverter.toDocumentProto(document).toBuilder();
+            String prefix = createPrefix(packageName, databaseName);
+            addPrefixToDocument(documentBuilder, prefix);
+
+            PutResultProto putResultProto = mIcingSearchEngineLocked.put(documentBuilder.build());
             addToMap(mNamespaceMapLocked, prefix, documentBuilder.getNamespace());
-            // The existing documents with same URI will be deleted, so there maybe some resources
-            // could be released after optimize().
-            checkForOptimizeLocked(/* force= */ false);
+
+            checkSuccess(putResultProto.getStatus());
         } finally {
             mReadWriteLock.writeLock().unlock();
         }
-        checkSuccess(putResultProto.getStatus());
     }
 
     /**
@@ -451,40 +491,42 @@
             @NonNull String uri,
             @NonNull Map<String, List<String>> typePropertyPaths)
             throws AppSearchException {
-        GetResultProto getResultProto;
-        List<TypePropertyMask> nonPrefixedPropertyMasks =
-                TypePropertyPathToProtoConverter.toTypePropertyMaskList(typePropertyPaths);
-        List<TypePropertyMask> prefixedPropertyMasks =
-                new ArrayList<>(nonPrefixedPropertyMasks.size());
-        for (int i = 0; i < nonPrefixedPropertyMasks.size(); ++i) {
-            TypePropertyMask typePropertyMask = nonPrefixedPropertyMasks.get(i);
-            String nonPrefixedType = typePropertyMask.getSchemaType();
-            String prefixedType =
-                    nonPrefixedType.equals(GetByUriRequest.PROJECTION_SCHEMA_TYPE_WILDCARD)
-                            ? nonPrefixedType
-                            : createPrefix(packageName, databaseName) + nonPrefixedType;
-            prefixedPropertyMasks.add(
-                    typePropertyMask.toBuilder().setSchemaType(prefixedType).build());
-        }
-        GetResultSpecProto getResultSpec =
-                GetResultSpecProto.newBuilder()
-                        .addAllTypePropertyMasks(prefixedPropertyMasks)
-                        .build();
         mReadWriteLock.readLock().lock();
         try {
-            getResultProto =
+            throwIfClosedLocked();
+
+            List<TypePropertyMask> nonPrefixedPropertyMasks =
+                    TypePropertyPathToProtoConverter.toTypePropertyMaskList(typePropertyPaths);
+            List<TypePropertyMask> prefixedPropertyMasks =
+                    new ArrayList<>(nonPrefixedPropertyMasks.size());
+            for (int i = 0; i < nonPrefixedPropertyMasks.size(); ++i) {
+                TypePropertyMask typePropertyMask = nonPrefixedPropertyMasks.get(i);
+                String nonPrefixedType = typePropertyMask.getSchemaType();
+                String prefixedType =
+                        nonPrefixedType.equals(GetByUriRequest.PROJECTION_SCHEMA_TYPE_WILDCARD)
+                                ? nonPrefixedType
+                                : createPrefix(packageName, databaseName) + nonPrefixedType;
+                prefixedPropertyMasks.add(
+                        typePropertyMask.toBuilder().setSchemaType(prefixedType).build());
+            }
+            GetResultSpecProto getResultSpec =
+                    GetResultSpecProto.newBuilder()
+                            .addAllTypePropertyMasks(prefixedPropertyMasks)
+                            .build();
+
+            GetResultProto getResultProto =
                     mIcingSearchEngineLocked.get(
                             createPrefix(packageName, databaseName) + namespace,
                             uri,
                             getResultSpec);
+            checkSuccess(getResultProto.getStatus());
+
+            DocumentProto.Builder documentBuilder = getResultProto.getDocument().toBuilder();
+            removePrefixesFromDocument(documentBuilder);
+            return GenericDocumentToProtoConverter.toGenericDocument(documentBuilder.build());
         } finally {
             mReadWriteLock.readLock().unlock();
         }
-        checkSuccess(getResultProto.getStatus());
-
-        DocumentProto.Builder documentBuilder = getResultProto.getDocument().toBuilder();
-        removePrefixesFromDocument(documentBuilder);
-        return GenericDocumentToProtoConverter.toGenericDocument(documentBuilder.build());
     }
 
     /**
@@ -507,17 +549,19 @@
             @NonNull String queryExpression,
             @NonNull SearchSpec searchSpec)
             throws AppSearchException {
-        List<String> filterPackageNames = searchSpec.getFilterPackageNames();
-        if (!filterPackageNames.isEmpty() && !filterPackageNames.contains(packageName)) {
-            // Client wanted to query over some packages that weren't its own. This isn't
-            // allowed through local query so we can return early with no results.
-            return new SearchResultPage(Bundle.EMPTY);
-        }
-
         mReadWriteLock.readLock().lock();
         try {
+            throwIfClosedLocked();
+
+            List<String> filterPackageNames = searchSpec.getFilterPackageNames();
+            if (!filterPackageNames.isEmpty() && !filterPackageNames.contains(packageName)) {
+                // Client wanted to query over some packages that weren't its own. This isn't
+                // allowed through local query so we can return early with no results.
+                return new SearchResultPage(Bundle.EMPTY);
+            }
+
             String prefix = createPrefix(packageName, databaseName);
-            Set<String> allowedPrefixedSchemas = getAllowedPrefixSchemas(prefix, searchSpec);
+            Set<String> allowedPrefixedSchemas = getAllowedPrefixSchemasLocked(prefix, searchSpec);
 
             return doQueryLocked(
                     Collections.singleton(createPrefix(packageName, databaseName)),
@@ -552,6 +596,8 @@
             throws AppSearchException {
         mReadWriteLock.readLock().lock();
         try {
+            throwIfClosedLocked();
+
             Set<String> packageFilters = new ArraySet<>(searchSpec.getFilterPackageNames());
             Set<String> prefixFilters = new ArraySet<>();
             Set<String> allPrefixes = mNamespaceMapLocked.keySet();
@@ -654,6 +700,8 @@
     public SearchResultPage getNextPage(long nextPageToken) throws AppSearchException {
         mReadWriteLock.readLock().lock();
         try {
+            throwIfClosedLocked();
+
             SearchResultProto searchResultProto =
                     mIcingSearchEngineLocked.getNextPage(nextPageToken);
             checkSuccess(searchResultProto.getStatus());
@@ -674,6 +722,8 @@
     public void invalidateNextPageToken(long nextPageToken) {
         mReadWriteLock.readLock().lock();
         try {
+            throwIfClosedLocked();
+
             mIcingSearchEngineLocked.invalidateNextPageToken(nextPageToken);
         } finally {
             mReadWriteLock.readLock().unlock();
@@ -688,16 +738,19 @@
             @NonNull String uri,
             long usageTimestampMillis)
             throws AppSearchException {
-        String prefixedNamespace = createPrefix(packageName, databaseName) + namespace;
-        UsageReport report =
-                UsageReport.newBuilder()
-                        .setDocumentNamespace(prefixedNamespace)
-                        .setDocumentUri(uri)
-                        .setUsageTimestampMs(usageTimestampMillis)
-                        .setUsageType(UsageReport.UsageType.USAGE_TYPE1)
-                        .build();
         mReadWriteLock.writeLock().lock();
         try {
+            throwIfClosedLocked();
+
+            String prefixedNamespace = createPrefix(packageName, databaseName) + namespace;
+            UsageReport report =
+                    UsageReport.newBuilder()
+                            .setDocumentNamespace(prefixedNamespace)
+                            .setDocumentUri(uri)
+                            .setUsageTimestampMs(usageTimestampMillis)
+                            .setUsageType(UsageReport.UsageType.USAGE_TYPE1)
+                            .build();
+
             ReportUsageResultProto result = mIcingSearchEngineLocked.reportUsage(report);
             checkSuccess(result.getStatus());
         } finally {
@@ -722,16 +775,18 @@
             @NonNull String namespace,
             @NonNull String uri)
             throws AppSearchException {
-        String prefixedNamespace = createPrefix(packageName, databaseName) + namespace;
-        DeleteResultProto deleteResultProto;
         mReadWriteLock.writeLock().lock();
         try {
-            deleteResultProto = mIcingSearchEngineLocked.delete(prefixedNamespace, uri);
-            checkForOptimizeLocked(/* force= */ false);
+            throwIfClosedLocked();
+
+            String prefixedNamespace = createPrefix(packageName, databaseName) + namespace;
+            DeleteResultProto deleteResultProto =
+                    mIcingSearchEngineLocked.delete(prefixedNamespace, uri);
+
+            checkSuccess(deleteResultProto.getStatus());
         } finally {
             mReadWriteLock.writeLock().unlock();
         }
-        checkSuccess(deleteResultProto.getStatus());
     }
 
     /**
@@ -751,22 +806,25 @@
             @NonNull String queryExpression,
             @NonNull SearchSpec searchSpec)
             throws AppSearchException {
-        List<String> filterPackageNames = searchSpec.getFilterPackageNames();
-        if (!filterPackageNames.isEmpty() && !filterPackageNames.contains(packageName)) {
-            // We're only removing documents within the parameter `packageName`. If we're not
-            // restricting our remove-query to this package name, then there's nothing for us to
-            // remove.
-            return;
-        }
-
-        SearchSpecProto searchSpecProto = SearchSpecToProtoConverter.toSearchSpecProto(searchSpec);
-        SearchSpecProto.Builder searchSpecBuilder =
-                searchSpecProto.toBuilder().setQuery(queryExpression);
-        DeleteByQueryResultProto deleteResultProto;
         mReadWriteLock.writeLock().lock();
         try {
+            throwIfClosedLocked();
+
+            List<String> filterPackageNames = searchSpec.getFilterPackageNames();
+            if (!filterPackageNames.isEmpty() && !filterPackageNames.contains(packageName)) {
+                // We're only removing documents within the parameter `packageName`. If we're not
+                // restricting our remove-query to this package name, then there's nothing for us to
+                // remove.
+                return;
+            }
+
+            SearchSpecProto searchSpecProto =
+                    SearchSpecToProtoConverter.toSearchSpecProto(searchSpec);
+            SearchSpecProto.Builder searchSpecBuilder =
+                    searchSpecProto.toBuilder().setQuery(queryExpression);
+
             String prefix = createPrefix(packageName, databaseName);
-            Set<String> allowedPrefixedSchemas = getAllowedPrefixSchemas(prefix, searchSpec);
+            Set<String> allowedPrefixedSchemas = getAllowedPrefixSchemasLocked(prefix, searchSpec);
 
             // rewriteSearchSpecForPrefixesLocked will return false if there is nothing to search
             // over given their search filters, so we can return early and skip sending request
@@ -775,15 +833,16 @@
                     searchSpecBuilder, Collections.singleton(prefix), allowedPrefixedSchemas)) {
                 return;
             }
-            deleteResultProto = mIcingSearchEngineLocked.deleteByQuery(searchSpecBuilder.build());
-            checkForOptimizeLocked(/* force= */ true);
+            DeleteByQueryResultProto deleteResultProto =
+                    mIcingSearchEngineLocked.deleteByQuery(searchSpecBuilder.build());
+
+            // It seems that the caller wants to get success if the data matching the query is
+            // not in the DB because it was not there or was successfully deleted.
+            checkCodeOneOf(
+                    deleteResultProto.getStatus(), StatusProto.Code.OK, StatusProto.Code.NOT_FOUND);
         } finally {
             mReadWriteLock.writeLock().unlock();
         }
-        // It seems that the caller wants to get success if the data matching the query is not in
-        // the DB because it was not there or was successfully deleted.
-        checkCodeOneOf(
-                deleteResultProto.getStatus(), StatusProto.Code.OK, StatusProto.Code.NOT_FOUND);
     }
 
     /**
@@ -799,9 +858,16 @@
      * @throws AppSearchException on any error that AppSearch persist data to disk.
      */
     public void persistToDisk() throws AppSearchException {
-        PersistToDiskResultProto persistToDiskResultProto =
-                mIcingSearchEngineLocked.persistToDisk();
-        checkSuccess(persistToDiskResultProto.getStatus());
+        mReadWriteLock.writeLock().lock();
+        try {
+            throwIfClosedLocked();
+
+            PersistToDiskResultProto persistToDiskResultProto =
+                    mIcingSearchEngineLocked.persistToDisk();
+            checkSuccess(persistToDiskResultProto.getStatus());
+        } finally {
+            mReadWriteLock.writeLock().unlock();
+        }
     }
 
     /**
@@ -814,21 +880,16 @@
      *
      * @throws AppSearchException on IcingSearchEngine error.
      */
-    private void reset() throws AppSearchException {
-        ResetResultProto resetResultProto;
-        mReadWriteLock.writeLock().lock();
-        try {
-            resetResultProto = mIcingSearchEngineLocked.reset();
-            mOptimizeIntervalCountLocked = 0;
-            mSchemaMapLocked.clear();
-            mNamespaceMapLocked.clear();
+    @GuardedBy("mReadWriteLock")
+    private void resetLocked() throws AppSearchException {
+        ResetResultProto resetResultProto = mIcingSearchEngineLocked.reset();
+        mOptimizeIntervalCountLocked = 0;
+        mSchemaMapLocked.clear();
+        mNamespaceMapLocked.clear();
 
-            // Must be called after everything else since VisibilityStore may repopulate
-            // IcingSearchEngine with an initial schema.
-            mVisibilityStoreLocked.handleReset();
-        } finally {
-            mReadWriteLock.writeLock().unlock();
-        }
+        // Must be called after everything else since VisibilityStore may repopulate
+        // IcingSearchEngine with an initial schema.
+        mVisibilityStoreLocked.handleReset();
         checkSuccess(resetResultProto.getStatus());
     }
 
@@ -1079,7 +1140,8 @@
      * <p>This only checks intersection of schema filters on the search spec with those that the
      * prefix owns itself. This does not check global query permissions.
      */
-    private Set<String> getAllowedPrefixSchemas(
+    @GuardedBy("mReadWriteLock")
+    private Set<String> getAllowedPrefixSchemasLocked(
             @NonNull String prefix, @NonNull SearchSpec searchSpec) {
         Set<String> allowedPrefixedSchemas = new ArraySet<>();
 
@@ -1295,35 +1357,72 @@
     /**
      * Checks whether {@link IcingSearchEngine#optimize()} should be called to release resources.
      *
-     * <p>This method should be only called in mutate methods and get the WRITE lock to keep thread
-     * safety.
+     * <p>This method should be only called after a mutation to local storage backend which deletes
+     * a mass of data and could release lots resources after {@link IcingSearchEngine#optimize()}.
+     *
+     * <p>This method will trigger {@link IcingSearchEngine#getOptimizeInfo()} to check resources
+     * that could be released for every {@link #CHECK_OPTIMIZE_INTERVAL} mutations.
      *
      * <p>{@link IcingSearchEngine#optimize()} should be called only if {@link
      * GetOptimizeInfoResultProto} shows there is enough resources could be released.
      *
-     * <p>{@link IcingSearchEngine#getOptimizeInfo()} should be called once per {@link
-     * #CHECK_OPTIMIZE_INTERVAL} of remove executions.
-     *
-     * @param force whether we should directly call {@link IcingSearchEngine#getOptimizeInfo()}.
+     * @param mutationSize The number of how many mutations have been executed for current request.
+     *     An inside counter will accumulates it. Once the counter reaches {@link
+     *     #CHECK_OPTIMIZE_INTERVAL}, {@link IcingSearchEngine#getOptimizeInfo()} will be triggered
+     *     and the counter will be reset.
      */
-    @GuardedBy("mReadWriteLock")
-    private void checkForOptimizeLocked(boolean force) throws AppSearchException {
-        ++mOptimizeIntervalCountLocked;
-        if (force || mOptimizeIntervalCountLocked >= CHECK_OPTIMIZE_INTERVAL) {
-            mOptimizeIntervalCountLocked = 0;
+    public void checkForOptimize(int mutationSize) throws AppSearchException {
+        mReadWriteLock.writeLock().lock();
+        try {
+            mOptimizeIntervalCountLocked += mutationSize;
+            if (mOptimizeIntervalCountLocked >= CHECK_OPTIMIZE_INTERVAL) {
+                checkForOptimize();
+            }
+        } finally {
+            mReadWriteLock.writeLock().unlock();
+        }
+    }
+
+    /**
+     * Checks whether {@link IcingSearchEngine#optimize()} should be called to release resources.
+     *
+     * <p>This method will directly trigger {@link IcingSearchEngine#getOptimizeInfo()} to check
+     * resources that could be released.
+     *
+     * <p>{@link IcingSearchEngine#optimize()} should be called only if {@link
+     * GetOptimizeInfoResultProto} shows there is enough resources could be released.
+     */
+    public void checkForOptimize() throws AppSearchException {
+        mReadWriteLock.writeLock().lock();
+        try {
             GetOptimizeInfoResultProto optimizeInfo = getOptimizeInfoResultLocked();
             checkSuccess(optimizeInfo.getStatus());
+            mOptimizeIntervalCountLocked = 0;
             // Second threshold, decide when to call optimize().
             if (optimizeInfo.getOptimizableDocs() >= OPTIMIZE_THRESHOLD_DOC_COUNT
                     || optimizeInfo.getEstimatedOptimizableBytes() >= OPTIMIZE_THRESHOLD_BYTES) {
-                // TODO(b/155939114): call optimize in the same thread will slow down api calls
-                //  significantly. Move this call to background.
-                OptimizeResultProto optimizeResultProto = mIcingSearchEngineLocked.optimize();
-                checkSuccess(optimizeResultProto.getStatus());
+                optimize();
             }
-            // TODO(b/147699081): Return OptimizeResultProto & log lost data detail once we add
-            //  a field to indicate lost_schema and lost_documents in OptimizeResultProto.
-            //  go/icing-library-apis.
+        } finally {
+            mReadWriteLock.writeLock().unlock();
+        }
+        // TODO(b/147699081): Return OptimizeResultProto & log lost data detail once we add
+        //  a field to indicate lost_schema and lost_documents in OptimizeResultProto.
+        //  go/icing-library-apis.
+    }
+
+    /**
+     * Triggers {@link IcingSearchEngine#optimize()} directly.
+     *
+     * <p>This method should be only called as a scheduled task in AppSearch Platform backend.
+     */
+    public void optimize() throws AppSearchException {
+        mReadWriteLock.writeLock().lock();
+        try {
+            OptimizeResultProto optimizeResultProto = mIcingSearchEngineLocked.optimize();
+            checkSuccess(optimizeResultProto.getStatus());
+        } finally {
+            mReadWriteLock.writeLock().unlock();
         }
     }
 
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchMigrationHelperImpl.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchMigrationHelperImpl.java
index b3f8264..a501e99 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchMigrationHelperImpl.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchMigrationHelperImpl.java
@@ -76,7 +76,7 @@
         int currentVersion = mCurrentVersionMap.get(schemaType);
         int finalVersion = mFinalVersionMap.get(schemaType);
         try (FileOutputStream outputStream = new FileOutputStream(mFile)) {
-            // TODO(b/151178558) change the output stream so that we can use it in platform
+            // TODO(b/177266929) change the output stream so that we can use it in platform
             CodedOutputStream codedOutputStream = CodedOutputStream.newInstance(outputStream);
             SearchResultPage searchResultPage =
                     mAppSearchImpl.query(
@@ -126,11 +126,13 @@
                 try {
                     mAppSearchImpl.putDocument(mPackageName, mDatabaseName, document);
                 } catch (Throwable t) {
-                    responseBuilder.setFailure(
-                            document.getSchemaType(),
-                            document.getNamespace(),
-                            document.getUri(),
-                            throwableToFailedResult(t));
+                    responseBuilder.addMigrationFailure(
+                            new SetSchemaResponse.MigrationFailure.Builder()
+                                    .setNamespace(document.getNamespace())
+                                    .setSchemaType(document.getSchemaType())
+                                    .setUri(document.getUri())
+                                    .setAppSearchResult(throwableToFailedResult(t))
+                                    .build());
                 }
             }
             mAppSearchImpl.persistToDisk();
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/converter/SetSchemaResultToProtoConverter.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/converter/SetSchemaResponseToProtoConverter.java
similarity index 70%
rename from apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/converter/SetSchemaResultToProtoConverter.java
rename to apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/converter/SetSchemaResponseToProtoConverter.java
index e1e7d46..a0f39ec 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/converter/SetSchemaResultToProtoConverter.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/converter/SetSchemaResponseToProtoConverter.java
@@ -17,45 +17,41 @@
 package com.android.server.appsearch.external.localstorage.converter;
 
 import android.annotation.NonNull;
-import android.app.appsearch.SetSchemaResult;
+import android.app.appsearch.SetSchemaResponse;
 
 import com.android.internal.util.Preconditions;
 
 import com.google.android.icing.proto.SetSchemaResultProto;
 
 /**
- * Translates a {@link SetSchemaResultProto} into {@link SetSchemaResult}.
+ * Translates a {@link SetSchemaResultProto} into {@link SetSchemaResponse}.
  *
  * @hide
  */
-public class SetSchemaResultToProtoConverter {
+public class SetSchemaResponseToProtoConverter {
 
-    private SetSchemaResultToProtoConverter() {}
+    private SetSchemaResponseToProtoConverter() {}
 
     /**
-     * Translate a {@link SetSchemaResultProto} into {@link SetSchemaResult}.
+     * Translate a {@link SetSchemaResultProto} into {@link SetSchemaResponse}.
      *
      * @param proto The {@link SetSchemaResultProto} containing results.
      * @param prefix The prefix need to removed from schemaTypes
-     * @return {@link SetSchemaResult} of results.
+     * @return The {@link SetSchemaResponse} object.
      */
     @NonNull
-    public static SetSchemaResult toSetSchemaResult(
+    public static SetSchemaResponse toSetSchemaResponse(
             @NonNull SetSchemaResultProto proto, @NonNull String prefix) {
         Preconditions.checkNotNull(proto);
         Preconditions.checkNotNull(prefix);
-        SetSchemaResult.Builder builder =
-                new SetSchemaResult.Builder()
-                        .setResultCode(
-                                ResultCodeToProtoConverter.toResultCode(
-                                        proto.getStatus().getCode()));
+        SetSchemaResponse.Builder builder = new SetSchemaResponse.Builder();
 
         for (int i = 0; i < proto.getDeletedSchemaTypesCount(); i++) {
-            builder.addDeletedSchemaType(proto.getDeletedSchemaTypes(i).substring(prefix.length()));
+            builder.addDeletedType(proto.getDeletedSchemaTypes(i).substring(prefix.length()));
         }
 
         for (int i = 0; i < proto.getIncompatibleSchemaTypesCount(); i++) {
-            builder.addIncompatibleSchemaType(
+            builder.addIncompatibleType(
                     proto.getIncompatibleSchemaTypes(i).substring(prefix.length()));
         }
 
diff --git a/apex/appsearch/synced_jetpack_changeid.txt b/apex/appsearch/synced_jetpack_changeid.txt
index 6ba5572..d076db3 100644
--- a/apex/appsearch/synced_jetpack_changeid.txt
+++ b/apex/appsearch/synced_jetpack_changeid.txt
@@ -1 +1 @@
-I2bf8bd9db1b71b7da4ab50dd7480e4529678413a
+Ia9a8daef1a6d7d9432f7808d440abd64f4797701
diff --git a/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/AppSearchSessionShim.java b/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/AppSearchSessionShim.java
index b2ffd5b..ea21e19 100644
--- a/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/AppSearchSessionShim.java
+++ b/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/AppSearchSessionShim.java
@@ -77,10 +77,14 @@
      *       android.app.appsearch.exceptions.AppSearchException} with the {@link
      *       AppSearchResult#RESULT_INVALID_SCHEMA} error code, all documents which are not
      *       compatible with the new schema will be deleted and the incompatible schema will be
-     *       applied.
+     *       applied. Incompatible types and deleted types will be set into {@link
+     *       SetSchemaResponse#getIncompatibleTypes()} and {@link
+     *       SetSchemaResponse#getDeletedTypes()}, respectively.
      *   <li>Add a {@link android.app.appsearch.AppSearchSchema.Migrator} for each incompatible type
      *       and make no deletion. The migrator will migrate documents from it's old schema version
-     *       to the new version. See the migration section below.
+     *       to the new version. Migrated types will be set into both {@link
+     *       SetSchemaResponse#getIncompatibleTypes()} and {@link
+     *       SetSchemaResponse#getMigratedTypes()}. See the migration section below.
      * </ul>
      *
      * <p>It is a no-op to set the same schema as has been previously set; this is handled
@@ -109,10 +113,8 @@
      * backwards-compatible, and stored documents won't have any observable changes.
      *
      * @param request The schema update request.
-     * @return The pending {@link SetSchemaResponse} of performing this operation. Success if the
-     *     the schema has been set and any migrations has been done. Otherwise, the failure {@link
-     *     android.app.appsearch.SetSchemaResponse.MigrationFailure} indicates which document is
-     *     fail to be migrated.
+     * @return A {@link ListenableFuture} with exception if we hit any error. Or the pending {@link
+     *     SetSchemaResponse} of performing this operation, if the schema has been successfully set.
      * @see android.app.appsearch.AppSearchSchema.Migrator
      * @see android.app.appsearch.AppSearchMigrationHelper.Transformer
      */
@@ -144,57 +146,79 @@
     ListenableFuture<AppSearchBatchResult<String, Void>> put(@NonNull PutDocumentsRequest request);
 
     /**
-     * Retrieves {@link GenericDocument}s by URI.
+     * Gets {@link GenericDocument} objects by URIs and namespace from the {@link
+     * AppSearchSessionShim} database.
      *
-     * @param request {@link GetByUriRequest} containing URIs to be retrieved.
-     * @return The pending result of performing this operation. The keys of the returned {@link
-     *     AppSearchBatchResult} are the input URIs. The values are the returned {@link
-     *     GenericDocument}s on success, or a failed {@link AppSearchResult} otherwise. URIs that
-     *     are not found will return a failed {@link AppSearchResult} with a result code of {@link
-     *     AppSearchResult#RESULT_NOT_FOUND}.
+     * @param request a request containing URIs and namespace to get documents for.
+     * @return A {@link ListenableFuture} which resolves to an {@link AppSearchBatchResult}. The
+     *     keys of the {@link AppSearchBatchResult} represent the input URIs from the {@link
+     *     GetByUriRequest} object. The values are either the corresponding {@link GenericDocument}
+     *     object for the URI on success, or an {@link AppSearchResult} object on failure. For
+     *     example, if a URI is not found, the value for that URI will be set to an {@link
+     *     AppSearchResult} object with result code: {@link AppSearchResult#RESULT_NOT_FOUND}.
      */
     @NonNull
     ListenableFuture<AppSearchBatchResult<String, GenericDocument>> getByUri(
             @NonNull GetByUriRequest request);
 
     /**
-     * Searches for documents based on a given query string.
+     * Retrieves documents from the open {@link AppSearchSessionShim} that match a given query
+     * string and type of search provided.
      *
-     * <p>Currently we support following features in the raw query format:
+     * <p>Query strings can be empty, contain one term with no operators, or contain multiple terms
+     * and operators.
+     *
+     * <p>For query strings that are empty, all documents that match the {@link SearchSpec} will be
+     * returned.
+     *
+     * <p>For query strings with a single term and no operators, documents that match the provided
+     * query string and {@link SearchSpec} will be returned.
+     *
+     * <p>The following operators are supported:
      *
      * <ul>
-     *   <li>AND
-     *       <p>AND joins (e.g. “match documents that have both the terms ‘dog’ and ‘cat’”).
-     *       Example: hello world matches documents that have both ‘hello’ and ‘world’
+     *   <li>AND (implicit)
+     *       <p>AND is an operator that matches documents that contain <i>all</i> provided terms.
+     *       <p><b>NOTE:</b> A space between terms is treated as an "AND" operator. Explicitly
+     *       including "AND" in a query string will treat "AND" as a term, returning documents that
+     *       also contain "AND".
+     *       <p>Example: "apple AND banana" matches documents that contain the terms "apple", "and",
+     *       "banana".
+     *       <p>Example: "apple banana" matches documents that contain both "apple" and "banana".
+     *       <p>Example: "apple banana cherry" matches documents that contain "apple", "banana", and
+     *       "cherry".
      *   <li>OR
-     *       <p>OR joins (e.g. “match documents that have either the term ‘dog’ or ‘cat’”). Example:
-     *       dog OR puppy
-     *   <li>Exclusion
-     *       <p>Exclude a term (e.g. “match documents that do not have the term ‘dog’”). Example:
-     *       -dog excludes the term ‘dog’
-     *   <li>Grouping terms
-     *       <p>Allow for conceptual grouping of subqueries to enable hierarchical structures (e.g.
-     *       “match documents that have either ‘dog’ or ‘puppy’, and either ‘cat’ or ‘kitten’”).
-     *       Example: (dog puppy) (cat kitten) two one group containing two terms.
-     *   <li>Property restricts
-     *       <p>Specifies which properties of a document to specifically match terms in (e.g. “match
-     *       documents where the ‘subject’ property contains ‘important’”). Example:
-     *       subject:important matches documents with the term ‘important’ in the ‘subject’ property
-     *   <li>Schema type restricts
-     *       <p>This is similar to property restricts, but allows for restricts on top-level
-     *       document fields, such as schema_type. Clients should be able to limit their query to
-     *       documents of a certain schema_type (e.g. “match documents that are of the ‘Email’
-     *       schema_type”). Example: { schema_type_filters: “Email”, “Video”,query: “dog” } will
-     *       match documents that contain the query term ‘dog’ and are of either the ‘Email’ schema
-     *       type or the ‘Video’ schema type.
+     *       <p>OR is an operator that matches documents that contain <i>any</i> provided term.
+     *       <p>Example: "apple OR banana" matches documents that contain either "apple" or
+     *       "banana".
+     *       <p>Example: "apple OR banana OR cherry" matches documents that contain any of "apple",
+     *       "banana", or "cherry".
+     *   <li>Exclusion (-)
+     *       <p>Exclusion (-) is an operator that matches documents that <i>do not</i> contain the
+     *       provided term.
+     *       <p>Example: "-apple" matches documents that do not contain "apple".
+     *   <li>Grouped Terms
+     *       <p>For queries that require multiple operators and terms, terms can be grouped into
+     *       subqueries. Subqueries are contained within an open "(" and close ")" parenthesis.
+     *       <p>Example: "(donut OR bagel) (coffee OR tea)" matches documents that contain either
+     *       "donut" or "bagel" and either "coffee" or "tea".
+     *   <li>Property Restricts
+     *       <p>For queries that require a term to match a specific {@link AppSearchSchema} property
+     *       of a document, a ":" must be included between the property name and the term.
+     *       <p>Example: "subject:important" matches documents that contain the term "important" in
+     *       the "subject" property.
      * </ul>
      *
+     * <p>Additional search specifications, such as filtering by {@link AppSearchSchema} type or
+     * adding projection, can be set by calling the corresponding {@link SearchSpec.Builder} setter.
+     *
      * <p>This method is lightweight. The heavy work will be done in {@link
      * SearchResultsShim#getNextPage()}.
      *
-     * @param queryExpression Query String to search.
-     * @param searchSpec Spec for setting filters, raw query etc.
-     * @return The search result of performing this operation.
+     * @param queryExpression query string to search.
+     * @param searchSpec spec for setting document filters, adding projection, setting term match
+     *     type, etc.
+     * @return a {@link SearchResultsShim} object for retrieved matched documents.
      */
     @NonNull
     SearchResultsShim search(@NonNull String queryExpression, @NonNull SearchSpec searchSpec);
@@ -219,13 +243,22 @@
     ListenableFuture<Void> reportUsage(@NonNull ReportUsageRequest request);
 
     /**
-     * Removes {@link GenericDocument}s from the index by URI.
+     * Removes {@link GenericDocument} objects by URIs and namespace from the {@link
+     * AppSearchSessionShim} database.
      *
-     * @param request Request containing URIs to be removed.
-     * @return The pending result of performing this operation. The keys of the returned {@link
-     *     AppSearchBatchResult} are the input URIs. The values are {@code null} on success, or a
-     *     failed {@link AppSearchResult} otherwise. URIs that are not found will return a failed
-     *     {@link AppSearchResult} with a result code of {@link AppSearchResult#RESULT_NOT_FOUND}.
+     * <p>Removed documents will no longer be surfaced by {@link #search} or {@link #getByUri}
+     * calls.
+     *
+     * <p><b>NOTE:</b>By default, documents are removed via a soft delete operation. Once the
+     * document crosses the count threshold or byte usage threshold, the documents will be removed
+     * from disk.
+     *
+     * @param request {@link RemoveByUriRequest} with URIs and namespace to remove from the index.
+     * @return a {@link ListenableFuture} which resolves to an {@link AppSearchBatchResult}. The
+     *     keys of the {@link AppSearchBatchResult} represent the input URIs from the {@link
+     *     RemoveByUriRequest} object. The values are either {@code null} on success, or a failed
+     *     {@link AppSearchResult} otherwise. URIs that are not found will return a failed {@link
+     *     AppSearchResult} with a result code of {@link AppSearchResult#RESULT_NOT_FOUND}.
      */
     @NonNull
     ListenableFuture<AppSearchBatchResult<String, Void>> remove(
diff --git a/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/AppSearchTestUtils.java b/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/AppSearchTestUtils.java
index 44d5180..37717d6 100644
--- a/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/AppSearchTestUtils.java
+++ b/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/AppSearchTestUtils.java
@@ -25,7 +25,6 @@
 import android.app.appsearch.GetByUriRequest;
 import android.app.appsearch.SearchResult;
 import android.app.appsearch.SearchResultsShim;
-import android.app.appsearch.SetSchemaResponse;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -43,15 +42,6 @@
         return result;
     }
 
-    // TODO(b/151178558) check setSchemaResponse.xxxtypes for the test need to verify.
-    public static void checkIsSetSchemaResponseSuccess(Future<SetSchemaResponse> future)
-            throws Exception {
-        SetSchemaResponse setSchemaResponse = future.get();
-        assertWithMessage("SetSchemaResponse not successful.")
-                .that(setSchemaResponse.isSuccess())
-                .isTrue();
-    }
-
     public static List<GenericDocument> doGet(
             AppSearchSessionShim session, String namespace, String... uris) throws Exception {
         AppSearchBatchResult<String, GenericDocument> result =
diff --git a/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/GlobalSearchSessionShim.java b/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/GlobalSearchSessionShim.java
index b96f99e..31c934f 100644
--- a/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/GlobalSearchSessionShim.java
+++ b/apex/appsearch/testing/java/com/android/server/appsearch/testing/external/GlobalSearchSessionShim.java
@@ -27,43 +27,26 @@
  */
 public interface GlobalSearchSessionShim extends Closeable {
     /**
-     * Searches across all documents in the storage based on a given query string.
+     * Retrieves documents from all AppSearch databases that the querying application has access to.
      *
-     * <p>Currently we support following features in the raw query format:
+     * <p>Applications can be granted access to documents by specifying {@link
+     * SetSchemaRequest.Builder#setSchemaTypeVisibilityForPackage}, or {@link
+     * SetSchemaRequest.Builder#setDocumentClassVisibilityForPackage} when building a schema.
      *
-     * <ul>
-     *   <li>AND
-     *       <p>AND joins (e.g. “match documents that have both the terms ‘dog’ and ‘cat’”).
-     *       Example: hello world matches documents that have both ‘hello’ and ‘world’
-     *   <li>OR
-     *       <p>OR joins (e.g. “match documents that have either the term ‘dog’ or ‘cat’”). Example:
-     *       dog OR puppy
-     *   <li>Exclusion
-     *       <p>Exclude a term (e.g. “match documents that do not have the term ‘dog’”). Example:
-     *       -dog excludes the term ‘dog’
-     *   <li>Grouping terms
-     *       <p>Allow for conceptual grouping of subqueries to enable hierarchical structures (e.g.
-     *       “match documents that have either ‘dog’ or ‘puppy’, and either ‘cat’ or ‘kitten’”).
-     *       Example: (dog puppy) (cat kitten) two one group containing two terms.
-     *   <li>Property restricts
-     *       <p>Specifies which properties of a document to specifically match terms in (e.g. “match
-     *       documents where the ‘subject’ property contains ‘important’”). Example:
-     *       subject:important matches documents with the term ‘important’ in the ‘subject’ property
-     *   <li>Schema type restricts
-     *       <p>This is similar to property restricts, but allows for restricts on top-level
-     *       document fields, such as schema_type. Clients should be able to limit their query to
-     *       documents of a certain schema_type (e.g. “match documents that are of the ‘Email’
-     *       schema_type”). Example: { schema_type_filters: “Email”, “Video”,query: “dog” } will
-     *       match documents that contain the query term ‘dog’ and are of either the ‘Email’ schema
-     *       type or the ‘Video’ schema type.
-     * </ul>
+     * <p>Document access can also be granted to system UIs by specifying {@link
+     * SetSchemaRequest.Builder#setSchemaTypeVisibilityForSystemUi}, or {@link
+     * SetSchemaRequest.Builder#setDocumentClassVisibilityForSystemUi} when building a schema.
+     *
+     * <p>See {@link AppSearchSessionShim#search(String, SearchSpec)} for a detailed explanation on
+     * forming a query string.
      *
      * <p>This method is lightweight. The heavy work will be done in {@link
      * SearchResultsShim#getNextPage()}.
      *
-     * @param queryExpression Query String to search.
-     * @param searchSpec Spec for setting filters, raw query etc.
-     * @return The search result of performing this operation.
+     * @param queryExpression query string to search.
+     * @param searchSpec spec for setting document filters, adding projection, setting term match
+     *     type, etc.
+     * @return a {@link SearchResultsShim} object for retrieved matched documents.
      */
     @NonNull
     SearchResultsShim search(@NonNull String queryExpression, @NonNull SearchSpec searchSpec);
diff --git a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java
index 9b6c723..73b0105 100644
--- a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java
+++ b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java
@@ -16,8 +16,6 @@
 
 package com.android.server.appsearch.external.localstorage;
 
-import static android.app.appsearch.AppSearchResult.RESULT_OK;
-
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.testng.Assert.expectThrows;
@@ -27,7 +25,7 @@
 import android.app.appsearch.SearchResult;
 import android.app.appsearch.SearchResultPage;
 import android.app.appsearch.SearchSpec;
-import android.app.appsearch.SetSchemaResult;
+import android.app.appsearch.SetSchemaResponse;
 import android.app.appsearch.exceptions.AppSearchException;
 import android.content.Context;
 import android.util.ArraySet;
@@ -66,14 +64,14 @@
     @Before
     public void setUp() throws Exception {
         Context context = ApplicationProvider.getApplicationContext();
+
         // Give ourselves global query permissions
         mAppSearchImpl =
                 AppSearchImpl.create(
                         mTemporaryFolder.newFolder(),
                         context,
                         VisibilityStore.NO_OP_USER_ID,
-                        /*globalQuerierPackage
-                        =*/ context.getPackageName());
+                        /*globalQuerierPackage=*/ context.getPackageName());
     }
 
     // TODO(b/175430168) add test to verify reset is working properly.
@@ -425,18 +423,24 @@
         GetOptimizeInfoResultProto optimizeInfo = mAppSearchImpl.getOptimizeInfoResultLocked();
         assertThat(optimizeInfo.getOptimizableDocs()).isEqualTo(0);
 
-        // delete 999 documents , we will reach the threshold to trigger optimize() in next
+        // delete 999 documents, we will reach the threshold to trigger optimize() in next
         // deletion.
         for (int i = 0; i < AppSearchImpl.OPTIMIZE_THRESHOLD_DOC_COUNT - 1; i++) {
             mAppSearchImpl.remove("package", "database", "namespace", "uri" + i);
         }
 
-        // optimize() still not be triggered since we are in the interval to call getOptimizeInfo()
+        // Updates the check for optimize counter, checkForOptimize() will be triggered since
+        // CHECK_OPTIMIZE_INTERVAL is reached but optimize() won't since
+        // OPTIMIZE_THRESHOLD_DOC_COUNT is not.
+        mAppSearchImpl.checkForOptimize(
+                /*mutateBatchSize=*/ AppSearchImpl.OPTIMIZE_THRESHOLD_DOC_COUNT - 1);
+
+        // Verify optimize() still not be triggered.
         optimizeInfo = mAppSearchImpl.getOptimizeInfoResultLocked();
         assertThat(optimizeInfo.getOptimizableDocs())
                 .isEqualTo(AppSearchImpl.OPTIMIZE_THRESHOLD_DOC_COUNT - 1);
 
-        // Keep delete docs, will reach the interval this time and trigger optimize().
+        // Keep delete docs
         for (int i = AppSearchImpl.OPTIMIZE_THRESHOLD_DOC_COUNT;
                 i
                         < AppSearchImpl.OPTIMIZE_THRESHOLD_DOC_COUNT
@@ -444,6 +448,9 @@
                 i++) {
             mAppSearchImpl.remove("package", "database", "namespace", "uri" + i);
         }
+        // updates the check for optimize counter, will reach both CHECK_OPTIMIZE_INTERVAL and
+        // OPTIMIZE_THRESHOLD_DOC_COUNT this time and trigger a optimize().
+        mAppSearchImpl.checkForOptimize(/*mutateBatchSize*/ AppSearchImpl.CHECK_OPTIMIZE_INTERVAL);
 
         // Verify optimize() is triggered
         optimizeInfo = mAppSearchImpl.getOptimizeInfoResultLocked();
@@ -783,7 +790,7 @@
                 Collections.singletonList(new AppSearchSchema.Builder("Email").build());
 
         // set email incompatible and delete text
-        SetSchemaResult setSchemaResult =
+        SetSchemaResponse setSchemaResponse =
                 mAppSearchImpl.setSchema(
                         "package",
                         "database1",
@@ -791,9 +798,8 @@
                         /*schemasNotPlatformSurfaceable=*/ Collections.emptyList(),
                         /*schemasPackageAccessible=*/ Collections.emptyMap(),
                         /*forceOverride=*/ true);
-        assertThat(setSchemaResult.getDeletedSchemaTypes()).containsExactly("Text");
-        assertThat(setSchemaResult.getIncompatibleSchemaTypes()).containsExactly("Email");
-        assertThat(setSchemaResult.getResultCode()).isEqualTo(RESULT_OK);
+        assertThat(setSchemaResponse.getDeletedTypes()).containsExactly("Text");
+        assertThat(setSchemaResponse.getIncompatibleTypes()).containsExactly("Email");
     }
 
     @Test
@@ -836,7 +842,7 @@
 
         final List<AppSearchSchema> finalSchemas =
                 Collections.singletonList(new AppSearchSchema.Builder("Email").build());
-        SetSchemaResult setSchemaResult =
+        SetSchemaResponse setSchemaResponse =
                 mAppSearchImpl.setSchema(
                         "package",
                         "database1",
@@ -846,7 +852,7 @@
                         /*forceOverride=*/ false);
 
         // Check the Document type has been deleted.
-        assertThat(setSchemaResult.getDeletedSchemaTypes()).containsExactly("Document");
+        assertThat(setSchemaResponse.getDeletedTypes()).containsExactly("Document");
 
         // ForceOverride to delete.
         mAppSearchImpl.setSchema(
@@ -1046,4 +1052,136 @@
                                     strippedDocumentProto.build()));
         }
     }
+
+    @Test
+    public void testThrowsExceptionIfClosed() throws Exception {
+        Context context = ApplicationProvider.getApplicationContext();
+        AppSearchImpl appSearchImpl =
+                AppSearchImpl.create(
+                        mTemporaryFolder.newFolder(),
+                        context,
+                        VisibilityStore.NO_OP_USER_ID,
+                        /*globalQuerierPackage=*/ "");
+
+        // Initial check that we could do something at first.
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(new AppSearchSchema.Builder("type").build());
+        appSearchImpl.setSchema(
+                "package",
+                "database",
+                schemas,
+                /*schemasNotPlatformSurfaceable=*/ Collections.emptyList(),
+                /*schemasPackageAccessible=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false);
+
+        appSearchImpl.close();
+
+        // Check all our public APIs
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.setSchema(
+                            "package",
+                            "database",
+                            schemas,
+                            /*schemasNotPlatformSurfaceable=*/ Collections.emptyList(),
+                            /*schemasPackageAccessible=*/ Collections.emptyMap(),
+                            /*forceOverride=*/ false);
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.getSchema("package", "database");
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.putDocument(
+                            "package",
+                            "database",
+                            new GenericDocument.Builder<>("uri", "type")
+                                    .setNamespace("namespace")
+                                    .build());
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.getDocument(
+                            "package", "database", "namespace", "uri", Collections.emptyMap());
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.query(
+                            "package",
+                            "database",
+                            "query",
+                            new SearchSpec.Builder()
+                                    .setTermMatch(TermMatchType.Code.PREFIX_VALUE)
+                                    .build());
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.globalQuery(
+                            "query",
+                            new SearchSpec.Builder()
+                                    .setTermMatch(TermMatchType.Code.PREFIX_VALUE)
+                                    .build(),
+                            "package",
+                            /*callerUid=*/ 1);
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.getNextPage(/*nextPageToken=*/ 1L);
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.invalidateNextPageToken(/*nextPageToken=*/ 1L);
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.reportUsage(
+                            "package",
+                            "database",
+                            "namespace",
+                            "uri",
+                            /*usageTimestampMillis=*/ 1000L);
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.remove("package", "database", "namespace", "uri");
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.removeByQuery(
+                            "package",
+                            "database",
+                            "query",
+                            new SearchSpec.Builder()
+                                    .setTermMatch(TermMatchType.Code.PREFIX_VALUE)
+                                    .build());
+                });
+
+        expectThrows(
+                IllegalStateException.class,
+                () -> {
+                    appSearchImpl.persistToDisk();
+                });
+    }
 }