Merge "Remove cached services when the network is disconnected or unrequested" into main
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 8e4ec2f..0adb290 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -1938,6 +1938,11 @@
                         mContext, MdnsFeatureFlags.NSD_QUERY_WITH_KNOWN_ANSWER))
                 .setAvoidAdvertisingEmptyTxtRecords(mDeps.isTetheringFeatureNotChickenedOut(
                         mContext, MdnsFeatureFlags.NSD_AVOID_ADVERTISING_EMPTY_TXT_RECORDS))
+                .setIsCachedServicesRemovalEnabled(mDeps.isFeatureEnabled(
+                        mContext, MdnsFeatureFlags.NSD_CACHED_SERVICES_REMOVAL))
+                .setCachedServicesRetentionTime(mDeps.getDeviceConfigPropertyInt(
+                        MdnsFeatureFlags.NSD_CACHED_SERVICES_RETENTION_TIME,
+                        MdnsFeatureFlags.DEFAULT_CACHED_SERVICES_RETENTION_TIME_MILLISECONDS))
                 .setOverrideProvider(new MdnsFeatureFlags.FlagOverrideProvider() {
                     @Override
                     public boolean isForceEnabledForTest(@NonNull String flag) {
@@ -1947,10 +1952,9 @@
                     }
 
                     @Override
-                    public int getIntValueForTest(@NonNull String flag) {
+                    public int getIntValueForTest(@NonNull String flag, int defaultValue) {
                         return mDeps.getDeviceConfigPropertyInt(
-                                FORCE_ENABLE_FLAG_FOR_TEST_PREFIX + flag,
-                                -1 /* defaultValue */);
+                                FORCE_ENABLE_FLAG_FOR_TEST_PREFIX + flag, defaultValue);
                     }
                 })
                 .build();
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 a74bdf7..b16d8bd 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
@@ -301,6 +301,17 @@
                         serviceTypeClient.notifySocketDestroyed();
                         executorProvider.shutdownExecutorService(serviceTypeClient.getExecutor());
                         perSocketServiceTypeClients.remove(serviceTypeClient);
+                        // The cached services may not be reliable after the socket is disconnected,
+                        // the service type client won't receive any updates for them. Therefore,
+                        // remove these cached services after exceeding the retention time
+                        // (currently 10s) if no service type client requires them.
+                        if (mdnsFeatureFlags.isCachedServicesRemovalEnabled()) {
+                            final MdnsServiceCache.CacheKey cacheKey =
+                                    serviceTypeClient.getCacheKey();
+                            discoveryExecutor.executeDelayed(
+                                    () -> handleRemoveCachedServices(cacheKey),
+                                    mdnsFeatureFlags.getCachedServicesRetentionTime());
+                        }
                     }
                 });
     }
@@ -337,6 +348,42 @@
                 // of the service type clients.
                 executorProvider.shutdownExecutorService(serviceTypeClient.getExecutor());
                 perSocketServiceTypeClients.remove(serviceTypeClient);
