Check service ttl expiration

Now, every services are cached on MdnsServiceCache, so the
remaining TTL should be checked when retrieving services from
the MdnsServiceCache and have a callback to notify the
MdnsServiceTypeClient about expired services.

Bug: 265787401
Test: atest FrameworksNetTests CtsNetTestCases
Change-Id: I99da6cc79bdf5df3c899e642e067907501bc9d4f
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 43357e4..ee5f25b 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -1645,8 +1645,8 @@
                         mContext, MdnsFeatureFlags.NSD_FORCE_DISABLE_MDNS_OFFLOAD))
                 .setIncludeInetAddressRecordsInProbing(mDeps.isFeatureEnabled(
                         mContext, MdnsFeatureFlags.INCLUDE_INET_ADDRESS_RECORDS_IN_PROBING))
-                .setIsExpiredServicesRemovalEnabled(mDeps.isTrunkStableFeatureEnabled(
-                        MdnsFeatureFlags.NSD_EXPIRED_SERVICES_REMOVAL))
+                .setIsExpiredServicesRemovalEnabled(mDeps.isFeatureEnabled(
+                        mContext, MdnsFeatureFlags.NSD_EXPIRED_SERVICES_REMOVAL))
                 .setIsLabelCountLimitEnabled(mDeps.isTetheringFeatureNotChickenedOut(
                         mContext, MdnsFeatureFlags.NSD_LIMIT_LABEL_COUNT))
                 .build();
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 738c151..0a6d8c1 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java
@@ -86,7 +86,7 @@
         public Builder() {
             mIsMdnsOffloadFeatureEnabled = false;
             mIncludeInetAddressRecordsInProbing = false;
-            mIsExpiredServicesRemovalEnabled = true; // Default enabled.
+            mIsExpiredServicesRemovalEnabled = false;
             mIsLabelCountLimitEnabled = true; // Default enabled.
         }
 
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsResponse.java b/service-t/src/com/android/server/connectivity/mdns/MdnsResponse.java
index e2288c1..05ad1be 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsResponse.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsResponse.java
@@ -33,6 +33,7 @@
 
 /** An mDNS response. */
 public class MdnsResponse {
+    public static final long EXPIRATION_NEVER = Long.MAX_VALUE;
     private final List<MdnsRecord> records;
     private final List<MdnsPointerRecord> pointerRecords;
     private MdnsServiceRecord serviceRecord;
@@ -349,6 +350,21 @@
         return serviceName;
     }
 
+    /** Get the min remaining ttl time from received records */
+    public long getMinRemainingTtl(long now) {
+        long minRemainingTtl = EXPIRATION_NEVER;
+        // TODO: Check other records(A, AAAA, TXT) ttl time.
+        if (!hasServiceRecord()) {
+            return EXPIRATION_NEVER;
+        }
+        // Check ttl time.
+        long remainingTtl = serviceRecord.getRemainingTTL(now);
+        if (remainingTtl < minRemainingTtl) {
+            minRemainingTtl = remainingTtl;
+        }
+        return minRemainingTtl;
+    }
+
     /**
      * Tests if this response is a goodbye message. This will be true if a service record is present
      * and any of the records have a TTL of 0.
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 d3493c7..e9a41d1 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
@@ -16,16 +16,22 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsResponse.EXPIRATION_NEVER;
 import static com.android.server.connectivity.mdns.util.MdnsUtils.ensureRunningOnHandlerThread;
 import static com.android.server.connectivity.mdns.util.MdnsUtils.equalsIgnoreDnsCase;
 import static com.android.server.connectivity.mdns.util.MdnsUtils.toDnsLowerCase;
 
+import static java.lang.Math.min;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.os.Handler;
 import android.os.Looper;
 import android.util.ArrayMap;
 
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.connectivity.mdns.util.MdnsUtils;
+
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
@@ -67,8 +73,11 @@
         }
     }
     /**
-     * A map of cached services. Key is composed of service name, type and socket. Value is the
-     * service which use the service type to discover from each socket.
+     * A map of cached services. Key is composed of service type and socket. Value is the list of
+     * services which are discovered from the given CacheKey.
+     * When the MdnsFeatureFlags#NSD_EXPIRED_SERVICES_REMOVAL flag is enabled, the lists are sorted
+     * by expiration time, with the earliest entries appearing first. This sorting allows the
+     * removal process to progress through the expiration check efficiently.
      */
     @NonNull
     private final ArrayMap<CacheKey, List<MdnsResponse>> mCachedServices = new ArrayMap<>();
