Fix bug in the large batch query options.

Create new query options every time a query requests it from the
DatastoreHelper. This is needed to prevent one query's fetch options
from interfering with the fetch options of another.

Test: live on staging
Bug: 63339916
Change-Id: If45b6f2c345ae73e08c833065670501a1fb52191

7 files changed, 30 insertions, 17 deletions

diff --git a/src/main/java/com/android/vts/servlet/ShowGraphServlet.java b/src/main/java/com/android/vts/servlet/ShowGraphServlet.java
index 6f51911..4ef709b 100644
--- a/src/main/java/com/android/vts/servlet/ShowGraphServlet.java
+++ b/src/main/java/com/android/vts/servlet/ShowGraphServlet.java
@@ -162,7 +162,9 @@ public class ShowGraphServlet extends BaseServlet {
         // Process the test runs in the query
         List<Key> gets = new ArrayList<>();
         for (Entity testRun :
-                datastore.prepare(testRunQuery).asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                datastore
+                        .prepare(testRunQuery)
+                        .asIterable(DatastoreHelper.getLargeBatchOptions())) {
             gets.add(
                     KeyFactory.createKey(
                             testRun.getKey(), ProfilingPointRunEntity.KIND, profilingPointName));
@@ -191,7 +193,7 @@ public class ShowGraphServlet extends BaseServlet {
                         .setKeysOnly();
         gets = new ArrayList<>();
         for (Entity device :
-                datastore.prepare(deviceQuery).asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                datastore.prepare(deviceQuery).asIterable(DatastoreHelper.getLargeBatchOptions())) {
             if (testRunProfiling.containsKey(device.getParent())) {
                 gets.add(device.getKey());
             }
diff --git a/src/main/java/com/android/vts/servlet/ShowPlanReleaseServlet.java b/src/main/java/com/android/vts/servlet/ShowPlanReleaseServlet.java
index 23300e7..5e0f430 100644
--- a/src/main/java/com/android/vts/servlet/ShowPlanReleaseServlet.java
+++ b/src/main/java/com/android/vts/servlet/ShowPlanReleaseServlet.java
@@ -160,7 +160,8 @@ public class ShowPlanReleaseServlet extends BaseServlet {
                     datastore
                             .prepare(testPlanRunQuery)
                             .asIterable(
-                                    DatastoreHelper.LARGE_BATCH_OPTIONS.limit(MAX_RUNS_PER_PAGE))) {
+                                    DatastoreHelper.getLargeBatchOptions()
+                                            .limit(MAX_RUNS_PER_PAGE))) {
                 TestPlanRunEntity testPlanRun = TestPlanRunEntity.fromEntity(testPlanRunEntity);
                 if (testPlanRun == null) {
                     logger.log(
@@ -220,7 +221,7 @@ public class ShowPlanReleaseServlet extends BaseServlet {
             for (Entity device :
                     datastore
                             .prepare(deviceQuery)
-                            .asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                            .asIterable(DatastoreHelper.getLargeBatchOptions())) {
                 if (testPlanMap.containsKey(device.getParent())) {
                     deviceGets.add(device.getKey());
                 }
diff --git a/src/main/java/com/android/vts/servlet/ShowTreeServlet.java b/src/main/java/com/android/vts/servlet/ShowTreeServlet.java
index 9cbd786..453c974 100644
--- a/src/main/java/com/android/vts/servlet/ShowTreeServlet.java
+++ b/src/main/java/com/android/vts/servlet/ShowTreeServlet.java
@@ -184,8 +184,8 @@ public class ShowTreeServlet extends BaseServlet {
                     datastore
                             .prepare(testRunQuery)
                             .asIterable(
-                                    DatastoreHelper.LARGE_BATCH_OPTIONS.limit(
-                                            MAX_BUILD_IDS_PER_PAGE))) {
+                                    DatastoreHelper.getLargeBatchOptions()
+                                            .limit(MAX_BUILD_IDS_PER_PAGE))) {
                 TestRunEntity testRunEntity = TestRunEntity.fromEntity(testRun);
                 if (testRunEntity == null) {
                     continue;
@@ -247,7 +247,7 @@ public class ShowTreeServlet extends BaseServlet {
             for (Entity device :
                     datastore
                             .prepare(deviceQuery)
-                            .asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                            .asIterable(DatastoreHelper.getLargeBatchOptions())) {
                 if (metadataMap.containsKey(device.getParent())) {
                     deviceGets.add(device.getKey());
                 }
diff --git a/src/main/java/com/android/vts/util/DatastoreHelper.java b/src/main/java/com/android/vts/util/DatastoreHelper.java
index ec10194..b738df7 100644
--- a/src/main/java/com/android/vts/util/DatastoreHelper.java
+++ b/src/main/java/com/android/vts/util/DatastoreHelper.java
@@ -61,8 +61,15 @@ import java.util.logging.Logger;
 public class DatastoreHelper {
     protected static final Logger logger = Logger.getLogger(DatastoreHelper.class.getName());
     public static final int MAX_WRITE_RETRIES = 5;
-    public static final FetchOptions LARGE_BATCH_OPTIONS =
-            FetchOptions.Builder.withChunkSize(1000).prefetchSize(1000);
+
+    /**
+     * Get query fetch options for large batches of entities.
+     *
+     * @return FetchOptions with a large chunk and prefetch size.
+     */
+    public static FetchOptions getLargeBatchOptions() {
+        return FetchOptions.Builder.withChunkSize(1000).prefetchSize(1000);
+    }
 
     /**
      * Returns true if there are data points newer than lowerBound in the results table.
diff --git a/src/main/java/com/android/vts/util/FilterUtil.java b/src/main/java/com/android/vts/util/FilterUtil.java
index 69a91ee..d81d969 100644
--- a/src/main/java/com/android/vts/util/FilterUtil.java
+++ b/src/main/java/com/android/vts/util/FilterUtil.java
@@ -374,7 +374,7 @@ public class FilterUtil {
         Query testQuery =
                 new Query(kind).setAncestor(ancestorKey).setFilter(testFilter).setKeysOnly();
         for (Entity testRunKey :
-                datastore.prepare(testQuery).asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                datastore.prepare(testQuery).asIterable(DatastoreHelper.getLargeBatchOptions())) {
             matchingTestKeys.add(testRunKey.getKey());
         }
 
@@ -391,7 +391,7 @@ public class FilterUtil {
             for (Entity device :
                     datastore
                             .prepare(deviceQuery)
-                            .asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                            .asIterable(DatastoreHelper.getLargeBatchOptions())) {
                 if (matchingTestKeys.contains(device.getKey().getParent())) {
                     allMatchingKeys.add(device.getKey().getParent());
                 }
diff --git a/src/main/java/com/android/vts/util/PerformanceUtil.java b/src/main/java/com/android/vts/util/PerformanceUtil.java
index 92917b4..236dfa2 100644
--- a/src/main/java/com/android/vts/util/PerformanceUtil.java
+++ b/src/main/java/com/android/vts/util/PerformanceUtil.java
@@ -231,7 +231,9 @@ public class PerformanceUtil {
 
         Set<Key> testRunKeys = new HashSet<>();
         for (Entity e :
-                datastore.prepare(testRunQuery).asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                datastore
+                        .prepare(testRunQuery)
+                        .asIterable(DatastoreHelper.getLargeBatchOptions())) {
             testRunKeys.add(e.getKey());
         }
 
@@ -254,7 +256,7 @@ public class PerformanceUtil {
             for (Entity e :
                     datastore
                             .prepare(deviceQuery)
-                            .asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                            .asIterable(DatastoreHelper.getLargeBatchOptions())) {
                 if (testRunKeys.contains(e.getKey().getParent())) {
                     matchingTestRunKeys.add(e.getKey().getParent());
                 }
@@ -263,8 +265,7 @@ public class PerformanceUtil {
         }
 
         Filter profilingFilter =
-                FilterUtil.getProfilingTimeFilter(
-                        testKey, TestRunEntity.KIND, startTime, endTime);
+                FilterUtil.getProfilingTimeFilter(testKey, TestRunEntity.KIND, startTime, endTime);
         Query profilingQuery =
                 new Query(ProfilingPointRunEntity.KIND)
                         .setAncestor(testKey)
@@ -273,7 +274,9 @@ public class PerformanceUtil {
 
         List<Key> profilingKeys = new ArrayList<>();
         for (Entity e :
-                datastore.prepare(profilingQuery).asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                datastore
+                        .prepare(profilingQuery)
+                        .asIterable(DatastoreHelper.getLargeBatchOptions())) {
             if (testRunKeys.contains(e.getKey().getParent())) {
                 profilingKeys.add(e.getKey());
             }
diff --git a/src/main/java/com/android/vts/util/TestResults.java b/src/main/java/com/android/vts/util/TestResults.java
index 9d2cc79..9c759b2 100644
--- a/src/main/java/com/android/vts/util/TestResults.java
+++ b/src/main/java/com/android/vts/util/TestResults.java
@@ -199,7 +199,7 @@ public class TestResults {
                         .setKeysOnly();
         List<Key> deviceGets = new ArrayList<>();
         for (Entity device :
-                datastore.prepare(deviceQuery).asIterable(DatastoreHelper.LARGE_BATCH_OPTIONS)) {
+                datastore.prepare(deviceQuery).asIterable(DatastoreHelper.getLargeBatchOptions())) {
             if (testCaseRunMap.containsKey(device.getParent())) {
                 deviceGets.add(device.getKey());
             }