Add ability to enforce limits on docs per package and size of doc.

This is needed to prevent abuse of our service and to share icing
docids in framework, which are the resource we are most likely to
run out of.

Without this change, an app could use the platform backend to index
a very high number of documents and exhaust the docid limit, thereby
preventing any other app from using AppSearch.

Bug: 170371356
Test: New testcases added to AppSearchImplTest
Change-Id: I03ade35072bc69b84f8fcefed72b3c3bc2b8ee68

12 files changed, 1327 insertions, 61 deletions

diff --git a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchConfig.java b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchConfig.java
index d5271a6cb92e..689aa1fcd371 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchConfig.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchConfig.java
@@ -60,6 +60,11 @@ public final class AppSearchConfig implements AutoCloseable {
     @VisibleForTesting
     static final int DEFAULT_SAMPLING_INTERVAL = 10;
 
+    @VisibleForTesting
+    static final int DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES = 512 * 1024; // 512KiB
+    @VisibleForTesting
+    static final int DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_COUNT = 20_000;
+
     /*
      * Keys for ALL the flags stored in DeviceConfig.
      */
@@ -70,13 +75,19 @@ public final class AppSearchConfig implements AutoCloseable {
             "sampling_interval_for_batch_call_stats";
     public static final String KEY_SAMPLING_INTERVAL_FOR_PUT_DOCUMENT_STATS =
             "sampling_interval_for_put_document_stats";
+    public static final String KEY_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES =
+            "limit_config_max_document_size_bytes";
+    public static final String KEY_LIMIT_CONFIG_MAX_DOCUMENT_COUNT =
+            "limit_config_max_document_docunt";
 
     // Array contains all the corresponding keys for the cached values.
     private static final String[] KEYS_TO_ALL_CACHED_VALUES = {
             KEY_MIN_TIME_INTERVAL_BETWEEN_SAMPLES_MILLIS,
             KEY_SAMPLING_INTERVAL_DEFAULT,
             KEY_SAMPLING_INTERVAL_FOR_BATCH_CALL_STATS,
-            KEY_SAMPLING_INTERVAL_FOR_PUT_DOCUMENT_STATS
+            KEY_SAMPLING_INTERVAL_FOR_PUT_DOCUMENT_STATS,
+            KEY_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES,
+            KEY_LIMIT_CONFIG_MAX_DOCUMENT_COUNT,
     };
 
     // Lock needed for all the operations in this class.
@@ -222,6 +233,24 @@ public final class AppSearchConfig implements AutoCloseable {
         }
     }
 
+    /** Returns the maximum serialized size an indexed document can be, in bytes. */
+    public int getCachedLimitConfigMaxDocumentSizeBytes() {
+        synchronized (mLock) {
+            throwIfClosedLocked();
+            return mBundleLocked.getInt(KEY_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES,
+                    DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES);
+        }
+    }
+
+    /** Returns the maximum number of active docs allowed per package. */
+    public int getCachedLimitConfigMaxDocumentCount() {
+        synchronized (mLock) {
+            throwIfClosedLocked();
+            return mBundleLocked.getInt(KEY_LIMIT_CONFIG_MAX_DOCUMENT_COUNT,
+                    DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_COUNT);
+        }
+    }
+
     @GuardedBy("mLock")
     private void throwIfClosedLocked() {
         if (mIsClosedLocked) {
@@ -264,6 +293,20 @@ public final class AppSearchConfig implements AutoCloseable {
                     mBundleLocked.putInt(key, properties.getInt(key, DEFAULT_SAMPLING_INTERVAL));
                 }
                 break;
+            case KEY_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES:
+                synchronized (mLock) {
+                    mBundleLocked.putInt(
+                            key,
+                            properties.getInt(key, DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES));
+                }
+                break;
+            case KEY_LIMIT_CONFIG_MAX_DOCUMENT_COUNT:
+                synchronized (mLock) {
+                    mBundleLocked.putInt(
+                            key,
+                            properties.getInt(key, DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_COUNT));
+                }
+                break;
             default:
                 break;
         }
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchUserInstanceManager.java b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchUserInstanceManager.java
index e067d4bcdf72..d0d2e8964cf0 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchUserInstanceManager.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchUserInstanceManager.java
@@ -173,8 +173,11 @@ public final class AppSearchUserInstanceManager {
         File appSearchDir = getAppSearchDir(userHandle);
         File icingDir = new File(appSearchDir, "icing");
         Log.i(TAG, "Creating new AppSearch instance at: " + icingDir);
-        AppSearchImpl appSearchImpl =
-                AppSearchImpl.create(icingDir, initStatsBuilder, new FrameworkOptimizeStrategy());
+        AppSearchImpl appSearchImpl = AppSearchImpl.create(
+                icingDir,
+                new FrameworkLimitConfig(config),
+                initStatsBuilder,
+                new FrameworkOptimizeStrategy());
 
         long prepareVisibilityStoreLatencyStartMillis = SystemClock.elapsedRealtime();
         VisibilityStoreImpl visibilityStore =
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/FrameworkLimitConfig.java b/apex/appsearch/service/java/com/android/server/appsearch/FrameworkLimitConfig.java
new file mode 100644
index 000000000000..d16168a915d5
--- /dev/null
+++ b/apex/appsearch/service/java/com/android/server/appsearch/FrameworkLimitConfig.java
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 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 com.android.server.appsearch;
+
+import android.annotation.NonNull;
+
+import com.android.server.appsearch.external.localstorage.LimitConfig;
+
+import java.util.Objects;
+
+class FrameworkLimitConfig implements LimitConfig {
+    private final AppSearchConfig mAppSearchConfig;
+
+    FrameworkLimitConfig(@NonNull AppSearchConfig appSearchConfig) {
+        mAppSearchConfig = Objects.requireNonNull(appSearchConfig);
+    }
+
+    @Override
+    public int getMaxDocumentSizeBytes() {
+        return mAppSearchConfig.getCachedLimitConfigMaxDocumentSizeBytes();
+    }
+
+    @Override
+    public int getMaxDocumentCount() {
+        return mAppSearchConfig.getCachedLimitConfigMaxDocumentCount();
+    }
+}
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 9dee179bd6f2..a1b93ce12975 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
@@ -88,6 +88,7 @@ import com.google.android.icing.proto.SearchResultProto;
 import com.google.android.icing.proto.SearchSpecProto;
 import com.google.android.icing.proto.SetSchemaResultProto;
 import com.google.android.icing.proto.StatusProto;
+import com.google.android.icing.proto.StorageInfoProto;
 import com.google.android.icing.proto.StorageInfoResultProto;
 import com.google.android.icing.proto.TypePropertyMask;
 import com.google.android.icing.proto.UsageReport;
@@ -147,10 +148,9 @@ public final class AppSearchImpl implements Closeable {
     @VisibleForTesting static final int CHECK_OPTIMIZE_INTERVAL = 100;
 
     private final ReadWriteLock mReadWriteLock = new ReentrantReadWriteLock();
-
     private final LogUtil mLogUtil = new LogUtil(TAG);
-
     private final OptimizeStrategy mOptimizeStrategy;
+    private final LimitConfig mLimitConfig;
 
     @GuardedBy("mReadWriteLock")
     @VisibleForTesting
@@ -169,6 +169,10 @@ public final class AppSearchImpl implements Closeable {
     @GuardedBy("mReadWriteLock")
     private final Map<String, Set<String>> mNamespaceMapLocked = new HashMap<>();
 
+    /** Maps package name to active document count. */
+    @GuardedBy("mReadWriteLock")
+    private final Map<String, Integer> mDocumentCountMapLocked = new ArrayMap<>();
+
     /**
      * The counter to check when to call {@link #checkForOptimize}. The interval is {@link
      * #CHECK_OPTIMIZE_INTERVAL}.
@@ -196,19 +200,22 @@ public final class AppSearchImpl implements Closeable {
     @NonNull
     public static AppSearchImpl create(
             @NonNull File icingDir,
+            @NonNull LimitConfig limitConfig,
             @Nullable InitializeStats.Builder initStatsBuilder,
             @NonNull OptimizeStrategy optimizeStrategy)
             throws AppSearchException {
-        return new AppSearchImpl(icingDir, initStatsBuilder, optimizeStrategy);
+        return new AppSearchImpl(icingDir, limitConfig, initStatsBuilder, optimizeStrategy);
     }
 
     /** @param initStatsBuilder collects stats for initialization if provided. */
     private AppSearchImpl(
             @NonNull File icingDir,
+            @NonNull LimitConfig limitConfig,
             @Nullable InitializeStats.Builder initStatsBuilder,
             @NonNull OptimizeStrategy optimizeStrategy)
             throws AppSearchException {
         Objects.requireNonNull(icingDir);
+        mLimitConfig = Objects.requireNonNull(limitConfig);
         mOptimizeStrategy = Objects.requireNonNull(optimizeStrategy);
 
         mReadWriteLock.writeLock().lock();
@@ -244,9 +251,9 @@ public final class AppSearchImpl implements Closeable {
                     AppSearchLoggerHelper.copyNativeStats(
                             initializeResultProto.getInitializeStats(), initStatsBuilder);
                 }
-
                 checkSuccess(initializeResultProto.getStatus());
 
+                // Read all protos we need to construct AppSearchImpl's cache maps
                 long prepareSchemaAndNamespacesLatencyStartMillis = SystemClock.elapsedRealtime();
                 SchemaProto schemaProto = getSchemaProtoLocked();
 
@@ -258,6 +265,9 @@ public final class AppSearchImpl implements Closeable {
                         getAllNamespacesResultProto.getNamespacesCount(),
                         getAllNamespacesResultProto);
 
+                StorageInfoProto storageInfoProto = getRawStorageInfoProto();
+
+                // Log the time it took to read the data that goes into the cache maps
                 if (initStatsBuilder != null) {
                     initStatsBuilder
                             .setStatusCode(
@@ -268,20 +278,27 @@ public final class AppSearchImpl implements Closeable {
                                             (SystemClock.elapsedRealtime()
                                                     - prepareSchemaAndNamespacesLatencyStartMillis));
                 }
-
                 checkSuccess(getAllNamespacesResultProto.getStatus());
 
                 // Populate schema map
-                for (SchemaTypeConfigProto schema : schemaProto.getTypesList()) {
+                List<SchemaTypeConfigProto> schemaProtoTypesList = schemaProto.getTypesList();
+                for (int i = 0; i < schemaProtoTypesList.size(); i++) {
+                    SchemaTypeConfigProto schema = schemaProtoTypesList.get(i);
                     String prefixedSchemaType = schema.getSchemaType();
                     addToMap(mSchemaMapLocked, getPrefix(prefixedSchemaType), schema);
                 }
 
                 // Populate namespace map
-                for (String prefixedNamespace : getAllNamespacesResultProto.getNamespacesList()) {
+                List<String> prefixedNamespaceList =
+                        getAllNamespacesResultProto.getNamespacesList();
+                for (int i = 0; i < prefixedNamespaceList.size(); i++) {
+                    String prefixedNamespace = prefixedNamespaceList.get(i);
                     addToMap(mNamespaceMapLocked, getPrefix(prefixedNamespace), prefixedNamespace);
                 }
 
+                // Populate document count map
+                rebuildDocumentCountMapLocked(storageInfoProto);
+
                 // logging prepare_schema_and_namespaces latency
                 if (initStatsBuilder != null) {
                     initStatsBuilder.setPrepareSchemaAndNamespacesLatencyMillis(
@@ -596,10 +613,19 @@ public final class AppSearchImpl implements Closeable {
             long rewriteDocumentTypeEndTimeMillis = SystemClock.elapsedRealtime();
             DocumentProto finalDocument = documentBuilder.build();
 
+            // Check limits
+            int newDocumentCount =
+                    enforceLimitConfigLocked(
+                            packageName, finalDocument.getUri(), finalDocument.getSerializedSize());
+
+            // Insert document
             mLogUtil.piiTrace("putDocument, request", finalDocument.getUri(), finalDocument);
-            PutResultProto putResultProto = mIcingSearchEngineLocked.put(documentBuilder.build());
+            PutResultProto putResultProto = mIcingSearchEngineLocked.put(finalDocument);
             mLogUtil.piiTrace("putDocument, response", putResultProto.getStatus(), putResultProto);
-            addToMap(mNamespaceMapLocked, prefix, documentBuilder.getNamespace());
+
+            // Update caches
+            addToMap(mNamespaceMapLocked, prefix, finalDocument.getNamespace());
+            mDocumentCountMapLocked.put(packageName, newDocumentCount);
 
             // Logging stats
             if (pStatsBuilder != null) {
@@ -631,6 +657,71 @@ public final class AppSearchImpl implements Closeable {
     }
 
     /**
+     * Checks that a new document can be added to the given packageName with the given serialized
+     * size without violating our {@link LimitConfig}.
+     *
+     * @return the new count of documents for the given package, including the new document.
+     * @throws AppSearchException with a code of {@link AppSearchResult#RESULT_OUT_OF_SPACE} if the
+     *     limits are violated by the new document.
+     */
+    @GuardedBy("mReadWriteLock")
+    private int enforceLimitConfigLocked(String packageName, String newDocUri, int newDocSize)
+            throws AppSearchException {
+        // Limits check: size of document
+        if (newDocSize > mLimitConfig.getMaxDocumentSizeBytes()) {
+            throw new AppSearchException(
+                    AppSearchResult.RESULT_OUT_OF_SPACE,
+                    "Document \""
+                            + newDocUri
+                            + "\" for package \""
+                            + packageName
+                            + "\" serialized to "
+                            + newDocSize
+                            + " bytes, which exceeds "
+                            + "limit of "
+                            + mLimitConfig.getMaxDocumentSizeBytes()
+                            + " bytes");
+        }
+
+        // Limits check: number of documents
+        Integer oldDocumentCount = mDocumentCountMapLocked.get(packageName);
+        int newDocumentCount;
+        if (oldDocumentCount == null) {
+            newDocumentCount = 1;
+        } else {
+            newDocumentCount = oldDocumentCount + 1;
+        }
+        if (newDocumentCount > mLimitConfig.getMaxDocumentCount()) {
+            // Our management of mDocumentCountMapLocked doesn't account for document
+            // replacements, so our counter might have overcounted if the app has replaced docs.
+            // Rebuild the counter from StorageInfo in case this is so.
+            // TODO(b/170371356):  If Icing lib exposes something in the result which says
+            //  whether the document was a replacement, we could subtract 1 again after the put
+            //  to keep the count accurate. That would allow us to remove this code.
+            rebuildDocumentCountMapLocked(getRawStorageInfoProto());
+            oldDocumentCount = mDocumentCountMapLocked.get(packageName);
+            if (oldDocumentCount == null) {
+                newDocumentCount = 1;
+            } else {
+                newDocumentCount = oldDocumentCount + 1;
+            }
+        }
+        if (newDocumentCount > mLimitConfig.getMaxDocumentCount()) {
+            // Now we really can't fit it in, even accounting for replacements.
+            throw new AppSearchException(
+                    AppSearchResult.RESULT_OUT_OF_SPACE,
+                    "Package \""
+                            + packageName
+                            + "\" exceeded limit of "
+                            + mLimitConfig.getMaxDocumentCount()
+                            + " documents. Some documents "
+                            + "must be removed to index additional ones.");
+        }
+
+        return newDocumentCount;
+    }
+
+    /**
      * Retrieves a document from the AppSearch index by namespace and document ID.
      *
      * <p>This method belongs to query group.
@@ -1121,6 +1212,9 @@ public final class AppSearchImpl implements Closeable {
                         deleteResultProto.getDeleteStats(), removeStatsBuilder);
             }
             checkSuccess(deleteResultProto.getStatus());
+
+            // Update derived maps
+            updateDocumentCountAfterRemovalLocked(packageName, /*numDocumentsDeleted=*/ 1);
         } finally {
             mReadWriteLock.writeLock().unlock();
             if (removeStatsBuilder != null) {
@@ -1196,6 +1290,11 @@ public final class AppSearchImpl implements Closeable {
             // not in the DB because it was not there or was successfully deleted.
             checkCodeOneOf(
                     deleteResultProto.getStatus(), StatusProto.Code.OK, StatusProto.Code.NOT_FOUND);
+
+            // Update derived maps
+            int numDocumentsDeleted =
+                    deleteResultProto.getDeleteStats().getNumDocumentsDeleted();
+            updateDocumentCountAfterRemovalLocked(packageName, numDocumentsDeleted);
         } finally {
             mReadWriteLock.writeLock().unlock();
             if (removeStatsBuilder != null) {
@@ -1205,6 +1304,22 @@ public final class AppSearchImpl implements Closeable {
         }
     }
 
+    @GuardedBy("mReadWriteLock")
+    private void updateDocumentCountAfterRemovalLocked(
+            @NonNull String packageName, int numDocumentsDeleted) {
+        if (numDocumentsDeleted > 0) {
+            Integer oldDocumentCount = mDocumentCountMapLocked.get(packageName);
+            // This should always be true: how can we delete documents for a package without
+            // having seen that package during init? This is just a safeguard.
+            if (oldDocumentCount != null) {
+                // This should always be >0; how can we remove more documents than we've indexed?
+                // This is just a safeguard.
+                int newDocumentCount = Math.max(oldDocumentCount - numDocumentsDeleted, 0);
+                mDocumentCountMapLocked.put(packageName, newDocumentCount);
+            }
+        }
+    }
+
     /** Estimates the storage usage info for a specific package. */
     @NonNull
     public StorageInfo getStorageInfoForPackage(@NonNull String packageName)
@@ -1233,7 +1348,7 @@ public final class AppSearchImpl implements Closeable {
                 return new StorageInfo.Builder().build();
             }
 
-            return getStorageInfoForNamespacesLocked(wantedPrefixedNamespaces);
+            return getStorageInfoForNamespaces(getRawStorageInfoProto(), wantedPrefixedNamespaces);
         } finally {
             mReadWriteLock.readLock().unlock();
         }
@@ -1264,29 +1379,45 @@ public final class AppSearchImpl implements Closeable {
                 return new StorageInfo.Builder().build();
             }
 
-            return getStorageInfoForNamespacesLocked(wantedPrefixedNamespaces);
+            return getStorageInfoForNamespaces(getRawStorageInfoProto(), wantedPrefixedNamespaces);
         } finally {
             mReadWriteLock.readLock().unlock();
         }
     }
 
-    @GuardedBy("mReadWriteLock")
+    /**
+     * Returns the native storage info capsuled in {@link StorageInfoResultProto} directly from
+     * IcingSearchEngine.
+     */
     @NonNull
-    private StorageInfo getStorageInfoForNamespacesLocked(@NonNull Set<String> prefixedNamespaces)
-            throws AppSearchException {
-        mLogUtil.piiTrace("getStorageInfo, request");
-        StorageInfoResultProto storageInfoResult = mIcingSearchEngineLocked.getStorageInfo();
-        mLogUtil.piiTrace(
-                "getStorageInfo, response", storageInfoResult.getStatus(), storageInfoResult);
-        checkSuccess(storageInfoResult.getStatus());
-        if (!storageInfoResult.hasStorageInfo()
-                || !storageInfoResult.getStorageInfo().hasDocumentStorageInfo()) {
+    public StorageInfoProto getRawStorageInfoProto() throws AppSearchException {
+        mReadWriteLock.readLock().lock();
+        try {
+            throwIfClosedLocked();
+            mLogUtil.piiTrace("getStorageInfo, request");
+            StorageInfoResultProto storageInfoResult = mIcingSearchEngineLocked.getStorageInfo();
+            mLogUtil.piiTrace(
+                    "getStorageInfo, response", storageInfoResult.getStatus(), storageInfoResult);
+            checkSuccess(storageInfoResult.getStatus());
+            return storageInfoResult.getStorageInfo();
+        } finally {
+            mReadWriteLock.readLock().unlock();
+        }
+    }
+
+    /**
+     * Extracts and returns {@link StorageInfo} from {@link StorageInfoProto} based on prefixed
+     * namespaces.
+     */
+    @NonNull
+    private static StorageInfo getStorageInfoForNamespaces(
+            @NonNull StorageInfoProto storageInfoProto, @NonNull Set<String> prefixedNamespaces) {
+        if (!storageInfoProto.hasDocumentStorageInfo()) {
             return new StorageInfo.Builder().build();
         }
-        long totalStorageSize = storageInfoResult.getStorageInfo().getTotalStorageSize();
 
-        DocumentStorageInfoProto documentStorageInfo =
-                storageInfoResult.getStorageInfo().getDocumentStorageInfo();
+        long totalStorageSize = storageInfoProto.getTotalStorageSize();
+        DocumentStorageInfoProto documentStorageInfo = storageInfoProto.getDocumentStorageInfo();
         int totalDocuments =
                 documentStorageInfo.getNumAliveDocuments()
                         + documentStorageInfo.getNumExpiredDocuments();
@@ -1436,6 +1567,7 @@ public final class AppSearchImpl implements Closeable {
                 String packageName = entry.getKey();
                 Set<String> databaseNames = entry.getValue();
                 if (!installedPackages.contains(packageName) && databaseNames != null) {
+                    mDocumentCountMapLocked.remove(packageName);
                     for (String databaseName : databaseNames) {
                         String removedPrefix = createPrefix(packageName, databaseName);
                         mSchemaMapLocked.remove(removedPrefix);
@@ -1468,6 +1600,7 @@ public final class AppSearchImpl implements Closeable {
         mOptimizeIntervalCountLocked = 0;
         mSchemaMapLocked.clear();
         mNamespaceMapLocked.clear();
+        mDocumentCountMapLocked.clear();
         if (initStatsBuilder != null) {
             initStatsBuilder
                     .setHasReset(true)
@@ -1477,6 +1610,26 @@ public final class AppSearchImpl implements Closeable {
         checkSuccess(resetResultProto.getStatus());
     }
 
+    @GuardedBy("mReadWriteLock")
+    private void rebuildDocumentCountMapLocked(@NonNull StorageInfoProto storageInfoProto) {
+        mDocumentCountMapLocked.clear();
+        List<NamespaceStorageInfoProto> namespaceStorageInfoProtoList =
+                storageInfoProto.getDocumentStorageInfo().getNamespaceStorageInfoList();
+        for (int i = 0; i < namespaceStorageInfoProtoList.size(); i++) {
+            NamespaceStorageInfoProto namespaceStorageInfoProto =
+                    namespaceStorageInfoProtoList.get(i);
+            String packageName = getPackageName(namespaceStorageInfoProto.getNamespace());
+            Integer oldCount = mDocumentCountMapLocked.get(packageName);
+            int newCount;
+            if (oldCount == null) {
+                newCount = namespaceStorageInfoProto.getNumAliveDocuments();
+            } else {
+                newCount = oldCount + namespaceStorageInfoProto.getNumAliveDocuments();
+            }
+            mDocumentCountMapLocked.put(packageName, newCount);
+        }
+    }
+
     /** Wrapper around schema changes */
     @VisibleForTesting
     static class RewrittenSchemaResults {
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/LimitConfig.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/LimitConfig.java
new file mode 100644
index 000000000000..3f5723ee53e0
--- /dev/null
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/LimitConfig.java
@@ -0,0 +1,57 @@
+/*
+ * 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 com.android.server.appsearch.external.localstorage;
+
+
+/**
+ * Defines limits placed on users of AppSearch and enforced by {@link AppSearchImpl}.
+ *
+ * @hide
+ */
+public interface LimitConfig {
+    /**
+     * The maximum number of bytes a single document is allowed to be.
+     *
+     * <p>Enforced at the time of serializing the document into a proto.
+     *
+     * <p>This limit has two purposes:
+     *
+     * <ol>
+     *   <li>Prevent the system service from using too much memory during indexing or querying by
+     *       capping the size of the data structures it needs to buffer
+     *   <li>Prevent apps from using a very large amount of data by storing exceptionally large
+     *       documents.
+     * </ol>
+     */
+    int getMaxDocumentSizeBytes();
+
+    /**
+     * The maximum number of documents a single app is allowed to index.
+     *
+     * <p>Enforced at indexing time.
+     *
+     * <p>This limit has two purposes:
+     *
+     * <ol>
+     *   <li>Protect icing lib's docid space from being overwhelmed by a single app. The overall
+     *       docid limit is currently 2^20 (~1 million)
+     *   <li>Prevent apps from using a very large amount of data on the system by storing too many
+     *       documents.
+     * </ol>
+     */
+    int getMaxDocumentCount();
+}
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/UnlimitedLimitConfig.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/UnlimitedLimitConfig.java
new file mode 100644
index 000000000000..0fabab04048b
--- /dev/null
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/UnlimitedLimitConfig.java
@@ -0,0 +1,37 @@
+/*
+ * 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 com.android.server.appsearch.external.localstorage;
+
+
+/**
+ * In Jetpack, AppSearch doesn't enforce artificial limits on number of documents or size of
+ * documents, since the app is the only user of the Icing instance. Icing still enforces a docid
+ * limit of 1M docs.
+ *
+ * @hide
+ */
+public class UnlimitedLimitConfig implements LimitConfig {
+    @Override
+    public int getMaxDocumentSizeBytes() {
+        return Integer.MAX_VALUE;
+    }
+
+    @Override
+    public int getMaxDocumentCount() {
+        return Integer.MAX_VALUE;
+    }
+}
diff --git a/services/tests/mockingservicestests/src/com/android/server/appsearch/AppSearchConfigTest.java b/services/tests/mockingservicestests/src/com/android/server/appsearch/AppSearchConfigTest.java
index c4860287c567..ffb1dd9c7e78 100644
--- a/services/tests/mockingservicestests/src/com/android/server/appsearch/AppSearchConfigTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/appsearch/AppSearchConfigTest.java
@@ -50,6 +50,10 @@ public class AppSearchConfigTest {
                 AppSearchConfig.DEFAULT_SAMPLING_INTERVAL);
         assertThat(appSearchConfig.getCachedSamplingIntervalForPutDocumentStats()).isEqualTo(
                 AppSearchConfig.DEFAULT_SAMPLING_INTERVAL);
+        assertThat(appSearchConfig.getCachedLimitConfigMaxDocumentSizeBytes()).isEqualTo(
+                AppSearchConfig.DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES);
+        assertThat(appSearchConfig.getCachedLimitConfigMaxDocumentCount()).isEqualTo(
+                AppSearchConfig.DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_COUNT);
     }
 
     @Test
@@ -265,6 +269,22 @@ public class AppSearchConfigTest {
     }
 
     @Test
+    public void testCustomizedValue() {
+        DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
+                AppSearchConfig.KEY_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES,
+                Integer.toString(2001),
+                /*makeDefault=*/ false);
+        DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
+                AppSearchConfig.KEY_LIMIT_CONFIG_MAX_DOCUMENT_COUNT,
+                Integer.toString(2002),
+                /*makeDefault=*/ false);
+
+        AppSearchConfig appSearchConfig = AppSearchConfig.create(DIRECT_EXECUTOR);
+        assertThat(appSearchConfig.getCachedLimitConfigMaxDocumentSizeBytes()).isEqualTo(2001);
+        assertThat(appSearchConfig.getCachedLimitConfigMaxDocumentCount()).isEqualTo(2002);
+    }
+
+    @Test
     public void testNotUsable_afterClose() {
         AppSearchConfig appSearchConfig = AppSearchConfig.create(DIRECT_EXECUTOR);
 
diff --git a/services/tests/mockingservicestests/src/com/android/server/appsearch/FrameworkLimitConfigTest.java b/services/tests/mockingservicestests/src/com/android/server/appsearch/FrameworkLimitConfigTest.java
new file mode 100644
index 000000000000..088ed277aa80
--- /dev/null
+++ b/services/tests/mockingservicestests/src/com/android/server/appsearch/FrameworkLimitConfigTest.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 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 com.android.server.appsearch;
+
+import static com.android.internal.util.ConcurrentUtils.DIRECT_EXECUTOR;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.provider.DeviceConfig;
+
+import com.android.server.testables.TestableDeviceConfig;
+
+import org.junit.Rule;
+import org.junit.Test;
+
+/**
+ * Tests for {@link FrameworkLimitConfig}.
+ *
+ * <p>Build/Install/Run: atest FrameworksMockingServicesTests:AppSearchConfigTest
+ */
+public class FrameworkLimitConfigTest {
+    @Rule
+    public final TestableDeviceConfig.TestableDeviceConfigRule
+            mDeviceConfigRule = new TestableDeviceConfig.TestableDeviceConfigRule();
+
+    @Test
+    public void testDefaultValues() {
+        AppSearchConfig appSearchConfig = AppSearchConfig.create(DIRECT_EXECUTOR);
+        FrameworkLimitConfig config = new FrameworkLimitConfig(appSearchConfig);
+        assertThat(config.getMaxDocumentSizeBytes()).isEqualTo(
+                AppSearchConfig.DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES);
+        assertThat(appSearchConfig.getCachedLimitConfigMaxDocumentCount()).isEqualTo(
+                AppSearchConfig.DEFAULT_LIMIT_CONFIG_MAX_DOCUMENT_COUNT);
+    }
+
+    @Test
+    public void testCustomizedValues() {
+        AppSearchConfig appSearchConfig = AppSearchConfig.create(DIRECT_EXECUTOR);
+        FrameworkLimitConfig config = new FrameworkLimitConfig(appSearchConfig);
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_APPSEARCH,
+                AppSearchConfig.KEY_LIMIT_CONFIG_MAX_DOCUMENT_SIZE_BYTES,
+                "2001",
+                /*makeDefault=*/ false);
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_APPSEARCH,
+                AppSearchConfig.KEY_LIMIT_CONFIG_MAX_DOCUMENT_COUNT,
+                "2002",
+                /*makeDefault=*/ false);
+
+        assertThat(config.getMaxDocumentSizeBytes()).isEqualTo(2001);
+        assertThat(appSearchConfig.getCachedLimitConfigMaxDocumentCount()).isEqualTo(2002);
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java b/services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java
index 3c10789bc792..0d475c00569e 100644
--- a/services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java
+++ b/services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java
@@ -92,7 +92,10 @@ public class AppSearchImplPlatformTest {
 
         // Give ourselves global query permissions
         mAppSearchImpl = AppSearchImpl.create(
-                mTemporaryFolder.newFolder(), /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                mTemporaryFolder.newFolder(),
+                new UnlimitedLimitConfig(),
+                /*initStatsBuilder=*/ null,
+                ALWAYS_OPTIMIZE);
         mVisibilityStore = VisibilityStoreImpl.create(mAppSearchImpl, mContext);
         mGlobalQuerierUid =
                 mContext.getPackageManager().getPackageUid(mContext.getPackageName(), /*flags=*/ 0);
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 330b1a74d879..91f49224fde8 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
@@ -54,6 +54,7 @@ import com.android.server.appsearch.icing.proto.SchemaTypeConfigProto;
 import com.android.server.appsearch.icing.proto.SearchResultProto;
 import com.android.server.appsearch.icing.proto.SearchSpecProto;
 import com.android.server.appsearch.icing.proto.StatusProto;
+import com.android.server.appsearch.icing.proto.StorageInfoProto;
 import com.android.server.appsearch.icing.proto.StringIndexingConfig;
 import com.android.server.appsearch.icing.proto.TermMatchType;
 
@@ -85,7 +86,10 @@ public class AppSearchImplTest {
     public void setUp() throws Exception {
         mAppSearchImpl =
                 AppSearchImpl.create(
-                        mTemporaryFolder.newFolder(), /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                        mTemporaryFolder.newFolder(),
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
     }
 
     /**
@@ -468,7 +472,11 @@ public class AppSearchImplTest {
         Context context = ApplicationProvider.getApplicationContext();
         File appsearchDir = mTemporaryFolder.newFolder();
         AppSearchImpl appSearchImpl =
-                AppSearchImpl.create(appsearchDir, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
 
         // Insert schema
         List<AppSearchSchema> schemas =
@@ -529,7 +537,12 @@ public class AppSearchImplTest {
         // Initialize AppSearchImpl. This should cause a reset.
         InitializeStats.Builder initStatsBuilder = new InitializeStats.Builder();
         appSearchImpl.close();
-        appSearchImpl = AppSearchImpl.create(appsearchDir, initStatsBuilder, ALWAYS_OPTIMIZE);
+        appSearchImpl =
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        initStatsBuilder,
+                        ALWAYS_OPTIMIZE);
 
         // Check recovery state
         InitializeStats initStats = initStatsBuilder.build();
@@ -1688,7 +1701,10 @@ public class AppSearchImplTest {
     public void testThrowsExceptionIfClosed() throws Exception {
         AppSearchImpl appSearchImpl =
                 AppSearchImpl.create(
-                        mTemporaryFolder.newFolder(), /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                        mTemporaryFolder.newFolder(),
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
 
         // Initial check that we could do something at first.
         List<AppSearchSchema> schemas =
@@ -1816,7 +1832,11 @@ public class AppSearchImplTest {
         // Setup the index
         File appsearchDir = mTemporaryFolder.newFolder();
         AppSearchImpl appSearchImpl =
-                AppSearchImpl.create(appsearchDir, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
 
         List<AppSearchSchema> schemas =
                 Collections.singletonList(new AppSearchSchema.Builder("type").build());
@@ -1843,7 +1863,11 @@ public class AppSearchImplTest {
 
         // That document should be visible even from another instance.
         AppSearchImpl appSearchImpl2 =
-                AppSearchImpl.create(appsearchDir, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
         getResult =
                 appSearchImpl2.getDocument(
                         "package", "database", "namespace1", "id1", Collections.emptyMap());
@@ -1855,7 +1879,11 @@ public class AppSearchImplTest {
         // Setup the index
         File appsearchDir = mTemporaryFolder.newFolder();
         AppSearchImpl appSearchImpl =
-                AppSearchImpl.create(appsearchDir, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
 
         List<AppSearchSchema> schemas =
                 Collections.singletonList(new AppSearchSchema.Builder("type").build());
@@ -1906,7 +1934,11 @@ public class AppSearchImplTest {
 
         // Only the second document should be retrievable from another instance.
         AppSearchImpl appSearchImpl2 =
-                AppSearchImpl.create(appsearchDir, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
         assertThrows(
                 AppSearchException.class,
                 () ->
@@ -1927,7 +1959,11 @@ public class AppSearchImplTest {
         // Setup the index
         File appsearchDir = mTemporaryFolder.newFolder();
         AppSearchImpl appSearchImpl =
-                AppSearchImpl.create(appsearchDir, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
 
         List<AppSearchSchema> schemas =
                 Collections.singletonList(new AppSearchSchema.Builder("type").build());
@@ -1986,7 +2022,11 @@ public class AppSearchImplTest {
 
         // Only the second document should be retrievable from another instance.
         AppSearchImpl appSearchImpl2 =
-                AppSearchImpl.create(appsearchDir, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
         assertThrows(
                 AppSearchException.class,
                 () ->
@@ -2001,4 +2041,784 @@ public class AppSearchImplTest {
                         "package", "database", "namespace2", "id2", Collections.emptyMap());
         assertThat(getResult).isEqualTo(document2);
     }
+
+    @Test
+    public void testGetIcingSearchEngineStorageInfo() throws Exception {
+        // Setup the index
+        File appsearchDir = mTemporaryFolder.newFolder();
+        AppSearchImpl appSearchImpl =
+                AppSearchImpl.create(
+                        appsearchDir,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(new AppSearchSchema.Builder("type").build());
+        appSearchImpl.setSchema(
+                "package",
+                "database",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+
+        // Add two documents
+        GenericDocument document1 =
+                new GenericDocument.Builder<>("namespace1", "id1", "type").build();
+        appSearchImpl.putDocument("package", "database", document1, /*logger=*/ null);
+        GenericDocument document2 =
+                new GenericDocument.Builder<>("namespace1", "id2", "type").build();
+        appSearchImpl.putDocument("package", "database", document2, /*logger=*/ null);
+
+        StorageInfoProto storageInfo = appSearchImpl.getRawStorageInfoProto();
+
+        // Simple checks to verify if we can get correct StorageInfoProto from IcingSearchEngine
+        // No need to cover all the fields
+        assertThat(storageInfo.getTotalStorageSize()).isGreaterThan(0);
+        assertThat(storageInfo.getDocumentStorageInfo().getNumAliveDocuments()).isEqualTo(2);
+        assertThat(storageInfo.getSchemaStoreStorageInfo().getNumSchemaTypes()).isEqualTo(1);
+    }
+
+    @Test
+    public void testLimitConfig_DocumentSize() throws Exception {
+        // Create a new mAppSearchImpl with a lower limit
+        mAppSearchImpl.close();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        mTemporaryFolder.newFolder(),
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return 80;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 1;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Insert schema
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(new AppSearchSchema.Builder("type").build());
+        mAppSearchImpl.setSchema(
+                "package",
+                "database",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+
+        // Insert a document which is too large
+        GenericDocument document =
+                new GenericDocument.Builder<>(
+                                "this_namespace_is_long_to_make_the_doc_big", "id", "type")
+                        .build();
+        AppSearchException e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains(
+                        "Document \"id\" for package \"package\" serialized to 99 bytes, which"
+                            + " exceeds limit of 80 bytes");
+
+        // Make sure this failure didn't increase our document count. We should still be able to
+        // index 1 document.
+        GenericDocument document2 =
+                new GenericDocument.Builder<>("namespace", "id2", "type").build();
+        mAppSearchImpl.putDocument("package", "database", document2, /*logger=*/ null);
+
+        // Now we should get a failure
+        GenericDocument document3 =
+                new GenericDocument.Builder<>("namespace", "id3", "type").build();
+        e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document3, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 1 documents");
+    }
+
+    @Test
+    public void testLimitConfig_Init() throws Exception {
+        // Create a new mAppSearchImpl with a lower limit
+        mAppSearchImpl.close();
+        File tempFolder = mTemporaryFolder.newFolder();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        tempFolder,
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return 80;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 1;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Insert schema
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(new AppSearchSchema.Builder("type").build());
+        mAppSearchImpl.setSchema(
+                "package",
+                "database",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+
+        // Index a document
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id1", "type").build(),
+                /*logger=*/ null);
+
+        // Now we should get a failure
+        GenericDocument document2 =
+                new GenericDocument.Builder<>("namespace", "id2", "type").build();
+        AppSearchException e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document2, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 1 documents");
+
+        // Close and reinitialize AppSearchImpl
+        mAppSearchImpl.close();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        tempFolder,
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return 80;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 1;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Make sure the limit is maintained
+        e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document2, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 1 documents");
+    }
+
+    @Test
+    public void testLimitConfig_Remove() throws Exception {
+        // Create a new mAppSearchImpl with a lower limit
+        mAppSearchImpl.close();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        mTemporaryFolder.newFolder(),
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return Integer.MAX_VALUE;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 3;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Insert schema
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(new AppSearchSchema.Builder("type").build());
+        mAppSearchImpl.setSchema(
+                "package",
+                "database",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+
+        // Index 3 documents
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id1", "type").build(),
+                /*logger=*/ null);
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id2", "type").build(),
+                /*logger=*/ null);
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id3", "type").build(),
+                /*logger=*/ null);
+
+        // Now we should get a failure
+        GenericDocument document4 =
+                new GenericDocument.Builder<>("namespace", "id4", "type").build();
+        AppSearchException e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document4, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 3 documents");
+
+        // Remove a document that doesn't exist
+        assertThrows(
+                AppSearchException.class,
+                () ->
+                        mAppSearchImpl.remove(
+                                "package",
+                                "database",
+                                "namespace",
+                                "id4",
+                                /*removeStatsBuilder=*/ null));
+
+        // Should still fail
+        e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document4, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 3 documents");
+
+        // Remove a document that does exist
+        mAppSearchImpl.remove(
+                "package", "database", "namespace", "id2", /*removeStatsBuilder=*/ null);
+
+        // Now doc4 should work
+        mAppSearchImpl.putDocument("package", "database", document4, /*logger=*/ null);
+
+        // The next one should fail again
+        e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package",
+                                        "database",
+                                        new GenericDocument.Builder<>("namespace", "id5", "type")
+                                                .build(),
+                                        /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 3 documents");
+    }
+
+    @Test
+    public void testLimitConfig_DifferentPackages() throws Exception {
+        // Create a new mAppSearchImpl with a lower limit
+        mAppSearchImpl.close();
+        File tempFolder = mTemporaryFolder.newFolder();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        tempFolder,
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return Integer.MAX_VALUE;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 2;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Insert schema
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(new AppSearchSchema.Builder("type").build());
+        mAppSearchImpl.setSchema(
+                "package1",
+                "database1",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+        mAppSearchImpl.setSchema(
+                "package1",
+                "database2",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+        mAppSearchImpl.setSchema(
+                "package2",
+                "database1",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+        mAppSearchImpl.setSchema(
+                "package2",
+                "database2",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+
+        // Index documents in package1/database1
+        mAppSearchImpl.putDocument(
+                "package1",
+                "database1",
+                new GenericDocument.Builder<>("namespace", "id1", "type").build(),
+                /*logger=*/ null);
+        mAppSearchImpl.putDocument(
+                "package1",
+                "database2",
+                new GenericDocument.Builder<>("namespace", "id2", "type").build(),
+                /*logger=*/ null);
+
+        // Indexing a third doc into package1 should fail (here we use database3)
+        AppSearchException e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package1",
+                                        "database3",
+                                        new GenericDocument.Builder<>("namespace", "id3", "type")
+                                                .build(),
+                                        /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package1\" exceeded limit of 2 documents");
+
+        // Indexing a doc into package2 should succeed
+        mAppSearchImpl.putDocument(
+                "package2",
+                "database1",
+                new GenericDocument.Builder<>("namespace", "id1", "type").build(),
+                /*logger=*/ null);
+
+        // Reinitialize to make sure packages are parsed correctly on init
+        mAppSearchImpl.close();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        tempFolder,
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return Integer.MAX_VALUE;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 2;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // package1 should still be out of space
+        e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package1",
+                                        "database4",
+                                        new GenericDocument.Builder<>("namespace", "id4", "type")
+                                                .build(),
+                                        /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package1\" exceeded limit of 2 documents");
+
+        // package2 has room for one more
+        mAppSearchImpl.putDocument(
+                "package2",
+                "database2",
+                new GenericDocument.Builder<>("namespace", "id2", "type").build(),
+                /*logger=*/ null);
+
+        // now package2 really is out of space
+        e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package2",
+                                        "database3",
+                                        new GenericDocument.Builder<>("namespace", "id3", "type")
+                                                .build(),
+                                        /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package2\" exceeded limit of 2 documents");
+    }
+
+    @Test
+    public void testLimitConfig_RemoveByQyery() throws Exception {
+        // Create a new mAppSearchImpl with a lower limit
+        mAppSearchImpl.close();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        mTemporaryFolder.newFolder(),
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return Integer.MAX_VALUE;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 3;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Insert schema
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(
+                        new AppSearchSchema.Builder("type")
+                                .addProperty(
+                                        new AppSearchSchema.StringPropertyConfig.Builder("body")
+                                                .setIndexingType(
+                                                        AppSearchSchema.StringPropertyConfig
+                                                                .INDEXING_TYPE_PREFIXES)
+                                                .setTokenizerType(
+                                                        AppSearchSchema.StringPropertyConfig
+                                                                .TOKENIZER_TYPE_PLAIN)
+                                                .build())
+                                .build());
+        mAppSearchImpl.setSchema(
+                "package",
+                "database",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+
+        // Index 3 documents
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id1", "type")
+                        .setPropertyString("body", "tablet")
+                        .build(),
+                /*logger=*/ null);
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id2", "type")
+                        .setPropertyString("body", "tabby")
+                        .build(),
+                /*logger=*/ null);
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id3", "type")
+                        .setPropertyString("body", "grabby")
+                        .build(),
+                /*logger=*/ null);
+
+        // Now we should get a failure
+        GenericDocument document4 =
+                new GenericDocument.Builder<>("namespace", "id4", "type").build();
+        AppSearchException e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document4, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 3 documents");
+
+        // Run removebyquery, deleting nothing
+        mAppSearchImpl.removeByQuery(
+                "package",
+                "database",
+                "nothing",
+                new SearchSpec.Builder().build(),
+                /*removeStatsBuilder=*/ null);
+
+        // Should still fail
+        e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document4, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 3 documents");
+
+        // Remove "tab*"
+        mAppSearchImpl.removeByQuery(
+                "package",
+                "database",
+                "tab",
+                new SearchSpec.Builder().build(),
+                /*removeStatsBuilder=*/ null);
+
+        // Now doc4 and doc5 should work
+        mAppSearchImpl.putDocument("package", "database", document4, /*logger=*/ null);
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id5", "type").build(),
+                /*logger=*/ null);
+
+        // We only deleted 2 docs so the next one should fail again
+        e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package",
+                                        "database",
+                                        new GenericDocument.Builder<>("namespace", "id6", "type")
+                                                .build(),
+                                        /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 3 documents");
+    }
+
+    @Test
+    public void testLimitConfig_Replace() throws Exception {
+        // Create a new mAppSearchImpl with a lower limit
+        mAppSearchImpl.close();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        mTemporaryFolder.newFolder(),
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return Integer.MAX_VALUE;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 2;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Insert schema
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(
+                        new AppSearchSchema.Builder("type")
+                                .addProperty(
+                                        new AppSearchSchema.StringPropertyConfig.Builder("body")
+                                                .build())
+                                .build());
+        mAppSearchImpl.setSchema(
+                "package",
+                "database",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+
+        // Index a document
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id1", "type")
+                        .setPropertyString("body", "id1.orig")
+                        .build(),
+                /*logger=*/ null);
+        // Replace it with another doc
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id1", "type")
+                        .setPropertyString("body", "id1.new")
+                        .build(),
+                /*logger=*/ null);
+
+        // Index id2. This should pass but only because we check for replacements.
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id2", "type").build(),
+                /*logger=*/ null);
+
+        // Now we should get a failure on id3
+        GenericDocument document3 =
+                new GenericDocument.Builder<>("namespace", "id3", "type").build();
+        AppSearchException e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document3, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 2 documents");
+    }
+
+    @Test
+    public void testLimitConfig_ReplaceReinit() throws Exception {
+        // Create a new mAppSearchImpl with a lower limit
+        mAppSearchImpl.close();
+        File tempFolder = mTemporaryFolder.newFolder();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        tempFolder,
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return Integer.MAX_VALUE;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 2;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Insert schema
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(
+                        new AppSearchSchema.Builder("type")
+                                .addProperty(
+                                        new AppSearchSchema.StringPropertyConfig.Builder("body")
+                                                .build())
+                                .build());
+        mAppSearchImpl.setSchema(
+                "package",
+                "database",
+                schemas,
+                /*visibilityStore=*/ null,
+                /*schemasNotDisplayedBySystem=*/ Collections.emptyList(),
+                /*schemasVisibleToPackages=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false,
+                /*version=*/ 0);
+
+        // Index a document
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id1", "type")
+                        .setPropertyString("body", "id1.orig")
+                        .build(),
+                /*logger=*/ null);
+        // Replace it with another doc
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id1", "type")
+                        .setPropertyString("body", "id1.new")
+                        .build(),
+                /*logger=*/ null);
+
+        // Reinitialize to make sure replacements are correctly accounted for by init
+        mAppSearchImpl.close();
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        tempFolder,
+                        new LimitConfig() {
+                            @Override
+                            public int getMaxDocumentSizeBytes() {
+                                return Integer.MAX_VALUE;
+                            }
+
+                            @Override
+                            public int getMaxDocumentCount() {
+                                return 2;
+                            }
+                        },
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
+
+        // Index id2. This should pass but only because we check for replacements.
+        mAppSearchImpl.putDocument(
+                "package",
+                "database",
+                new GenericDocument.Builder<>("namespace", "id2", "type").build(),
+                /*logger=*/ null);
+
+        // Now we should get a failure on id3
+        GenericDocument document3 =
+                new GenericDocument.Builder<>("namespace", "id3", "type").build();
+        AppSearchException e =
+                assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        "package", "database", document3, /*logger=*/ null));
+        assertThat(e.getResultCode()).isEqualTo(AppSearchResult.RESULT_OUT_OF_SPACE);
+        assertThat(e)
+                .hasMessageThat()
+                .contains("Package \"package\" exceeded limit of 2 documents");
+    }
 }
diff --git a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchLoggerTest.java b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchLoggerTest.java
index 080c375ac7c6..7bacbb63f10c 100644
--- a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchLoggerTest.java
+++ b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchLoggerTest.java
@@ -67,7 +67,10 @@ public class AppSearchLoggerTest {
     public void setUp() throws Exception {
         mAppSearchImpl =
                 AppSearchImpl.create(
-                        mTemporaryFolder.newFolder(), /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                        mTemporaryFolder.newFolder(),
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
         mLogger = new TestLogger();
     }
 
@@ -290,7 +293,11 @@ public class AppSearchLoggerTest {
     public void testLoggingStats_initializeWithoutDocuments_success() throws Exception {
         // Create an unused AppSearchImpl to generated an InitializeStats.
         InitializeStats.Builder initStatsBuilder = new InitializeStats.Builder();
-        AppSearchImpl.create(mTemporaryFolder.newFolder(), initStatsBuilder, ALWAYS_OPTIMIZE);
+        AppSearchImpl.create(
+                mTemporaryFolder.newFolder(),
+                new UnlimitedLimitConfig(),
+                initStatsBuilder,
+                ALWAYS_OPTIMIZE);
         InitializeStats iStats = initStatsBuilder.build();
 
         assertThat(iStats).isNotNull();
@@ -314,7 +321,11 @@ public class AppSearchLoggerTest {
         final File folder = mTemporaryFolder.newFolder();
 
         AppSearchImpl appSearchImpl =
-                AppSearchImpl.create(folder, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        folder,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
         List<AppSearchSchema> schemas =
                 ImmutableList.of(
                         new AppSearchSchema.Builder("Type1").build(),
@@ -336,7 +347,7 @@ public class AppSearchLoggerTest {
 
         // Create another appsearchImpl on the same folder
         InitializeStats.Builder initStatsBuilder = new InitializeStats.Builder();
-        AppSearchImpl.create(folder, initStatsBuilder, ALWAYS_OPTIMIZE);
+        AppSearchImpl.create(folder, new UnlimitedLimitConfig(), initStatsBuilder, ALWAYS_OPTIMIZE);
         InitializeStats iStats = initStatsBuilder.build();
 
         assertThat(iStats).isNotNull();
@@ -360,7 +371,11 @@ public class AppSearchLoggerTest {
         final File folder = mTemporaryFolder.newFolder();
 
         AppSearchImpl appSearchImpl =
-                AppSearchImpl.create(folder, /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                AppSearchImpl.create(
+                        folder,
+                        new UnlimitedLimitConfig(),
+                        /*initStatsBuilder=*/ null,
+                        ALWAYS_OPTIMIZE);
 
         List<AppSearchSchema> schemas =
                 ImmutableList.of(
@@ -393,7 +408,7 @@ public class AppSearchLoggerTest {
 
         // Create another appsearchImpl on the same folder
         InitializeStats.Builder initStatsBuilder = new InitializeStats.Builder();
-        AppSearchImpl.create(folder, initStatsBuilder, ALWAYS_OPTIMIZE);
+        AppSearchImpl.create(folder, new UnlimitedLimitConfig(), initStatsBuilder, ALWAYS_OPTIMIZE);
         InitializeStats iStats = initStatsBuilder.build();
 
         // Some of other fields are already covered by AppSearchImplTest#testReset()
@@ -484,11 +499,13 @@ public class AppSearchLoggerTest {
                         .setPropertyString("nonExist", "testPut example1")
                         .build();
 
-        // We mainly want to check the status code in stats. So we don't need to inspect the
-        // exception here.
-        Assert.assertThrows(
-                AppSearchException.class,
-                () -> mAppSearchImpl.putDocument(testPackageName, testDatabase, document, mLogger));
+        AppSearchException exception =
+                Assert.assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.putDocument(
+                                        testPackageName, testDatabase, document, mLogger));
+        assertThat(exception.getResultCode()).isEqualTo(AppSearchResult.RESULT_NOT_FOUND);
 
         PutDocumentStats pStats = mLogger.mPutDocumentStats;
         assertThat(pStats).isNotNull();
@@ -676,17 +693,17 @@ public class AppSearchLoggerTest {
 
         RemoveStats.Builder rStatsBuilder = new RemoveStats.Builder(testPackageName, testDatabase);
 
-        // We mainly want to check the status code in stats. So we don't need to inspect the
-        // exception here.
-        Assert.assertThrows(
-                AppSearchException.class,
-                () ->
-                        mAppSearchImpl.remove(
-                                testPackageName,
-                                testDatabase,
-                                testNamespace,
-                                "invalidId",
-                                rStatsBuilder));
+        AppSearchException exception =
+                Assert.assertThrows(
+                        AppSearchException.class,
+                        () ->
+                                mAppSearchImpl.remove(
+                                        testPackageName,
+                                        testDatabase,
+                                        testNamespace,
+                                        "invalidId",
+                                        rStatsBuilder));
+        assertThat(exception.getResultCode()).isEqualTo(AppSearchResult.RESULT_NOT_FOUND);
 
         RemoveStats rStats = rStatsBuilder.build();
         assertThat(rStats.getPackageName()).isEqualTo(testPackageName);
diff --git a/services/tests/servicestests/src/com/android/server/appsearch/visibilitystore/VisibilityStoreImplTest.java b/services/tests/servicestests/src/com/android/server/appsearch/visibilitystore/VisibilityStoreImplTest.java
index 07a728bac2a5..374642b676d2 100644
--- a/services/tests/servicestests/src/com/android/server/appsearch/visibilitystore/VisibilityStoreImplTest.java
+++ b/services/tests/servicestests/src/com/android/server/appsearch/visibilitystore/VisibilityStoreImplTest.java
@@ -38,6 +38,7 @@ import androidx.test.core.app.ApplicationProvider;
 
 import com.android.server.appsearch.external.localstorage.AppSearchImpl;
 import com.android.server.appsearch.external.localstorage.OptimizeStrategy;
+import com.android.server.appsearch.external.localstorage.UnlimitedLimitConfig;
 import com.android.server.appsearch.external.localstorage.util.PrefixUtil;
 import com.android.server.appsearch.external.localstorage.visibilitystore.VisibilityStore;
 
@@ -88,7 +89,10 @@ public class VisibilityStoreImplTest {
 
         // Give ourselves global query permissions
         AppSearchImpl appSearchImpl = AppSearchImpl.create(
-                mTemporaryFolder.newFolder(), /*initStatsBuilder=*/ null, ALWAYS_OPTIMIZE);
+                mTemporaryFolder.newFolder(),
+                new UnlimitedLimitConfig(),
+                /*initStatsBuilder=*/ null,
+                ALWAYS_OPTIMIZE);
         mVisibilityStore = VisibilityStoreImpl.create(appSearchImpl, mContext);
         mUid = mContext.getPackageManager().getPackageUid(mContext.getPackageName(), /*flags=*/ 0);
     }