Merge "Ensure that search results sync only runs for the current search query" into main

13 files changed, 222 insertions, 42 deletions

diff --git a/photopicker/src/com/android/photopicker/data/MediaProviderClient.kt b/photopicker/src/com/android/photopicker/data/MediaProviderClient.kt
index 4ecba6077..d7dc60428 100644
--- a/photopicker/src/com/android/photopicker/data/MediaProviderClient.kt
+++ b/photopicker/src/com/android/photopicker/data/MediaProviderClient.kt
@@ -907,6 +907,67 @@ open class MediaProviderClient {
         resolver: ContentResolver,
         config: PhotopickerConfiguration,
     ): Int {
+        val extras: Bundle =
+            prepareSearchResultsExtras(
+                searchRequest = searchRequest,
+                providers = providers,
+                config = config,
+            )
+
+        val result: Bundle? =
+            resolver.call(
+                MEDIA_PROVIDER_AUTHORITY,
+                SEARCH_REQUEST_INIT_CALL_METHOD,
+                /* arg */ null,
+                extras,
+            )
+        return checkNotNull(result?.getInt(SEARCH_REQUEST_ID)) {
+            "Search request ID cannot be null"
+        }
+    }
+
+    /**
+     * Notifies the Data Source that the previously known search query is performed again by the
+     * user in the same session.
+     *
+     * This call lets [MediaProvider] know that the user has triggered a known search request again
+     * and the backend should prepare to handle search results queries for the given search request.
+     */
+    suspend fun ensureSearchResults(
+        searchRequest: SearchRequest,
+        searchRequestId: Int,
+        providers: List<Provider>,
+        resolver: ContentResolver,
+        config: PhotopickerConfiguration,
+    ) {
+        val extras: Bundle =
+            prepareSearchResultsExtras(
+                searchRequest = searchRequest,
+                searchRequestId = searchRequestId,
+                providers = providers,
+                config = config,
+            )
+
+        resolver.call(
+            MEDIA_PROVIDER_AUTHORITY,
+            SEARCH_REQUEST_INIT_CALL_METHOD,
+            /* arg */ null,
+            extras,
+        )
+    }
+
+    /**
+     * Creates an extras [Bundle] with the required args for MediaProvider's
+     * [SEARCH_REQUEST_INIT_CALL_METHOD].
+     *
+     * See [createSearchRequest] and [ensureSearchResults].
+     */
+    private fun prepareSearchResultsExtras(
+        searchRequest: SearchRequest,
+        searchRequestId: Int? = null,
+        providers: List<Provider>,
+        config: PhotopickerConfiguration,
+    ): Bundle {
         val extras =
             bundleOf(
                 EXTRA_MIME_TYPES to config.mimeTypes,
@@ -917,6 +978,10 @@ open class MediaProviderClient {
                     },
             )
 
+        if (searchRequestId != null) {
+            extras.putInt(SEARCH_REQUEST_ID, searchRequestId)
+        }
+
         when (searchRequest) {
             is SearchRequest.SearchTextRequest ->
                 extras.putString(SearchRequestInitRequest.SEARCH_TEXT.key, searchRequest.searchText)
@@ -940,16 +1005,7 @@ open class MediaProviderClient {
             }
         }
 
-        val result: Bundle? =
-            resolver.call(
-                MEDIA_PROVIDER_AUTHORITY,
-                SEARCH_REQUEST_INIT_CALL_METHOD,
-                /* arg */ null,
-                extras,
-            )
-        return checkNotNull(result?.getInt(SEARCH_REQUEST_ID)) {
-            "Search request ID cannot be null"
-        }
+        return extras
     }
 
     /**
diff --git a/photopicker/src/com/android/photopicker/features/categorygrid/data/CategoryDataServiceImpl.kt b/photopicker/src/com/android/photopicker/features/categorygrid/data/CategoryDataServiceImpl.kt
index 03e2c78d5..a67fc5386 100644
--- a/photopicker/src/com/android/photopicker/features/categorygrid/data/CategoryDataServiceImpl.kt
+++ b/photopicker/src/com/android/photopicker/features/categorygrid/data/CategoryDataServiceImpl.kt
@@ -304,8 +304,7 @@ class CategoryDataServiceImpl(
                             cancellationSignal = cancellationSignal,
                         )
                     // Ensure that cancellation get propagated to the data source when the paging
-                    // source
-                    // is invalidated.
+                    // source is invalidated.
                     pagingSource.registerInvalidatedCallback { cancellationSignal?.cancel() }
 
                     Log.v(
diff --git a/photopicker/src/com/android/photopicker/features/categorygrid/paging/CategoryAndAlbumPagingSource.kt b/photopicker/src/com/android/photopicker/features/categorygrid/paging/CategoryAndAlbumPagingSource.kt
index fa1d195d7..b6d5b2016 100644
--- a/photopicker/src/com/android/photopicker/features/categorygrid/paging/CategoryAndAlbumPagingSource.kt
+++ b/photopicker/src/com/android/photopicker/features/categorygrid/paging/CategoryAndAlbumPagingSource.kt
@@ -52,6 +52,7 @@ class CategoryAndAlbumPagingSource(
     override suspend fun load(params: LoadParams<GroupPageKey>): LoadResult<GroupPageKey, Group> {
         val pageKey = params.key ?: GroupPageKey()
         val pageSize = params.loadSize
+
         // Switch to the background thread from the main thread using [withContext].
         val result =
             withContext(dispatcher) {
diff --git a/photopicker/src/com/android/photopicker/features/search/Search.kt b/photopicker/src/com/android/photopicker/features/search/Search.kt
index 191927d5f..46dae64a8 100644
--- a/photopicker/src/com/android/photopicker/features/search/Search.kt
+++ b/photopicker/src/com/android/photopicker/features/search/Search.kt
@@ -895,7 +895,7 @@ private fun ResultMediaGrid(
                 delay(1000)
                 if (items.itemCount == 0) {
                     resultsState = ResultsState.LOADING_WITH_INDICATOR
-                    delay(4000)
+                    delay(10000)
                     if (resultsState == ResultsState.LOADING_WITH_INDICATOR)
                         resultsState = ResultsState.EMPTY
                 }
diff --git a/photopicker/src/com/android/photopicker/features/search/data/SearchDataServiceImpl.kt b/photopicker/src/com/android/photopicker/features/search/data/SearchDataServiceImpl.kt
index 9aacde80f..013de1dd9 100644
--- a/photopicker/src/com/android/photopicker/features/search/data/SearchDataServiceImpl.kt
+++ b/photopicker/src/com/android/photopicker/features/search/data/SearchDataServiceImpl.kt
@@ -83,7 +83,7 @@ class SearchDataServiceImpl(
 ) : SearchDataService {
     companion object {
         // Timeout for receiving suggestions from the data source in milli seconds.
-        private const val SUGGESTIONS_TIMEOUT: Long = 1500
+        private const val SUGGESTIONS_TIMEOUT: Long = 3000
     }
 
     // An internal lock to allow thread-safe updates to the search request and results cache.
@@ -190,7 +190,7 @@ class SearchDataServiceImpl(
             } catch (e: RuntimeException) {
                 Log.w(
                     SearchDataService.TAG,
-                    "An error occurred while fetching " + "search suggestions for prefix $prefix",
+                    "An error occurred while fetching search suggestions for prefix $prefix",
                     e,
                 )
 
@@ -276,7 +276,7 @@ class SearchDataServiceImpl(
                     }
 
                     Log.d(
-                        DataService.TAG,
+                        SearchDataService.TAG,
                         "Created a search results paging source that queries $availableProviders " +
                             "for search request id $searchRequestId",
                     )
@@ -307,9 +307,12 @@ class SearchDataServiceImpl(
      * Checks if this is a new search request in the current session.
      * 1. If this is a new search requests, [MediaProvider] is notified with the new search request
      *    and it creates and returns a search request id.
-     * 2. If this is not a new search request, previously caches search request id is returned.
+     * 2. If this is not a new search request, previously cached search request id is returned.
+     *
+     * In both scenarios, this notifies the backend to refresh search results cache for the given
+     * search request id.
      */
