Merge "Fix validation crash for nextPageToken." into sc-dev
diff --git a/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java b/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
index 18e1ef5..15916cc 100644
--- a/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
+++ b/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
@@ -146,6 +146,9 @@
 public final class AppSearchImpl implements Closeable {
     private static final String TAG = "AppSearchImpl";
 
+    /** A value 0 means that there're no more pages in the search results. */
+    private static final long EMPTY_PAGE_TOKEN = 0;
+
     @VisibleForTesting static final int CHECK_OPTIMIZE_INTERVAL = 100;
 
     private final ReadWriteLock mReadWriteLock = new ReentrantReadWriteLock();
@@ -1139,6 +1142,16 @@
                     searchResultProto.getResultsCount(),
                     searchResultProto);
             checkSuccess(searchResultProto.getStatus());
+            if (nextPageToken != EMPTY_PAGE_TOKEN
+                    && searchResultProto.getNextPageToken() == EMPTY_PAGE_TOKEN) {
+                // At this point, we're guaranteed that this nextPageToken exists for this package,
+                // otherwise checkNextPageToken would've thrown an exception.
+                // Since the new token is 0, this is the last page. We should remove the old token
+                // from our cache since it no longer refers to this query.
+                synchronized (mNextPageTokensLocked) {
+                    mNextPageTokensLocked.get(packageName).remove(nextPageToken);
+                }
+            }
             return rewriteSearchResultProto(searchResultProto, mSchemaMapLocked);
         } finally {
             mReadWriteLock.readLock().unlock();
@@ -2059,6 +2072,10 @@
     }
 
     private void addNextPageToken(String packageName, long nextPageToken) {
+        if (nextPageToken == EMPTY_PAGE_TOKEN) {
+            // There is no more pages. No need to add it.
+            return;
+        }
         synchronized (mNextPageTokensLocked) {
             Set<Long> tokens = mNextPageTokensLocked.get(packageName);
             if (tokens == null) {
@@ -2071,6 +2088,11 @@
 
     private void checkNextPageToken(String packageName, long nextPageToken)
             throws AppSearchException {
+        if (nextPageToken == EMPTY_PAGE_TOKEN) {
+            // Swallow the check for empty page token, token = 0 means there is no more page and it
+            // won't return anything from Icing.
+            return;
+        }
         synchronized (mNextPageTokensLocked) {
             Set<Long> nextPageTokens = mNextPageTokensLocked.get(packageName);
             if (nextPageTokens == null || !nextPageTokens.contains(nextPageToken)) {
diff --git a/synced_jetpack_changeid.txt b/synced_jetpack_changeid.txt
index 6555107..a81d7d8 100644
--- a/synced_jetpack_changeid.txt
+++ b/synced_jetpack_changeid.txt
@@ -1 +1 @@
-c7387d9b58726a23a0608a9365fb3a1b57b7b8a1
+Ie04f1ecc033faae8085afcb51eb9e40a298998d5
diff --git a/testing/coretests/src/android/app/appsearch/AppSearchSessionUnitTest.java b/testing/coretests/src/android/app/appsearch/AppSearchSessionUnitTest.java
index d51004c..07e4333 100644
--- a/testing/coretests/src/android/app/appsearch/AppSearchSessionUnitTest.java
+++ b/testing/coretests/src/android/app/appsearch/AppSearchSessionUnitTest.java
@@ -16,6 +16,8 @@
 
 package android.app.appsearch;
 
+import static android.app.appsearch.SearchSpec.TERM_MATCH_PREFIX;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.testng.Assert.expectThrows;
@@ -30,6 +32,8 @@
 import org.junit.Before;
 import org.junit.Test;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -100,4 +104,127 @@
                 .isEqualTo(AppSearchResult.RESULT_INTERNAL_ERROR);
         assertThat(appSearchException.getMessage()).startsWith("NullPointerException");
     }
+
+    @Test
+    public void testGetEmptyNextPage() throws Exception {
+        // Set the schema.
+        CompletableFuture<AppSearchResult<SetSchemaResponse>> schemaFuture =
+                new CompletableFuture<>();
+        mSearchSession.setSchema(
+                new SetSchemaRequest.Builder()
+                        .addSchemas(new AppSearchSchema.Builder("schema1").build())
+                        .setForceOverride(true).build(),
+                mExecutor, mExecutor, schemaFuture::complete);
+        schemaFuture.get().getResultValue();
+
+        // Create a document and index it.
+        GenericDocument document1 = new GenericDocument.Builder<>("namespace", "id1",
+                "schema1").build();
+        CompletableFuture<AppSearchBatchResult<String, Void>> putDocumentsFuture =
+                new CompletableFuture<>();
+        mSearchSession.put(
+                new PutDocumentsRequest.Builder().addGenericDocuments(document1).build(),
+                mExecutor, new BatchResultCallback<String, Void>() {
+                    @Override
+                    public void onResult(AppSearchBatchResult<String, Void> result) {
+                        putDocumentsFuture.complete(result);
+                    }
+
+                    @Override
+                    public void onSystemError(Throwable throwable) {
+                        putDocumentsFuture.completeExceptionally(throwable);
+                    }
+                });
+        putDocumentsFuture.get();
+
+        // Search and get the first page.
+        SearchSpec searchSpec = new SearchSpec.Builder()
+                .setTermMatch(TERM_MATCH_PREFIX)
+                .setResultCountPerPage(1)
+                .build();
+        SearchResults searchResults = mSearchSession.search("", searchSpec);
+
+        CompletableFuture<AppSearchResult<List<SearchResult>>> getNextPageFuture =
+                new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        List<SearchResult> results = getNextPageFuture.get().getResultValue();
+        assertThat(results).hasSize(1);
+        assertThat(results.get(0).getGenericDocument()).isEqualTo(document1);
+
+        // We get all documents, and it shouldn't fail if we keep calling getNextPage().
+        getNextPageFuture = new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        results = getNextPageFuture.get().getResultValue();
+        assertThat(results).isEmpty();
+    }
+
+    @Test
+    public void testGetEmptyNextPage_multiPages() throws Exception {
+        // Set the schema.
+        CompletableFuture<AppSearchResult<SetSchemaResponse>> schemaFuture =
+                new CompletableFuture<>();
+        mSearchSession.setSchema(
+                new SetSchemaRequest.Builder()
+                        .addSchemas(new AppSearchSchema.Builder("schema1").build())
+                        .setForceOverride(true).build(),
+                mExecutor, mExecutor, schemaFuture::complete);
+        schemaFuture.get().getResultValue();
+
+        // Create a document and insert 3 package1 documents
+        GenericDocument document1 = new GenericDocument.Builder<>("namespace", "id1",
+                "schema1").build();
+        GenericDocument document2 = new GenericDocument.Builder<>("namespace", "id2",
+                "schema1").build();
+        GenericDocument document3 = new GenericDocument.Builder<>("namespace", "id3",
+                "schema1").build();
+        CompletableFuture<AppSearchBatchResult<String, Void>> putDocumentsFuture =
+                new CompletableFuture<>();
+        mSearchSession.put(
+                new PutDocumentsRequest.Builder()
+                        .addGenericDocuments(document1, document2, document3).build(),
+                mExecutor, new BatchResultCallback<String, Void>() {
+                    @Override
+                    public void onResult(AppSearchBatchResult<String, Void> result) {
+                        putDocumentsFuture.complete(result);
+                    }
+
+                    @Override
+                    public void onSystemError(Throwable throwable) {
+                        putDocumentsFuture.completeExceptionally(throwable);
+                    }
+                });
+        putDocumentsFuture.get();
+
+        // Search for only 2 result per page
+        SearchSpec searchSpec = new SearchSpec.Builder()
+                .setTermMatch(TERM_MATCH_PREFIX)
+                .setResultCountPerPage(2)
+                .build();
+        SearchResults searchResults = mSearchSession.search("", searchSpec);
+
+        // Get the first page, it contains 2 results.
+        List<GenericDocument> outDocs = new ArrayList<>();
+        CompletableFuture<AppSearchResult<List<SearchResult>>> getNextPageFuture =
+                new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        List<SearchResult> results = getNextPageFuture.get().getResultValue();
+        assertThat(results).hasSize(2);
+        outDocs.add(results.get(0).getGenericDocument());
+        outDocs.add(results.get(1).getGenericDocument());
+
+        // Get the second page, it contains only 1 result.
+        getNextPageFuture = new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        results = getNextPageFuture.get().getResultValue();
+        assertThat(results).hasSize(1);
+        outDocs.add(results.get(0).getGenericDocument());
+
+        assertThat(outDocs).containsExactly(document1, document2, document3);
+
+        // We get all documents, and it shouldn't fail if we keep calling getNextPage().
+        getNextPageFuture = new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        results = getNextPageFuture.get().getResultValue();
+        assertThat(results).isEmpty();
+    }
 }