+                // The cached services may not be reliable after the socket is disconnected, the
+                // service type client won't receive any updates for them. Therefore, remove these
+                // cached services after exceeding the retention time (currently 10s) if no service
+                // type client requires them.
+                // Note: This removal is only called if the requested socket is still active for
+                // other requests. If the requested socket is no longer needed after the listener
+                // is unregistered, SocketCreationCallback#onSocketDestroyed callback will remove
+                // both the service type client and cached services there.
+                //
+                // List some multiple listener cases for the cached service removal flow.
+                //
+                // Case 1 - Same service type, different network requests
+                //  - Register Listener A (service type X, requesting all networks: Y and Z)
+                //  - Create service type clients X-Y and X-Z
+                //  - Register Listener B (service type X, requesting network Y)
+                //  - Reuse service type client X-Y
+                //  - Unregister Listener A
+                //  - Socket destroyed on network Z; remove the X-Z client. Unregister the listener
+                //    from the X-Y client and keep it, as it's still being used by Listener B.
+                //  - Remove cached services associated with the X-Z client after 10 seconds.
+                //
+                // Case 2 - Different service types, same network request
+                //  - Register Listener A (service type X, requesting network Y)
+                //  - Create service type client X-Y
+                //  - Register Listener B (service type Z, requesting network Y)
+                //  - Create service type client Z-Y
+                //  - Unregister Listener A
+                //  - No socket is destroyed because network Y is still being used by Listener B.
+                //  - Unregister the listener from the X-Y client, then remove it.
+                //  - Remove cached services associated with the X-Y client after 10 seconds.
+                if (mdnsFeatureFlags.isCachedServicesRemovalEnabled()) {
+                    final MdnsServiceCache.CacheKey cacheKey = serviceTypeClient.getCacheKey();
+                    discoveryExecutor.executeDelayed(
+                            () -> handleRemoveCachedServices(cacheKey),
+                            mdnsFeatureFlags.getCachedServicesRetentionTime());
+                }
             }
         }
         if (perSocketServiceTypeClients.isEmpty()) {
@@ -381,6 +428,26 @@
         }
     }
 
+    private void handleRemoveCachedServices(@NonNull MdnsServiceCache.CacheKey cacheKey) {
+        // Check if there is an active service type client that requires the cached services. If so,
+        // do not remove associated services from cache.
+        for (MdnsServiceTypeClient client : getMdnsServiceTypeClient(cacheKey.mSocketKey)) {
+            if (client.getCacheKey().equals(cacheKey)) {
+                // Found a client that has same CacheKey.
+                return;
+            }
+        }
+        sharedLog.log("Remove cached services for " + cacheKey);
+        // No client has same CacheKey. Remove associated services.
+        getServiceCache().removeServices(cacheKey);
+    }
+
+    @VisibleForTesting
+    @NonNull
+    MdnsServiceCache getServiceCache() {
+        return serviceCache;
+    }
+
     @VisibleForTesting
     MdnsServiceTypeClient createServiceTypeClient(@NonNull String serviceType,
             @NonNull SocketKey socketKey) {
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java b/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java
index b2be6ce..4e27fef 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java
@@ -73,6 +73,22 @@
     public static final String NSD_AVOID_ADVERTISING_EMPTY_TXT_RECORDS =
             "nsd_avoid_advertising_empty_txt_records";
 
+    /**
+     * A feature flag to control whether the cached services removal should be enabled.
+     * The removal will be triggered if the retention time has elapsed after all listeners have been
+     * unregistered from the service type client or the interface has been destroyed.
+     */
+    public static final String NSD_CACHED_SERVICES_REMOVAL = "nsd_cached_services_removal";
+
+    /**
+     * A feature flag to control the retention time for cached services.
+     *
+     * <p> Making the retention time configurable allows for testing and future adjustments.
+     */
+    public static final String NSD_CACHED_SERVICES_RETENTION_TIME =
+            "nsd_cached_services_retention_time";
+    public static final int DEFAULT_CACHED_SERVICES_RETENTION_TIME_MILLISECONDS = 10000;
+
     // Flag for offload feature
     public final boolean mIsMdnsOffloadFeatureEnabled;
 
@@ -100,6 +116,12 @@
     // Flag for avoiding advertising empty TXT records
     public final boolean mAvoidAdvertisingEmptyTxtRecords;
 
+    // Flag for cached services removal
+    public final boolean mIsCachedServicesRemovalEnabled;
+
+    // Retention Time for cached services
+    public final long mCachedServicesRetentionTime;
+
     @Nullable
     private final FlagOverrideProvider mOverrideProvider;
 
@@ -116,7 +138,7 @@
         /**
          * Get the int value of the flag for testing purposes.
          */
-        int getIntValueForTest(@NonNull String flag);
+        int getIntValueForTest(@NonNull String flag, int defaultValue);
     }
 
     /**
@@ -129,13 +151,14 @@
     /**
      * Get the int value of the flag for testing purposes.
      *
-     * @return the test int value, or -1 if it is unset or the OverrideProvider doesn't exist.
+     * @return the test int value, or given default value if it is unset or the OverrideProvider
+     * doesn't exist.
      */
