Merge "Remove the ServiceTypeClient after socket destroyed"

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,18 +135,36 @@ 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
-                MdnsServiceTypeClient serviceTypeClient =
+                    public void onSocketCreated(@Nullable Network network) {
-                        perNetworkServiceTypeClients.get(serviceType, network);
+                        synchronized (MdnsDiscoveryManager.this) {
-                if (serviceTypeClient == null) {
+                            // All listeners of the same service types shares the same
-                    serviceTypeClient = createServiceTypeClient(serviceType, network);
+                            // MdnsServiceTypeClient.
-                    perNetworkServiceTypeClients.put(serviceType, network, serviceTypeClient);
+                            MdnsServiceTypeClient serviceTypeClient =
-                }
+                                    perNetworkServiceTypeClients.get(serviceType, network);
-                serviceTypeClient.startSendAndReceive(listener, searchOptions);
+                            if (serviceTypeClient == null) {
-            }
+                                serviceTypeClient = createServiceTypeClient(serviceType, network);
-        });
+                                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));
     }