Remove the ServiceTypeClient after socket destroyed

The ServiceTypeClient should be removed after socket destroyed
because it's no longer used by any request. The requests will be
matched to the newly created ServiceTypeClient.

Bug: 278635632
Test: atest FrameworksNetTests CtsNetTestCases
(cherry picked from https://android-review.googlesource.com/q/commit:f2a51ac5f8cc6b1cfaaa0f8fd5c28af522806ca0)
Merged-In: Ia917b14d3666f3bfe8e874606a34800a4ce65c5a
Change-Id: Ia917b14d3666f3bfe8e874606a34800a4ce65c5a

framework-t/src/android/net/nsd/NsdManager.java

@@ -429,6 +429,10 @@ public final class NsdManager {
             private final DiscoveryListener mWrapped;
             private final Executor mWrappedExecutor;
             private final ArraySet<TrackedNsdInfo> mFoundInfo = new ArraySet<>();
+            // When this flag is set to true, no further service found or lost callbacks should be
+            // handled. This flag indicates that the network for this DelegatingDiscoveryListener is
+            // lost, and any further callbacks would be redundant.
+            private boolean mAllServicesLost = false;
 
             private DelegatingDiscoveryListener(Network network, DiscoveryListener listener,
                     Executor executor) {
@@ -445,6 +449,7 @@ public final class NsdManager {
                     serviceInfo.setNetwork(mNetwork);
                     mWrappedExecutor.execute(() -> mWrapped.onServiceLost(serviceInfo));
                 }
+                mAllServicesLost = true;
             }
 
             @Override
@@ -486,12 +491,22 @@ public final class NsdManager {
 
             @Override
             public void onServiceFound(NsdServiceInfo serviceInfo) {
+                if (mAllServicesLost) {
+                    // This DelegatingDiscoveryListener no longer has a network connection. Ignore
+                    // the callback.
+                    return;
+                }
                 mFoundInfo.add(new TrackedNsdInfo(serviceInfo));
                 mWrappedExecutor.execute(() -> mWrapped.onServiceFound(serviceInfo));
             }
 
             @Override
             public void onServiceLost(NsdServiceInfo serviceInfo) {
+                if (mAllServicesLost) {
+                    // This DelegatingDiscoveryListener no longer has a network connection. Ignore
+                    // the callback.
+                    return;
+                }
                 mFoundInfo.remove(new TrackedNsdInfo(serviceInfo));
                 mWrappedExecutor.execute(() -> mWrapped.onServiceLost(serviceInfo));
             }

service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java

@@ -135,17 +135,35 @@ public class MdnsDiscoveryManager implements MdnsSocketClientBase.Callback {
             }
         }
         // Request the network for discovery.
-        socketClient.notifyNetworkRequested(listener, searchOptions.getNetwork(), network -> {
+        socketClient.notifyNetworkRequested(listener, searchOptions.getNetwork(),
-            synchronized (this) {
+                new MdnsSocketClientBase.SocketCreationCallback() {
-                // All listeners of the same service types shares the same MdnsServiceTypeClient.
+                    @Override
+                    public void onSocketCreated(@Nullable Network network) {
+                        synchronized (MdnsDiscoveryManager.this) {
+                            // All listeners of the same service types shares the same
+                            // MdnsServiceTypeClient.
                             MdnsServiceTypeClient serviceTypeClient =
                                     perNetworkServiceTypeClients.get(serviceType, network);
                             if (serviceTypeClient == null) {
                                 serviceTypeClient = createServiceTypeClient(serviceType, network);
-                    perNetworkServiceTypeClients.put(serviceType, network, serviceTypeClient);
+                                perNetworkServiceTypeClients.put(serviceType, network,
+                                        serviceTypeClient);
                             }
                             serviceTypeClient.startSendAndReceive(listener, searchOptions);
                         }
+                    }
+
+                    @Override
+                    public void onSocketDestroyed(@Nullable Network network) {
+                        synchronized (MdnsDiscoveryManager.this) {
+                            final MdnsServiceTypeClient serviceTypeClient =
+                                    perNetworkServiceTypeClients.get(serviceType, network);
+                            if (serviceTypeClient == null) return;
+                            // Notify all listeners that all services are removed from this socket.
+                            serviceTypeClient.notifyAllServicesRemoved();
+                            perNetworkServiceTypeClients.remove(serviceTypeClient);
+                        }
+                    }
                 });
     }
 

service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java

@@ -90,6 +90,7 @@ public class MdnsMultinetworkSocketClient implements MdnsSocketClientBase {
                 @NonNull MdnsInterfaceSocket socket) {
             mSocketPacketHandlers.remove(socket);
             mActiveNetworkSockets.remove(socket);
+            mSocketCreationCallback.onSocketDestroyed(network);
         }
     }
 

service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java

@@ -263,6 +263,28 @@ public class MdnsServiceTypeClient {
         }
     }
 
+    /** Notify all services are removed because the socket is destroyed. */
+    public void notifyAllServicesRemoved() {
+        synchronized (lock) {
+            for (MdnsResponse response : instanceNameToResponse.values()) {
+                final String name = response.getServiceInstanceName();
+                if (name == null) continue;
+                for (int i = 0; i < listeners.size(); i++) {
+                    if (!responseMatchesOptions(response, listeners.valueAt(i))) continue;
+                    final MdnsServiceBrowserListener listener = listeners.keyAt(i);
+                    final MdnsServiceInfo serviceInfo =
+                            buildMdnsServiceInfoFromResponse(response, serviceTypeLabels);
+                    if (response.isComplete()) {
+                        sharedLog.log("Socket destroyed. onServiceRemoved: " + name);
+                        listener.onServiceRemoved(serviceInfo);
+                    }
+                    sharedLog.log("Socket destroyed. onServiceNameRemoved: " + name);
+                    listener.onServiceNameRemoved(serviceInfo);
+                }
+            }
+        }
+    }
+
     private void onResponseModified(@NonNull MdnsResponse response) {
         final String serviceInstanceName = response.getServiceInstanceName();
         final MdnsResponse currentResponse =

service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java

@@ -86,5 +86,8 @@ public interface MdnsSocketClientBase {
     interface SocketCreationCallback {
         /*** Notify requested socket is created */
         void onSocketCreated(@Nullable Network network);
+
+        /*** Notify requested socket is destroyed */
+        void onSocketDestroyed(@Nullable Network network);
     }
 }

tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java

@@ -20,6 +20,7 @@ import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -193,6 +194,60 @@ public class MdnsDiscoveryManagerTests {
         verify(mockServiceTypeClientTwo2).processResponse(responseForSubtype, ifIndex, NETWORK_2);
     }
 
+    @Test
+    public void testSocketCreatedAndDestroyed() throws IOException {
+        // Create a ServiceTypeClient for SERVICE_TYPE_1 and NETWORK_1
+        final MdnsSearchOptions options1 =
+                MdnsSearchOptions.newBuilder().setNetwork(NETWORK_1).build();
+        final SocketCreationCallback callback = expectSocketCreationCallback(
+                SERVICE_TYPE_1, mockListenerOne, options1);
+        callback.onSocketCreated(NETWORK_1);
+        verify(mockServiceTypeClientOne1).startSendAndReceive(mockListenerOne, options1);
+
+        // Create a ServiceTypeClient for SERVICE_TYPE_2 and NETWORK_2
+        final MdnsSearchOptions options2 =
+                MdnsSearchOptions.newBuilder().setNetwork(NETWORK_2).build();
+        final SocketCreationCallback callback2 = expectSocketCreationCallback(
+                SERVICE_TYPE_2, mockListenerTwo, options2);
+        callback2.onSocketCreated(NETWORK_2);
+        verify(mockServiceTypeClientTwo2).startSendAndReceive(mockListenerTwo, options2);
+
+        // Receive a response, it should be processed on both clients.
+        final MdnsPacket response = createMdnsPacket(SERVICE_TYPE_1);
+        final int ifIndex = 1;
+        discoveryManager.onResponseReceived(response, ifIndex, null /* network */);
+        verify(mockServiceTypeClientOne1).processResponse(response, ifIndex, null /* network */);
+        verify(mockServiceTypeClientTwo2).processResponse(response, ifIndex, null /* network */);
+
+        // The client for NETWORK_1 receives the callback that the NETWORK_1 has been destroyed,
+        // mockServiceTypeClientOne1 should send service removed notifications and remove from the
+        // list of clients.
+        callback.onSocketDestroyed(NETWORK_1);
+        verify(mockServiceTypeClientOne1).notifyAllServicesRemoved();
+
+        // Receive a response again, it should be processed only on mockServiceTypeClientTwo2.
+        // Because the mockServiceTypeClientOne1 is removed from the list of clients, it is no
+        // longer able to process responses.
+        discoveryManager.onResponseReceived(response, ifIndex, null /* network */);
+        verify(mockServiceTypeClientOne1, times(1))
+                .processResponse(response, ifIndex, null /* network */);
+        verify(mockServiceTypeClientTwo2, times(2))
+                .processResponse(response, ifIndex, null /* network */);
+
+        // The client for NETWORK_2 receives the callback that the NETWORK_1 has been destroyed,
+        // mockServiceTypeClientTwo2 shouldn't send any notifications.
+        callback2.onSocketDestroyed(NETWORK_1);
+        verify(mockServiceTypeClientTwo2, never()).notifyAllServicesRemoved();
+
+        // Receive a response again, mockServiceTypeClientTwo2 is still in the list of clients, it's
+        // still able to process responses.
+        discoveryManager.onResponseReceived(response, ifIndex, null /* network */);
+        verify(mockServiceTypeClientOne1, times(1))
+                .processResponse(response, ifIndex, null /* network */);
+        verify(mockServiceTypeClientTwo2, times(3))
+                .processResponse(response, ifIndex, null /* network */);
+    }
+
     private MdnsPacket createMdnsPacket(String serviceType) {
         final String[] type = TextUtils.split(serviceType, "\\.");
         final ArrayList<String> name = new ArrayList<>(type.length + 1);

tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java

@@ -1056,6 +1056,61 @@ public class MdnsServiceTypeClientTests {
         inOrder.verify(mockListenerTwo).onServiceRemoved(matchServiceName(otherInstance));
     }
 
+    @Test
+    public void testNotifyAllServicesRemoved() {
+        client = new MdnsServiceTypeClient(
+                SERVICE_TYPE, mockSocketClient, currentThreadExecutor, mockNetwork, mockSharedLog);
+
+        final String requestedInstance = "instance1";
+        final String otherInstance = "instance2";
+        final String ipV4Address = "192.0.2.0";
+
+        final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
+                // Use different case in the options
+                .setResolveInstanceName("Instance1").build();
+
+        client.startSendAndReceive(mockListenerOne, resolveOptions);
+        client.startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+
+        // Complete response from instanceName
+        client.processResponse(createResponse(
+                        requestedInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
+                        Collections.emptyMap() /* textAttributes */, TEST_TTL),
+                INTERFACE_INDEX, mockNetwork);
+
+        // Complete response from otherInstanceName
+        client.processResponse(createResponse(
+                        otherInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
+                        Collections.emptyMap() /* textAttributes */, TEST_TTL),
+                INTERFACE_INDEX, mockNetwork);
+
+        client.notifyAllServicesRemoved();
+
+        // mockListenerOne gets notified for the requested instance
+        final InOrder inOrder1 = inOrder(mockListenerOne);
+        inOrder1.verify(mockListenerOne).onServiceNameDiscovered(
+                matchServiceName(requestedInstance));
+        inOrder1.verify(mockListenerOne).onServiceFound(matchServiceName(requestedInstance));
+        inOrder1.verify(mockListenerOne).onServiceRemoved(matchServiceName(requestedInstance));
+        inOrder1.verify(mockListenerOne).onServiceNameRemoved(matchServiceName(requestedInstance));
+        verify(mockListenerOne, never()).onServiceFound(matchServiceName(otherInstance));
+        verify(mockListenerOne, never()).onServiceNameDiscovered(matchServiceName(otherInstance));
+        verify(mockListenerOne, never()).onServiceRemoved(matchServiceName(otherInstance));
+        verify(mockListenerOne, never()).onServiceNameRemoved(matchServiceName(otherInstance));
+
+        // mockListenerTwo gets notified for both though
+        final InOrder inOrder2 = inOrder(mockListenerTwo);
+        inOrder2.verify(mockListenerTwo).onServiceNameDiscovered(
+                matchServiceName(requestedInstance));
+        inOrder2.verify(mockListenerTwo).onServiceFound(matchServiceName(requestedInstance));
+        inOrder2.verify(mockListenerTwo).onServiceNameDiscovered(matchServiceName(otherInstance));
+        inOrder2.verify(mockListenerTwo).onServiceFound(matchServiceName(otherInstance));
+        inOrder2.verify(mockListenerTwo).onServiceRemoved(matchServiceName(otherInstance));
+        inOrder2.verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(otherInstance));
+        inOrder2.verify(mockListenerTwo).onServiceRemoved(matchServiceName(requestedInstance));
+        inOrder2.verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(requestedInstance));
+    }
+
     private static MdnsServiceInfo matchServiceName(String name) {
         return argThat(info -> info.getServiceInstanceName().equals(name));
     }