-    private int getIntValueForTest(@NonNull String flag) {
+    private int getIntValueForTest(@NonNull String flag, int defaultValue) {
         if (mOverrideProvider == null) {
-            return -1;
+            return defaultValue;
         }
-        return mOverrideProvider.getIntValueForTest(flag);
+        return mOverrideProvider.getIntValueForTest(flag, defaultValue);
     }
 
     /**
@@ -178,6 +201,23 @@
     }
 
     /**
+     * Indicates whether {@link #NSD_CACHED_SERVICES_REMOVAL} is enabled, including for testing.
+     */
+    public boolean isCachedServicesRemovalEnabled() {
+        return mIsCachedServicesRemovalEnabled
+                || isForceEnabledForTest(NSD_CACHED_SERVICES_REMOVAL);
+    }
+
+    /**
+     * Get the value which is set to {@link #NSD_CACHED_SERVICES_RETENTION_TIME}, including for
+     * testing.
+     */
+    public long getCachedServicesRetentionTime() {
+        return getIntValueForTest(
+                NSD_CACHED_SERVICES_RETENTION_TIME, (int) mCachedServicesRetentionTime);
+    }
+
+    /**
      * The constructor for {@link MdnsFeatureFlags}.
      */
     public MdnsFeatureFlags(boolean isOffloadFeatureEnabled,
@@ -189,6 +229,8 @@
             boolean isAggressiveQueryModeEnabled,
             boolean isQueryWithKnownAnswerEnabled,
             boolean avoidAdvertisingEmptyTxtRecords,
+            boolean isCachedServicesRemovalEnabled,
+            long cachedServicesRetentionTime,
             @Nullable FlagOverrideProvider overrideProvider) {
         mIsMdnsOffloadFeatureEnabled = isOffloadFeatureEnabled;
         mIncludeInetAddressRecordsInProbing = includeInetAddressRecordsInProbing;
@@ -199,6 +241,8 @@
         mIsAggressiveQueryModeEnabled = isAggressiveQueryModeEnabled;
         mIsQueryWithKnownAnswerEnabled = isQueryWithKnownAnswerEnabled;
         mAvoidAdvertisingEmptyTxtRecords = avoidAdvertisingEmptyTxtRecords;
+        mIsCachedServicesRemovalEnabled = isCachedServicesRemovalEnabled;
+        mCachedServicesRetentionTime = cachedServicesRetentionTime;
         mOverrideProvider = overrideProvider;
     }
 
@@ -220,6 +264,8 @@
         private boolean mIsAggressiveQueryModeEnabled;
         private boolean mIsQueryWithKnownAnswerEnabled;
         private boolean mAvoidAdvertisingEmptyTxtRecords;
+        private boolean mIsCachedServicesRemovalEnabled;
+        private long mCachedServicesRetentionTime;
         private FlagOverrideProvider mOverrideProvider;
 
         /**
@@ -235,6 +281,8 @@
             mIsAggressiveQueryModeEnabled = false;
             mIsQueryWithKnownAnswerEnabled = false;
             mAvoidAdvertisingEmptyTxtRecords = true; // Default enabled.
+            mIsCachedServicesRemovalEnabled = false;
+            mCachedServicesRetentionTime = DEFAULT_CACHED_SERVICES_RETENTION_TIME_MILLISECONDS;
             mOverrideProvider = null;
         }
 
@@ -341,6 +389,26 @@
         }
 
         /**
+         * Set whether the cached services removal is enabled.
+         *
+         * @see #NSD_CACHED_SERVICES_REMOVAL
+         */
+        public Builder setIsCachedServicesRemovalEnabled(boolean isCachedServicesRemovalEnabled) {
+            mIsCachedServicesRemovalEnabled = isCachedServicesRemovalEnabled;
+            return this;
+        }
+
+        /**
+         * Set cached services retention time.
+         *
+         * @see #NSD_CACHED_SERVICES_RETENTION_TIME
+         */
+        public Builder setCachedServicesRetentionTime(long cachedServicesRetentionTime) {
+            mCachedServicesRetentionTime = cachedServicesRetentionTime;
+            return this;
+        }
+
+        /**
          * Builds a {@link MdnsFeatureFlags} with the arguments supplied to this builder.
          */
         public MdnsFeatureFlags build() {
@@ -353,6 +421,8 @@
                     mIsAggressiveQueryModeEnabled,
                     mIsQueryWithKnownAnswerEnabled,
                     mAvoidAdvertisingEmptyTxtRecords,
+                    mIsCachedServicesRemovalEnabled,
+                    mCachedServicesRetentionTime,
                     mOverrideProvider);
         }
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
index 591ed8b..22f7a03 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
@@ -49,7 +49,7 @@
  *  to their default value (0, false or null).
  */
 public class MdnsServiceCache {
-    static class CacheKey {
+    public static class CacheKey {
         @NonNull final String mUpperCaseServiceType;
         @NonNull final SocketKey mSocketKey;
 
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 4b55ea9..a5dd536 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -456,6 +456,14 @@
         return executor;
     }
 
+    /**
+     * Get the cache key for this service type client.
+     */
+    @NonNull
+    public MdnsServiceCache.CacheKey getCacheKey() {
+        return cacheKey;
+    }
+
     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 ec47618..d801fba 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
@@ -102,6 +102,7 @@
     @Mock MdnsServiceBrowserListener mockListenerOne;
     @Mock MdnsServiceBrowserListener mockListenerTwo;
     @Mock SharedLog sharedLog;
+    @Mock MdnsServiceCache mockServiceCache;
     private MdnsDiscoveryManager discoveryManager;
     private HandlerThread thread;
     private Handler handler;
@@ -145,7 +146,9 @@
                         return null;
                     }
                 };