@@ -82,10 +91,20 @@
     private final Handler mHandler;
     @NonNull
     private final MdnsFeatureFlags mMdnsFeatureFlags;
+    @NonNull
+    private final MdnsUtils.Clock mClock;
+    private long mNextExpirationTime = EXPIRATION_NEVER;
 
     public MdnsServiceCache(@NonNull Looper looper, @NonNull MdnsFeatureFlags mdnsFeatureFlags) {
+        this(looper, mdnsFeatureFlags, new MdnsUtils.Clock());
+    }
+
+    @VisibleForTesting
+    MdnsServiceCache(@NonNull Looper looper, @NonNull MdnsFeatureFlags mdnsFeatureFlags,
+            @NonNull MdnsUtils.Clock clock) {
         mHandler = new Handler(looper);
         mMdnsFeatureFlags = mdnsFeatureFlags;
+        mClock = clock;
     }
 
     /**
@@ -97,6 +116,9 @@
     @NonNull
     public List<MdnsResponse> getCachedServices(@NonNull CacheKey cacheKey) {
         ensureRunningOnHandlerThread(mHandler);
+        if (mMdnsFeatureFlags.mIsExpiredServicesRemovalEnabled) {
+            maybeRemoveExpiredServices(cacheKey, mClock.elapsedRealtime());
+        }
         return mCachedServices.containsKey(cacheKey)
                 ? Collections.unmodifiableList(new ArrayList<>(mCachedServices.get(cacheKey)))
                 : Collections.emptyList();
@@ -129,6 +151,9 @@
     @Nullable
     public MdnsResponse getCachedService(@NonNull String serviceName, @NonNull CacheKey cacheKey) {
         ensureRunningOnHandlerThread(mHandler);
+        if (mMdnsFeatureFlags.mIsExpiredServicesRemovalEnabled) {
+            maybeRemoveExpiredServices(cacheKey, mClock.elapsedRealtime());
+        }
         final List<MdnsResponse> responses = mCachedServices.get(cacheKey);
         if (responses == null) {
             return null;
@@ -137,6 +162,16 @@
         return response != null ? new MdnsResponse(response) : null;
     }
 
+    static void insertResponseAndSortList(
+            List<MdnsResponse> responses, MdnsResponse response, long now) {
+        // binarySearch returns "the index of the search key, if it is contained in the list;
+        // otherwise, (-(insertion point) - 1)"
+        final int searchRes = Collections.binarySearch(responses, response,
+                // Sort the list by ttl.
+                (o1, o2) -> Long.compare(o1.getMinRemainingTtl(now), o2.getMinRemainingTtl(now)));
+        responses.add(searchRes >= 0 ? searchRes : (-searchRes - 1), response);
+    }
+
     /**
      * Add or update a service.
      *
@@ -151,7 +186,15 @@
         final MdnsResponse existing =
                 findMatchedResponse(responses, response.getServiceInstanceName());
         responses.remove(existing);
-        responses.add(response);
+        if (mMdnsFeatureFlags.mIsExpiredServicesRemovalEnabled) {
+            final long now = mClock.elapsedRealtime();
+            // Insert and sort service
+            insertResponseAndSortList(responses, response, now);
+            // Update the next expiration check time when a new service is added.
+            mNextExpirationTime = getNextExpirationTime(now);
+        } else {
+            responses.add(response);
+        }
     }
 
     /**
@@ -168,14 +211,25 @@
             return null;
         }
         final Iterator<MdnsResponse> iterator = responses.iterator();
+        MdnsResponse removedResponse = null;
         while (iterator.hasNext()) {
             final MdnsResponse response = iterator.next();
             if (equalsIgnoreDnsCase(serviceName, response.getServiceInstanceName())) {
                 iterator.remove();
-                return response;
+                removedResponse = response;
+                break;
             }
         }
-        return null;
+
+        if (mMdnsFeatureFlags.mIsExpiredServicesRemovalEnabled) {
+            // Remove the serviceType if no response.
+            if (responses.isEmpty()) {
+                mCachedServices.remove(cacheKey);
+            }
+            // Update the next expiration check time when a service is removed.
+            mNextExpirationTime = getNextExpirationTime(mClock.elapsedRealtime());
+        }
+        return removedResponse;
     }
 
     /**
@@ -203,6 +257,87 @@
         mCallbacks.remove(cacheKey);
     }
 
+    private void notifyServiceExpired(@NonNull CacheKey cacheKey,
+            @NonNull MdnsResponse previousResponse, @Nullable MdnsResponse newResponse) {
+        final ServiceExpiredCallback callback = mCallbacks.get(cacheKey);
+        if (callback == null) {
+            // The cached service is no listener.
+            return;
+        }
+        mHandler.post(()-> callback.onServiceRecordExpired(previousResponse, newResponse));
+    }
+
+    static List<MdnsResponse> removeExpiredServices(@NonNull List<MdnsResponse> responses,
+            long now) {
+        final List<MdnsResponse> removedResponses = new ArrayList<>();
+        final Iterator<MdnsResponse> iterator = responses.iterator();
+        while (iterator.hasNext()) {
+            final MdnsResponse response = iterator.next();
+            // TODO: Check other records (A, AAAA, TXT) ttl time and remove the record if it's
+            //  expired. Then send service update notification.
+            if (!response.hasServiceRecord() || response.getMinRemainingTtl(now) > 0) {
+                // The responses are sorted by the service record ttl time. Break out of loop
+                // early if service is not expired or no service record.
+                break;
+            }
+            // Remove the ttl expired service.
+            iterator.remove();
+            removedResponses.add(response);
+        }
+        return removedResponses;
+    }
+
+    private long getNextExpirationTime(long now) {
+        if (mCachedServices.isEmpty()) {
+            return EXPIRATION_NEVER;
+        }
+
+        long minRemainingTtl = EXPIRATION_NEVER;
+        for (int i = 0; i < mCachedServices.size(); i++) {
+            minRemainingTtl = min(minRemainingTtl,
+                    // The empty lists are not kept in the map, so there's always at least one
+                    // element in the list. Therefore, it's fine to get the first element without a
+                    // null check.
+                    mCachedServices.valueAt(i).get(0).getMinRemainingTtl(now));
+        }
+        return minRemainingTtl == EXPIRATION_NEVER ? EXPIRATION_NEVER : now + minRemainingTtl;
+    }
+
+    /**
+     * Check whether the ttl time is expired on each service and notify to the listeners
+     */
+    private void maybeRemoveExpiredServices(CacheKey cacheKey, long now) {
+        ensureRunningOnHandlerThread(mHandler);
+        if (now < mNextExpirationTime) {
+            // Skip the check if ttl time is not expired.
+            return;
+        }
+
+        final List<MdnsResponse> responses = mCachedServices.get(cacheKey);
+        if (responses == null) {
+            // No such services.
+            return;
+        }
+
+        final List<MdnsResponse> removedResponses = removeExpiredServices(responses, now);
+        if (removedResponses.isEmpty()) {
+            // No expired services.
+            return;
+        }
+
+        for (MdnsResponse previousResponse : removedResponses) {
+            notifyServiceExpired(cacheKey, previousResponse, null /* newResponse */);
+        }
+
+        // Remove the serviceType if no response.
+        if (responses.isEmpty()) {
+            mCachedServices.remove(cacheKey);
+        }
+
+        // Update next expiration time.
+        mNextExpirationTime = getNextExpirationTime(now);
+    }
+
     /*** Callbacks for listening service expiration */
     public interface ServiceExpiredCallback {
         /*** Notify the service is expired */
@@ -210,5 +345,5 @@
                 @Nullable MdnsResponse newResponse);
     }
 
-    // TODO: check ttl expiration for each service and notify to the clients.
+    // TODO: Schedule a job to check ttl expiration for all services and notify to the clients.
 }
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 0a03186..32f604e 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -312,8 +312,7 @@
         this.searchOptions = searchOptions;
         boolean hadReply = false;
         if (listeners.put(listener, searchOptions) == null) {
-            for (MdnsResponse existingResponse :
-                    serviceCache.getCachedServices(cacheKey)) {
+            for (MdnsResponse existingResponse : serviceCache.getCachedServices(cacheKey)) {
                 if (!responseMatchesOptions(existingResponse, searchOptions)) continue;
                 final MdnsServiceInfo info =
                         buildMdnsServiceInfoFromResponse(existingResponse, serviceTypeLabels);