From 58bb0a129a4eb1f19904265104af2a14db248cb3 Mon Sep 17 00:00:00 2001
From: Remi NGUYEN VAN <reminv@google.com>
Date: Fri, 12 May 2023 16:10:25 +0900
Subject: [PATCH] Do not immediately send new queries on cache hit

If a listener request is fulfilled immediately from cache, do not
restart queries immediately, but just schedule the next run.

This avoids doubling the number of queries when discovery, then resolve
is requested.

Bug: 281793453
Test: atest
Change-Id: Ibe41ee20427068a1ccdbc1f04525988a89a58899
---
 .../mdns/MdnsServiceTypeClient.java           | 31 ++++++++++++-------
 .../mdns/MdnsServiceTypeClientTests.java      | 29 +++++++++++++++++
 2 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
index 35f9b0475d..298e632535 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -175,6 +175,7 @@ public class MdnsServiceTypeClient {
             @NonNull MdnsSearchOptions searchOptions) {
         synchronized (lock) {
             this.searchOptions = searchOptions;
+            boolean hadReply = false;
             if (listeners.put(listener, searchOptions) == null) {
                 for (MdnsResponse existingResponse : instanceNameToResponse.values()) {
                     if (!responseMatchesOptions(existingResponse, searchOptions)) continue;
@@ -183,6 +184,7 @@ public class MdnsServiceTypeClient {
                     listener.onServiceNameDiscovered(info);
                     if (existingResponse.isComplete()) {
                         listener.onServiceFound(info);
+                        hadReply = true;
                     }
                 }
             }
@@ -192,14 +194,16 @@ public class MdnsServiceTypeClient {
             }
             // Keep tracking the ScheduledFuture for the task so we can cancel it if caller is not
             // interested anymore.
-            requestTaskFuture =
-                    executor.submit(
-                            new QueryTask(
-                                    new QueryTaskConfig(
-                                            searchOptions.getSubtypes(),
-                                            searchOptions.isPassiveMode(),
-                                            ++currentSessionId,
-                                            searchOptions.getNetwork())));
+            final QueryTaskConfig taskConfig = new QueryTaskConfig(
+                    searchOptions.getSubtypes(),
+                    searchOptions.isPassiveMode(),
+                    ++currentSessionId,
+                    searchOptions.getNetwork());
+            if (hadReply) {
+                requestTaskFuture = scheduleNextRunLocked(taskConfig);
+            } else {
+                requestTaskFuture = executor.submit(new QueryTask(taskConfig));
+            }
         }
     }
 
@@ -574,11 +578,14 @@ public class MdnsServiceTypeClient {
                         }
                     }
                 }
-                QueryTaskConfig config = this.config.getConfigForNextRun();
-                requestTaskFuture =
-                        executor.schedule(
-                                new QueryTask(config), config.timeToRunNextTaskInMs, MILLISECONDS);
+                requestTaskFuture = scheduleNextRunLocked(this.config);
             }
         }
     }
+
+    @NonNull
+    private Future<?> scheduleNextRunLocked(@NonNull QueryTaskConfig lastRunConfig) {
+        QueryTaskConfig config = lastRunConfig.getConfigForNextRun();
+        return executor.schedule(new QueryTask(config), config.timeToRunNextTaskInMs, MILLISECONDS);
+    }
 }
\ No newline at end of file
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 2fcdff24d6..98cea627db 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -21,6 +21,7 @@ import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
@@ -422,6 +423,34 @@ public class MdnsServiceTypeClientTests {
         assertNull(currentThreadExecutor.getAndClearLastScheduledRunnable());
     }
 
+    @Test
+    public void testQueryScheduledWhenAnsweredFromCache() {
+        final MdnsSearchOptions searchOptions = MdnsSearchOptions.getDefaultOptions();
+        client.startSendAndReceive(mockListenerOne, searchOptions);
+        assertNotNull(currentThreadExecutor.getAndClearSubmittedRunnable());
+
+        client.processResponse(createResponse(
+                "service-instance-1", "192.0.2.123", 5353,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), TEST_TTL), /* interfaceIndex= */ 20, mockNetwork);
+
+        verify(mockListenerOne).onServiceNameDiscovered(any());
+        verify(mockListenerOne).onServiceFound(any());
+
+        // File another identical query
+        client.startSendAndReceive(mockListenerTwo, searchOptions);
+
+        verify(mockListenerTwo).onServiceNameDiscovered(any());
+        verify(mockListenerTwo).onServiceFound(any());
+
+        // This time no query is submitted, only scheduled
+        assertNull(currentThreadExecutor.getAndClearSubmittedRunnable());
+        assertNotNull(currentThreadExecutor.getAndClearLastScheduledRunnable());
+        // This just skips the first query of the first burst
+        assertEquals(MdnsConfigs.timeBetweenQueriesInBurstMs(),
+                currentThreadExecutor.getAndClearLastScheduledDelayInMs());
+    }
+
     private static void verifyServiceInfo(MdnsServiceInfo serviceInfo, String serviceName,
             String[] serviceType, List<String> ipv4Addresses, List<String> ipv6Addresses, int port,
             List<String> subTypes, Map<String, String> attributes, int interfaceIndex,