Merge "Move query scheduling time to MdnsQueryScheduler" into main
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsQueryScheduler.java b/service-t/src/com/android/server/connectivity/mdns/MdnsQueryScheduler.java
index 356b738..7495aec 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsQueryScheduler.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsQueryScheduler.java
@@ -16,9 +16,14 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.AGGRESSIVE_QUERY_MODE;
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.PASSIVE_QUERY_MODE;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 /**
  * The query scheduler class for calculating next query tasks parameters.
  * <p>
@@ -26,6 +31,25 @@
  */
 public class MdnsQueryScheduler {
 
+    @VisibleForTesting
+    // RFC 6762 5.2: The interval between the first two queries MUST be at least one second.
+    static final int INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS = 1000;
+    private static final int INITIAL_TIME_BETWEEN_BURSTS_MS =
+            (int) MdnsConfigs.initialTimeBetweenBurstsMs();
+    private static final int MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS =
+            (int) MdnsConfigs.timeBetweenBurstsMs();
+    private static final int QUERIES_PER_BURST = (int) MdnsConfigs.queriesPerBurst();
+    private static final int TIME_BETWEEN_QUERIES_IN_BURST_MS =
+            (int) MdnsConfigs.timeBetweenQueriesInBurstMs();
+    private static final int QUERIES_PER_BURST_PASSIVE_MODE =
+            (int) MdnsConfigs.queriesPerBurstPassive();
+    @VisibleForTesting
+    // Basically this tries to send one query per typical DTIM interval 100ms, to maximize the
+    // chances that a query will be received if devices are using a DTIM multiplier (in which case
+    // they only listen once every [multiplier] DTIM intervals).
+    static final int TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS = 100;
+    static final int MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS = 60000;
+
     /**
      * The argument for tracking the query tasks status.
      */
@@ -72,19 +96,21 @@
         if (mLastScheduledQueryTaskArgs == null) {
             return null;
         }
-        if (!mLastScheduledQueryTaskArgs.config.shouldUseQueryBackoff(numOfQueriesBeforeBackoff)) {
+        final QueryTaskConfig lastConfig = mLastScheduledQueryTaskArgs.config;
+        if (!shouldUseQueryBackoff(lastConfig.queryIndex, lastConfig.queryMode,
+                numOfQueriesBeforeBackoff)) {
             return null;
         }
 
         final long timeToRun = calculateTimeToRun(mLastScheduledQueryTaskArgs,
-                mLastScheduledQueryTaskArgs.config, now, minRemainingTtl, lastSentTime,
+                lastConfig.queryIndex, lastConfig.queryMode, now, minRemainingTtl, lastSentTime,
                 numOfQueriesBeforeBackoff, false /* forceEnableBackoff */);
 
         if (timeToRun <= mLastScheduledQueryTaskArgs.timeToRun) {
             return null;
         }
 
-        mLastScheduledQueryTaskArgs = new ScheduledQueryTaskArgs(mLastScheduledQueryTaskArgs.config,
+        mLastScheduledQueryTaskArgs = new ScheduledQueryTaskArgs(lastConfig,
                 timeToRun,
                 minRemainingTtl + now,
                 sessionId);
@@ -104,17 +130,19 @@
             int queryMode,
             int numOfQueriesBeforeBackoff,
             boolean forceEnableBackoff) {
-        final QueryTaskConfig nextRunConfig = currentConfig.getConfigForNextRun(queryMode);
+        final int newQueryIndex = currentConfig.getConfigForNextRun(queryMode).queryIndex;
         long timeToRun;
         if (mLastScheduledQueryTaskArgs == null && !forceEnableBackoff) {
-            timeToRun = now + nextRunConfig.getDelayBeforeTaskWithoutBackoff();
+            timeToRun = now + getDelayBeforeTaskWithoutBackoff(
+                    newQueryIndex, queryMode);
         } else {
-            timeToRun = calculateTimeToRun(mLastScheduledQueryTaskArgs,
-                    nextRunConfig, now, minRemainingTtl, lastSentTime, numOfQueriesBeforeBackoff,
-                    forceEnableBackoff);
+            timeToRun = calculateTimeToRun(mLastScheduledQueryTaskArgs, newQueryIndex,
+                    queryMode, now, minRemainingTtl, lastSentTime,
+                    numOfQueriesBeforeBackoff, forceEnableBackoff);
         }
-        mLastScheduledQueryTaskArgs = new ScheduledQueryTaskArgs(nextRunConfig, timeToRun,
-                minRemainingTtl + now,
+        mLastScheduledQueryTaskArgs = new ScheduledQueryTaskArgs(
+                currentConfig.getConfigForNextRun(queryMode),
+                timeToRun, minRemainingTtl + now,
                 sessionId);
         return mLastScheduledQueryTaskArgs;
     }
@@ -131,11 +159,11 @@
     }
 
     private static long calculateTimeToRun(@Nullable ScheduledQueryTaskArgs taskArgs,
-            QueryTaskConfig queryTaskConfig, long now, long minRemainingTtl, long lastSentTime,
+            int queryIndex, int queryMode, long now, long minRemainingTtl, long lastSentTime,
             int numOfQueriesBeforeBackoff, boolean forceEnableBackoff) {
-        final long baseDelayInMs = queryTaskConfig.getDelayBeforeTaskWithoutBackoff();
+        final long baseDelayInMs = getDelayBeforeTaskWithoutBackoff(queryIndex, queryMode);
         if (!(forceEnableBackoff
-                || queryTaskConfig.shouldUseQueryBackoff(numOfQueriesBeforeBackoff))) {
+                || shouldUseQueryBackoff(queryIndex, queryMode, numOfQueriesBeforeBackoff))) {
             return lastSentTime + baseDelayInMs;
         }
         if (minRemainingTtl <= 0) {
@@ -152,4 +180,93 @@
         }
         return Math.max(now + (long) (0.8 * minRemainingTtl), lastSentTime + baseDelayInMs);
     }
+
+    private static int getBurstIndex(int queryIndex, int queryMode) {
+        if (queryMode == PASSIVE_QUERY_MODE && queryIndex >= QUERIES_PER_BURST) {
+            // In passive mode, after the first burst of QUERIES_PER_BURST queries, subsequent
+            // bursts have QUERIES_PER_BURST_PASSIVE_MODE queries.
+            final int queryIndexAfterFirstBurst = queryIndex - QUERIES_PER_BURST;
+            return 1 + (queryIndexAfterFirstBurst / QUERIES_PER_BURST_PASSIVE_MODE);
+        } else {
+            return queryIndex / QUERIES_PER_BURST;
+        }
+    }
+
+    private static int getQueryIndexInBurst(int queryIndex, int queryMode) {
+        if (queryMode == PASSIVE_QUERY_MODE && queryIndex >= QUERIES_PER_BURST) {
+            final int queryIndexAfterFirstBurst = queryIndex - QUERIES_PER_BURST;
+            return queryIndexAfterFirstBurst % QUERIES_PER_BURST_PASSIVE_MODE;
+        } else {
+            return queryIndex % QUERIES_PER_BURST;
+        }
+    }
+
+    private static boolean isFirstBurst(int queryIndex, int queryMode) {
+        return getBurstIndex(queryIndex, queryMode) == 0;
+    }
+
+    static boolean isFirstQueryInBurst(int queryIndex, int queryMode) {
+        return getQueryIndexInBurst(queryIndex, queryMode) == 0;
+    }
+
+    private static long getDelayBeforeTaskWithoutBackoff(int queryIndex, int queryMode) {
+        final int burstIndex = getBurstIndex(queryIndex, queryMode);
+        final int queryIndexInBurst = getQueryIndexInBurst(queryIndex, queryMode);
+        if (queryIndexInBurst == 0) {
+            return getTimeToBurstMs(burstIndex, queryMode);
+        } else if (queryIndexInBurst == 1 && queryMode == AGGRESSIVE_QUERY_MODE) {
+            // In aggressive mode, the first 2 queries are sent without delay.
+            return 0;
+        }
+        return queryMode == AGGRESSIVE_QUERY_MODE
+                ? TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS
+                : TIME_BETWEEN_QUERIES_IN_BURST_MS;
+    }
+
+    /**
+     * Shifts a value left by the specified number of bits, coercing to at most maxValue.
+     *
+     * <p>This allows calculating min(value*2^shift, maxValue) without overflow.
+     */
+    private static int boundedLeftShift(int value, int shift, int maxValue) {
+        // There must be at least one leading zero for positive values, so the maximum left shift
+        // without overflow is the number of leading zeros minus one.
+        final int maxShift = Integer.numberOfLeadingZeros(value) - 1;
+        if (shift > maxShift) {
+            // The shift would overflow positive integers, so is greater than maxValue.
+            return maxValue;
+        }
+        return Math.min(value << shift, maxValue);
+    }
+
+    private static int getTimeToBurstMs(int burstIndex, int queryMode) {
+        if (burstIndex == 0) {
+            // No delay before the first burst
+            return 0;
+        }
+        switch (queryMode) {
+            case PASSIVE_QUERY_MODE:
+                return MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS;
+            case AGGRESSIVE_QUERY_MODE:
+                return boundedLeftShift(INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS,
+                        burstIndex - 1,
+                        MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS);
+            default: // ACTIVE_QUERY_MODE
+                return boundedLeftShift(INITIAL_TIME_BETWEEN_BURSTS_MS,
+                        burstIndex - 1,
+                        MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS);
+        }
+    }
+
+    /**
+     * Determine if the query backoff should be used.
+     */
+    public static boolean shouldUseQueryBackoff(int queryIndex, int queryMode,
+            int numOfQueriesBeforeBackoff) {
+        // Don't enable backoff mode during the burst or in the first burst
+        if (!isFirstQueryInBurst(queryIndex, queryMode) || isFirstBurst(queryIndex, queryMode)) {
+            return false;
+        }
+        return queryIndex > numOfQueriesBeforeBackoff;
+    }
 }
