Merge "Add ability to enforce limits on docs per package and size of doc." into sc-dev

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;
         }

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 appSearchImpl = AppSearchImpl.create(
-                AppSearchImpl.create(icingDir, initStatsBuilder, new FrameworkOptimizeStrategy());
+                icingDir,
+                new FrameworkLimitConfig(config),
+                initStatsBuilder,
+                new FrameworkOptimizeStrategy());
 
         long prepareVisibilityStoreLatencyStartMillis = SystemClock.elapsedRealtime();
         VisibilityStoreImpl visibilityStore =

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();
+    }
+}

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) {
@@ -630,6 +656,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.
      *
@@ -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)
+    public StorageInfoProto getRawStorageInfoProto() throws AppSearchException {
-            throws AppSearchException {
+        mReadWriteLock.readLock().lock();
-        mLogUtil.piiTrace("getStorageInfo, request");
+        try {
-        StorageInfoResultProto storageInfoResult = mIcingSearchEngineLocked.getStorageInfo();
+            throwIfClosedLocked();
-        mLogUtil.piiTrace(
+            mLogUtil.piiTrace("getStorageInfo, request");
-                "getStorageInfo, response", storageInfoResult.getStatus(), storageInfoResult);
+            StorageInfoResultProto storageInfoResult = mIcingSearchEngineLocked.getStorageInfo();
-        checkSuccess(storageInfoResult.getStatus());
+            mLogUtil.piiTrace(
-        if (!storageInfoResult.hasStorageInfo()
+                    "getStorageInfo, response", storageInfoResult.getStatus(), storageInfoResult);
-                || !storageInfoResult.getStorageInfo().hasDocumentStorageInfo()) {
+            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 =
+        long totalStorageSize = storageInfoProto.getTotalStorageSize();
-                storageInfoResult.getStorageInfo().getDocumentStorageInfo();
+        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 {

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();
+}

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;
+    }
+}

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
@@ -264,6 +268,22 @@ public class AppSearchConfigTest {
                 samplingIntervalBatchCallStats);
     }
 
+    @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);

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);
+    }
+}

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);

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");
+    }
 }

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
+        AppSearchException exception =
-        // exception here.
+                Assert.assertThrows(
-        Assert.assertThrows(
+                        AppSearchException.class,
-                AppSearchException.class,
+                        () ->
-                () -> mAppSearchImpl.putDocument(testPackageName, testDatabase, document, mLogger));
+                                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
+        AppSearchException exception =
-        // exception here.
+                Assert.assertThrows(
-        Assert.assertThrows(
+                        AppSearchException.class,
-                AppSearchException.class,
+                        () ->
-                () ->
+                                mAppSearchImpl.remove(
-                        mAppSearchImpl.remove(
+                                        testPackageName,
-                                testPackageName,
+                                        testDatabase,
-                                testDatabase,
+                                        testNamespace,
-                                testNamespace,
+                                        "invalidId",
-                                "invalidId",
+                                        rStatsBuilder));
-                                rStatsBuilder));
+        assertThat(exception.getResultCode()).isEqualTo(AppSearchResult.RESULT_NOT_FOUND);
 
         RemoveStats rStats = rStatsBuilder.build();
         assertThat(rStats.getPackageName()).isEqualTo(testPackageName);

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);
     }