1 /*
2  * Copyright 2020 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 com.android.server.appsearch.external.localstorage;
18 
19 import static com.android.server.appsearch.external.localstorage.util.PrefixUtil.addPrefixToDocument;
20 import static com.android.server.appsearch.external.localstorage.util.PrefixUtil.createPrefix;
21 import static com.android.server.appsearch.external.localstorage.util.PrefixUtil.getDatabaseName;
22 import static com.android.server.appsearch.external.localstorage.util.PrefixUtil.getPackageName;
23 import static com.android.server.appsearch.external.localstorage.util.PrefixUtil.getPrefix;
24 import static com.android.server.appsearch.external.localstorage.util.PrefixUtil.removePrefix;
25 import static com.android.server.appsearch.external.localstorage.util.PrefixUtil.removePrefixesFromDocument;
26 
27 import android.annotation.NonNull;
28 import android.annotation.Nullable;
29 import android.annotation.WorkerThread;
30 import android.app.appsearch.AppSearchResult;
31 import android.app.appsearch.AppSearchSchema;
32 import android.app.appsearch.GenericDocument;
33 import android.app.appsearch.GetByDocumentIdRequest;
34 import android.app.appsearch.GetSchemaResponse;
35 import android.app.appsearch.PackageIdentifier;
36 import android.app.appsearch.SearchResultPage;
37 import android.app.appsearch.SearchSpec;
38 import android.app.appsearch.SetSchemaResponse;
39 import android.app.appsearch.StorageInfo;
40 import android.app.appsearch.exceptions.AppSearchException;
41 import android.app.appsearch.util.LogUtil;
42 import android.os.Bundle;
43 import android.os.SystemClock;
44 import android.util.ArrayMap;
45 import android.util.ArraySet;
46 import android.util.Log;
47 
48 import com.android.internal.annotations.GuardedBy;
49 import com.android.internal.annotations.VisibleForTesting;
50 import com.android.server.appsearch.external.localstorage.converter.GenericDocumentToProtoConverter;
51 import com.android.server.appsearch.external.localstorage.converter.ResultCodeToProtoConverter;
52 import com.android.server.appsearch.external.localstorage.converter.SchemaToProtoConverter;
53 import com.android.server.appsearch.external.localstorage.converter.SearchResultToProtoConverter;
54 import com.android.server.appsearch.external.localstorage.converter.SearchSpecToProtoConverter;
55 import com.android.server.appsearch.external.localstorage.converter.SetSchemaResponseToProtoConverter;
56 import com.android.server.appsearch.external.localstorage.converter.TypePropertyPathToProtoConverter;
57 import com.android.server.appsearch.external.localstorage.stats.InitializeStats;
58 import com.android.server.appsearch.external.localstorage.stats.OptimizeStats;
59 import com.android.server.appsearch.external.localstorage.stats.PutDocumentStats;
60 import com.android.server.appsearch.external.localstorage.stats.RemoveStats;
61 import com.android.server.appsearch.external.localstorage.stats.SearchStats;
62 import com.android.server.appsearch.external.localstorage.stats.SetSchemaStats;
63 import com.android.server.appsearch.external.localstorage.visibilitystore.VisibilityStore;
64 
65 import com.google.android.icing.IcingSearchEngine;
66 import com.google.android.icing.proto.DeleteByQueryResultProto;
67 import com.google.android.icing.proto.DeleteResultProto;
68 import com.google.android.icing.proto.DocumentProto;
69 import com.google.android.icing.proto.DocumentStorageInfoProto;
70 import com.google.android.icing.proto.GetAllNamespacesResultProto;
71 import com.google.android.icing.proto.GetOptimizeInfoResultProto;
72 import com.google.android.icing.proto.GetResultProto;
73 import com.google.android.icing.proto.GetResultSpecProto;
74 import com.google.android.icing.proto.GetSchemaResultProto;
75 import com.google.android.icing.proto.IcingSearchEngineOptions;
76 import com.google.android.icing.proto.InitializeResultProto;
77 import com.google.android.icing.proto.NamespaceStorageInfoProto;
78 import com.google.android.icing.proto.OptimizeResultProto;
79 import com.google.android.icing.proto.PersistToDiskResultProto;
80 import com.google.android.icing.proto.PersistType;
81 import com.google.android.icing.proto.PropertyConfigProto;
82 import com.google.android.icing.proto.PutResultProto;
83 import com.google.android.icing.proto.ReportUsageResultProto;
84 import com.google.android.icing.proto.ResetResultProto;
85 import com.google.android.icing.proto.ResultSpecProto;
86 import com.google.android.icing.proto.SchemaProto;
87 import com.google.android.icing.proto.SchemaTypeConfigProto;
88 import com.google.android.icing.proto.ScoringSpecProto;
89 import com.google.android.icing.proto.SearchResultProto;
90 import com.google.android.icing.proto.SearchSpecProto;
91 import com.google.android.icing.proto.SetSchemaResultProto;
92 import com.google.android.icing.proto.StatusProto;
93 import com.google.android.icing.proto.StorageInfoProto;
94 import com.google.android.icing.proto.StorageInfoResultProto;
95 import com.google.android.icing.proto.TypePropertyMask;
96 import com.google.android.icing.proto.UsageReport;
97 
98 import java.io.Closeable;
99 import java.io.File;
100 import java.util.ArrayList;
101 import java.util.Collections;
102 import java.util.HashMap;
103 import java.util.Iterator;
104 import java.util.List;
105 import java.util.Map;
106 import java.util.Objects;
107 import java.util.Set;
108 import java.util.concurrent.locks.ReadWriteLock;
109 import java.util.concurrent.locks.ReentrantReadWriteLock;
110 
111 /**
112  * Manages interaction with the native IcingSearchEngine and other components to implement AppSearch
113  * functionality.
114  *
115  * <p>Never create two instances using the same folder.
116  *
117  * <p>A single instance of {@link AppSearchImpl} can support all packages and databases. This is
118  * done by combining the package and database name into a unique prefix and prefixing the schemas
119  * and documents stored under that owner. Schemas and documents are physically saved together in
120  * {@link IcingSearchEngine}, but logically isolated:
121  *
122  * <ul>
123  *   <li>Rewrite SchemaType in SchemaProto by adding the package-database prefix and save into
124  *       SchemaTypes set in {@link #setSchema}.
125  *   <li>Rewrite namespace and SchemaType in DocumentProto by adding package-database prefix and
126  *       save to namespaces set in {@link #putDocument}.
127  *   <li>Remove package-database prefix when retrieving documents in {@link #getDocument} and {@link
128  *       #query}.
129  *   <li>Rewrite filters in {@link SearchSpecProto} to have all namespaces and schema types of the
130  *       queried database when user using empty filters in {@link #query}.
131  * </ul>
132  *
133  * <p>Methods in this class belong to two groups, the query group and the mutate group.
134  *
135  * <ul>
136  *   <li>All methods are going to modify global parameters and data in Icing are executed under
137  *       WRITE lock to keep thread safety.
138  *   <li>All methods are going to access global parameters or query data from Icing are executed
139  *       under READ lock to improve query performance.
140  * </ul>
141  *
142  * <p>This class is thread safe.
143  *
144  * @hide
145  */
146 @WorkerThread
147 public final class AppSearchImpl implements Closeable {
148     private static final String TAG = "AppSearchImpl";
149 
150     /** A value 0 means that there're no more pages in the search results. */
151     private static final long EMPTY_PAGE_TOKEN = 0;
152 
153     @VisibleForTesting static final int CHECK_OPTIMIZE_INTERVAL = 100;
154 
155     private final ReadWriteLock mReadWriteLock = new ReentrantReadWriteLock();
156     private final LogUtil mLogUtil = new LogUtil(TAG);
157     private final OptimizeStrategy mOptimizeStrategy;
158     private final LimitConfig mLimitConfig;
159 
160     @GuardedBy("mReadWriteLock")
161     @VisibleForTesting
162     final IcingSearchEngine mIcingSearchEngineLocked;
163 
164     // This map contains schema types and SchemaTypeConfigProtos for all package-database
165     // prefixes. It maps each package-database prefix to an inner-map. The inner-map maps each
166     // prefixed schema type to its respective SchemaTypeConfigProto.
167     @GuardedBy("mReadWriteLock")
168     private final Map<String, Map<String, SchemaTypeConfigProto>> mSchemaMapLocked =
169             new ArrayMap<>();
170 
171     // This map contains namespaces for all package-database prefixes. All values in the map are
172     // prefixed with the package-database prefix.
173     // TODO(b/172360376): Check if this can be replaced with an ArrayMap
174     @GuardedBy("mReadWriteLock")
175     private final Map<String, Set<String>> mNamespaceMapLocked = new HashMap<>();
176 
177     /** Maps package name to active document count. */
178     @GuardedBy("mReadWriteLock")
179     private final Map<String, Integer> mDocumentCountMapLocked = new ArrayMap<>();
180 
181     // Maps packages to the set of valid nextPageTokens that the package can manipulate. A token
182     // is unique and constant per query (i.e. the same token '123' is used to iterate through
183     // pages of search results). The tokens themselves are generated and tracked by
184     // IcingSearchEngine. IcingSearchEngine considers a token valid and won't be reused
185     // until we call invalidateNextPageToken on the token.
186     //
187     // Note that we synchronize on itself because the nextPageToken cache is checked at
188     // query-time, and queries are done in parallel with a read lock. Ideally, this would be
189     // guarded by the normal mReadWriteLock.writeLock, but ReentrantReadWriteLocks can't upgrade
190     // read to write locks. This lock should be acquired at the smallest scope possible.
191     // mReadWriteLock is a higher-level lock, so calls shouldn't be made out
192     // to any functions that grab the lock.
193     @GuardedBy("mNextPageTokensLocked")
194     private final Map<String, Set<Long>> mNextPageTokensLocked = new ArrayMap<>();
195 
196     /**
197      * The counter to check when to call {@link #checkForOptimize}. The interval is {@link
198      * #CHECK_OPTIMIZE_INTERVAL}.
199      */
200     @GuardedBy("mReadWriteLock")
201     private int mOptimizeIntervalCountLocked = 0;
202 
203     /** Whether this instance has been closed, and therefore unusable. */
204     @GuardedBy("mReadWriteLock")
205     private boolean mClosedLocked = false;
206 
207     /**
208      * Creates and initializes an instance of {@link AppSearchImpl} which writes data to the given
209      * folder.
210      *
211      * <p>Clients can pass a {@link AppSearchLogger} here through their AppSearchSession, but it
212      * can't be saved inside {@link AppSearchImpl}, because the impl will be shared by all the
213      * sessions for the same package in JetPack.
214      *
215      * <p>Instead, logger instance needs to be passed to each individual method, like create, query
216      * and putDocument.
217      *
218      * @param initStatsBuilder collects stats for initialization if provided.
219      */
220     @NonNull
create( @onNull File icingDir, @NonNull LimitConfig limitConfig, @Nullable InitializeStats.Builder initStatsBuilder, @NonNull OptimizeStrategy optimizeStrategy)221     public static AppSearchImpl create(
222             @NonNull File icingDir,
223             @NonNull LimitConfig limitConfig,
224             @Nullable InitializeStats.Builder initStatsBuilder,
225             @NonNull OptimizeStrategy optimizeStrategy)
226             throws AppSearchException {
227         return new AppSearchImpl(icingDir, limitConfig, initStatsBuilder, optimizeStrategy);
228     }
229 
230     /** @param initStatsBuilder collects stats for initialization if provided. */
AppSearchImpl( @onNull File icingDir, @NonNull LimitConfig limitConfig, @Nullable InitializeStats.Builder initStatsBuilder, @NonNull OptimizeStrategy optimizeStrategy)231     private AppSearchImpl(
232             @NonNull File icingDir,
233             @NonNull LimitConfig limitConfig,
234             @Nullable InitializeStats.Builder initStatsBuilder,
235             @NonNull OptimizeStrategy optimizeStrategy)
236             throws AppSearchException {
237         Objects.requireNonNull(icingDir);
238         mLimitConfig = Objects.requireNonNull(limitConfig);
239         mOptimizeStrategy = Objects.requireNonNull(optimizeStrategy);
240 
241         mReadWriteLock.writeLock().lock();
242         try {
243             // We synchronize here because we don't want to call IcingSearchEngine.initialize() more
244             // than once. It's unnecessary and can be a costly operation.
245             IcingSearchEngineOptions options =
246                     IcingSearchEngineOptions.newBuilder()
247                             .setBaseDir(icingDir.getAbsolutePath())
248                             .build();
249             mLogUtil.piiTrace("Constructing IcingSearchEngine, request", options);
250             mIcingSearchEngineLocked = new IcingSearchEngine(options);
251             mLogUtil.piiTrace(
252                     "Constructing IcingSearchEngine, response",
253                     Objects.hashCode(mIcingSearchEngineLocked));
254 
255             // The core initialization procedure. If any part of this fails, we bail into
256             // resetLocked(), deleting all data (but hopefully allowing AppSearchImpl to come up).
257             try {
258                 mLogUtil.piiTrace("icingSearchEngine.initialize, request");
259                 InitializeResultProto initializeResultProto = mIcingSearchEngineLocked.initialize();
260                 mLogUtil.piiTrace(
261                         "icingSearchEngine.initialize, response",
262                         initializeResultProto.getStatus(),
263                         initializeResultProto);
264 
265                 if (initStatsBuilder != null) {
266                     initStatsBuilder
267                             .setStatusCode(
268                                     statusProtoToResultCode(initializeResultProto.getStatus()))
269                             // TODO(b/173532925) how to get DeSyncs value
270                             .setHasDeSync(false);
271                     AppSearchLoggerHelper.copyNativeStats(
272                             initializeResultProto.getInitializeStats(), initStatsBuilder);
273                 }
274                 checkSuccess(initializeResultProto.getStatus());
275 
276                 // Read all protos we need to construct AppSearchImpl's cache maps
277                 long prepareSchemaAndNamespacesLatencyStartMillis = SystemClock.elapsedRealtime();
278                 SchemaProto schemaProto = getSchemaProtoLocked();
279 
280                 mLogUtil.piiTrace("init:getAllNamespaces, request");
281                 GetAllNamespacesResultProto getAllNamespacesResultProto =
282                         mIcingSearchEngineLocked.getAllNamespaces();
283                 mLogUtil.piiTrace(
284                         "init:getAllNamespaces, response",
285                         getAllNamespacesResultProto.getNamespacesCount(),
286                         getAllNamespacesResultProto);
287 
288                 StorageInfoProto storageInfoProto = getRawStorageInfoProto();
289 
290                 // Log the time it took to read the data that goes into the cache maps
291                 if (initStatsBuilder != null) {
292                     // In case there is some error for getAllNamespaces, we can still
293                     // set the latency for preparation.
294                     // If there is no error, the value will be overridden by the actual one later.
295                     initStatsBuilder
296                             .setStatusCode(
297                                     statusProtoToResultCode(
298                                             getAllNamespacesResultProto.getStatus()))
299                             .setPrepareSchemaAndNamespacesLatencyMillis(
300                                     (int)
301                                             (SystemClock.elapsedRealtime()
302                                                     - prepareSchemaAndNamespacesLatencyStartMillis));
303                 }
304                 checkSuccess(getAllNamespacesResultProto.getStatus());
305 
306                 // Populate schema map
307                 List<SchemaTypeConfigProto> schemaProtoTypesList = schemaProto.getTypesList();
308                 for (int i = 0; i < schemaProtoTypesList.size(); i++) {
309                     SchemaTypeConfigProto schema = schemaProtoTypesList.get(i);
310                     String prefixedSchemaType = schema.getSchemaType();
311                     addToMap(mSchemaMapLocked, getPrefix(prefixedSchemaType), schema);
312                 }
313 
314                 // Populate namespace map
315                 List<String> prefixedNamespaceList =
316                         getAllNamespacesResultProto.getNamespacesList();
317                 for (int i = 0; i < prefixedNamespaceList.size(); i++) {
318                     String prefixedNamespace = prefixedNamespaceList.get(i);
319                     addToMap(mNamespaceMapLocked, getPrefix(prefixedNamespace), prefixedNamespace);
320                 }
321 
322                 // Populate document count map
323                 rebuildDocumentCountMapLocked(storageInfoProto);
324 
325                 // logging prepare_schema_and_namespaces latency
326                 if (initStatsBuilder != null) {
327                     initStatsBuilder.setPrepareSchemaAndNamespacesLatencyMillis(
328                             (int)
329                                     (SystemClock.elapsedRealtime()
330                                             - prepareSchemaAndNamespacesLatencyStartMillis));
331                 }
332 
333                 mLogUtil.piiTrace("Init completed successfully");
334 
335             } catch (AppSearchException e) {
336                 // Some error. Reset and see if it fixes it.
337                 Log.e(TAG, "Error initializing, resetting IcingSearchEngine.", e);
338                 if (initStatsBuilder != null) {
339                     initStatsBuilder.setStatusCode(e.getResultCode());
340                 }
341                 resetLocked(initStatsBuilder);
342             }
343 
344         } finally {
345             mReadWriteLock.writeLock().unlock();
346         }
347     }
348 
349     @GuardedBy("mReadWriteLock")
throwIfClosedLocked()350     private void throwIfClosedLocked() {
351         if (mClosedLocked) {
352             throw new IllegalStateException("Trying to use a closed AppSearchImpl instance.");
353         }
354     }
355 
356     /**
357      * Persists data to disk and closes the instance.
358      *
359      * <p>This instance is no longer usable after it's been closed. Call {@link #create} to create a
360      * new, usable instance.
361      */
362     @Override
close()363     public void close() {
364         mReadWriteLock.writeLock().lock();
365         try {
366             if (mClosedLocked) {
367                 return;
368             }
369             persistToDisk(PersistType.Code.FULL);
370             mLogUtil.piiTrace("icingSearchEngine.close, request");
371             mIcingSearchEngineLocked.close();
372             mLogUtil.piiTrace("icingSearchEngine.close, response");
373             mClosedLocked = true;
374         } catch (AppSearchException e) {
375             Log.w(TAG, "Error when closing AppSearchImpl.", e);
376         } finally {
377             mReadWriteLock.writeLock().unlock();
378         }
379     }
380 
381     /**
382      * Updates the AppSearch schema for this app.
383      *
384      * <p>This method belongs to mutate group.
385      *
386      * @param packageName The package name that owns the schemas.
387      * @param databaseName The name of the database where this schema lives.
388      * @param schemas Schemas to set for this app.
389      * @param visibilityStore If set, {@code schemasNotDisplayedBySystem} and {@code
390      *     schemasVisibleToPackages} will be saved here if the schema is successfully applied.
391      * @param schemasNotDisplayedBySystem Schema types that should not be surfaced on platform
392      *     surfaces.
393      * @param schemasVisibleToPackages Schema types that are visible to the specified packages.
394      * @param forceOverride Whether to force-apply the schema even if it is incompatible. Documents
395      *     which do not comply with the new schema will be deleted.
396      * @param version The overall version number of the request.
397      * @param setSchemaStatsBuilder Builder for {@link SetSchemaStats} to hold stats for setSchema
398      * @return The response contains deleted schema types and incompatible schema types of this
399      *     call.
400      * @throws AppSearchException On IcingSearchEngine error. If the status code is
401      *     FAILED_PRECONDITION for the incompatible change, the exception will be converted to the
402      *     SetSchemaResponse.
403      */
404     @NonNull
setSchema( @onNull String packageName, @NonNull String databaseName, @NonNull List<AppSearchSchema> schemas, @Nullable VisibilityStore visibilityStore, @NonNull List<String> schemasNotDisplayedBySystem, @NonNull Map<String, List<PackageIdentifier>> schemasVisibleToPackages, boolean forceOverride, int version, @Nullable SetSchemaStats.Builder setSchemaStatsBuilder)405     public SetSchemaResponse setSchema(
406             @NonNull String packageName,
407             @NonNull String databaseName,
408             @NonNull List<AppSearchSchema> schemas,
409             @Nullable VisibilityStore visibilityStore,
410             @NonNull List<String> schemasNotDisplayedBySystem,
411             @NonNull Map<String, List<PackageIdentifier>> schemasVisibleToPackages,
412             boolean forceOverride,
413             int version,
414             @Nullable SetSchemaStats.Builder setSchemaStatsBuilder)
415             throws AppSearchException {
416         mReadWriteLock.writeLock().lock();
417         try {
418             throwIfClosedLocked();
419 
420             SchemaProto.Builder existingSchemaBuilder = getSchemaProtoLocked().toBuilder();
421 
422             SchemaProto.Builder newSchemaBuilder = SchemaProto.newBuilder();
423             for (int i = 0; i < schemas.size(); i++) {
424                 AppSearchSchema schema = schemas.get(i);
425                 SchemaTypeConfigProto schemaTypeProto =
426                         SchemaToProtoConverter.toSchemaTypeConfigProto(schema, version);
427                 newSchemaBuilder.addTypes(schemaTypeProto);
428             }
429 
430             String prefix = createPrefix(packageName, databaseName);
431             // Combine the existing schema (which may have types from other prefixes) with this
432             // prefix's new schema. Modifies the existingSchemaBuilder.
433             RewrittenSchemaResults rewrittenSchemaResults =
434                     rewriteSchema(prefix, existingSchemaBuilder, newSchemaBuilder.build());
435 
436             // Apply schema
437             SchemaProto finalSchema = existingSchemaBuilder.build();
438             mLogUtil.piiTrace("setSchema, request", finalSchema.getTypesCount(), finalSchema);
439             SetSchemaResultProto setSchemaResultProto =
440                     mIcingSearchEngineLocked.setSchema(finalSchema, forceOverride);
441             mLogUtil.piiTrace(
442                     "setSchema, response", setSchemaResultProto.getStatus(), setSchemaResultProto);
443 
444             if (setSchemaStatsBuilder != null) {
445                 setSchemaStatsBuilder.setStatusCode(
446                         statusProtoToResultCode(setSchemaResultProto.getStatus()));
447                 AppSearchLoggerHelper.copyNativeStats(setSchemaResultProto, setSchemaStatsBuilder);
448             }
449 
450             // Determine whether it succeeded.
451             try {
452                 checkSuccess(setSchemaResultProto.getStatus());
453             } catch (AppSearchException e) {
454                 // Swallow the exception for the incompatible change case. We will propagate
455                 // those deleted schemas and incompatible types to the SetSchemaResponse.
456                 boolean isFailedPrecondition =
457                         setSchemaResultProto.getStatus().getCode()
458                                 == StatusProto.Code.FAILED_PRECONDITION;
459                 boolean isIncompatible =
460                         setSchemaResultProto.getDeletedSchemaTypesCount() > 0
461                                 || setSchemaResultProto.getIncompatibleSchemaTypesCount() > 0;
462                 if (isFailedPrecondition && isIncompatible) {
463                     return SetSchemaResponseToProtoConverter.toSetSchemaResponse(
464                             setSchemaResultProto, prefix);
465                 } else {
466                     throw e;
467                 }
468             }
469 
470             // Update derived data structures.
471             for (SchemaTypeConfigProto schemaTypeConfigProto :
472                     rewrittenSchemaResults.mRewrittenPrefixedTypes.values()) {
473                 addToMap(mSchemaMapLocked, prefix, schemaTypeConfigProto);
474             }
475 
476             for (String schemaType : rewrittenSchemaResults.mDeletedPrefixedTypes) {
477                 removeFromMap(mSchemaMapLocked, prefix, schemaType);
478             }
479 
480             if (visibilityStore != null) {
481                 Set<String> prefixedSchemasNotDisplayedBySystem =
482                         new ArraySet<>(schemasNotDisplayedBySystem.size());
483                 for (int i = 0; i < schemasNotDisplayedBySystem.size(); i++) {
484                     prefixedSchemasNotDisplayedBySystem.add(
485                             prefix + schemasNotDisplayedBySystem.get(i));
486                 }
487 
488                 Map<String, List<PackageIdentifier>> prefixedSchemasVisibleToPackages =
489                         new ArrayMap<>(schemasVisibleToPackages.size());
490                 for (Map.Entry<String, List<PackageIdentifier>> entry :
491                         schemasVisibleToPackages.entrySet()) {
492                     prefixedSchemasVisibleToPackages.put(prefix + entry.getKey(), entry.getValue());
493                 }
494 
495                 visibilityStore.setVisibility(
496                         packageName,
497                         databaseName,
498                         prefixedSchemasNotDisplayedBySystem,
499                         prefixedSchemasVisibleToPackages);
500             }
501 
502             return SetSchemaResponseToProtoConverter.toSetSchemaResponse(
503                     setSchemaResultProto, prefix);
504         } finally {
505             mReadWriteLock.writeLock().unlock();
506         }
507     }
508 
509     /**
510      * Retrieves the AppSearch schema for this package name, database.
511      *
512      * <p>This method belongs to query group.
513      *
514      * @param packageName Package name that owns this schema
515      * @param databaseName The name of the database where this schema lives.
516      * @throws AppSearchException on IcingSearchEngine error.
517      */
518     @NonNull
getSchema(@onNull String packageName, @NonNull String databaseName)519     public GetSchemaResponse getSchema(@NonNull String packageName, @NonNull String databaseName)
520             throws AppSearchException {
521         mReadWriteLock.readLock().lock();
522         try {
523             throwIfClosedLocked();
524 
525             SchemaProto fullSchema = getSchemaProtoLocked();
526 
527             String prefix = createPrefix(packageName, databaseName);
528             GetSchemaResponse.Builder responseBuilder = new GetSchemaResponse.Builder();
529             for (int i = 0; i < fullSchema.getTypesCount(); i++) {
530                 String typePrefix = getPrefix(fullSchema.getTypes(i).getSchemaType());
531                 if (!prefix.equals(typePrefix)) {
532                     continue;
533                 }
534                 // Rewrite SchemaProto.types.schema_type
535                 SchemaTypeConfigProto.Builder typeConfigBuilder =
536                         fullSchema.getTypes(i).toBuilder();
537                 String newSchemaType = typeConfigBuilder.getSchemaType().substring(prefix.length());
538                 typeConfigBuilder.setSchemaType(newSchemaType);
539 
540                 // Rewrite SchemaProto.types.properties.schema_type
541                 for (int propertyIdx = 0;
542                         propertyIdx < typeConfigBuilder.getPropertiesCount();
543                         propertyIdx++) {
544                     PropertyConfigProto.Builder propertyConfigBuilder =
545                             typeConfigBuilder.getProperties(propertyIdx).toBuilder();
546                     if (!propertyConfigBuilder.getSchemaType().isEmpty()) {
547                         String newPropertySchemaType =
548                                 propertyConfigBuilder.getSchemaType().substring(prefix.length());
549                         propertyConfigBuilder.setSchemaType(newPropertySchemaType);
550                         typeConfigBuilder.setProperties(propertyIdx, propertyConfigBuilder);
551                     }
552                 }
553 
554                 AppSearchSchema schema =
555                         SchemaToProtoConverter.toAppSearchSchema(typeConfigBuilder);
556 
557                 // TODO(b/183050495) find a place to store the version for the database, rather
558                 // than read from a schema.
559                 responseBuilder.setVersion(fullSchema.getTypes(i).getVersion());
560                 responseBuilder.addSchema(schema);
561             }
562             return responseBuilder.build();
563         } finally {
564             mReadWriteLock.readLock().unlock();
565         }
566     }
567 
568     /**
569      * Retrieves the list of namespaces with at least one document for this package name, database.
570      *
571      * <p>This method belongs to query group.
572      *
573      * @param packageName Package name that owns this schema
574      * @param databaseName The name of the database where this schema lives.
575      * @throws AppSearchException on IcingSearchEngine error.
576      */
577     @NonNull
getNamespaces(@onNull String packageName, @NonNull String databaseName)578     public List<String> getNamespaces(@NonNull String packageName, @NonNull String databaseName)
579             throws AppSearchException {
580         mReadWriteLock.readLock().lock();
581         try {
582             throwIfClosedLocked();
583             mLogUtil.piiTrace("getAllNamespaces, request");
584             // We can't just use mNamespaceMap here because we have no way to prune namespaces from
585             // mNamespaceMap when they have no more documents (e.g. after setting schema to empty or
586             // using deleteByQuery).
587             GetAllNamespacesResultProto getAllNamespacesResultProto =
588                     mIcingSearchEngineLocked.getAllNamespaces();
589             mLogUtil.piiTrace(
590                     "getAllNamespaces, response",
591                     getAllNamespacesResultProto.getNamespacesCount(),
592                     getAllNamespacesResultProto);
593             checkSuccess(getAllNamespacesResultProto.getStatus());
594             String prefix = createPrefix(packageName, databaseName);
595             List<String> results = new ArrayList<>();
596             for (int i = 0; i < getAllNamespacesResultProto.getNamespacesCount(); i++) {
597                 String prefixedNamespace = getAllNamespacesResultProto.getNamespaces(i);
598                 if (prefixedNamespace.startsWith(prefix)) {
599                     results.add(prefixedNamespace.substring(prefix.length()));
600                 }
601             }
602             return results;
603         } finally {
604             mReadWriteLock.readLock().unlock();
605         }
606     }
607 
608     /**
609      * Adds a document to the AppSearch index.
610      *
611      * <p>This method belongs to mutate group.
612      *
613      * @param packageName The package name that owns this document.
614      * @param databaseName The databaseName this document resides in.
615      * @param document The document to index.
616      * @throws AppSearchException on IcingSearchEngine error.
617      */
putDocument( @onNull String packageName, @NonNull String databaseName, @NonNull GenericDocument document, @Nullable AppSearchLogger logger)618     public void putDocument(
619             @NonNull String packageName,
620             @NonNull String databaseName,
621             @NonNull GenericDocument document,
622             @Nullable AppSearchLogger logger)
623             throws AppSearchException {
624         PutDocumentStats.Builder pStatsBuilder = null;
625         if (logger != null) {
626             pStatsBuilder = new PutDocumentStats.Builder(packageName, databaseName);
627         }
628         long totalStartTimeMillis = SystemClock.elapsedRealtime();
629 
630         mReadWriteLock.writeLock().lock();
631         try {
632             throwIfClosedLocked();
633 
634             // Generate Document Proto
635             long generateDocumentProtoStartTimeMillis = SystemClock.elapsedRealtime();
636             DocumentProto.Builder documentBuilder =
637                     GenericDocumentToProtoConverter.toDocumentProto(document).toBuilder();
638             long generateDocumentProtoEndTimeMillis = SystemClock.elapsedRealtime();
639 
640             // Rewrite Document Type
641             long rewriteDocumentTypeStartTimeMillis = SystemClock.elapsedRealtime();
642             String prefix = createPrefix(packageName, databaseName);
643             addPrefixToDocument(documentBuilder, prefix);
644             long rewriteDocumentTypeEndTimeMillis = SystemClock.elapsedRealtime();
645             DocumentProto finalDocument = documentBuilder.build();
646 
647             // Check limits
648             int newDocumentCount =
649                     enforceLimitConfigLocked(
650                             packageName, finalDocument.getUri(), finalDocument.getSerializedSize());
651 
652             // Insert document
653             mLogUtil.piiTrace("putDocument, request", finalDocument.getUri(), finalDocument);
654             PutResultProto putResultProto = mIcingSearchEngineLocked.put(finalDocument);
655             mLogUtil.piiTrace("putDocument, response", putResultProto.getStatus(), putResultProto);
656 
657             // Update caches
658             addToMap(mNamespaceMapLocked, prefix, finalDocument.getNamespace());
659             mDocumentCountMapLocked.put(packageName, newDocumentCount);
660 
661             // Logging stats
662             if (pStatsBuilder != null) {
663                 pStatsBuilder
664                         .setStatusCode(statusProtoToResultCode(putResultProto.getStatus()))
665                         .setGenerateDocumentProtoLatencyMillis(
666                                 (int)
667                                         (generateDocumentProtoEndTimeMillis
668                                                 - generateDocumentProtoStartTimeMillis))
669                         .setRewriteDocumentTypesLatencyMillis(
670                                 (int)
671                                         (rewriteDocumentTypeEndTimeMillis
672                                                 - rewriteDocumentTypeStartTimeMillis));
673                 AppSearchLoggerHelper.copyNativeStats(
674                         putResultProto.getPutDocumentStats(), pStatsBuilder);
675             }
676 
677             checkSuccess(putResultProto.getStatus());
678         } finally {
679             mReadWriteLock.writeLock().unlock();
680 
681             if (logger != null) {
682                 long totalEndTimeMillis = SystemClock.elapsedRealtime();
683                 pStatsBuilder.setTotalLatencyMillis(
684                         (int) (totalEndTimeMillis - totalStartTimeMillis));
685                 logger.logStats(pStatsBuilder.build());
686             }
687         }
688     }
689 
690     /**
691      * Checks that a new document can be added to the given packageName with the given serialized
692      * size without violating our {@link LimitConfig}.
693      *
694      * @return the new count of documents for the given package, including the new document.
695      * @throws AppSearchException with a code of {@link AppSearchResult#RESULT_OUT_OF_SPACE} if the
696      *     limits are violated by the new document.
697      */
698     @GuardedBy("mReadWriteLock")
enforceLimitConfigLocked(String packageName, String newDocUri, int newDocSize)699     private int enforceLimitConfigLocked(String packageName, String newDocUri, int newDocSize)
700             throws AppSearchException {
701         // Limits check: size of document
702         if (newDocSize > mLimitConfig.getMaxDocumentSizeBytes()) {
703             throw new AppSearchException(
704                     AppSearchResult.RESULT_OUT_OF_SPACE,
705                     "Document \""
706                             + newDocUri
707                             + "\" for package \""
708                             + packageName
709                             + "\" serialized to "
710                             + newDocSize
711                             + " bytes, which exceeds "
712                             + "limit of "
713                             + mLimitConfig.getMaxDocumentSizeBytes()
714                             + " bytes");
715         }
716 
717         // Limits check: number of documents
718         Integer oldDocumentCount = mDocumentCountMapLocked.get(packageName);
719         int newDocumentCount;
720         if (oldDocumentCount == null) {
721             newDocumentCount = 1;
722         } else {
723             newDocumentCount = oldDocumentCount + 1;
724         }
725         if (newDocumentCount > mLimitConfig.getMaxDocumentCount()) {
726             // Our management of mDocumentCountMapLocked doesn't account for document
727             // replacements, so our counter might have overcounted if the app has replaced docs.
728             // Rebuild the counter from StorageInfo in case this is so.
729             // TODO(b/170371356):  If Icing lib exposes something in the result which says
730             //  whether the document was a replacement, we could subtract 1 again after the put
731             //  to keep the count accurate. That would allow us to remove this code.
732             rebuildDocumentCountMapLocked(getRawStorageInfoProto());
733             oldDocumentCount = mDocumentCountMapLocked.get(packageName);
734             if (oldDocumentCount == null) {
735                 newDocumentCount = 1;
736             } else {
737                 newDocumentCount = oldDocumentCount + 1;
738             }
739         }
740         if (newDocumentCount > mLimitConfig.getMaxDocumentCount()) {
741             // Now we really can't fit it in, even accounting for replacements.
742             throw new AppSearchException(
743                     AppSearchResult.RESULT_OUT_OF_SPACE,
744                     "Package \""
745                             + packageName
746                             + "\" exceeded limit of "
747                             + mLimitConfig.getMaxDocumentCount()
748                             + " documents. Some documents "
749                             + "must be removed to index additional ones.");
750         }
751 
752         return newDocumentCount;
753     }
754 
755     /**
756      * Retrieves a document from the AppSearch index by namespace and document ID.
757      *
758      * <p>This method belongs to query group.
759      *
760      * @param packageName The package that owns this document.
761      * @param databaseName The databaseName this document resides in.
762      * @param namespace The namespace this document resides in.
763      * @param id The ID of the document to get.
764      * @param typePropertyPaths A map of schema type to a list of property paths to return in the
765      *     result.
766      * @return The Document contents
767      * @throws AppSearchException on IcingSearchEngine error.
768      */
769     @NonNull
getDocument( @onNull String packageName, @NonNull String databaseName, @NonNull String namespace, @NonNull String id, @NonNull Map<String, List<String>> typePropertyPaths)770     public GenericDocument getDocument(
771             @NonNull String packageName,
772             @NonNull String databaseName,
773             @NonNull String namespace,
774             @NonNull String id,
775             @NonNull Map<String, List<String>> typePropertyPaths)
776             throws AppSearchException {
777         mReadWriteLock.readLock().lock();
778         try {
779             throwIfClosedLocked();
780             String prefix = createPrefix(packageName, databaseName);
781             List<TypePropertyMask> nonPrefixedPropertyMasks =
782                     TypePropertyPathToProtoConverter.toTypePropertyMaskList(typePropertyPaths);
783             List<TypePropertyMask> prefixedPropertyMasks =
784                     new ArrayList<>(nonPrefixedPropertyMasks.size());
785             for (int i = 0; i < nonPrefixedPropertyMasks.size(); ++i) {
786                 TypePropertyMask typePropertyMask = nonPrefixedPropertyMasks.get(i);
787                 String nonPrefixedType = typePropertyMask.getSchemaType();
788                 String prefixedType =
789                         nonPrefixedType.equals(
790                                         GetByDocumentIdRequest.PROJECTION_SCHEMA_TYPE_WILDCARD)
791                                 ? nonPrefixedType
792                                 : prefix + nonPrefixedType;
793                 prefixedPropertyMasks.add(
794                         typePropertyMask.toBuilder().setSchemaType(prefixedType).build());
795             }
796             GetResultSpecProto getResultSpec =
797                     GetResultSpecProto.newBuilder()
798                             .addAllTypePropertyMasks(prefixedPropertyMasks)
799                             .build();
800 
801             String finalNamespace = createPrefix(packageName, databaseName) + namespace;
802             if (mLogUtil.isPiiTraceEnabled()) {
803                 mLogUtil.piiTrace(
804                         "getDocument, request", finalNamespace + ", " + id + "," + getResultSpec);
805             }
806             GetResultProto getResultProto =
807                     mIcingSearchEngineLocked.get(finalNamespace, id, getResultSpec);
808             mLogUtil.piiTrace("getDocument, response", getResultProto.getStatus(), getResultProto);
809             checkSuccess(getResultProto.getStatus());
810 
811             // The schema type map cannot be null at this point. It could only be null if no
812             // schema had ever been set for that prefix. Given we have retrieved a document from
813             // the index, we know a schema had to have been set.
814             Map<String, SchemaTypeConfigProto> schemaTypeMap = mSchemaMapLocked.get(prefix);
815             DocumentProto.Builder documentBuilder = getResultProto.getDocument().toBuilder();
816             removePrefixesFromDocument(documentBuilder);
817             return GenericDocumentToProtoConverter.toGenericDocument(
818                     documentBuilder.build(), prefix, schemaTypeMap);
819         } finally {
820             mReadWriteLock.readLock().unlock();
821         }
822     }
823 
824     /**
825      * Executes a query against the AppSearch index and returns results.
826      *
827      * <p>This method belongs to query group.
828      *
829      * @param packageName The package name that is performing the query.
830      * @param databaseName The databaseName this query for.
831      * @param queryExpression Query String to search.
832      * @param searchSpec Spec for setting filters, raw query etc.
833      * @param logger logger to collect query stats
834      * @return The results of performing this search. It may contain an empty list of results if no
835      *     documents matched the query.
836      * @throws AppSearchException on IcingSearchEngine error.
837      */
838     @NonNull
query( @onNull String packageName, @NonNull String databaseName, @NonNull String queryExpression, @NonNull SearchSpec searchSpec, @Nullable AppSearchLogger logger)839     public SearchResultPage query(
840             @NonNull String packageName,
841             @NonNull String databaseName,
842             @NonNull String queryExpression,
843             @NonNull SearchSpec searchSpec,
844             @Nullable AppSearchLogger logger)
845             throws AppSearchException {
846         long totalLatencyStartMillis = SystemClock.elapsedRealtime();
847         SearchStats.Builder sStatsBuilder = null;
848         if (logger != null) {
849             sStatsBuilder =
850                     new SearchStats.Builder(SearchStats.VISIBILITY_SCOPE_LOCAL, packageName)
851                             .setDatabase(databaseName);
852         }
853 
854         mReadWriteLock.readLock().lock();
855         try {
856             throwIfClosedLocked();
857 
858             List<String> filterPackageNames = searchSpec.getFilterPackageNames();
859             if (!filterPackageNames.isEmpty() && !filterPackageNames.contains(packageName)) {
860                 // Client wanted to query over some packages that weren't its own. This isn't
861                 // allowed through local query so we can return early with no results.
862                 if (logger != null) {
863                     sStatsBuilder.setStatusCode(AppSearchResult.RESULT_SECURITY_ERROR);
864                 }
865                 return new SearchResultPage(Bundle.EMPTY);
866             }
867 
868             String prefix = createPrefix(packageName, databaseName);
869             Set<String> allowedPrefixedSchemas = getAllowedPrefixSchemasLocked(prefix, searchSpec);
870 
871             SearchResultPage searchResultPage =
872                     doQueryLocked(
873                             Collections.singleton(createPrefix(packageName, databaseName)),
874                             allowedPrefixedSchemas,
875                             queryExpression,
876                             searchSpec,
877                             sStatsBuilder);
878             addNextPageToken(packageName, searchResultPage.getNextPageToken());
879             return searchResultPage;
880         } finally {
881             mReadWriteLock.readLock().unlock();
882             if (logger != null) {
883                 sStatsBuilder.setTotalLatencyMillis(
884                         (int) (SystemClock.elapsedRealtime() - totalLatencyStartMillis));
885                 logger.logStats(sStatsBuilder.build());
886             }
887         }
888     }
889 
890     /**
891      * Executes a global query, i.e. over all permitted prefixes, against the AppSearch index and
892      * returns results.
893      *
894      * <p>This method belongs to query group.
895      *
896      * @param queryExpression Query String to search.
897      * @param searchSpec Spec for setting filters, raw query etc.
898      * @param callerPackageName Package name of the caller, should belong to the {@code
899      *     callerUserHandle}.
900      * @param visibilityStore Optional visibility store to obtain system and package visibility
901      *     settings from
902      * @param callerUid UID of the client making the globalQuery call.
903      * @param callerHasSystemAccess Whether the caller has been positively identified as having
904      *     access to schemas marked system surfaceable.
905      * @param logger logger to collect globalQuery stats
906      * @return The results of performing this search. It may contain an empty list of results if no
907      *     documents matched the query.
908      * @throws AppSearchException on IcingSearchEngine error.
909      */
910     @NonNull
globalQuery( @onNull String queryExpression, @NonNull SearchSpec searchSpec, @NonNull String callerPackageName, @Nullable VisibilityStore visibilityStore, int callerUid, boolean callerHasSystemAccess, @Nullable AppSearchLogger logger)911     public SearchResultPage globalQuery(
912             @NonNull String queryExpression,
913             @NonNull SearchSpec searchSpec,
914             @NonNull String callerPackageName,
915             @Nullable VisibilityStore visibilityStore,
916             int callerUid,
917             boolean callerHasSystemAccess,
918             @Nullable AppSearchLogger logger)
919             throws AppSearchException {
920         long totalLatencyStartMillis = SystemClock.elapsedRealtime();
921         SearchStats.Builder sStatsBuilder = null;
922         if (logger != null) {
923             sStatsBuilder =
924                     new SearchStats.Builder(SearchStats.VISIBILITY_SCOPE_GLOBAL, callerPackageName);
925         }
926 
927         mReadWriteLock.readLock().lock();
928         try {
929             throwIfClosedLocked();
930 
931             // Convert package filters to prefix filters
932             Set<String> packageFilters = new ArraySet<>(searchSpec.getFilterPackageNames());
933             Set<String> prefixFilters = new ArraySet<>();
934             if (packageFilters.isEmpty()) {
935                 // Client didn't restrict their search over packages. Try to query over all
936                 // packages/prefixes
937                 prefixFilters = mNamespaceMapLocked.keySet();
938             } else {
939                 // Client did restrict their search over packages. Only include the prefixes that
940                 // belong to the specified packages.
941                 for (String prefix : mNamespaceMapLocked.keySet()) {
942                     String packageName = getPackageName(prefix);
943                     if (packageFilters.contains(packageName)) {
944                         prefixFilters.add(prefix);
945                     }
946                 }
947             }
948 
949             // Convert schema filters to prefixed schema filters
950             ArraySet<String> prefixedSchemaFilters = new ArraySet<>();
951             for (String prefix : prefixFilters) {
952                 List<String> schemaFilters = searchSpec.getFilterSchemas();
953                 if (schemaFilters.isEmpty()) {
954                     // Client didn't specify certain schemas to search over, check all schemas
955                     prefixedSchemaFilters.addAll(mSchemaMapLocked.get(prefix).keySet());
956                 } else {
957                     // Client specified some schemas to search over, check each one
958                     for (int i = 0; i < schemaFilters.size(); i++) {
959                         prefixedSchemaFilters.add(prefix + schemaFilters.get(i));
960                     }
961                 }
962             }
963 
964             // Remove the schemas the client is not allowed to search over
965             Iterator<String> prefixedSchemaIt = prefixedSchemaFilters.iterator();
966             while (prefixedSchemaIt.hasNext()) {
967                 String prefixedSchema = prefixedSchemaIt.next();
968                 String packageName = getPackageName(prefixedSchema);
969 
970                 boolean allow;
971                 if (packageName.equals(callerPackageName)) {
972                     // Callers can always retrieve their own data
973                     allow = true;
974                 } else if (visibilityStore == null) {
975                     // If there's no visibility store, there's no extra access
976                     allow = false;
977                 } else {
978                     String databaseName = getDatabaseName(prefixedSchema);
979                     allow =
980                             visibilityStore.isSchemaSearchableByCaller(
981                                     packageName,
982                                     databaseName,
983                                     prefixedSchema,
984                                     callerUid,
985                                     callerHasSystemAccess);
986                 }
987 
988                 if (!allow) {
989                     prefixedSchemaIt.remove();
990                 }
991             }
992 
993             SearchResultPage searchResultPage =
994                     doQueryLocked(
995                             prefixFilters,
996                             prefixedSchemaFilters,
997                             queryExpression,
998                             searchSpec,
999                             sStatsBuilder);
1000             addNextPageToken(callerPackageName, searchResultPage.getNextPageToken());
1001             return searchResultPage;
1002         } finally {
1003             mReadWriteLock.readLock().unlock();
1004 
1005             if (logger != null) {
1006                 sStatsBuilder.setTotalLatencyMillis(
1007                         (int) (SystemClock.elapsedRealtime() - totalLatencyStartMillis));
1008                 logger.logStats(sStatsBuilder.build());
1009             }
1010         }
1011     }
1012 
1013     /**
1014      * Returns a mapping of package names to all the databases owned by that package.
1015      *
1016      * <p>This method is inefficient to call repeatedly.
1017      */
1018     @NonNull
getPackageToDatabases()1019     public Map<String, Set<String>> getPackageToDatabases() {
1020         mReadWriteLock.readLock().lock();
1021         try {
1022             Map<String, Set<String>> packageToDatabases = new ArrayMap<>();
1023             for (String prefix : mSchemaMapLocked.keySet()) {
1024                 String packageName = getPackageName(prefix);
1025 
1026                 Set<String> databases = packageToDatabases.get(packageName);
1027                 if (databases == null) {
1028                     databases = new ArraySet<>();
1029                     packageToDatabases.put(packageName, databases);
1030                 }
1031 
1032                 String databaseName = getDatabaseName(prefix);
1033                 databases.add(databaseName);
1034             }
1035 
1036             return packageToDatabases;
1037         } finally {
1038             mReadWriteLock.readLock().unlock();
1039         }
1040     }
1041 
1042     @GuardedBy("mReadWriteLock")
doQueryLocked( @onNull Set<String> prefixes, @NonNull Set<String> allowedPrefixedSchemas, @NonNull String queryExpression, @NonNull SearchSpec searchSpec, @Nullable SearchStats.Builder sStatsBuilder)1043     private SearchResultPage doQueryLocked(
1044             @NonNull Set<String> prefixes,
1045             @NonNull Set<String> allowedPrefixedSchemas,
1046             @NonNull String queryExpression,
1047             @NonNull SearchSpec searchSpec,
1048             @Nullable SearchStats.Builder sStatsBuilder)
1049             throws AppSearchException {
1050         long rewriteSearchSpecLatencyStartMillis = SystemClock.elapsedRealtime();
1051 
1052         SearchSpecProto.Builder searchSpecBuilder =
1053                 SearchSpecToProtoConverter.toSearchSpecProto(searchSpec).toBuilder()
1054                         .setQuery(queryExpression);
1055         // rewriteSearchSpecForPrefixesLocked will return false if there is nothing to search
1056         // over given their search filters, so we can return an empty SearchResult and skip
1057         // sending request to Icing.
1058         if (!rewriteSearchSpecForPrefixesLocked(
1059                 searchSpecBuilder, prefixes, allowedPrefixedSchemas)) {
1060             if (sStatsBuilder != null) {
1061                 sStatsBuilder.setRewriteSearchSpecLatencyMillis(
1062                         (int)
1063                                 (SystemClock.elapsedRealtime()
1064                                         - rewriteSearchSpecLatencyStartMillis));
1065             }
1066             return new SearchResultPage(Bundle.EMPTY);
1067         }
1068 
1069         // rewriteSearchSpec, rewriteResultSpec and convertScoringSpec are all counted in
1070         // rewriteSearchSpecLatencyMillis
1071         ResultSpecProto.Builder resultSpecBuilder =
1072                 SearchSpecToProtoConverter.toResultSpecProto(searchSpec).toBuilder();
1073 
1074         int groupingType = searchSpec.getResultGroupingTypeFlags();
1075         if ((groupingType & SearchSpec.GROUPING_TYPE_PER_PACKAGE) != 0
1076                 && (groupingType & SearchSpec.GROUPING_TYPE_PER_NAMESPACE) != 0) {
1077             addPerPackagePerNamespaceResultGroupingsLocked(
1078                     resultSpecBuilder, prefixes, searchSpec.getResultGroupingLimit());
1079         } else if ((groupingType & SearchSpec.GROUPING_TYPE_PER_PACKAGE) != 0) {
1080             addPerPackageResultGroupingsLocked(
1081                     resultSpecBuilder, prefixes, searchSpec.getResultGroupingLimit());
1082         } else if ((groupingType & SearchSpec.GROUPING_TYPE_PER_NAMESPACE) != 0) {
1083             addPerNamespaceResultGroupingsLocked(
1084                     resultSpecBuilder, prefixes, searchSpec.getResultGroupingLimit());
1085         }
1086 
1087         rewriteResultSpecForPrefixesLocked(resultSpecBuilder, prefixes, allowedPrefixedSchemas);
1088         ScoringSpecProto scoringSpec = SearchSpecToProtoConverter.toScoringSpecProto(searchSpec);
1089         SearchSpecProto finalSearchSpec = searchSpecBuilder.build();
1090         ResultSpecProto finalResultSpec = resultSpecBuilder.build();
1091 
1092         long rewriteSearchSpecLatencyEndMillis = SystemClock.elapsedRealtime();
1093 
1094         if (mLogUtil.isPiiTraceEnabled()) {
1095             mLogUtil.piiTrace(
1096                     "search, request",
1097                     finalSearchSpec.getQuery(),
1098                     finalSearchSpec + ", " + scoringSpec + ", " + finalResultSpec);
1099         }
1100         SearchResultProto searchResultProto =
1101                 mIcingSearchEngineLocked.search(finalSearchSpec, scoringSpec, finalResultSpec);
1102         mLogUtil.piiTrace(
1103                 "search, response", searchResultProto.getResultsCount(), searchResultProto);
1104 
1105         if (sStatsBuilder != null) {
1106             sStatsBuilder
1107                     .setStatusCode(statusProtoToResultCode(searchResultProto.getStatus()))
1108                     .setRewriteSearchSpecLatencyMillis(
1109                             (int)
1110                                     (rewriteSearchSpecLatencyEndMillis
1111                                             - rewriteSearchSpecLatencyStartMillis));
1112             AppSearchLoggerHelper.copyNativeStats(searchResultProto.getQueryStats(), sStatsBuilder);
1113         }
1114 
1115         checkSuccess(searchResultProto.getStatus());
1116 
1117         long rewriteSearchResultLatencyStartMillis = SystemClock.elapsedRealtime();
1118         SearchResultPage resultPage = rewriteSearchResultProto(searchResultProto, mSchemaMapLocked);
1119         if (sStatsBuilder != null) {
1120             sStatsBuilder.setRewriteSearchResultLatencyMillis(
1121                     (int) (SystemClock.elapsedRealtime() - rewriteSearchResultLatencyStartMillis));
1122         }
1123 
1124         return resultPage;
1125     }
1126 
1127     /**
1128      * Fetches the next page of results of a previously executed query. Results can be empty if
1129      * next-page token is invalid or all pages have been returned.
1130      *
1131      * <p>This method belongs to query group.
1132      *
1133      * @param packageName Package name of the caller.
1134      * @param nextPageToken The token of pre-loaded results of previously executed query.
1135      * @return The next page of results of previously executed query.
1136      * @throws AppSearchException on IcingSearchEngine error or if can't advance on nextPageToken.
1137      */
1138     @NonNull
getNextPage( @onNull String packageName, long nextPageToken, @Nullable SearchStats.Builder statsBuilder)1139     public SearchResultPage getNextPage(
1140             @NonNull String packageName,
1141             long nextPageToken,
1142             @Nullable SearchStats.Builder statsBuilder)
1143             throws AppSearchException {
1144         long totalLatencyStartMillis = SystemClock.elapsedRealtime();
1145 
1146         mReadWriteLock.readLock().lock();
1147         try {
1148             throwIfClosedLocked();
1149 
1150             mLogUtil.piiTrace("getNextPage, request", nextPageToken);
1151             checkNextPageToken(packageName, nextPageToken);
1152             SearchResultProto searchResultProto =
1153                     mIcingSearchEngineLocked.getNextPage(nextPageToken);
1154 
1155             if (statsBuilder != null) {
1156                 statsBuilder.setStatusCode(statusProtoToResultCode(searchResultProto.getStatus()));
1157                 AppSearchLoggerHelper.copyNativeStats(
1158                         searchResultProto.getQueryStats(), statsBuilder);
1159             }
1160 
1161             mLogUtil.piiTrace(
1162                     "getNextPage, response",
1163                     searchResultProto.getResultsCount(),
1164                     searchResultProto);
1165             checkSuccess(searchResultProto.getStatus());
1166             if (nextPageToken != EMPTY_PAGE_TOKEN
1167                     && searchResultProto.getNextPageToken() == EMPTY_PAGE_TOKEN) {
1168                 // At this point, we're guaranteed that this nextPageToken exists for this package,
1169                 // otherwise checkNextPageToken would've thrown an exception.
1170                 // Since the new token is 0, this is the last page. We should remove the old token
1171                 // from our cache since it no longer refers to this query.
1172                 synchronized (mNextPageTokensLocked) {
1173                     mNextPageTokensLocked.get(packageName).remove(nextPageToken);
1174                 }
1175             }
1176             long rewriteSearchResultLatencyStartMillis = SystemClock.elapsedRealtime();
1177             SearchResultPage resultPage =
1178                     rewriteSearchResultProto(searchResultProto, mSchemaMapLocked);
1179             if (statsBuilder != null) {
1180                 statsBuilder.setRewriteSearchResultLatencyMillis(
1181                         (int)
1182                                 (SystemClock.elapsedRealtime()
1183                                         - rewriteSearchResultLatencyStartMillis));
1184             }
1185             return resultPage;
1186         } finally {
1187             mReadWriteLock.readLock().unlock();
1188             if (statsBuilder != null) {
1189                 statsBuilder.setTotalLatencyMillis(
1190                         (int) (SystemClock.elapsedRealtime() - totalLatencyStartMillis));
1191             }
1192         }
1193     }
1194 
1195     /**
1196      * Invalidates the next-page token so that no more results of the related query can be returned.
1197      *
1198      * <p>This method belongs to query group.
1199      *
1200      * @param packageName Package name of the caller.
1201      * @param nextPageToken The token of pre-loaded results of previously executed query to be
1202      *     Invalidated.
1203      * @throws AppSearchException if nextPageToken is unusable.
1204      */
invalidateNextPageToken(@onNull String packageName, long nextPageToken)1205     public void invalidateNextPageToken(@NonNull String packageName, long nextPageToken)
1206             throws AppSearchException {
1207         mReadWriteLock.readLock().lock();
1208         try {
1209             throwIfClosedLocked();
1210 
1211             mLogUtil.piiTrace("invalidateNextPageToken, request", nextPageToken);
1212             checkNextPageToken(packageName, nextPageToken);
1213             mIcingSearchEngineLocked.invalidateNextPageToken(nextPageToken);
1214 
1215             synchronized (mNextPageTokensLocked) {
1216                 // At this point, we're guaranteed that this nextPageToken exists for this package,
1217                 // otherwise checkNextPageToken would've thrown an exception.
1218                 mNextPageTokensLocked.get(packageName).remove(nextPageToken);
1219             }
1220         } finally {
1221             mReadWriteLock.readLock().unlock();
1222         }
1223     }
1224 
1225     /** Reports a usage of the given document at the given timestamp. */
reportUsage( @onNull String packageName, @NonNull String databaseName, @NonNull String namespace, @NonNull String documentId, long usageTimestampMillis, boolean systemUsage)1226     public void reportUsage(
1227             @NonNull String packageName,
1228             @NonNull String databaseName,
1229             @NonNull String namespace,
1230             @NonNull String documentId,
1231             long usageTimestampMillis,
1232             boolean systemUsage)
1233             throws AppSearchException {
1234         mReadWriteLock.writeLock().lock();
1235         try {
1236             throwIfClosedLocked();
1237 
1238             String prefixedNamespace = createPrefix(packageName, databaseName) + namespace;
1239             UsageReport.UsageType usageType =
1240                     systemUsage
1241                             ? UsageReport.UsageType.USAGE_TYPE2
1242                             : UsageReport.UsageType.USAGE_TYPE1;
1243             UsageReport report =
1244                     UsageReport.newBuilder()
1245                             .setDocumentNamespace(prefixedNamespace)
1246                             .setDocumentUri(documentId)
1247                             .setUsageTimestampMs(usageTimestampMillis)
1248                             .setUsageType(usageType)
1249                             .build();
1250 
1251             mLogUtil.piiTrace("reportUsage, request", report.getDocumentUri(), report);
1252             ReportUsageResultProto result = mIcingSearchEngineLocked.reportUsage(report);
1253             mLogUtil.piiTrace("reportUsage, response", result.getStatus(), result);
1254             checkSuccess(result.getStatus());
1255         } finally {
1256             mReadWriteLock.writeLock().unlock();
1257         }
1258     }
1259 
1260     /**
1261      * Removes the given document by id.
1262      *
1263      * <p>This method belongs to mutate group.
1264      *
1265      * @param packageName The package name that owns the document.
1266      * @param databaseName The databaseName the document is in.
1267      * @param namespace Namespace of the document to remove.
1268      * @param id ID of the document to remove.
1269      * @param removeStatsBuilder builder for {@link RemoveStats} to hold stats for remove
1270      * @throws AppSearchException on IcingSearchEngine error.
1271      */
remove( @onNull String packageName, @NonNull String databaseName, @NonNull String namespace, @NonNull String id, @Nullable RemoveStats.Builder removeStatsBuilder)1272     public void remove(
1273             @NonNull String packageName,
1274             @NonNull String databaseName,
1275             @NonNull String namespace,
1276             @NonNull String id,
1277             @Nullable RemoveStats.Builder removeStatsBuilder)
1278             throws AppSearchException {
1279         long totalLatencyStartTimeMillis = SystemClock.elapsedRealtime();
1280         mReadWriteLock.writeLock().lock();
1281         try {
1282             throwIfClosedLocked();
1283 
1284             String prefixedNamespace = createPrefix(packageName, databaseName) + namespace;
1285             if (mLogUtil.isPiiTraceEnabled()) {
1286                 mLogUtil.piiTrace("removeById, request", prefixedNamespace + ", " + id);
1287             }
1288             DeleteResultProto deleteResultProto =
1289                     mIcingSearchEngineLocked.delete(prefixedNamespace, id);
1290             mLogUtil.piiTrace(
1291                     "removeById, response", deleteResultProto.getStatus(), deleteResultProto);
1292 
1293             if (removeStatsBuilder != null) {
1294                 removeStatsBuilder.setStatusCode(
1295                         statusProtoToResultCode(deleteResultProto.getStatus()));
1296                 AppSearchLoggerHelper.copyNativeStats(
1297                         deleteResultProto.getDeleteStats(), removeStatsBuilder);
1298             }
1299             checkSuccess(deleteResultProto.getStatus());
1300 
1301             // Update derived maps
1302             updateDocumentCountAfterRemovalLocked(packageName, /*numDocumentsDeleted=*/ 1);
1303         } finally {
1304             mReadWriteLock.writeLock().unlock();
1305             if (removeStatsBuilder != null) {
1306                 removeStatsBuilder.setTotalLatencyMillis(
1307                         (int) (SystemClock.elapsedRealtime() - totalLatencyStartTimeMillis));
1308             }
1309         }
1310     }
1311 
1312     /**
1313      * Removes documents by given query.
1314      *
1315      * <p>This method belongs to mutate group.
1316      *
1317      * @param packageName The package name that owns the documents.
1318      * @param databaseName The databaseName the document is in.
1319      * @param queryExpression Query String to search.
1320      * @param searchSpec Defines what and how to remove
1321      * @param removeStatsBuilder builder for {@link RemoveStats} to hold stats for remove
1322      * @throws AppSearchException on IcingSearchEngine error.
1323      */
removeByQuery( @onNull String packageName, @NonNull String databaseName, @NonNull String queryExpression, @NonNull SearchSpec searchSpec, @Nullable RemoveStats.Builder removeStatsBuilder)1324     public void removeByQuery(
1325             @NonNull String packageName,
1326             @NonNull String databaseName,
1327             @NonNull String queryExpression,
1328             @NonNull SearchSpec searchSpec,
1329             @Nullable RemoveStats.Builder removeStatsBuilder)
1330             throws AppSearchException {
1331         long totalLatencyStartTimeMillis = SystemClock.elapsedRealtime();
1332         mReadWriteLock.writeLock().lock();
1333         try {
1334             throwIfClosedLocked();
1335 
1336             List<String> filterPackageNames = searchSpec.getFilterPackageNames();
1337             if (!filterPackageNames.isEmpty() && !filterPackageNames.contains(packageName)) {
1338                 // We're only removing documents within the parameter `packageName`. If we're not
1339                 // restricting our remove-query to this package name, then there's nothing for us to
1340                 // remove.
1341                 return;
1342             }
1343 
1344             SearchSpecProto searchSpecProto =
1345                     SearchSpecToProtoConverter.toSearchSpecProto(searchSpec);
1346             SearchSpecProto.Builder searchSpecBuilder =
1347                     searchSpecProto.toBuilder().setQuery(queryExpression);
1348 
1349             String prefix = createPrefix(packageName, databaseName);
1350             Set<String> allowedPrefixedSchemas = getAllowedPrefixSchemasLocked(prefix, searchSpec);
1351 
1352             // rewriteSearchSpecForPrefixesLocked will return false if there is nothing to search
1353             // over given their search filters, so we can return early and skip sending request
1354             // to Icing.
1355             if (!rewriteSearchSpecForPrefixesLocked(
1356                     searchSpecBuilder, Collections.singleton(prefix), allowedPrefixedSchemas)) {
1357                 return;
1358             }
1359             SearchSpecProto finalSearchSpec = searchSpecBuilder.build();
1360             mLogUtil.piiTrace("removeByQuery, request", finalSearchSpec);
1361             DeleteByQueryResultProto deleteResultProto =
1362                     mIcingSearchEngineLocked.deleteByQuery(finalSearchSpec);
1363             mLogUtil.piiTrace(
1364                     "removeByQuery, response", deleteResultProto.getStatus(), deleteResultProto);
1365 
1366             if (removeStatsBuilder != null) {
1367                 removeStatsBuilder.setStatusCode(
1368                         statusProtoToResultCode(deleteResultProto.getStatus()));
1369                 // TODO(b/187206766) also log query stats here once IcingLib returns it
1370                 AppSearchLoggerHelper.copyNativeStats(
1371                         deleteResultProto.getDeleteByQueryStats(), removeStatsBuilder);
1372             }
1373 
1374             // It seems that the caller wants to get success if the data matching the query is
1375             // not in the DB because it was not there or was successfully deleted.
1376             checkCodeOneOf(
1377                     deleteResultProto.getStatus(), StatusProto.Code.OK, StatusProto.Code.NOT_FOUND);
1378 
1379             // Update derived maps
1380             int numDocumentsDeleted =
1381                     deleteResultProto.getDeleteByQueryStats().getNumDocumentsDeleted();
1382             updateDocumentCountAfterRemovalLocked(packageName, numDocumentsDeleted);
1383         } finally {
1384             mReadWriteLock.writeLock().unlock();
1385             if (removeStatsBuilder != null) {
1386                 removeStatsBuilder.setTotalLatencyMillis(
1387                         (int) (SystemClock.elapsedRealtime() - totalLatencyStartTimeMillis));
1388             }
1389         }
1390     }
1391 
1392     @GuardedBy("mReadWriteLock")
updateDocumentCountAfterRemovalLocked( @onNull String packageName, int numDocumentsDeleted)1393     private void updateDocumentCountAfterRemovalLocked(
1394             @NonNull String packageName, int numDocumentsDeleted) {
1395         if (numDocumentsDeleted > 0) {
1396             Integer oldDocumentCount = mDocumentCountMapLocked.get(packageName);
1397             // This should always be true: how can we delete documents for a package without
1398             // having seen that package during init? This is just a safeguard.
1399             if (oldDocumentCount != null) {
1400                 // This should always be >0; how can we remove more documents than we've indexed?
1401                 // This is just a safeguard.
1402                 int newDocumentCount = Math.max(oldDocumentCount - numDocumentsDeleted, 0);
1403                 mDocumentCountMapLocked.put(packageName, newDocumentCount);
1404             }
1405         }
1406     }
1407 
1408     /** Estimates the storage usage info for a specific package. */
1409     @NonNull
getStorageInfoForPackage(@onNull String packageName)1410     public StorageInfo getStorageInfoForPackage(@NonNull String packageName)
1411             throws AppSearchException {
1412         mReadWriteLock.readLock().lock();
1413         try {
1414             throwIfClosedLocked();
1415 
1416             Map<String, Set<String>> packageToDatabases = getPackageToDatabases();
1417             Set<String> databases = packageToDatabases.get(packageName);
1418             if (databases == null) {
1419                 // Package doesn't exist, no storage info to report
1420                 return new StorageInfo.Builder().build();
1421             }
1422 
1423             // Accumulate all the namespaces we're interested in.
1424             Set<String> wantedPrefixedNamespaces = new ArraySet<>();
1425             for (String database : databases) {
1426                 Set<String> prefixedNamespaces =
1427                         mNamespaceMapLocked.get(createPrefix(packageName, database));
1428                 if (prefixedNamespaces != null) {
1429                     wantedPrefixedNamespaces.addAll(prefixedNamespaces);
1430                 }
1431             }
1432             if (wantedPrefixedNamespaces.isEmpty()) {
1433                 return new StorageInfo.Builder().build();
1434             }
1435 
1436             return getStorageInfoForNamespaces(getRawStorageInfoProto(), wantedPrefixedNamespaces);
1437         } finally {
1438             mReadWriteLock.readLock().unlock();
1439         }
1440     }
1441 
1442     /** Estimates the storage usage info for a specific database in a package. */
1443     @NonNull
getStorageInfoForDatabase( @onNull String packageName, @NonNull String databaseName)1444     public StorageInfo getStorageInfoForDatabase(
1445             @NonNull String packageName, @NonNull String databaseName) throws AppSearchException {
1446         mReadWriteLock.readLock().lock();
1447         try {
1448             throwIfClosedLocked();
1449 
1450             Map<String, Set<String>> packageToDatabases = getPackageToDatabases();
1451             Set<String> databases = packageToDatabases.get(packageName);
1452             if (databases == null) {
1453                 // Package doesn't exist, no storage info to report
1454                 return new StorageInfo.Builder().build();
1455             }
1456             if (!databases.contains(databaseName)) {
1457                 // Database doesn't exist, no storage info to report
1458                 return new StorageInfo.Builder().build();
1459             }
1460 
1461             Set<String> wantedPrefixedNamespaces =
1462                     mNamespaceMapLocked.get(createPrefix(packageName, databaseName));
1463             if (wantedPrefixedNamespaces == null || wantedPrefixedNamespaces.isEmpty()) {
1464                 return new StorageInfo.Builder().build();
1465             }
1466 
1467             return getStorageInfoForNamespaces(getRawStorageInfoProto(), wantedPrefixedNamespaces);
1468         } finally {
1469             mReadWriteLock.readLock().unlock();
1470         }
1471     }
1472 
1473     /**
1474      * Returns the native storage info capsuled in {@link StorageInfoResultProto} directly from
1475      * IcingSearchEngine.
1476      */
1477     @NonNull
getRawStorageInfoProto()1478     public StorageInfoProto getRawStorageInfoProto() throws AppSearchException {
1479         mReadWriteLock.readLock().lock();
1480         try {
1481             throwIfClosedLocked();
1482             mLogUtil.piiTrace("getStorageInfo, request");
1483             StorageInfoResultProto storageInfoResult = mIcingSearchEngineLocked.getStorageInfo();
1484             mLogUtil.piiTrace(
1485                     "getStorageInfo, response", storageInfoResult.getStatus(), storageInfoResult);
1486             checkSuccess(storageInfoResult.getStatus());
1487             return storageInfoResult.getStorageInfo();
1488         } finally {
1489             mReadWriteLock.readLock().unlock();
1490         }
1491     }
1492 
1493     /**
1494      * Extracts and returns {@link StorageInfo} from {@link StorageInfoProto} based on prefixed
1495      * namespaces.
1496      */
1497     @NonNull
getStorageInfoForNamespaces( @onNull StorageInfoProto storageInfoProto, @NonNull Set<String> prefixedNamespaces)1498     private static StorageInfo getStorageInfoForNamespaces(
1499             @NonNull StorageInfoProto storageInfoProto, @NonNull Set<String> prefixedNamespaces) {
1500         if (!storageInfoProto.hasDocumentStorageInfo()) {
1501             return new StorageInfo.Builder().build();
1502         }
1503 
1504         long totalStorageSize = storageInfoProto.getTotalStorageSize();
1505         DocumentStorageInfoProto documentStorageInfo = storageInfoProto.getDocumentStorageInfo();
1506         int totalDocuments =
1507                 documentStorageInfo.getNumAliveDocuments()
1508                         + documentStorageInfo.getNumExpiredDocuments();
1509 
1510         if (totalStorageSize == 0 || totalDocuments == 0) {
1511             // Maybe we can exit early and also avoid a divide by 0 error.
1512             return new StorageInfo.Builder().build();
1513         }
1514 
1515         // Accumulate stats across the package's namespaces.
1516         int aliveDocuments = 0;
1517         int expiredDocuments = 0;
1518         int aliveNamespaces = 0;
1519         List<NamespaceStorageInfoProto> namespaceStorageInfos =
1520                 documentStorageInfo.getNamespaceStorageInfoList();
1521         for (int i = 0; i < namespaceStorageInfos.size(); i++) {
1522             NamespaceStorageInfoProto namespaceStorageInfo = namespaceStorageInfos.get(i);
1523             // The namespace from icing lib is already the prefixed format
1524             if (prefixedNamespaces.contains(namespaceStorageInfo.getNamespace())) {
1525                 if (namespaceStorageInfo.getNumAliveDocuments() > 0) {
1526                     aliveNamespaces++;
1527                     aliveDocuments += namespaceStorageInfo.getNumAliveDocuments();
1528                 }
1529                 expiredDocuments += namespaceStorageInfo.getNumExpiredDocuments();
1530             }
1531         }
1532         int namespaceDocuments = aliveDocuments + expiredDocuments;
1533 
1534         // Since we don't have the exact size of all the documents, we do an estimation. Note
1535         // that while the total storage takes into account schema, index, etc. in addition to
1536         // documents, we'll only calculate the percentage based on number of documents a
1537         // client has.
1538         return new StorageInfo.Builder()
1539                 .setSizeBytes((long) (namespaceDocuments * 1.0 / totalDocuments * totalStorageSize))
1540                 .setAliveDocumentsCount(aliveDocuments)
1541                 .setAliveNamespacesCount(aliveNamespaces)
1542                 .build();
1543     }
1544 
1545     /**
1546      * Persists all update/delete requests to the disk.
1547      *
1548      * <p>If the app crashes after a call to PersistToDisk with {@link PersistType.Code#FULL}, Icing
1549      * would be able to fully recover all data written up to this point without a costly recovery
1550      * process.
1551      *
1552      * <p>If the app crashes after a call to PersistToDisk with {@link PersistType.Code#LITE}, Icing
1553      * would trigger a costly recovery process in next initialization. After that, Icing would still
1554      * be able to recover all written data - excepting Usage data. Usage data is only guaranteed to
1555      * be safe after a call to PersistToDisk with {@link PersistType.Code#FULL}
1556      *
1557      * <p>If the app crashes after an update/delete request has been made, but before any call to
1558      * PersistToDisk, then all data in Icing will be lost.
1559      *
1560      * @param persistType the amount of data to persist. {@link PersistType.Code#LITE} will only
1561      *     persist the minimal amount of data to ensure all data can be recovered. {@link
1562      *     PersistType.Code#FULL} will persist all data necessary to prevent data loss without
1563      *     needing data recovery.
1564      * @throws AppSearchException on any error that AppSearch persist data to disk.
1565      */
persistToDisk(@onNull PersistType.Code persistType)1566     public void persistToDisk(@NonNull PersistType.Code persistType) throws AppSearchException {
1567         mReadWriteLock.writeLock().lock();
1568         try {
1569             throwIfClosedLocked();
1570 
1571             mLogUtil.piiTrace("persistToDisk, request", persistType);
1572             PersistToDiskResultProto persistToDiskResultProto =
1573                     mIcingSearchEngineLocked.persistToDisk(persistType);
1574             mLogUtil.piiTrace(
1575                     "persistToDisk, response",
1576                     persistToDiskResultProto.getStatus(),
1577                     persistToDiskResultProto);
1578             checkSuccess(persistToDiskResultProto.getStatus());
1579         } finally {
1580             mReadWriteLock.writeLock().unlock();
1581         }
1582     }
1583 
1584     /**
1585      * Remove all {@link AppSearchSchema}s and {@link GenericDocument}s under the given package.
1586      *
1587      * @param packageName The name of package to be removed.
1588      * @throws AppSearchException if we cannot remove the data.
1589      */
clearPackageData(@onNull String packageName)1590     public void clearPackageData(@NonNull String packageName) throws AppSearchException {
1591         mReadWriteLock.writeLock().lock();
1592         try {
1593             throwIfClosedLocked();
1594             Set<String> existingPackages = getPackageToDatabases().keySet();
1595             if (existingPackages.contains(packageName)) {
1596                 existingPackages.remove(packageName);
1597                 prunePackageData(existingPackages);
1598             }
1599         } finally {
1600             mReadWriteLock.writeLock().unlock();
1601         }
1602     }
1603 
1604     /**
1605      * Remove all {@link AppSearchSchema}s and {@link GenericDocument}s that doesn't belong to any
1606      * of the given installed packages
1607      *
1608      * @param installedPackages The name of all installed package.
1609      * @throws AppSearchException if we cannot remove the data.
1610      */
prunePackageData(@onNull Set<String> installedPackages)1611     public void prunePackageData(@NonNull Set<String> installedPackages) throws AppSearchException {
1612         mReadWriteLock.writeLock().lock();
1613         try {
1614             throwIfClosedLocked();
1615             Map<String, Set<String>> packageToDatabases = getPackageToDatabases();
1616             if (installedPackages.containsAll(packageToDatabases.keySet())) {
1617                 // No package got removed. We are good.
1618                 return;
1619             }
1620 
1621             // Prune schema proto
1622             SchemaProto existingSchema = getSchemaProtoLocked();
1623             SchemaProto.Builder newSchemaBuilder = SchemaProto.newBuilder();
1624             for (int i = 0; i < existingSchema.getTypesCount(); i++) {
1625                 String packageName = getPackageName(existingSchema.getTypes(i).getSchemaType());
1626                 if (installedPackages.contains(packageName)) {
1627                     newSchemaBuilder.addTypes(existingSchema.getTypes(i));
1628                 }
1629             }
1630 
1631             SchemaProto finalSchema = newSchemaBuilder.build();
1632 
1633             // Apply schema, set force override to true to remove all schemas and documents that
1634             // doesn't belong to any of these installed packages.
1635             mLogUtil.piiTrace(
1636                     "clearPackageData.setSchema, request",
1637                     finalSchema.getTypesCount(),
1638                     finalSchema);
1639             SetSchemaResultProto setSchemaResultProto =
1640                     mIcingSearchEngineLocked.setSchema(
1641                             finalSchema, /*ignoreErrorsAndDeleteDocuments=*/ true);
1642             mLogUtil.piiTrace(
1643                     "clearPackageData.setSchema, response",
1644                     setSchemaResultProto.getStatus(),
1645                     setSchemaResultProto);
1646 
1647             // Determine whether it succeeded.
1648             checkSuccess(setSchemaResultProto.getStatus());
1649 
1650             // Prune cached maps
1651             for (Map.Entry<String, Set<String>> entry : packageToDatabases.entrySet()) {
1652                 String packageName = entry.getKey();
1653                 Set<String> databaseNames = entry.getValue();
1654                 if (!installedPackages.contains(packageName) && databaseNames != null) {
1655                     mDocumentCountMapLocked.remove(packageName);
1656                     synchronized (mNextPageTokensLocked) {
1657                         mNextPageTokensLocked.remove(packageName);
1658                     }
1659                     for (String databaseName : databaseNames) {
1660                         String removedPrefix = createPrefix(packageName, databaseName);
1661                         mSchemaMapLocked.remove(removedPrefix);
1662                         mNamespaceMapLocked.remove(removedPrefix);
1663                     }
1664                 }
1665             }
1666             // TODO(b/145759910) clear visibility setting for package.
1667         } finally {
1668             mReadWriteLock.writeLock().unlock();
1669         }
1670     }
1671 
1672     /**
1673      * Clears documents and schema across all packages and databaseNames.
1674      *
1675      * <p>This method belongs to mutate group.
1676      *
1677      * @throws AppSearchException on IcingSearchEngine error.
1678      */
1679     @GuardedBy("mReadWriteLock")
resetLocked(@ullable InitializeStats.Builder initStatsBuilder)1680     private void resetLocked(@Nullable InitializeStats.Builder initStatsBuilder)
1681             throws AppSearchException {
1682         mLogUtil.piiTrace("icingSearchEngine.reset, request");
1683         ResetResultProto resetResultProto = mIcingSearchEngineLocked.reset();
1684         mLogUtil.piiTrace(
1685                 "icingSearchEngine.reset, response",
1686                 resetResultProto.getStatus(),
1687                 resetResultProto);
1688         mOptimizeIntervalCountLocked = 0;
1689         mSchemaMapLocked.clear();
1690         mNamespaceMapLocked.clear();
1691         mDocumentCountMapLocked.clear();
1692         synchronized (mNextPageTokensLocked) {
1693             mNextPageTokensLocked.clear();
1694         }
1695         if (initStatsBuilder != null) {
1696             initStatsBuilder
1697                     .setHasReset(true)
1698                     .setResetStatusCode(statusProtoToResultCode(resetResultProto.getStatus()));
1699         }
1700 
1701         checkSuccess(resetResultProto.getStatus());
1702     }
1703 
1704     @GuardedBy("mReadWriteLock")
rebuildDocumentCountMapLocked(@onNull StorageInfoProto storageInfoProto)1705     private void rebuildDocumentCountMapLocked(@NonNull StorageInfoProto storageInfoProto) {
1706         mDocumentCountMapLocked.clear();
1707         List<NamespaceStorageInfoProto> namespaceStorageInfoProtoList =
1708                 storageInfoProto.getDocumentStorageInfo().getNamespaceStorageInfoList();
1709         for (int i = 0; i < namespaceStorageInfoProtoList.size(); i++) {
1710             NamespaceStorageInfoProto namespaceStorageInfoProto =
1711                     namespaceStorageInfoProtoList.get(i);
1712             String packageName = getPackageName(namespaceStorageInfoProto.getNamespace());
1713             Integer oldCount = mDocumentCountMapLocked.get(packageName);
1714             int newCount;
1715             if (oldCount == null) {
1716                 newCount = namespaceStorageInfoProto.getNumAliveDocuments();
1717             } else {
1718                 newCount = oldCount + namespaceStorageInfoProto.getNumAliveDocuments();
1719             }
1720             mDocumentCountMapLocked.put(packageName, newCount);
1721         }
1722     }
1723 
1724     /** Wrapper around schema changes */
1725     @VisibleForTesting
1726     static class RewrittenSchemaResults {
1727         // Any prefixed types that used to exist in the schema, but are deleted in the new one.
1728         final Set<String> mDeletedPrefixedTypes = new ArraySet<>();
1729 
1730         // Map of prefixed schema types to SchemaTypeConfigProtos that were part of the new schema.
1731         final Map<String, SchemaTypeConfigProto> mRewrittenPrefixedTypes = new ArrayMap<>();
1732     }
1733 
1734     /**
1735      * Rewrites all types mentioned in the given {@code newSchema} to prepend {@code prefix}.
1736      * Rewritten types will be added to the {@code existingSchema}.
1737      *
1738      * @param prefix The full prefix to prepend to the schema.
1739      * @param existingSchema A schema that may contain existing types from across all prefixes. Will
1740      *     be mutated to contain the properly rewritten schema types from {@code newSchema}.
1741      * @param newSchema Schema with types to add to the {@code existingSchema}.
1742      * @return a RewrittenSchemaResults that contains all prefixed schema type names in the given
1743      *     prefix as well as a set of schema types that were deleted.
1744      */
1745     @VisibleForTesting
rewriteSchema( @onNull String prefix, @NonNull SchemaProto.Builder existingSchema, @NonNull SchemaProto newSchema)1746     static RewrittenSchemaResults rewriteSchema(
1747             @NonNull String prefix,
1748             @NonNull SchemaProto.Builder existingSchema,
1749             @NonNull SchemaProto newSchema)
1750             throws AppSearchException {
1751         HashMap<String, SchemaTypeConfigProto> newTypesToProto = new HashMap<>();
1752         // Rewrite the schema type to include the typePrefix.
1753         for (int typeIdx = 0; typeIdx < newSchema.getTypesCount(); typeIdx++) {
1754             SchemaTypeConfigProto.Builder typeConfigBuilder =
1755                     newSchema.getTypes(typeIdx).toBuilder();
1756 
1757             // Rewrite SchemaProto.types.schema_type
1758             String newSchemaType = prefix + typeConfigBuilder.getSchemaType();
1759             typeConfigBuilder.setSchemaType(newSchemaType);
1760 
1761             // Rewrite SchemaProto.types.properties.schema_type
1762             for (int propertyIdx = 0;
1763                     propertyIdx < typeConfigBuilder.getPropertiesCount();
1764                     propertyIdx++) {
1765                 PropertyConfigProto.Builder propertyConfigBuilder =
1766                         typeConfigBuilder.getProperties(propertyIdx).toBuilder();
1767                 if (!propertyConfigBuilder.getSchemaType().isEmpty()) {
1768                     String newPropertySchemaType = prefix + propertyConfigBuilder.getSchemaType();
1769                     propertyConfigBuilder.setSchemaType(newPropertySchemaType);
1770                     typeConfigBuilder.setProperties(propertyIdx, propertyConfigBuilder);
1771                 }
1772             }
1773 
1774             newTypesToProto.put(newSchemaType, typeConfigBuilder.build());
1775         }
1776 
1777         // newTypesToProto is modified below, so we need a copy first
1778         RewrittenSchemaResults rewrittenSchemaResults = new RewrittenSchemaResults();
1779         rewrittenSchemaResults.mRewrittenPrefixedTypes.putAll(newTypesToProto);
1780 
1781         // Combine the existing schema (which may have types from other prefixes) with this
1782         // prefix's new schema. Modifies the existingSchemaBuilder.
1783         // Check if we need to replace any old schema types with the new ones.
1784         for (int i = 0; i < existingSchema.getTypesCount(); i++) {
1785             String schemaType = existingSchema.getTypes(i).getSchemaType();
1786             SchemaTypeConfigProto newProto = newTypesToProto.remove(schemaType);
1787             if (newProto != null) {
1788                 // Replacement
1789                 existingSchema.setTypes(i, newProto);
1790             } else if (prefix.equals(getPrefix(schemaType))) {
1791                 // All types existing before but not in newSchema should be removed.
1792                 existingSchema.removeTypes(i);
1793                 --i;
1794                 rewrittenSchemaResults.mDeletedPrefixedTypes.add(schemaType);
1795             }
1796         }
1797         // We've been removing existing types from newTypesToProto, so everything that remains is
1798         // new.
1799         existingSchema.addAllTypes(newTypesToProto.values());
1800 
1801         return rewrittenSchemaResults;
1802     }
1803 
1804     /**
1805      * Rewrites the search spec filters with {@code prefixes}.
1806      *
1807      * <p>This method should be only called in query methods and get the READ lock to keep thread
1808      * safety.
1809      *
1810      * @param searchSpecBuilder Client-provided SearchSpec
1811      * @param prefixes Prefixes that we should prepend to all our filters
1812      * @param allowedPrefixedSchemas Prefixed schemas that the client is allowed to query over. This
1813      *     supersedes the schema filters that may exist on the {@code searchSpecBuilder}.
1814      * @return false if none there would be nothing to search over.
1815      */
1816     @VisibleForTesting
1817     @GuardedBy("mReadWriteLock")
rewriteSearchSpecForPrefixesLocked( @onNull SearchSpecProto.Builder searchSpecBuilder, @NonNull Set<String> prefixes, @NonNull Set<String> allowedPrefixedSchemas)1818     boolean rewriteSearchSpecForPrefixesLocked(
1819             @NonNull SearchSpecProto.Builder searchSpecBuilder,
1820             @NonNull Set<String> prefixes,
1821             @NonNull Set<String> allowedPrefixedSchemas) {
1822         // Create a copy since retainAll() modifies the original set.
1823         Set<String> existingPrefixes = new ArraySet<>(mNamespaceMapLocked.keySet());
1824         existingPrefixes.retainAll(prefixes);
1825 
1826         if (existingPrefixes.isEmpty()) {
1827             // None of the prefixes exist, empty query.
1828             return false;
1829         }
1830 
1831         if (allowedPrefixedSchemas.isEmpty()) {
1832             // Not allowed to search over any schemas, empty query.
1833             return false;
1834         }
1835 
1836         // Clear the schema type filters since we'll be rewriting them with the
1837         // allowedPrefixedSchemas.
1838         searchSpecBuilder.clearSchemaTypeFilters();
1839         searchSpecBuilder.addAllSchemaTypeFilters(allowedPrefixedSchemas);
1840 
1841         // Cache the namespaces before clearing everything.
1842         List<String> namespaceFilters = searchSpecBuilder.getNamespaceFiltersList();
1843         searchSpecBuilder.clearNamespaceFilters();
1844 
1845         // Rewrite non-schema filters to include a prefix.
1846         for (String prefix : existingPrefixes) {
1847             // TODO(b/169883602): We currently grab every namespace for every prefix. We can
1848             //  optimize this by checking if a prefix has any allowedSchemaTypes. If not, that
1849             //  means we don't want to query over anything in that prefix anyways, so we don't
1850             //  need to grab its namespaces either.
1851 
1852             // Empty namespaces on the search spec means to query over all namespaces.
1853             Set<String> existingNamespaces = mNamespaceMapLocked.get(prefix);
1854             if (existingNamespaces != null) {
1855                 if (namespaceFilters.isEmpty()) {
1856                     // Include all namespaces
1857                     searchSpecBuilder.addAllNamespaceFilters(existingNamespaces);
1858                 } else {
1859                     // Prefix the given namespaces.
1860                     for (int i = 0; i < namespaceFilters.size(); i++) {
1861                         String prefixedNamespace = prefix + namespaceFilters.get(i);
1862                         if (existingNamespaces.contains(prefixedNamespace)) {
1863                             searchSpecBuilder.addNamespaceFilters(prefixedNamespace);
1864                         }
1865                     }
1866                 }
1867             }
1868         }
1869 
1870         return true;
1871     }
1872 
1873     /**
1874      * Returns the set of allowed prefixed schemas that the {@code prefix} can query while taking
1875      * into account the {@code searchSpec} schema filters.
1876      *
1877      * <p>This only checks intersection of schema filters on the search spec with those that the
1878      * prefix owns itself. This does not check global query permissions.
1879      */
1880     @GuardedBy("mReadWriteLock")
getAllowedPrefixSchemasLocked( @onNull String prefix, @NonNull SearchSpec searchSpec)1881     private Set<String> getAllowedPrefixSchemasLocked(
1882             @NonNull String prefix, @NonNull SearchSpec searchSpec) {
1883         Set<String> allowedPrefixedSchemas = new ArraySet<>();
1884 
1885         // Add all the schema filters the client specified.
1886         List<String> schemaFilters = searchSpec.getFilterSchemas();
1887         for (int i = 0; i < schemaFilters.size(); i++) {
1888             allowedPrefixedSchemas.add(prefix + schemaFilters.get(i));
1889         }
1890 
1891         if (allowedPrefixedSchemas.isEmpty()) {
1892             // If the client didn't specify any schema filters, search over all of their schemas
1893             Map<String, SchemaTypeConfigProto> prefixedSchemaMap = mSchemaMapLocked.get(prefix);
1894             if (prefixedSchemaMap != null) {
1895                 allowedPrefixedSchemas.addAll(prefixedSchemaMap.keySet());
1896             }
1897         }
1898         return allowedPrefixedSchemas;
1899     }
1900 
1901     /**
1902      * Rewrites the typePropertyMasks that exist in {@code prefixes}.
1903      *
1904      * <p>This method should be only called in query methods and get the READ lock to keep thread
1905      * safety.
1906      *
1907      * @param resultSpecBuilder ResultSpecs as specified by client
1908      * @param prefixes Prefixes that we should prepend to all our filters
1909      * @param allowedPrefixedSchemas Prefixed schemas that the client is allowed to query over.
1910      */
1911     @VisibleForTesting
1912     @GuardedBy("mReadWriteLock")
rewriteResultSpecForPrefixesLocked( @onNull ResultSpecProto.Builder resultSpecBuilder, @NonNull Set<String> prefixes, @NonNull Set<String> allowedPrefixedSchemas)1913     void rewriteResultSpecForPrefixesLocked(
1914             @NonNull ResultSpecProto.Builder resultSpecBuilder,
1915             @NonNull Set<String> prefixes,
1916             @NonNull Set<String> allowedPrefixedSchemas) {
1917         // Create a copy since retainAll() modifies the original set.
1918         Set<String> existingPrefixes = new ArraySet<>(mNamespaceMapLocked.keySet());
1919         existingPrefixes.retainAll(prefixes);
1920 
1921         List<TypePropertyMask> prefixedTypePropertyMasks = new ArrayList<>();
1922         // Rewrite filters to include a database prefix.
1923         for (String prefix : existingPrefixes) {
1924             // Qualify the given schema types
1925             for (TypePropertyMask typePropertyMask : resultSpecBuilder.getTypePropertyMasksList()) {
1926                 String unprefixedType = typePropertyMask.getSchemaType();
1927                 boolean isWildcard =
1928                         unprefixedType.equals(SearchSpec.PROJECTION_SCHEMA_TYPE_WILDCARD);
1929                 String prefixedType = isWildcard ? unprefixedType : prefix + unprefixedType;
1930                 if (isWildcard || allowedPrefixedSchemas.contains(prefixedType)) {
1931                     prefixedTypePropertyMasks.add(
1932                             typePropertyMask.toBuilder().setSchemaType(prefixedType).build());
1933                 }
1934             }
1935         }
1936         resultSpecBuilder
1937                 .clearTypePropertyMasks()
1938                 .addAllTypePropertyMasks(prefixedTypePropertyMasks);
1939     }
1940 
1941     /**
1942      * Adds result groupings for each namespace in each package being queried for.
1943      *
1944      * <p>This method should be only called in query methods and get the READ lock to keep thread
1945      * safety.
1946      *
1947      * @param resultSpecBuilder ResultSpecs as specified by client
1948      * @param prefixes Prefixes that we should prepend to all our filters
1949      * @param maxNumResults The maximum number of results for each grouping to support.
1950      */
1951     @GuardedBy("mReadWriteLock")
addPerPackagePerNamespaceResultGroupingsLocked( @onNull ResultSpecProto.Builder resultSpecBuilder, @NonNull Set<String> prefixes, int maxNumResults)1952     private void addPerPackagePerNamespaceResultGroupingsLocked(
1953             @NonNull ResultSpecProto.Builder resultSpecBuilder,
1954             @NonNull Set<String> prefixes,
1955             int maxNumResults) {
1956         Set<String> existingPrefixes = new ArraySet<>(mNamespaceMapLocked.keySet());
1957         existingPrefixes.retainAll(prefixes);
1958 
1959         // Create a map for package+namespace to prefixedNamespaces. This is NOT necessarily the
1960         // same as the list of namespaces. If one package has multiple databases, each with the same
1961         // namespace, then those should be grouped together.
1962         Map<String, List<String>> packageAndNamespaceToNamespaces = new ArrayMap<>();
1963         for (String prefix : existingPrefixes) {
1964             Set<String> prefixedNamespaces = mNamespaceMapLocked.get(prefix);
1965             if (prefixedNamespaces == null) {
1966                 continue;
1967             }
1968             String packageName = getPackageName(prefix);
1969             // Create a new prefix without the database name. This will allow us to group namespaces
1970             // that have the same name and package but a different database name together.
1971             String emptyDatabasePrefix = createPrefix(packageName, /*databaseName*/ "");
1972             for (String prefixedNamespace : prefixedNamespaces) {
1973                 String namespace;
1974                 try {
1975                     namespace = removePrefix(prefixedNamespace);
1976                 } catch (AppSearchException e) {
1977                     // This should never happen. Skip this namespace if it does.
1978                     Log.e(TAG, "Prefixed namespace " + prefixedNamespace + " is malformed.");
1979                     continue;
1980                 }
1981                 String emptyDatabasePrefixedNamespace = emptyDatabasePrefix + namespace;
1982                 List<String> namespaceList =
1983                         packageAndNamespaceToNamespaces.get(emptyDatabasePrefixedNamespace);
1984                 if (namespaceList == null) {
1985                     namespaceList = new ArrayList<>();
1986                     packageAndNamespaceToNamespaces.put(
1987                             emptyDatabasePrefixedNamespace, namespaceList);
1988                 }
1989                 namespaceList.add(prefixedNamespace);
1990             }
1991         }
1992 
1993         for (List<String> namespaces : packageAndNamespaceToNamespaces.values()) {
1994             resultSpecBuilder.addResultGroupings(
1995                     ResultSpecProto.ResultGrouping.newBuilder()
1996                             .addAllNamespaces(namespaces)
1997                             .setMaxResults(maxNumResults));
1998         }
1999     }
2000 
2001     /**
2002      * Adds result groupings for each package being queried for.
2003      *
2004      * <p>This method should be only called in query methods and get the READ lock to keep thread
2005      * safety.
2006      *
2007      * @param resultSpecBuilder ResultSpecs as specified by client
2008      * @param prefixes Prefixes that we should prepend to all our filters
2009      * @param maxNumResults The maximum number of results for each grouping to support.
2010      */
2011     @GuardedBy("mReadWriteLock")
addPerPackageResultGroupingsLocked( @onNull ResultSpecProto.Builder resultSpecBuilder, @NonNull Set<String> prefixes, int maxNumResults)2012     private void addPerPackageResultGroupingsLocked(
2013             @NonNull ResultSpecProto.Builder resultSpecBuilder,
2014             @NonNull Set<String> prefixes,
2015             int maxNumResults) {
2016         Set<String> existingPrefixes = new ArraySet<>(mNamespaceMapLocked.keySet());
2017         existingPrefixes.retainAll(prefixes);
2018 
2019         // Build up a map of package to namespaces.
2020         Map<String, List<String>> packageToNamespacesMap = new ArrayMap<>();
2021         for (String prefix : existingPrefixes) {
2022             Set<String> prefixedNamespaces = mNamespaceMapLocked.get(prefix);
2023             if (prefixedNamespaces == null) {
2024                 continue;
2025             }
2026             String packageName = getPackageName(prefix);
2027             List<String> packageNamespaceList = packageToNamespacesMap.get(packageName);
2028             if (packageNamespaceList == null) {
2029                 packageNamespaceList = new ArrayList<>();
2030                 packageToNamespacesMap.put(packageName, packageNamespaceList);
2031             }
2032             packageNamespaceList.addAll(prefixedNamespaces);
2033         }
2034 
2035         for (List<String> prefixedNamespaces : packageToNamespacesMap.values()) {
2036             resultSpecBuilder.addResultGroupings(
2037                     ResultSpecProto.ResultGrouping.newBuilder()
2038                             .addAllNamespaces(prefixedNamespaces)
2039                             .setMaxResults(maxNumResults));
2040         }
2041     }
2042 
2043     /**
2044      * Adds result groupings for each namespace being queried for.
2045      *
2046      * <p>This method should be only called in query methods and get the READ lock to keep thread
2047      * safety.
2048      *
2049      * @param resultSpecBuilder ResultSpecs as specified by client
2050      * @param prefixes Prefixes that we should prepend to all our filters
2051      * @param maxNumResults The maximum number of results for each grouping to support.
2052      */
2053     @GuardedBy("mReadWriteLock")
addPerNamespaceResultGroupingsLocked( @onNull ResultSpecProto.Builder resultSpecBuilder, @NonNull Set<String> prefixes, int maxNumResults)2054     private void addPerNamespaceResultGroupingsLocked(
2055             @NonNull ResultSpecProto.Builder resultSpecBuilder,
2056             @NonNull Set<String> prefixes,
2057             int maxNumResults) {
2058         Set<String> existingPrefixes = new ArraySet<>(mNamespaceMapLocked.keySet());
2059         existingPrefixes.retainAll(prefixes);
2060 
2061         // Create a map of namespace to prefixedNamespaces. This is NOT necessarily the
2062         // same as the list of namespaces. If a namespace exists under different packages and/or
2063         // different databases, they should still be grouped together.
2064         Map<String, List<String>> namespaceToPrefixedNamespaces = new ArrayMap<>();
2065         for (String prefix : existingPrefixes) {
2066             Set<String> prefixedNamespaces = mNamespaceMapLocked.get(prefix);
2067             if (prefixedNamespaces == null) {
2068                 continue;
2069             }
2070             for (String prefixedNamespace : prefixedNamespaces) {
2071                 String namespace;
2072                 try {
2073                     namespace = removePrefix(prefixedNamespace);
2074                 } catch (AppSearchException e) {
2075                     // This should never happen. Skip this namespace if it does.
2076                     Log.e(TAG, "Prefixed namespace " + prefixedNamespace + " is malformed.");
2077                     continue;
2078                 }
2079                 List<String> groupedPrefixedNamespaces =
2080                         namespaceToPrefixedNamespaces.get(namespace);
2081                 if (groupedPrefixedNamespaces == null) {
2082                     groupedPrefixedNamespaces = new ArrayList<>();
2083                     namespaceToPrefixedNamespaces.put(namespace, groupedPrefixedNamespaces);
2084                 }
2085                 groupedPrefixedNamespaces.add(prefixedNamespace);
2086             }
2087         }
2088 
2089         for (List<String> namespaces : namespaceToPrefixedNamespaces.values()) {
2090             resultSpecBuilder.addResultGroupings(
2091                     ResultSpecProto.ResultGrouping.newBuilder()
2092                             .addAllNamespaces(namespaces)
2093                             .setMaxResults(maxNumResults));
2094         }
2095     }
2096 
2097     @VisibleForTesting
2098     @GuardedBy("mReadWriteLock")
getSchemaProtoLocked()2099     SchemaProto getSchemaProtoLocked() throws AppSearchException {
2100         mLogUtil.piiTrace("getSchema, request");
2101         GetSchemaResultProto schemaProto = mIcingSearchEngineLocked.getSchema();
2102         mLogUtil.piiTrace("getSchema, response", schemaProto.getStatus(), schemaProto);
2103         // TODO(b/161935693) check GetSchemaResultProto is success or not. Call reset() if it's not.
2104         // TODO(b/161935693) only allow GetSchemaResultProto NOT_FOUND on first run
2105         checkCodeOneOf(schemaProto.getStatus(), StatusProto.Code.OK, StatusProto.Code.NOT_FOUND);
2106         return schemaProto.getSchema();
2107     }
2108 
addNextPageToken(String packageName, long nextPageToken)2109     private void addNextPageToken(String packageName, long nextPageToken) {
2110         if (nextPageToken == EMPTY_PAGE_TOKEN) {
2111             // There is no more pages. No need to add it.
2112             return;
2113         }
2114         synchronized (mNextPageTokensLocked) {
2115             Set<Long> tokens = mNextPageTokensLocked.get(packageName);
2116             if (tokens == null) {
2117                 tokens = new ArraySet<>();
2118                 mNextPageTokensLocked.put(packageName, tokens);
2119             }
2120             tokens.add(nextPageToken);
2121         }
2122     }
2123 
checkNextPageToken(String packageName, long nextPageToken)2124     private void checkNextPageToken(String packageName, long nextPageToken)
2125             throws AppSearchException {
2126         if (nextPageToken == EMPTY_PAGE_TOKEN) {
2127             // Swallow the check for empty page token, token = 0 means there is no more page and it
2128             // won't return anything from Icing.
2129             return;
2130         }
2131         synchronized (mNextPageTokensLocked) {
2132             Set<Long> nextPageTokens = mNextPageTokensLocked.get(packageName);
2133             if (nextPageTokens == null || !nextPageTokens.contains(nextPageToken)) {
2134                 throw new AppSearchException(
2135                         AppSearchResult.RESULT_SECURITY_ERROR,
2136                         "Package \""
2137                                 + packageName
2138                                 + "\" cannot use nextPageToken: "
2139                                 + nextPageToken);
2140             }
2141         }
2142     }
2143 
addToMap( Map<String, Set<String>> map, String prefix, String prefixedValue)2144     private static void addToMap(
2145             Map<String, Set<String>> map, String prefix, String prefixedValue) {
2146         Set<String> values = map.get(prefix);
2147         if (values == null) {
2148             values = new ArraySet<>();
2149             map.put(prefix, values);
2150         }
2151         values.add(prefixedValue);
2152     }
2153 
addToMap( Map<String, Map<String, SchemaTypeConfigProto>> map, String prefix, SchemaTypeConfigProto schemaTypeConfigProto)2154     private static void addToMap(
2155             Map<String, Map<String, SchemaTypeConfigProto>> map,
2156             String prefix,
2157             SchemaTypeConfigProto schemaTypeConfigProto) {
2158         Map<String, SchemaTypeConfigProto> schemaTypeMap = map.get(prefix);
2159         if (schemaTypeMap == null) {
2160             schemaTypeMap = new ArrayMap<>();
2161             map.put(prefix, schemaTypeMap);
2162         }
2163         schemaTypeMap.put(schemaTypeConfigProto.getSchemaType(), schemaTypeConfigProto);
2164     }
2165 
removeFromMap( Map<String, Map<String, SchemaTypeConfigProto>> map, String prefix, String schemaType)2166     private static void removeFromMap(
2167             Map<String, Map<String, SchemaTypeConfigProto>> map, String prefix, String schemaType) {
2168         Map<String, SchemaTypeConfigProto> schemaTypeMap = map.get(prefix);
2169         if (schemaTypeMap != null) {
2170             schemaTypeMap.remove(schemaType);
2171         }
2172     }
2173 
2174     /**
2175      * Checks the given status code and throws an {@link AppSearchException} if code is an error.
2176      *
2177      * @throws AppSearchException on error codes.
2178      */
checkSuccess(StatusProto statusProto)2179     private static void checkSuccess(StatusProto statusProto) throws AppSearchException {
2180         checkCodeOneOf(statusProto, StatusProto.Code.OK);
2181     }
2182 
2183     /**
2184      * Checks the given status code is one of the provided codes, and throws an {@link
2185      * AppSearchException} if it is not.
2186      */
checkCodeOneOf(StatusProto statusProto, StatusProto.Code... codes)2187     private static void checkCodeOneOf(StatusProto statusProto, StatusProto.Code... codes)
2188             throws AppSearchException {
2189         for (int i = 0; i < codes.length; i++) {
2190             if (codes[i] == statusProto.getCode()) {
2191                 // Everything's good
2192                 return;
2193             }
2194         }
2195 
2196         if (statusProto.getCode() == StatusProto.Code.WARNING_DATA_LOSS) {
2197             // TODO: May want to propagate WARNING_DATA_LOSS up to AppSearchSession so they can
2198             //  choose to log the error or potentially pass it on to clients.
2199             Log.w(TAG, "Encountered WARNING_DATA_LOSS: " + statusProto.getMessage());
2200             return;
2201         }
2202 
2203         throw new AppSearchException(
2204                 ResultCodeToProtoConverter.toResultCode(statusProto.getCode()),
2205                 statusProto.getMessage());
2206     }
2207 
2208     /**
2209      * Checks whether {@link IcingSearchEngine#optimize()} should be called to release resources.
2210      *
2211      * <p>This method should be only called after a mutation to local storage backend which deletes
2212      * a mass of data and could release lots resources after {@link IcingSearchEngine#optimize()}.
2213      *
2214      * <p>This method will trigger {@link IcingSearchEngine#getOptimizeInfo()} to check resources
2215      * that could be released for every {@link #CHECK_OPTIMIZE_INTERVAL} mutations.
2216      *
2217      * <p>{@link IcingSearchEngine#optimize()} should be called only if {@link
2218      * GetOptimizeInfoResultProto} shows there is enough resources could be released.
2219      *
2220      * @param mutationSize The number of how many mutations have been executed for current request.
2221      *     An inside counter will accumulates it. Once the counter reaches {@link
2222      *     #CHECK_OPTIMIZE_INTERVAL}, {@link IcingSearchEngine#getOptimizeInfo()} will be triggered
2223      *     and the counter will be reset.
2224      */
checkForOptimize(int mutationSize, @Nullable OptimizeStats.Builder builder)2225     public void checkForOptimize(int mutationSize, @Nullable OptimizeStats.Builder builder)
2226             throws AppSearchException {
2227         mReadWriteLock.writeLock().lock();
2228         try {
2229             mOptimizeIntervalCountLocked += mutationSize;
2230             if (mOptimizeIntervalCountLocked >= CHECK_OPTIMIZE_INTERVAL) {
2231                 checkForOptimize(builder);
2232             }
2233         } finally {
2234             mReadWriteLock.writeLock().unlock();
2235         }
2236     }
2237 
2238     /**
2239      * Checks whether {@link IcingSearchEngine#optimize()} should be called to release resources.
2240      *
2241      * <p>This method will directly trigger {@link IcingSearchEngine#getOptimizeInfo()} to check
2242      * resources that could be released.
2243      *
2244      * <p>{@link IcingSearchEngine#optimize()} should be called only if {@link
2245      * OptimizeStrategy#shouldOptimize(GetOptimizeInfoResultProto)} return true.
2246      */
checkForOptimize(@ullable OptimizeStats.Builder builder)2247     public void checkForOptimize(@Nullable OptimizeStats.Builder builder)
2248             throws AppSearchException {
2249         mReadWriteLock.writeLock().lock();
2250         try {
2251             GetOptimizeInfoResultProto optimizeInfo = getOptimizeInfoResultLocked();
2252             checkSuccess(optimizeInfo.getStatus());
2253             mOptimizeIntervalCountLocked = 0;
2254             if (mOptimizeStrategy.shouldOptimize(optimizeInfo)) {
2255                 optimize(builder);
2256             }
2257         } finally {
2258             mReadWriteLock.writeLock().unlock();
2259         }
2260         // TODO(b/147699081): Return OptimizeResultProto & log lost data detail once we add
2261         //  a field to indicate lost_schema and lost_documents in OptimizeResultProto.
2262         //  go/icing-library-apis.
2263     }
2264 
2265     /** Triggers {@link IcingSearchEngine#optimize()} directly. */
optimize(@ullable OptimizeStats.Builder builder)2266     public void optimize(@Nullable OptimizeStats.Builder builder) throws AppSearchException {
2267         mReadWriteLock.writeLock().lock();
2268         try {
2269             mLogUtil.piiTrace("optimize, request");
2270             OptimizeResultProto optimizeResultProto = mIcingSearchEngineLocked.optimize();
2271             mLogUtil.piiTrace(
2272                     "optimize, response", optimizeResultProto.getStatus(), optimizeResultProto);
2273             if (builder != null) {
2274                 builder.setStatusCode(statusProtoToResultCode(optimizeResultProto.getStatus()));
2275                 AppSearchLoggerHelper.copyNativeStats(
2276                         optimizeResultProto.getOptimizeStats(), builder);
2277             }
2278             checkSuccess(optimizeResultProto.getStatus());
2279         } finally {
2280             mReadWriteLock.writeLock().unlock();
2281         }
2282     }
2283 
2284     /** Remove the rewritten schema types from any result documents. */
2285     @NonNull
2286     @VisibleForTesting
rewriteSearchResultProto( @onNull SearchResultProto searchResultProto, @NonNull Map<String, Map<String, SchemaTypeConfigProto>> schemaMap)2287     static SearchResultPage rewriteSearchResultProto(
2288             @NonNull SearchResultProto searchResultProto,
2289             @NonNull Map<String, Map<String, SchemaTypeConfigProto>> schemaMap)
2290             throws AppSearchException {
2291         // Parallel array of package names for each document search result.
2292         List<String> packageNames = new ArrayList<>(searchResultProto.getResultsCount());
2293 
2294         // Parallel array of database names for each document search result.
2295         List<String> databaseNames = new ArrayList<>(searchResultProto.getResultsCount());
2296 
2297         SearchResultProto.Builder resultsBuilder = searchResultProto.toBuilder();
2298         for (int i = 0; i < searchResultProto.getResultsCount(); i++) {
2299             SearchResultProto.ResultProto.Builder resultBuilder =
2300                     searchResultProto.getResults(i).toBuilder();
2301             DocumentProto.Builder documentBuilder = resultBuilder.getDocument().toBuilder();
2302             String prefix = removePrefixesFromDocument(documentBuilder);
2303             packageNames.add(getPackageName(prefix));
2304             databaseNames.add(getDatabaseName(prefix));
2305             resultBuilder.setDocument(documentBuilder);
2306             resultsBuilder.setResults(i, resultBuilder);
2307         }
2308         return SearchResultToProtoConverter.toSearchResultPage(
2309                 resultsBuilder, packageNames, databaseNames, schemaMap);
2310     }
2311 
2312     @GuardedBy("mReadWriteLock")
2313     @VisibleForTesting
getOptimizeInfoResultLocked()2314     GetOptimizeInfoResultProto getOptimizeInfoResultLocked() {
2315         mLogUtil.piiTrace("getOptimizeInfo, request");
2316         GetOptimizeInfoResultProto result = mIcingSearchEngineLocked.getOptimizeInfo();
2317         mLogUtil.piiTrace("getOptimizeInfo, response", result.getStatus(), result);
2318         return result;
2319     }
2320 
2321     /**
2322      * Converts an erroneous status code from the Icing status enums to the AppSearchResult enums.
2323      *
2324      * <p>Callers should ensure that the status code is not OK or WARNING_DATA_LOSS.
2325      *
2326      * @param statusProto StatusProto with error code to translate into an {@link AppSearchResult}
2327      *     code.
2328      * @return {@link AppSearchResult} error code
2329      */
statusProtoToResultCode( @onNull StatusProto statusProto)2330     private static @AppSearchResult.ResultCode int statusProtoToResultCode(
2331             @NonNull StatusProto statusProto) {
2332         return ResultCodeToProtoConverter.toResultCode(statusProto.getCode());
2333     }
2334 }
2335