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

6 files changed, 361 insertions, 27 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;
+    }
+}