1 /*
2  * Copyright 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package androidx.appsearch.localstorage;
18 
19 import static androidx.appsearch.localstorage.util.PrefixUtil.getPackageName;
20 
21 import androidx.annotation.RestrictTo;
22 import androidx.appsearch.app.AppSearchResult;
23 import androidx.appsearch.exceptions.AppSearchException;
24 import androidx.appsearch.flags.Flags;
25 import androidx.appsearch.localstorage.util.MapUtil;
26 import androidx.collection.ArrayMap;
27 import androidx.core.util.Preconditions;
28 
29 import com.google.android.icing.proto.NamespaceStorageInfoProto;
30 
31 import org.jspecify.annotations.NonNull;
32 
33 import java.util.List;
34 import java.util.Map;
35 import java.util.concurrent.Callable;
36 
37 /**
38  * A class that encapsulates per-package document count tracking and limit enforcement.
39  *
40  * This class is configured with a {@link #mDocumentLimitStartThreshold}. While the total number of
41  * documents in the system is below that threshold, all packages will be allowed to put as many
42  * documents into the index as they wish. Once the total number of documents exceed
43  * {@link #mDocumentLimitStartThreshold}, then each package will be limited to no more than
44  * {@link #mPerPackageDocumentCountLimit} documents.
45  *
46  *  <p>This class is not thread safe.
47  *
48  * @exportToFramework:hide
49  */
50 @RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
51 public class DocumentLimiter {
52     private final int mDocumentLimitStartThreshold;
53     private final int mPerPackageDocumentCountLimit;
54     private int mTotalDocumentCount;
55     private final Map<String, Integer> mDocumentCountMap;
56 
57     /**
58      * @param documentLimitStartThreshold the total number of documents in the system at which the
59      *                                    limiter should begin applying the
60      *                                    perPackageDocumentCountLimit limit.
61      * @param perPackageDocumentCountLimit the maximum number of documents that each package is
62      *                                   allowed to have once the total number of documents in the
63      *                                   system exceeds documentLimitStartThreshold.
64      * @param namespaceStorageInfoProtoList a list of NamespaceStorageInfoProtos that reflects the
65      *                                     state of the index when this DocumentLimiter was created.
66      */
DocumentLimiter(int documentLimitStartThreshold, int perPackageDocumentCountLimit, @NonNull List<NamespaceStorageInfoProto> namespaceStorageInfoProtoList)67     public DocumentLimiter(int documentLimitStartThreshold, int perPackageDocumentCountLimit,
68             @NonNull List<NamespaceStorageInfoProto> namespaceStorageInfoProtoList) {
69         mDocumentLimitStartThreshold = documentLimitStartThreshold;
70         mPerPackageDocumentCountLimit = perPackageDocumentCountLimit;
71         mTotalDocumentCount = 0;
72         mDocumentCountMap = new ArrayMap<>(namespaceStorageInfoProtoList.size());
73         buildDocumentCountMap(Preconditions.checkNotNull(namespaceStorageInfoProtoList));
74     }
75 
76     /**
77      * Checks whether the package identified by packageName should be allowed to add another
78      * document.
79      *
80      * @param packageName the name of the package attempting to add the document
81      *
82      * @param namespaceStorageInfoProducer a callable that returns an up-to-date list of
83      *                                    NamespaceStorageInfoProtos.
84      *
85      * @throws AppSearchException if the document limit is in force (because the total number of
86      * documents in the system exceeds {@link #mDocumentLimitStartThreshold}) and the package
87      * identified by packageName has already added more documents than
88      * {@link #mPerPackageDocumentCountLimit}.
89      */
enforceDocumentCountLimit( @onNull String packageName, @NonNull Callable<List<NamespaceStorageInfoProto>> namespaceStorageInfoProducer)90     public void enforceDocumentCountLimit(
91             @NonNull String packageName,
92             @NonNull Callable<List<NamespaceStorageInfoProto>> namespaceStorageInfoProducer)
93             throws AppSearchException {
94         Preconditions.checkNotNull(packageName);
95         if (mTotalDocumentCount < mDocumentLimitStartThreshold) {
96             return;
97         }
98         Integer newDocumentCount = MapUtil.getOrDefault(mDocumentCountMap, packageName, 0) + 1;
99         if (!Flags.enableDocumentLimiterReplaceTracking()
100                 && newDocumentCount > mPerPackageDocumentCountLimit) {
101             // Our management of mDocumentCountMap doesn't account for document
102             // replacements, so our counter might have overcounted if the app has replaced docs.
103             // Rebuild the counter from StorageInfo in case this is so.
104             refreshDocumentCount(namespaceStorageInfoProducer);
105             newDocumentCount = MapUtil.getOrDefault(mDocumentCountMap, packageName, 0) + 1;
106         }
107         if (newDocumentCount > mPerPackageDocumentCountLimit) {
108             // Now we really can't fit it in, even accounting for replacements.
109             throw new AppSearchException(
110                     AppSearchResult.RESULT_OUT_OF_SPACE,
111                     "Package \"" + packageName + "\" exceeded limit of "
112                             + mPerPackageDocumentCountLimit + " documents. Some documents "
113                             + "must be removed to index additional ones.");
114         }
115     }
116 
117     /**
118      * Informs the DocumentLimiter that another document has been added for the package identified
119      * by package.
120      *
121      * @param packageName the name of the package that owns the added document.
122      */
reportDocumentAdded( @onNull String packageName, @NonNull Callable<List<NamespaceStorageInfoProto>> namespaceStorageInfoProducer)123     public void reportDocumentAdded(
124             @NonNull String packageName,
125             @NonNull Callable<List<NamespaceStorageInfoProto>> namespaceStorageInfoProducer)
126             throws AppSearchException {
127         Preconditions.checkNotNull(packageName);
128         ++mTotalDocumentCount;
129         Integer newDocumentCount = MapUtil.getOrDefault(mDocumentCountMap, packageName, 0) + 1;
130         mDocumentCountMap.put(packageName, newDocumentCount);
131         if (!Flags.enableDocumentLimiterReplaceTracking()
132                 && mTotalDocumentCount == mDocumentLimitStartThreshold) {
133             // We just hit the document limit start threshold. If
134             // Flags.enableDocumentLimiterReplaceTracking is false, then it's possible that we're
135             // over-counting documents. So we refresh the document count to make sure that we have
136             // the right count.
137             refreshDocumentCount(namespaceStorageInfoProducer);
138         }
139     }
140 
141     /**
142      * Informs the DocumentLimiter that numDocumentsDeleted documents, owned by the package
143      * identified by packageName, have been deleted.
144      *
145      * @param packageName the name of the package that owns the deleted documents.
146      * @param numDocumentsDeleted the number of documents that were deleted.
147      */
reportDocumentsRemoved(@onNull String packageName, int numDocumentsDeleted)148     public void reportDocumentsRemoved(@NonNull String packageName, int numDocumentsDeleted) {
149         Preconditions.checkNotNull(packageName);
150         if (numDocumentsDeleted <= 0) {
151             return;
152         }
153         mTotalDocumentCount -= numDocumentsDeleted;
154         Integer oldDocumentCount = mDocumentCountMap.get(packageName);
155         // This should always be true: how can we delete documents for a package without
156         // having seen that package during init? This is just a safeguard.
157         if (oldDocumentCount != null) {
158             if (numDocumentsDeleted >= oldDocumentCount) {
159                 mDocumentCountMap.remove(packageName);
160             } else {
161                 mDocumentCountMap.put(packageName, oldDocumentCount - numDocumentsDeleted);
162             }
163         }
164     }
165 
166     /**
167      * Informs the DocumentLimiter that the package identified by packageName has been removed from
168      * the system entirely.
169      *
170      * @param packageName the name of the package that was removed.
171      */
reportPackageRemoved(@onNull String packageName)172     public void reportPackageRemoved(@NonNull String packageName) {
173         Preconditions.checkNotNull(packageName);
174         Integer oldDocumentCount = mDocumentCountMap.remove(packageName);
175         if (oldDocumentCount != null) {
176             // This should always be true: how can we remove a package without having seen that
177             // package during init? This is just a safeguard.
178             mTotalDocumentCount -= oldDocumentCount;
179         }
180     }
181 
refreshDocumentCount( @onNull Callable<List<NamespaceStorageInfoProto>> namespaceStorageInfoProducer)182     private void refreshDocumentCount(
183             @NonNull Callable<List<NamespaceStorageInfoProto>> namespaceStorageInfoProducer)
184             throws AppSearchException {
185         try {
186             List<NamespaceStorageInfoProto> namespaceStorageInfos =
187                     namespaceStorageInfoProducer.call();
188             buildDocumentCountMap(namespaceStorageInfos);
189         } catch (AppSearchException e) {
190             throw e;
191         } catch (Exception e) {
192             // This should never happen.
193             throw new AppSearchException(
194                     AppSearchResult.RESULT_UNKNOWN_ERROR,
195                     "Encountered unexpected exception when retrieving namespace storage info.",
196                     e);
197         }
198     }
199 
buildDocumentCountMap( @onNull List<NamespaceStorageInfoProto> namespaceStorageInfoProtoList)200     private void buildDocumentCountMap(
201             @NonNull List<NamespaceStorageInfoProto> namespaceStorageInfoProtoList) {
202         mDocumentCountMap.clear();
203         mTotalDocumentCount = 0;
204         for (int i = 0; i < namespaceStorageInfoProtoList.size(); i++) {
205             NamespaceStorageInfoProto namespaceStorageInfoProto =
206                     namespaceStorageInfoProtoList.get(i);
207             mTotalDocumentCount += namespaceStorageInfoProto.getNumAliveDocuments();
208             String packageName = getPackageName(namespaceStorageInfoProto.getNamespace());
209             Integer newCount =
210                     MapUtil.getOrDefault(mDocumentCountMap, packageName, 0)
211                             + namespaceStorageInfoProto.getNumAliveDocuments();
212             mDocumentCountMap.put(packageName, newCount);
213         }
214     }
215 }
216