+        discoveryManager = makeDiscoveryManager(MdnsFeatureFlags.newBuilder().build());
         doReturn(mockExecutorService).when(mockServiceTypeClientType1NullNetwork).getExecutor();
+        doReturn(mockExecutorService).when(mockServiceTypeClientType1Network1).getExecutor();
     }
 
     @After
@@ -156,6 +159,40 @@
         }
     }
 
+    private MdnsDiscoveryManager makeDiscoveryManager(@NonNull MdnsFeatureFlags featureFlags) {
+        return new MdnsDiscoveryManager(executorProvider, socketClient, sharedLog, featureFlags) {
+            @Override
+            MdnsServiceTypeClient createServiceTypeClient(@NonNull String serviceType,
+                    @NonNull SocketKey socketKey) {
+                createdServiceTypeClientCount++;
+                final Pair<String, SocketKey> perSocketServiceType =
+                        Pair.create(serviceType, socketKey);
+                if (perSocketServiceType.equals(PER_SOCKET_SERVICE_TYPE_1_NULL_NETWORK)) {
+                    return mockServiceTypeClientType1NullNetwork;
+                } else if (perSocketServiceType.equals(
+                        PER_SOCKET_SERVICE_TYPE_1_NETWORK_1)) {
+                    return mockServiceTypeClientType1Network1;
+                } else if (perSocketServiceType.equals(
+                        PER_SOCKET_SERVICE_TYPE_2_NULL_NETWORK)) {
+                    return mockServiceTypeClientType2NullNetwork;
+                } else if (perSocketServiceType.equals(
+                        PER_SOCKET_SERVICE_TYPE_2_NETWORK_1)) {
+                    return mockServiceTypeClientType2Network1;
+                } else if (perSocketServiceType.equals(
+                        PER_SOCKET_SERVICE_TYPE_2_NETWORK_2)) {
+                    return mockServiceTypeClientType2Network2;
+                }
+                fail("Unexpected perSocketServiceType: " + perSocketServiceType);
+                return null;
+            }
+
+            @Override
+            MdnsServiceCache getServiceCache() {
+                return mockServiceCache;
+            }
+        };
+    }
+
     private void runOnHandler(Runnable r) {
         handler.post(r);
         HandlerUtils.waitForIdle(handler, DEFAULT_TIMEOUT);
@@ -438,6 +475,57 @@
         }
     }
 
