Merge "Ensure the ServiceTypeClient Executor threads can be shutdown" into main
diff --git a/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java b/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
index 0eebc61..161669b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
+++ b/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
@@ -16,6 +16,7 @@
 
 package com.android.server.connectivity.mdns;
 
+import android.annotation.NonNull;
 import android.util.ArraySet;
 
 import java.util.Set;
@@ -47,5 +48,17 @@
             }
             executor.shutdownNow();
         }
+        serviceTypeClientSchedulerExecutors.clear();
+    }
+
+    /**
+     * Shutdown one executor service and remove the executor service from the set.
+     * @param executorService the executorService to be shutdown
+     */
+    public void shutdownExecutorService(@NonNull ScheduledExecutorService executorService) {
+        if (!executorService.isShutdown()) {
+            executorService.shutdownNow();
+        }
+        serviceTypeClientSchedulerExecutors.remove(executorService);
     }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
index f386dd4..d55098c 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
@@ -204,6 +204,7 @@
                         if (serviceTypeClient == null) return;
                         // Notify all listeners that all services are removed from this socket.
                         serviceTypeClient.notifySocketDestroyed();
+                        executorProvider.shutdownExecutorService(serviceTypeClient.getExecutor());
                         perSocketServiceTypeClients.remove(serviceTypeClient);
                     }
                 });
@@ -238,6 +239,7 @@
             if (serviceTypeClient.stopSendAndReceive(listener)) {
                 // No listener is registered for the service type anymore, remove it from the list
                 // of the service type clients.
+                executorProvider.shutdownExecutorService(serviceTypeClient.getExecutor());
                 perSocketServiceTypeClients.remove(serviceTypeClient);
             }
         }
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 0752c01..1f76fbd 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -330,6 +330,13 @@
         }
     }
 
+    /**
+     * Get the executor service.
+     */
+    public ScheduledExecutorService getExecutor() {
+        return executor;
+    }
+
     private void removeScheduledTask() {
         dependencies.removeMessages(handler, EVENT_START_QUERYTASK);
         sharedLog.log("Remove EVENT_START_QUERYTASK"
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
index 1a4ae5d..e869b91 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
@@ -53,6 +53,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
 
 /** Tests for {@link MdnsDiscoveryManager}. */
 @RunWith(DevSdkIgnoreRunner.class)
@@ -80,6 +81,7 @@
     private static final Pair<String, SocketKey> PER_SOCKET_SERVICE_TYPE_2_NETWORK_2 =
             Pair.create(SERVICE_TYPE_2, SOCKET_KEY_NETWORK_2);
     @Mock private ExecutorProvider executorProvider;
+    @Mock private ScheduledExecutorService mockExecutorService;
     @Mock private MdnsSocketClientBase socketClient;
     @Mock private MdnsServiceTypeClient mockServiceTypeClientType1NullNetwork;
     @Mock private MdnsServiceTypeClient mockServiceTypeClientType1Network1;
@@ -128,6 +130,7 @@
                         return null;
                     }
                 };
+        doReturn(mockExecutorService).when(mockServiceTypeClientType1NullNetwork).getExecutor();
     }
 
     @After
@@ -165,11 +168,25 @@
         when(mockServiceTypeClientType1NullNetwork.stopSendAndReceive(mockListenerOne))
                 .thenReturn(true);
         runOnHandler(() -> discoveryManager.unregisterListener(SERVICE_TYPE_1, mockListenerOne));
+        verify(executorProvider).shutdownExecutorService(mockExecutorService);
         verify(mockServiceTypeClientType1NullNetwork).stopSendAndReceive(mockListenerOne);
         verify(socketClient).stopDiscovery();
     }
 
     @Test
+    public void onSocketDestroy_shutdownExecutorService() throws IOException {
+        final MdnsSearchOptions options =
+                MdnsSearchOptions.newBuilder().setNetwork(null /* network */).build();
+        final SocketCreationCallback callback = expectSocketCreationCallback(
+                SERVICE_TYPE_1, mockListenerOne, options);
+        runOnHandler(() -> callback.onSocketCreated(SOCKET_KEY_NULL_NETWORK));
+        verify(mockServiceTypeClientType1NullNetwork).startSendAndReceive(mockListenerOne, options);
+
+        runOnHandler(() -> callback.onSocketDestroyed(SOCKET_KEY_NULL_NETWORK));
+        verify(executorProvider).shutdownExecutorService(mockExecutorService);
+    }
+
+    @Test
     public void registerMultipleListeners() throws IOException {
         final MdnsSearchOptions options =
                 MdnsSearchOptions.newBuilder().setNetwork(null /* network */).build();