diff --git a/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java b/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
index dd4073f..d193e14 100644
--- a/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
+++ b/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
@@ -17,7 +17,6 @@
 package com.android.server.connectivity.mdns;
 
 import static com.android.server.connectivity.mdns.MdnsSearchOptions.AGGRESSIVE_QUERY_MODE;
-import static com.android.server.connectivity.mdns.MdnsSearchOptions.PASSIVE_QUERY_MODE;
 
 import com.android.internal.annotations.VisibleForTesting;
 
@@ -26,136 +25,25 @@
  * Call to getConfigForNextRun returns a config that can be used to build the next query task.
  */
 public class QueryTaskConfig {
-
-    private static final int INITIAL_TIME_BETWEEN_BURSTS_MS =
-            (int) MdnsConfigs.initialTimeBetweenBurstsMs();
-    private static final int MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS =
-            (int) MdnsConfigs.timeBetweenBurstsMs();
-    private static final int QUERIES_PER_BURST = (int) MdnsConfigs.queriesPerBurst();
-    private static final int TIME_BETWEEN_QUERIES_IN_BURST_MS =
-            (int) MdnsConfigs.timeBetweenQueriesInBurstMs();
-    private static final int QUERIES_PER_BURST_PASSIVE_MODE =
-            (int) MdnsConfigs.queriesPerBurstPassive();
     private static final int UNSIGNED_SHORT_MAX_VALUE = 65536;
-    @VisibleForTesting
-    // RFC 6762 5.2: The interval between the first two queries MUST be at least one second.
-    static final int INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS = 1000;
-    @VisibleForTesting
-    // Basically this tries to send one query per typical DTIM interval 100ms, to maximize the
-    // chances that a query will be received if devices are using a DTIM multiplier (in which case
-    // they only listen once every [multiplier] DTIM intervals).
-    static final int TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS = 100;
-    static final int MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS = 60000;
     private final boolean alwaysAskForUnicastResponse =
             MdnsConfigs.alwaysAskForUnicastResponseInEachBurst();
     @VisibleForTesting
     final int transactionId;
     @VisibleForTesting
     final boolean expectUnicastResponse;
-    private final int queryIndex;
-    private final int queryMode;
+    final int queryIndex;
+    final int queryMode;
 
-    QueryTaskConfig(int queryMode, int queryIndex, int transactionId,
-            boolean expectUnicastResponse) {
+    QueryTaskConfig(int queryMode, int queryIndex, int transactionId) {
         this.queryMode = queryMode;
         this.transactionId = transactionId;
         this.queryIndex = queryIndex;
-        this.expectUnicastResponse = expectUnicastResponse;
+        this.expectUnicastResponse = getExpectUnicastResponse();
     }
 
     QueryTaskConfig(int queryMode) {
-        this(queryMode, 0, 1, true);
-    }
-
-    private static int getBurstIndex(int queryIndex, int queryMode) {
-        if (queryMode == PASSIVE_QUERY_MODE && queryIndex >= QUERIES_PER_BURST) {
-            // In passive mode, after the first burst of QUERIES_PER_BURST queries, subsequent
-            // bursts have QUERIES_PER_BURST_PASSIVE_MODE queries.
-            final int queryIndexAfterFirstBurst = queryIndex - QUERIES_PER_BURST;
-            return 1 + (queryIndexAfterFirstBurst / QUERIES_PER_BURST_PASSIVE_MODE);
-        } else {
-            return queryIndex / QUERIES_PER_BURST;
-        }
-    }
-
-    private static int getQueryIndexInBurst(int queryIndex, int queryMode) {
-        if (queryMode == PASSIVE_QUERY_MODE && queryIndex >= QUERIES_PER_BURST) {
-            final int queryIndexAfterFirstBurst = queryIndex - QUERIES_PER_BURST;
-            return queryIndexAfterFirstBurst % QUERIES_PER_BURST_PASSIVE_MODE;
-        } else {
-            return queryIndex % QUERIES_PER_BURST;
-        }
-    }
-
-    private static boolean isFirstBurst(int queryIndex, int queryMode) {
-        return getBurstIndex(queryIndex, queryMode) == 0;
-    }
-
-    private static boolean isFirstQueryInBurst(int queryIndex, int queryMode) {
-        return getQueryIndexInBurst(queryIndex, queryMode) == 0;
-    }
-
-    // TODO: move delay calculations to MdnsQueryScheduler
-    long getDelayBeforeTaskWithoutBackoff() {
-        return getDelayBeforeTaskWithoutBackoff(queryIndex, queryMode);
-    }
-
-    private static long getDelayBeforeTaskWithoutBackoff(int queryIndex, int queryMode) {
-        final int burstIndex = getBurstIndex(queryIndex, queryMode);
-        final int queryIndexInBurst = getQueryIndexInBurst(queryIndex, queryMode);
-        if (queryIndexInBurst == 0) {
-            return getTimeToBurstMs(burstIndex, queryMode);
-        } else if (queryIndexInBurst == 1 && queryMode == AGGRESSIVE_QUERY_MODE) {
-            // In aggressive mode, the first 2 queries are sent without delay.
-            return 0;
-        }
-        return queryMode == AGGRESSIVE_QUERY_MODE
-                ? TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS
-                : TIME_BETWEEN_QUERIES_IN_BURST_MS;
-    }
-
-    private boolean getExpectUnicastResponse(int queryIndex, int queryMode) {
-        if (queryMode == AGGRESSIVE_QUERY_MODE) {
-            if (isFirstQueryInBurst(queryIndex, queryMode)) {
-                return true;
-            }
-        }
-        return alwaysAskForUnicastResponse;
-    }
-
-    /**
-     * Shifts a value left by the specified number of bits, coercing to at most maxValue.
-     *
-     * <p>This allows calculating min(value*2^shift, maxValue) without overflow.
-     */
-    private static int boundedLeftShift(int value, int shift, int maxValue) {
-        // There must be at least one leading zero for positive values, so the maximum left shift
-        // without overflow is the number of leading zeros minus one.
-        final int maxShift = Integer.numberOfLeadingZeros(value) - 1;
-        if (shift > maxShift) {
-            // The shift would overflow positive integers, so is greater than maxValue.
-            return maxValue;
-        }
-        return Math.min(value << shift, maxValue);
-    }
-
-    private static int getTimeToBurstMs(int burstIndex, int queryMode) {
-        if (burstIndex == 0) {
-            // No delay before the first burst
-            return 0;
-        }
-        switch (queryMode) {
-            case PASSIVE_QUERY_MODE:
-                return MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS;
-            case AGGRESSIVE_QUERY_MODE:
-                return boundedLeftShift(INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS,
-                        burstIndex - 1,
-                        MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS);
-            default: // ACTIVE_QUERY_MODE
-                return boundedLeftShift(INITIAL_TIME_BETWEEN_BURSTS_MS,
-                        burstIndex - 1,
-                        MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS);
-        }
+        this(queryMode, 0, 1);
     }
 
     /**
@@ -168,18 +56,15 @@
             newTransactionId = 1;
         }
 
-        return new QueryTaskConfig(queryMode, newQueryIndex, newTransactionId,
-                getExpectUnicastResponse(newQueryIndex, queryMode));
+        return new QueryTaskConfig(queryMode, newQueryIndex, newTransactionId);
     }
 
-    /**
-     * Determine if the query backoff should be used.
-     */
-    public boolean shouldUseQueryBackoff(int numOfQueriesBeforeBackoff) {
-        // Don't enable backoff mode during the burst or in the first burst
-        if (!isFirstQueryInBurst(queryIndex, queryMode) || isFirstBurst(queryIndex, queryMode)) {
-            return false;
+    private boolean getExpectUnicastResponse() {
+        if (queryMode == AGGRESSIVE_QUERY_MODE) {
+            if (MdnsQueryScheduler.isFirstQueryInBurst(queryIndex, queryMode)) {
+                return true;
+            }
         }
-        return queryIndex > numOfQueriesBeforeBackoff;
+        return queryIndex == 0 || alwaysAskForUnicastResponse;
     }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
index ed95e4b..b8a9707 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -16,13 +16,13 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsQueryScheduler.INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS;
+import static com.android.server.connectivity.mdns.MdnsQueryScheduler.MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS;
+import static com.android.server.connectivity.mdns.MdnsQueryScheduler.TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS;
 import static com.android.server.connectivity.mdns.MdnsSearchOptions.ACTIVE_QUERY_MODE;
 import static com.android.server.connectivity.mdns.MdnsSearchOptions.AGGRESSIVE_QUERY_MODE;
 import static com.android.server.connectivity.mdns.MdnsSearchOptions.PASSIVE_QUERY_MODE;
 import static com.android.server.connectivity.mdns.MdnsServiceTypeClient.EVENT_START_QUERYTASK;
-import static com.android.server.connectivity.mdns.QueryTaskConfig.INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS;
-import static com.android.server.connectivity.mdns.QueryTaskConfig.MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS;
-import static com.android.server.connectivity.mdns.QueryTaskConfig.TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS;
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.junit.Assert.assertArrayEquals;