+    @Test
+    public void testRemoveServicesAfterAllListenersUnregistered() throws IOException {
+        final MdnsFeatureFlags mdnsFeatureFlags = MdnsFeatureFlags.newBuilder()
+                .setIsCachedServicesRemovalEnabled(true)
+                .setCachedServicesRetentionTime(0L)
+                .build();
+        discoveryManager = makeDiscoveryManager(mdnsFeatureFlags);
+
+        final MdnsSearchOptions options =
+                MdnsSearchOptions.newBuilder().setNetwork(NETWORK_1).build();
+        final SocketCreationCallback callback = expectSocketCreationCallback(
+                SERVICE_TYPE_1, mockListenerOne, options);
+        runOnHandler(() -> callback.onSocketCreated(SOCKET_KEY_NETWORK_1));
+        verify(mockServiceTypeClientType1Network1).startSendAndReceive(mockListenerOne, options);
+
+        final MdnsServiceCache.CacheKey cacheKey =
+                new MdnsServiceCache.CacheKey(SERVICE_TYPE_1, SOCKET_KEY_NETWORK_1);
+        doReturn(cacheKey).when(mockServiceTypeClientType1Network1).getCacheKey();
+        doReturn(true).when(mockServiceTypeClientType1Network1)
+                .stopSendAndReceive(mockListenerOne);
+        runOnHandler(() -> discoveryManager.unregisterListener(SERVICE_TYPE_1, mockListenerOne));
+        verify(executorProvider).shutdownExecutorService(mockExecutorService);
+        verify(mockServiceTypeClientType1Network1).stopSendAndReceive(mockListenerOne);
+        verify(socketClient).stopDiscovery();
+        verify(mockServiceCache).removeServices(cacheKey);
+    }
+
+    @Test
+    public void testRemoveServicesAfterSocketDestroyed() throws IOException {
+        final MdnsFeatureFlags mdnsFeatureFlags = MdnsFeatureFlags.newBuilder()
+                .setIsCachedServicesRemovalEnabled(true)
+                .setCachedServicesRetentionTime(0L)
+                .build();
+        discoveryManager = makeDiscoveryManager(mdnsFeatureFlags);
+
+        final MdnsSearchOptions options =
+                MdnsSearchOptions.newBuilder().setNetwork(NETWORK_1).build();
+        final SocketCreationCallback callback = expectSocketCreationCallback(
+                SERVICE_TYPE_1, mockListenerOne, options);
+        runOnHandler(() -> callback.onSocketCreated(SOCKET_KEY_NETWORK_1));
+        verify(mockServiceTypeClientType1Network1).startSendAndReceive(mockListenerOne, options);
+
+        final MdnsServiceCache.CacheKey cacheKey =
+                new MdnsServiceCache.CacheKey(SERVICE_TYPE_1, SOCKET_KEY_NETWORK_1);
+        doReturn(cacheKey).when(mockServiceTypeClientType1Network1).getCacheKey();
+        runOnHandler(() -> callback.onSocketDestroyed(SOCKET_KEY_NETWORK_1));
+        verify(mockServiceTypeClientType1Network1).notifySocketDestroyed();
+        verify(executorProvider).shutdownExecutorService(mockExecutorService);
+        verify(mockServiceCache).removeServices(cacheKey);
+    }
+
     private MdnsPacket createMdnsPacket(String serviceType) {
         final String[] type = TextUtils.split(serviceType, "\\.");
         final ArrayList<String> name = new ArrayList<>(type.length + 1);