-    private fun getSearchRequestId(
+    private suspend fun getSearchRequestId(
         searchRequest: SearchRequest,
         availableProviders: List<Provider>,
         contentResolver: ContentResolver,
@@ -322,14 +325,40 @@ class SearchDataServiceImpl(
                     "Search request id is available for search request $searchRequest. " +
                         "Not creating a new search request id.",
                 )
-                searchRequestIdMap[searchRequest]!!
+
+                val searchRequestId = searchRequestIdMap[searchRequest]!!
+
+                try {
+                    // Ensure search results data in data source is ready for the search query.
+                    mediaProviderClient.ensureSearchResults(
+                        searchRequest,
+                        searchRequestId,
+                        availableProviders,
+                        contentResolver,
+                        config,
+                    )
+                } catch (e: RuntimeException) {
+                    Log.e(SearchDataService.TAG, "Could not ensure search results", e)
+                }
+
+                searchRequestId
             } else {
-                mediaProviderClient.createSearchRequest(
-                    searchRequest,
-                    availableProviders,
-                    contentResolver,
-                    config,
+                Log.d(
+                    SearchDataService.TAG,
+                    "Search request id is not available for search request $searchRequest. " +
+                        "Creating a new search request id.",
                 )
+
+                val newSearchRequestId =
+                    mediaProviderClient.createSearchRequest(
+                        searchRequest,
+                        availableProviders,
+                        contentResolver,
+                        config,
+                    )
+
+                searchRequestIdMap[searchRequest] = newSearchRequestId
+                newSearchRequestId
             }
         }
     }
diff --git a/src/com/android/providers/media/MediaProvider.java b/src/com/android/providers/media/MediaProvider.java
index 21d625106..a220c19c6 100644
--- a/src/com/android/providers/media/MediaProvider.java
+++ b/src/com/android/providers/media/MediaProvider.java
@@ -7767,7 +7767,7 @@ public class MediaProvider extends ContentProvider {
             throw new SecurityException(
                     getSecurityExceptionMessage("Picker search media init"));
         }
-        return PickerDataLayerV2.handleNewSearchRequest(getContext(), extras);
+        return PickerDataLayerV2.handleSearchResultsInit(getContext(), extras);
     }
 
     private void initMediaSets(@NonNull Bundle extras) {
diff --git a/src/com/android/providers/media/photopicker/sync/PickerSyncManager.java b/src/com/android/providers/media/photopicker/sync/PickerSyncManager.java
index b3076b0b7..6d879a306 100644
--- a/src/com/android/providers/media/photopicker/sync/PickerSyncManager.java
+++ b/src/com/android/providers/media/photopicker/sync/PickerSyncManager.java
@@ -17,6 +17,7 @@
 package com.android.providers.media.photopicker.sync;
 
 import static com.android.providers.media.photopicker.sync.SyncTrackerRegistry.markAlbumMediaSyncAsComplete;
+import static com.android.providers.media.photopicker.sync.SyncTrackerRegistry.markAllSearchResultsSyncAsComplete;
 import static com.android.providers.media.photopicker.sync.SyncTrackerRegistry.markMediaInMediaSetSyncAsComplete;
 import static com.android.providers.media.photopicker.sync.SyncTrackerRegistry.markMediaSetsSyncAsComplete;
 import static com.android.providers.media.photopicker.sync.SyncTrackerRegistry.markSearchResultsSyncAsComplete;
@@ -509,6 +510,10 @@ public class PickerSyncManager {
         final OneTimeWorkRequest syncRequest =
                 buildOneTimeWorkerRequest(SearchResultsSyncWorker.class, inputData);
 
+        // Clear all existing requests since there can be only one unique work running and our
+        // new sync work will replace the existing work (if any).
+        markAllSearchResultsSyncAsComplete(syncSource);
+
         // Track the new sync request
         trackNewSearchResultsSyncRequests(syncSource, syncRequest.getId());
 
@@ -519,7 +524,7 @@ public class PickerSyncManager {
         try {
             final Operation enqueueOperation = mWorkManager.enqueueUniqueWork(
                     workName,
-                    ExistingWorkPolicy.APPEND_OR_REPLACE,
+                    ExistingWorkPolicy.REPLACE,
                     syncRequest);
 
             // Check that the request has been successfully enqueued.
diff --git a/src/com/android/providers/media/photopicker/sync/SearchResultsSyncWorker.java b/src/com/android/providers/media/photopicker/sync/SearchResultsSyncWorker.java
index 6efd7ed0a..f05da24be 100644
--- a/src/com/android/providers/media/photopicker/sync/SearchResultsSyncWorker.java
+++ b/src/com/android/providers/media/photopicker/sync/SearchResultsSyncWorker.java
@@ -166,7 +166,7 @@ public class SearchResultsSyncWorker extends Worker {
         final Pair<String, String> resumeKey = getResumeKey(searchRequest, syncSource);
 
         if (SYNC_COMPLETE_RESUME_KEY.equals(resumeKey.first)) {
-            Log.i(TAG, "Sync has already been completed.");
+            Log.i(TAG, "Sync was already complete.");
             return;
         }
 
diff --git a/src/com/android/providers/media/photopicker/sync/SyncTracker.java b/src/com/android/providers/media/photopicker/sync/SyncTracker.java
index 87d3a204f..02f0bebb9 100644
--- a/src/com/android/providers/media/photopicker/sync/SyncTracker.java
+++ b/src/com/android/providers/media/photopicker/sync/SyncTracker.java
@@ -88,6 +88,22 @@ public class SyncTracker {
     }
 
     /**
+     * Use this method to mark all picker sync futures as complete.
+     *
+     * This is useful when using {@link ExistingWorkPolicy.REPLACE} to enqueue Unique Work.
+     * It clears the tracker of the work that will get cancelled by the REPLACE policy.
+     */
+    public void markAllSyncsCompleted() {
+        synchronized (mFutureMap) {
+            for (CompletableFuture future : mFutureMap.values()) {
+                future.complete(FUTURE_RESULT);
+            }
+            mFutureMap.clear();
+            Log.i(TAG, String.format("Marked all sync futures as complete"));
+        }
+    }
+
+    /**
      * Use this method to check if any sync request is still pending.
      * @return a {@link Collection} of {@link CompletableFuture} of pending syncs. This can be
      * used to track when all pending are complete.
diff --git a/src/com/android/providers/media/photopicker/sync/SyncTrackerRegistry.java b/src/com/android/providers/media/photopicker/sync/SyncTrackerRegistry.java
index 0cd43a7ab..3c2898b96 100644
--- a/src/com/android/providers/media/photopicker/sync/SyncTrackerRegistry.java
+++ b/src/com/android/providers/media/photopicker/sync/SyncTrackerRegistry.java
@@ -351,6 +351,19 @@ public class SyncTrackerRegistry {
     }
 
     /**
+     * Mark all the pending futures as complete.
+     */
+    public static void markAllSearchResultsSyncAsComplete(
+            @PickerSyncManager.SyncSource int syncSource) {
+        if (syncSource == SYNC_LOCAL_ONLY || syncSource == SYNC_LOCAL_AND_CLOUD) {
+            getLocalSearchSyncTracker().markAllSyncsCompleted();
+        }
+        if (syncSource == SYNC_CLOUD_ONLY || syncSource == SYNC_LOCAL_AND_CLOUD) {
+            getCloudSearchSyncTracker().markAllSyncsCompleted();
+        }
+    }
+
+    /**
      * Mark the required futures as complete for existing media set sync requests.
      */
     public static void markMediaSetsSyncAsComplete(
diff --git a/src/com/android/providers/media/photopicker/v2/PickerDataLayerV2.java b/src/com/android/providers/media/photopicker/v2/PickerDataLayerV2.java
index 5fc0a3c4d..c5693c357 100644
--- a/src/com/android/providers/media/photopicker/v2/PickerDataLayerV2.java
+++ b/src/com/android/providers/media/photopicker/v2/PickerDataLayerV2.java
@@ -1081,7 +1081,10 @@ public class PickerDataLayerV2 {
             return null;
         }
 
+        Log.d(TAG, "Fetching albums from CMP " + cloudAuthority);
         final Cursor cursor = getAlbumsCursorFromProvider(appContext, query, cloudAuthority);
+
+        Log.d(TAG, "Received albums from CMP " + cloudAuthority);
         return cursor == null
                 ? null
                 : new AlbumsCursorWrapper(cursor, cloudAuthority, localAuthority);
@@ -1308,24 +1311,39 @@ public class PickerDataLayerV2 {
     }
 
     /**
-     * Handle Picker application's request to create a new search request and return a Bundle with
-     * the search request Id.
+     * Handle Picker application's request to initialize search request media. If a new search
+     * request id needs to be created, return a Bundle with the search request Id.
+     *
      * Also trigger search results sync with the providers and saves the incoming search request in
-     * the search history table.
+     * the search history table if this is a new request.
+     *
+     * By default use ForkJoinPool.commonPool() for small background tasks to reduce resource
+     * usage instead of creating a custom pool. Its threads are slowly reclaimed during periods
+     * of non-use, and reinstated upon subsequent use.
      *
      * @param appContext Application context.
      * @param extras Bundle with input parameters.
      * @return a response Bundle.
      */
     @NonNull
-    public static Bundle handleNewSearchRequest(
+    public static Bundle handleSearchResultsInit(
             @NonNull Context appContext,
             @NonNull Bundle extras) {
-        // By default use ForkJoinPool.commonPool() to reduce resource usage instead of creating a
-        // custom pool. Its threads are slowly reclaimed during periods of non-use, and reinstated
-        // upon subsequent use.
-        return handleNewSearchRequest(appContext, extras, ForkJoinPool.commonPool(),
-                getWorkManager(appContext));
+        final int defaultSearchRequestId = -1;
+        final int inputSearchRequestId = extras.getInt("search_request_id", defaultSearchRequestId);
+        final WorkManager workManager = getWorkManager(appContext);
+
+        if (inputSearchRequestId == defaultSearchRequestId) {
+            Log.d(TAG, "New search request id needs to be created");
+
+            return handleNewSearchRequest(appContext, extras,
+                    ForkJoinPool.commonPool(), workManager);
+        } else {
+            Log.d(TAG, "Search request id already exists. Ensure that sync is complete for the "
+                    + "search request id.");
+
+            return ensureSearchResultsSynced(appContext, extras, workManager);
+        }
     }
 
     /**
@@ -1346,7 +1364,7 @@ public class PickerDataLayerV2 {
                                                 @NonNull Executor executor,
                                                 @NonNull WorkManager workManager) {
         requireNonNull(extras);
-        Log.d(TAG, "Received a search request: " + extras);
+        Log.d(TAG, "Received a new search request: " + extras);
 
         final SearchRequest searchRequest = SearchRequest.create(extras);
         final SQLiteDatabase database = PickerSyncController.getInstanceOrThrow().getDbFacade()
@@ -1365,7 +1383,9 @@ public class PickerDataLayerV2 {
                     executor);
 
             // Schedule search results sync
-            scheduleSearchResultsSync(appContext, searchRequest, searchRequestId, extras,
+            final Set<String> providers = new HashSet<>(
+                    Objects.requireNonNull(extras.getStringArrayList("providers")));
+            scheduleSearchResultsSync(appContext, searchRequest, searchRequestId, providers,
                     workManager);
 
             Log.d(TAG, "Returning search request id: " + searchRequestId);
@@ -1374,6 +1394,36 @@ public class PickerDataLayerV2 {
     }
 
     /**
+     * Ensure that the search results are synced for the given search request id. If not,
+     * start or resume the sync for the search results. Both of these aspects are taken care of by
+     * scheduling a work request.
+     *
+     * @param appContext Application context.
+     * @param extras Bundle with input parameters.
+     * @param workManager An instance of {@link WorkManager}
+     * @return a response Bundle.
+     */
+    @NonNull
+    public static Bundle ensureSearchResultsSynced(@NonNull Context appContext,
+                                                @NonNull Bundle extras,
+                                                @NonNull WorkManager workManager) {
+        requireNonNull(extras);
+        Log.d(TAG, "Received a previously known search request again: " + extras);
+
+        final int searchRequestId = extras.getInt("search_request_id", -1);
+        final SearchRequest searchRequest = SearchRequest.create(extras);
+
+        // Schedule search results sync with REPLACE policy. This takes care of cancelling any
+        // existing search results sync that might be obsolete.
+        final Set<String> providers = new HashSet<>(
+                Objects.requireNonNull(extras.getStringArrayList("providers")));
+        scheduleSearchResultsSync(appContext, searchRequest, searchRequestId, providers,
+                workManager);
+
+        return getSearchRequestInitResponse(searchRequestId);
+    }
+
+    /**
      * Handles Photopicker's request to trigger a sync for media items in a media set
      * based on whether the provider implements search categories and media sets
      * @param extras Bundle with all input parameters
@@ -1564,11 +1614,9 @@ public class PickerDataLayerV2 {
             @NonNull Context appContext,
             @NonNull SearchRequest searchRequest,
             int searchRequestId,
-            @NonNull Bundle extras,
+            @NonNull Set<String> providers,
             WorkManager workManager) {
         final PickerSyncManager syncManager = new PickerSyncManager(workManager, appContext);
-        final Set<String> providers = new HashSet<>(
-                Objects.requireNonNull(extras.getStringArrayList("providers")));
 
         final boolean localSyncWasScheduled = scheduleSearchSyncWithLocalProvider(
                 searchRequest, searchRequestId, syncManager, providers);
diff --git a/src/com/android/providers/media/photopicker/v2/sqlite/SearchSuggestionsDatabaseUtils.java b/src/com/android/providers/media/photopicker/v2/sqlite/SearchSuggestionsDatabaseUtils.java
index 38b8aaff1..25c0208fd 100644
--- a/src/com/android/providers/media/photopicker/v2/sqlite/SearchSuggestionsDatabaseUtils.java
+++ b/src/com/android/providers/media/photopicker/v2/sqlite/SearchSuggestionsDatabaseUtils.java
@@ -270,7 +270,7 @@ public class SearchSuggestionsDatabaseUtils {
             } while (cursor.moveToNext());
         }
 
-        Log.d(TAG, "Extracted suggestions from cursor: " + searchSuggestions);
+        Log.d(TAG, "Extracted suggestions from cursor: " + searchSuggestions.size());
         return searchSuggestions;
     }
 
diff --git a/tests/src/com/android/providers/media/photopicker/sync/SyncTrackerTests.java b/tests/src/com/android/providers/media/photopicker/sync/SyncTrackerTests.java
index 1eac73a78..574580963 100644
--- a/tests/src/com/android/providers/media/photopicker/sync/SyncTrackerTests.java
+++ b/tests/src/com/android/providers/media/photopicker/sync/SyncTrackerTests.java
@@ -50,6 +50,19 @@ public class SyncTrackerTests {
     }
 
     @Test
+    public void testMarkAllSyncsComplete() {
+        SyncTracker syncTracker = new SyncTracker();
+        syncTracker.createSyncFuture(UUID.randomUUID());
+        syncTracker.createSyncFuture(UUID.randomUUID());
+        syncTracker.createSyncFuture(UUID.randomUUID());
+        Collection<CompletableFuture<Object>> futures = syncTracker.pendingSyncFutures();
+        assertThat(futures.size()).isEqualTo(3);
+
+        syncTracker.markAllSyncsCompleted();
+        assertThat(futures.size()).isEqualTo(0);
+    }
+
+    @Test
     public void testCompleteOnTimeoutSyncFuture()
             throws InterruptedException, ExecutionException, TimeoutException {
         SyncTracker syncTracker = new SyncTracker();