diff --git a/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt b/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt
index 0885f4f..ece4a34 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt
+++ b/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt
@@ -58,11 +58,6 @@
     @After
     @Throws(Exception::class)
     fun tearDown() {
-        // cancel active requests to enable engine shutdown.
-        stream?.let {
-            it.cancel()
-            callback.blockForDone()
-        }
         httpEngine.shutdown()
     }
 
diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
index c5f9631..ec63e41 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
@@ -119,7 +119,14 @@
       // the session and delegates writing. The corresponding handler will write
       // with the setting specified in the trace config.
       NetworkTraceHandler::Trace([&](NetworkTraceHandler::TraceContext ctx) {
-        ctx.GetDataSourceLocked()->Write(packets, ctx);
+        perfetto::LockedHandle<NetworkTraceHandler> handle =
+            ctx.GetDataSourceLocked();
+        // The underlying handle can be invalidated between when Trace starts
+        // and GetDataSourceLocked is called, but not while the LockedHandle
+        // exists and holds the lock. Check validity prior to use.
+        if (handle.valid()) {
+          handle->Write(packets, ctx);
+        }
       });
     });
 
diff --git a/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java b/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
index bd4ec20..13f6dac 100644
--- a/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
+++ b/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
@@ -17,7 +17,6 @@
 package com.android.server.connectivity.mdns;
 
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.text.TextUtils;
 import android.util.Log;
 import android.util.Pair;
@@ -106,12 +105,11 @@
     // Incompatible return type for override of Callable#call().
     @SuppressWarnings("nullness:override.return.invalid")
     @Override
-    @Nullable
     public Pair<Integer, List<String>> call() {
         try {
             MdnsSocketClientBase requestSender = weakRequestSender.get();
             if (requestSender == null) {
-                return null;
+                return Pair.create(-1, new ArrayList<>());
             }
 
             int numQuestions = 0;
@@ -158,7 +156,7 @@
 
             if (numQuestions == 0) {
                 // No query to send
-                return null;
+                return Pair.create(-1, new ArrayList<>());
             }
 
             // Header.
@@ -197,7 +195,7 @@
         } catch (IOException e) {
             LOGGER.e(String.format("Failed to create mDNS packet for subtype: %s.",
                     TextUtils.join(",", subtypes)), e);
-            return null;
+            return Pair.create(-1, new ArrayList<>());
         }
     }
 
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java b/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java
index f5e7790..8cb3e96 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java
@@ -54,10 +54,6 @@
         return true;
     }
 
-    public static boolean shouldCancelScanTaskWhenFutureIsNull() {
-        return false;
-    }
-
     public static long sleepTimeForSocketThreadMs() {
         return 20_000L;
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsConstants.java b/service-t/src/com/android/server/connectivity/mdns/MdnsConstants.java
index ce5f540..0c32cf1 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsConstants.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsConstants.java
@@ -16,18 +16,13 @@
 
 package com.android.server.connectivity.mdns;
 
-import static com.android.internal.annotations.VisibleForTesting.Visibility.PACKAGE;
-
 import static java.nio.charset.StandardCharsets.UTF_8;
 
-import com.android.internal.annotations.VisibleForTesting;
-
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.nio.charset.Charset;
 
 /** mDNS-related constants. */
-@VisibleForTesting(visibility = PACKAGE)
 public final class MdnsConstants {
     public static final int MDNS_PORT = 5353;
     // Flags word format is:
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 dfaec75..f386dd4 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
@@ -51,6 +51,7 @@
     @NonNull private final PerSocketServiceTypeClients perSocketServiceTypeClients;
     @NonNull private final Handler handler;
     @Nullable private final HandlerThread handlerThread;
+    @NonNull private final MdnsServiceCache serviceCache;
 
     private static class PerSocketServiceTypeClients {
         private final ArrayMap<Pair<String, SocketKey>, MdnsServiceTypeClient> clients =
@@ -119,10 +120,12 @@
         if (socketClient.getLooper() != null) {
             this.handlerThread = null;
             this.handler = new Handler(socketClient.getLooper());
+            this.serviceCache = new MdnsServiceCache(socketClient.getLooper());
         } else {
             this.handlerThread = new HandlerThread(MdnsDiscoveryManager.class.getSimpleName());
             this.handlerThread.start();
             this.handler = new Handler(handlerThread.getLooper());
+            this.serviceCache = new MdnsServiceCache(handlerThread.getLooper());
         }
     }
 
@@ -289,6 +292,6 @@
         return new MdnsServiceTypeClient(
                 serviceType, socketClient,
                 executorProvider.newServiceTypeClientSchedulerExecutor(), socketKey,
-                sharedLog.forSubComponent(tag), handler.getLooper());
+                sharedLog.forSubComponent(tag), handler.getLooper(), serviceCache);
     }
 }
\ No newline at end of file
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 dc99e49..ec6af9b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
@@ -96,7 +96,14 @@
                 : Collections.emptyList();
     }
 
-    private MdnsResponse findMatchedResponse(@NonNull List<MdnsResponse> responses,
+    /**
+     * Find a matched response for given service name
+     *
+     * @param responses the responses to be searched.
+     * @param serviceName the target service name
+     * @return the response which matches the given service name or null if not found.
+     */
+    public static MdnsResponse findMatchedResponse(@NonNull List<MdnsResponse> responses,
             @NonNull String serviceName) {
         for (MdnsResponse response : responses) {
             if (equalsIgnoreDnsCase(serviceName, response.getServiceInstanceName())) {
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 9c49b8f..b15defd 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -16,20 +16,19 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsServiceCache.findMatchedResponse;
 import static com.android.server.connectivity.mdns.util.MdnsUtils.ensureRunningOnHandlerThread;
 
-import static java.util.concurrent.TimeUnit.MILLISECONDS;
-
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.Message;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Pair;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.SharedLog;
@@ -41,11 +40,8 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
-import java.util.Map;
-import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 
 /**
@@ -56,6 +52,9 @@
 
     private static final String TAG = MdnsServiceTypeClient.class.getSimpleName();
     private static final int DEFAULT_MTU = 1500;
+    @VisibleForTesting
+    static final int EVENT_START_QUERYTASK = 1;
+    static final int EVENT_QUERY_RESULT = 2;
 
     private final String serviceType;
     private final String[] serviceTypeLabels;
@@ -65,12 +64,13 @@
     @NonNull private final SocketKey socketKey;
     @NonNull private final SharedLog sharedLog;
     @NonNull private final Handler handler;
-    private final Object lock = new Object();
+    @NonNull private final Dependencies dependencies;
+    /**
+     * The service caches for each socket. It should be accessed from looper thread only.
+     */
+    @NonNull private final MdnsServiceCache serviceCache;
     private final ArrayMap<MdnsServiceBrowserListener, MdnsSearchOptions> listeners =
             new ArrayMap<>();
-    // TODO: change instanceNameToResponse to TreeMap with case insensitive comparator.
-    @GuardedBy("lock")
-    private final Map<String, MdnsResponse> instanceNameToResponse = new HashMap<>();
     private final boolean removeServiceAfterTtlExpires =
             MdnsConfigs.removeServiceAfterTtlExpires();
     private final MdnsResponseDecoder.Clock clock;
@@ -82,17 +82,103 @@
     // new subtypes. It stays the same between packets for same subtypes.
     private long currentSessionId = 0;
 
-    @GuardedBy("lock")
     @Nullable
-    private Future<?> nextQueryTaskFuture;
-
-    @GuardedBy("lock")
-    @Nullable
-    private QueryTask lastScheduledTask;
-
-    @GuardedBy("lock")
+    private ScheduledQueryTaskArgs lastScheduledQueryTaskArgs;
     private long lastSentTime;
 
+    private class QueryTaskHandler extends Handler {
+        QueryTaskHandler(Looper looper) {
+            super(looper);
+        }
+
+        @Override
+        @SuppressWarnings("FutureReturnValueIgnored")
+        public void handleMessage(Message msg) {
+            switch (msg.what) {
+                case EVENT_START_QUERYTASK: {
+                    final ScheduledQueryTaskArgs taskArgs = (ScheduledQueryTaskArgs) msg.obj;
+                    // QueryTask should be run immediately after being created (not be scheduled in
+                    // advance). Because the result of "makeResponsesForResolve" depends on answers
+                    // that were received before it is called, so to take into account all answers
+                    // before sending the query, it needs to be called just before sending it.
+                    final List<MdnsResponse> servicesToResolve = makeResponsesForResolve(socketKey);
+                    final QueryTask queryTask = new QueryTask(taskArgs, servicesToResolve,
+                            servicesToResolve.size() < listeners.size() /* sendDiscoveryQueries */);
+                    executor.submit(queryTask);
+                    break;
+                }
+                case EVENT_QUERY_RESULT: {
+                    final QuerySentResult sentResult = (QuerySentResult) msg.obj;
+                    if (MdnsConfigs.useSessionIdToScheduleMdnsTask()) {
+                        // In case that the task is not canceled successfully, use session ID to
+                        // check if this task should continue to schedule more.
+                        if (sentResult.taskArgs.sessionId != currentSessionId) {
+                            break;
+                        }
+                    }
+
+                    if ((sentResult.transactionId != -1)) {
+                        for (int i = 0; i < listeners.size(); i++) {
+                            listeners.keyAt(i).onDiscoveryQuerySent(
+                                    sentResult.subTypes, sentResult.transactionId);
+                        }
+                    }
+
+                    tryRemoveServiceAfterTtlExpires();
+
+                    final QueryTaskConfig nextRunConfig =
+                            sentResult.taskArgs.config.getConfigForNextRun();
+                    final long now = clock.elapsedRealtime();
+                    lastSentTime = now;
+                    final long minRemainingTtl = getMinRemainingTtl(now);
+                    final long timeToRun = calculateTimeToRun(lastScheduledQueryTaskArgs,
+                            nextRunConfig, now, minRemainingTtl, lastSentTime);
+                    scheduleNextRun(nextRunConfig, minRemainingTtl, now, timeToRun,
+                            lastScheduledQueryTaskArgs.sessionId);
+                    break;
+                }
+                default:
+                    sharedLog.e("Unrecognized event " + msg.what);
+                    break;
+            }
+        }
+    }
+
+    /**
+     * Dependencies of MdnsServiceTypeClient, for injection in tests.
+     */
+    @VisibleForTesting
+    public static class Dependencies {
+        /**
+         * @see Handler#sendMessageDelayed(Message, long)
+         */
+        public void sendMessageDelayed(@NonNull Handler handler, @NonNull Message message,
+                long delayMillis) {
+            handler.sendMessageDelayed(message, delayMillis);
+        }
+
+        /**
+         * @see Handler#removeMessages(int)
+         */
+        public void removeMessages(@NonNull Handler handler, int what) {
+            handler.removeMessages(what);
+        }
+
+        /**
+         * @see Handler#hasMessages(int)
+         */
+        public boolean hasMessages(@NonNull Handler handler, int what) {
+            return handler.hasMessages(what);
+        }
+
+        /**
+         * @see Handler#post(Runnable)
+         */
+        public void sendMessage(@NonNull Handler handler, @NonNull Message message) {
+            handler.sendMessage(message);
+        }
+    }
+
     /**
      * Constructor of {@link MdnsServiceTypeClient}.
      *
@@ -105,9 +191,10 @@
             @NonNull ScheduledExecutorService executor,
             @NonNull SocketKey socketKey,
             @NonNull SharedLog sharedLog,
-            @NonNull Looper looper) {
+            @NonNull Looper looper,
+            @NonNull MdnsServiceCache serviceCache) {
         this(serviceType, socketClient, executor, new MdnsResponseDecoder.Clock(), socketKey,
-                sharedLog, looper);
+                sharedLog, looper, new Dependencies(), serviceCache);
     }
 
     @VisibleForTesting
@@ -118,7 +205,9 @@
             @NonNull MdnsResponseDecoder.Clock clock,
             @NonNull SocketKey socketKey,
             @NonNull SharedLog sharedLog,
-            @NonNull Looper looper) {
+            @NonNull Looper looper,
+            @NonNull Dependencies dependencies,
+            @NonNull MdnsServiceCache serviceCache) {
         this.serviceType = serviceType;
         this.socketClient = socketClient;
         this.executor = executor;
@@ -127,7 +216,9 @@
         this.clock = clock;
         this.socketKey = socketKey;
         this.sharedLog = sharedLog;
-        this.handler = new Handler(looper);
+        this.handler = new QueryTaskHandler(looper);
+        this.dependencies = dependencies;
+        this.serviceCache = serviceCache;
     }
 
     private static MdnsServiceInfo buildMdnsServiceInfoFromResponse(
@@ -187,69 +278,63 @@
      * @param listener      The {@link MdnsServiceBrowserListener} to register.
      * @param searchOptions {@link MdnsSearchOptions} contains the list of subtypes to discover.
      */
+    @SuppressWarnings("FutureReturnValueIgnored")
     public void startSendAndReceive(
             @NonNull MdnsServiceBrowserListener listener,
             @NonNull MdnsSearchOptions searchOptions) {
         ensureRunningOnHandlerThread(handler);
-        synchronized (lock) {
-            this.searchOptions = searchOptions;
-            boolean hadReply = false;
-            if (listeners.put(listener, searchOptions) == null) {
-                for (MdnsResponse existingResponse : instanceNameToResponse.values()) {
-                    if (!responseMatchesOptions(existingResponse, searchOptions)) continue;
-                    final MdnsServiceInfo info =
-                            buildMdnsServiceInfoFromResponse(existingResponse, serviceTypeLabels);
-                    listener.onServiceNameDiscovered(info);
-                    if (existingResponse.isComplete()) {
-                        listener.onServiceFound(info);
-                        hadReply = true;
-                    }
+        this.searchOptions = searchOptions;
+        boolean hadReply = false;
+        if (listeners.put(listener, searchOptions) == null) {
+            for (MdnsResponse existingResponse :
+                    serviceCache.getCachedServices(serviceType, socketKey)) {
+                if (!responseMatchesOptions(existingResponse, searchOptions)) continue;
+                final MdnsServiceInfo info =
+                        buildMdnsServiceInfoFromResponse(existingResponse, serviceTypeLabels);
+                listener.onServiceNameDiscovered(info);
+                if (existingResponse.isComplete()) {
+                    listener.onServiceFound(info);
+                    hadReply = true;
                 }
             }
-            // Cancel the next scheduled periodical task.
-            if (nextQueryTaskFuture != null) {
-                cancelRequestTaskLocked();
-            }
-            // Keep tracking the ScheduledFuture for the task so we can cancel it if caller is not
-            // interested anymore.
-            final QueryTaskConfig taskConfig = new QueryTaskConfig(
-                    searchOptions.getSubtypes(),
-                    searchOptions.isPassiveMode(),
-                    searchOptions.onlyUseIpv6OnIpv6OnlyNetworks(),
-                    searchOptions.numOfQueriesBeforeBackoff(),
-                    socketKey);
-            final long now = clock.elapsedRealtime();
-            if (lastSentTime == 0) {
-                lastSentTime = now;
-            }
-            if (hadReply) {
-                final QueryTaskConfig queryTaskConfig = taskConfig.getConfigForNextRun();
-                final long minRemainingTtl = getMinRemainingTtlLocked(now);
-                final long timeToRun = now + queryTaskConfig.delayUntilNextTaskWithoutBackoffMs;
-                nextQueryTaskFuture = scheduleNextRunLocked(queryTaskConfig,
-                        minRemainingTtl, now, timeToRun, currentSessionId);
-            } else {
-                lastScheduledTask = new QueryTask(taskConfig,
-                        now /* timeToRun */,
-                        now + getMinRemainingTtlLocked(now)/* minTtlExpirationTimeWhenScheduled */,
-                        currentSessionId);
-                nextQueryTaskFuture = executor.submit(lastScheduledTask);
-            }
+        }
+        // Remove the next scheduled periodical task.
+        removeScheduledTask();
+        // Keep tracking the ScheduledFuture for the task so we can cancel it if caller is not
+        // interested anymore.
+        final QueryTaskConfig taskConfig = new QueryTaskConfig(
+                searchOptions.getSubtypes(),
+                searchOptions.isPassiveMode(),
+                searchOptions.onlyUseIpv6OnIpv6OnlyNetworks(),
+                searchOptions.numOfQueriesBeforeBackoff(),
+                socketKey);
+        final long now = clock.elapsedRealtime();
+        if (lastSentTime == 0) {
+            lastSentTime = now;
+        }
+        if (hadReply) {
+            final QueryTaskConfig queryTaskConfig = taskConfig.getConfigForNextRun();
+            final long minRemainingTtl = getMinRemainingTtl(now);
+            final long timeToRun = now + queryTaskConfig.delayUntilNextTaskWithoutBackoffMs;
+            scheduleNextRun(
+                    queryTaskConfig, minRemainingTtl, now, timeToRun, currentSessionId);
+        } else {
+            final List<MdnsResponse> servicesToResolve = makeResponsesForResolve(socketKey);
+            lastScheduledQueryTaskArgs = new ScheduledQueryTaskArgs(taskConfig, now /* timeToRun */,
+                    now + getMinRemainingTtl(now)/* minTtlExpirationTimeWhenScheduled */,
+                    currentSessionId);
+            final QueryTask queryTask = new QueryTask(lastScheduledQueryTaskArgs, servicesToResolve,
+                    servicesToResolve.size() < listeners.size() /* sendDiscoveryQueries */);
+            executor.submit(queryTask);
         }
     }
 
-    @GuardedBy("lock")
-    private void cancelRequestTaskLocked() {
-        final boolean canceled = nextQueryTaskFuture.cancel(true);
-        sharedLog.log("task canceled:" + canceled + ", current session: " + currentSessionId
-                + " task hashcode: " + getHexString(nextQueryTaskFuture));
+    private void removeScheduledTask() {
+        dependencies.removeMessages(handler, EVENT_START_QUERYTASK);
+        sharedLog.log("Remove EVENT_START_QUERYTASK"
+                + ", current session: " + currentSessionId);
         ++currentSessionId;
-        nextQueryTaskFuture = null;
-        lastScheduledTask = null;
-    }
-
-    private static String getHexString(Object o) {
-        return Integer.toHexString(System.identityHashCode(o));
+        lastScheduledQueryTaskArgs = null;
     }
 
     private boolean responseMatchesOptions(@NonNull MdnsResponse response,
@@ -281,15 +366,13 @@
      */
     public boolean stopSendAndReceive(@NonNull MdnsServiceBrowserListener listener) {
         ensureRunningOnHandlerThread(handler);
-        synchronized (lock) {
-            if (listeners.remove(listener) == null) {
-                return listeners.isEmpty();
-            }
-            if (listeners.isEmpty() && nextQueryTaskFuture != null) {
-                cancelRequestTaskLocked();
-            }
+        if (listeners.remove(listener) == null) {
             return listeners.isEmpty();
         }
+        if (listeners.isEmpty()) {
+            removeScheduledTask();
+        }
+        return listeners.isEmpty();
     }
 
     /**
@@ -298,50 +381,51 @@
     public synchronized void processResponse(@NonNull MdnsPacket packet,
             @NonNull SocketKey socketKey) {
         ensureRunningOnHandlerThread(handler);
-        synchronized (lock) {
-            // Augment the list of current known responses, and generated responses for resolve
-            // requests if there is no known response
-            final List<MdnsResponse> currentList = new ArrayList<>(instanceNameToResponse.values());
-            List<MdnsResponse> additionalResponses = makeResponsesForResolve(socketKey);
-            for (MdnsResponse additionalResponse : additionalResponses) {
-                if (!instanceNameToResponse.containsKey(
-                        additionalResponse.getServiceInstanceName())) {
-                    currentList.add(additionalResponse);
-                }
+        // Augment the list of current known responses, and generated responses for resolve
+        // requests if there is no known response
+        final List<MdnsResponse> cachedList =
+                serviceCache.getCachedServices(serviceType, socketKey);
+        final List<MdnsResponse> currentList = new ArrayList<>(cachedList);
+        List<MdnsResponse> additionalResponses = makeResponsesForResolve(socketKey);
+        for (MdnsResponse additionalResponse : additionalResponses) {
+            if (findMatchedResponse(
+                    cachedList, additionalResponse.getServiceInstanceName()) == null) {
+                currentList.add(additionalResponse);
             }
-            final Pair<ArraySet<MdnsResponse>, ArrayList<MdnsResponse>> augmentedResult =
-                    responseDecoder.augmentResponses(packet, currentList,
-                            socketKey.getInterfaceIndex(), socketKey.getNetwork());
+        }
+        final Pair<ArraySet<MdnsResponse>, ArrayList<MdnsResponse>> augmentedResult =
+                responseDecoder.augmentResponses(packet, currentList,
+                        socketKey.getInterfaceIndex(), socketKey.getNetwork());
 
-            final ArraySet<MdnsResponse> modifiedResponse = augmentedResult.first;
-            final ArrayList<MdnsResponse> allResponses = augmentedResult.second;
+        final ArraySet<MdnsResponse> modifiedResponse = augmentedResult.first;
+        final ArrayList<MdnsResponse> allResponses = augmentedResult.second;
 
-            for (MdnsResponse response : allResponses) {
-                if (modifiedResponse.contains(response)) {
-                    if (response.isGoodbye()) {
-                        onGoodbyeReceivedLocked(response.getServiceInstanceName());
-                    } else {
-                        onResponseModifiedLocked(response);
-                    }
-                } else if (instanceNameToResponse.containsKey(response.getServiceInstanceName())) {
-                    // If the response is not modified and already in the cache. The cache will
-                    // need to be updated to refresh the last receipt time.
-                    instanceNameToResponse.put(response.getServiceInstanceName(), response);
+        for (MdnsResponse response : allResponses) {
+            final String serviceInstanceName = response.getServiceInstanceName();
+            if (modifiedResponse.contains(response)) {
+                if (response.isGoodbye()) {
+                    onGoodbyeReceived(serviceInstanceName);
+                } else {
+                    onResponseModified(response);
                 }
+            } else if (findMatchedResponse(cachedList, serviceInstanceName) != null) {
+                // If the response is not modified and already in the cache. The cache will
+                // need to be updated to refresh the last receipt time.
+                serviceCache.addOrUpdateService(serviceType, socketKey, response);
             }
-            if (nextQueryTaskFuture != null && lastScheduledTask != null
-                    && lastScheduledTask.config.shouldUseQueryBackoff()) {
-                final long now = clock.elapsedRealtime();
-                final long minRemainingTtl = getMinRemainingTtlLocked(now);
-                final long timeToRun = calculateTimeToRun(lastScheduledTask,
-                        lastScheduledTask.config, now,
-                        minRemainingTtl, lastSentTime);
-                if (timeToRun > lastScheduledTask.timeToRun) {
-                    QueryTaskConfig lastTaskConfig = lastScheduledTask.config;
-                    cancelRequestTaskLocked();
-                    nextQueryTaskFuture = scheduleNextRunLocked(lastTaskConfig, minRemainingTtl,
-                            now, timeToRun, currentSessionId);
-                }
+        }
+        if (dependencies.hasMessages(handler, EVENT_START_QUERYTASK)
+                && lastScheduledQueryTaskArgs != null
+                && lastScheduledQueryTaskArgs.config.shouldUseQueryBackoff()) {
+            final long now = clock.elapsedRealtime();
+            final long minRemainingTtl = getMinRemainingTtl(now);
+            final long timeToRun = calculateTimeToRun(lastScheduledQueryTaskArgs,
+                    lastScheduledQueryTaskArgs.config, now,
+                    minRemainingTtl, lastSentTime);
+            if (timeToRun > lastScheduledQueryTaskArgs.timeToRun) {
+                QueryTaskConfig lastTaskConfig = lastScheduledQueryTaskArgs.config;
+                removeScheduledTask();
+                scheduleNextRun(lastTaskConfig, minRemainingTtl, now, timeToRun, currentSessionId);
             }
         }
     }
@@ -356,46 +440,40 @@
     /** Notify all services are removed because the socket is destroyed. */
     public void notifySocketDestroyed() {
         ensureRunningOnHandlerThread(handler);
-        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);
+        for (MdnsResponse response : serviceCache.getCachedServices(serviceType, socketKey)) {
+            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);
                 }
-            }
-
-            if (nextQueryTaskFuture != null) {
-                cancelRequestTaskLocked();
+                sharedLog.log("Socket destroyed. onServiceNameRemoved: " + name);
+                listener.onServiceNameRemoved(serviceInfo);
             }
         }
+        removeScheduledTask();
     }
 
-    @GuardedBy("lock")
-    private void onResponseModifiedLocked(@NonNull MdnsResponse response) {
+    private void onResponseModified(@NonNull MdnsResponse response) {
         final String serviceInstanceName = response.getServiceInstanceName();
         final MdnsResponse currentResponse =
-                instanceNameToResponse.get(serviceInstanceName);
+                serviceCache.getCachedService(serviceInstanceName, serviceType, socketKey);
 
         boolean newServiceFound = false;
         boolean serviceBecomesComplete = false;
         if (currentResponse == null) {
             newServiceFound = true;
             if (serviceInstanceName != null) {
-                instanceNameToResponse.put(serviceInstanceName, response);
+                serviceCache.addOrUpdateService(serviceType, socketKey, response);
             }
         } else {
             boolean before = currentResponse.isComplete();
-            instanceNameToResponse.put(serviceInstanceName, response);
+            serviceCache.addOrUpdateService(serviceType, socketKey, response);
             boolean after = response.isComplete();
             serviceBecomesComplete = !before && after;
         }
@@ -427,9 +505,9 @@
         }
     }
 
-    @GuardedBy("lock")
-    private void onGoodbyeReceivedLocked(@Nullable String serviceInstanceName) {
-        final MdnsResponse response = instanceNameToResponse.remove(serviceInstanceName);
+    private void onGoodbyeReceived(@Nullable String serviceInstanceName) {
+        final MdnsResponse response =
+                serviceCache.removeService(serviceInstanceName, serviceType, socketKey);
         if (response == null) {
             return;
         }
@@ -605,7 +683,8 @@
             if (resolveName == null) {
                 continue;
             }
-            MdnsResponse knownResponse = instanceNameToResponse.get(resolveName);
+            MdnsResponse knownResponse =
+                    serviceCache.getCachedService(resolveName, serviceType, socketKey);
             if (knownResponse == null) {
                 final ArrayList<String> instanceFullName = new ArrayList<>(
                         serviceTypeLabels.length + 1);
@@ -620,34 +699,82 @@
         return resolveResponses;
     }
 
-    // A FutureTask that enqueues a single query, and schedule a new FutureTask for the next task.
-    private class QueryTask implements Runnable {
+    private void tryRemoveServiceAfterTtlExpires() {
+        if (!shouldRemoveServiceAfterTtlExpires()) return;
 
+        Iterator<MdnsResponse> iter =
+                serviceCache.getCachedServices(serviceType, socketKey).iterator();
+        while (iter.hasNext()) {
+            MdnsResponse existingResponse = iter.next();
+            final String serviceInstanceName = existingResponse.getServiceInstanceName();
+            if (existingResponse.hasServiceRecord()
+                    && existingResponse.getServiceRecord()
+                    .getRemainingTTL(clock.elapsedRealtime()) == 0) {
+                serviceCache.removeService(serviceInstanceName, serviceType, socketKey);
+                for (int i = 0; i < listeners.size(); i++) {
+                    if (!responseMatchesOptions(existingResponse, listeners.valueAt(i))) {
+                        continue;
+                    }
+                    final MdnsServiceBrowserListener listener = listeners.keyAt(i);
+                    if (serviceInstanceName != null) {
+                        final MdnsServiceInfo serviceInfo = buildMdnsServiceInfoFromResponse(
+                                existingResponse, serviceTypeLabels);
+                        if (existingResponse.isComplete()) {
+                            sharedLog.log("TTL expired. onServiceRemoved: " + serviceInfo);
+                            listener.onServiceRemoved(serviceInfo);
+                        }
+                        sharedLog.log("TTL expired. onServiceNameRemoved: " + serviceInfo);
+                        listener.onServiceNameRemoved(serviceInfo);
+                    }
+                }
+            }
+        }
+    }
+
+    private static class ScheduledQueryTaskArgs {
         private final QueryTaskConfig config;
         private final long timeToRun;
         private final long minTtlExpirationTimeWhenScheduled;
         private final long sessionId;
 
-        QueryTask(@NonNull QueryTaskConfig config, long timeToRun,
-                long minTtlExpirationTimeWhenScheduled,
-                long sessionId) {
+        ScheduledQueryTaskArgs(@NonNull QueryTaskConfig config, long timeToRun,
+                long minTtlExpirationTimeWhenScheduled, long sessionId) {
             this.config = config;
             this.timeToRun = timeToRun;
             this.minTtlExpirationTimeWhenScheduled = minTtlExpirationTimeWhenScheduled;
             this.sessionId = sessionId;
         }
+    }
+
+    private static class QuerySentResult {
+        private final int transactionId;
+        private final List<String> subTypes = new ArrayList<>();
+        private final ScheduledQueryTaskArgs taskArgs;
+
+        QuerySentResult(int transactionId, @NonNull List<String> subTypes,
+                @NonNull ScheduledQueryTaskArgs taskArgs) {
+            this.transactionId = transactionId;
+            this.subTypes.addAll(subTypes);
+            this.taskArgs = taskArgs;
+        }
+    }
+
+    // A FutureTask that enqueues a single query, and schedule a new FutureTask for the next task.
+    private class QueryTask implements Runnable {
+
+        private final ScheduledQueryTaskArgs taskArgs;
+        private final List<MdnsResponse> servicesToResolve = new ArrayList<>();
+        private final boolean sendDiscoveryQueries;
+
+        QueryTask(@NonNull ScheduledQueryTaskArgs taskArgs,
+                @NonNull List<MdnsResponse> servicesToResolve, boolean sendDiscoveryQueries) {
+            this.taskArgs = taskArgs;
+            this.servicesToResolve.addAll(servicesToResolve);
+            this.sendDiscoveryQueries = sendDiscoveryQueries;
+        }
 
         @Override
         public void run() {
-            final List<MdnsResponse> servicesToResolve;
-            final boolean sendDiscoveryQueries;
-            synchronized (lock) {
-                // The listener is requesting to resolve a service that has no info in
-                // cache. Use the provided name to generate a minimal response, so other records are
-                // queried to complete it.
-                servicesToResolve = makeResponsesForResolve(config.socketKey);
-                sendDiscoveryQueries = servicesToResolve.size() < listeners.size();
-            }
             Pair<Integer, List<String>> result;
             try {
                 result =
@@ -655,88 +782,27 @@
                                 socketClient,
                                 createMdnsPacketWriter(),
                                 serviceType,
-                                config.subtypes,
-                                config.expectUnicastResponse,
-                                config.transactionId,
-                                config.socketKey,
-                                config.onlyUseIpv6OnIpv6OnlyNetworks,
+                                taskArgs.config.subtypes,
+                                taskArgs.config.expectUnicastResponse,
+                                taskArgs.config.transactionId,
+                                taskArgs.config.socketKey,
+                                taskArgs.config.onlyUseIpv6OnIpv6OnlyNetworks,
                                 sendDiscoveryQueries,
                                 servicesToResolve,
                                 clock)
                                 .call();
             } catch (RuntimeException e) {
                 sharedLog.e(String.format("Failed to run EnqueueMdnsQueryCallable for subtype: %s",
-                        TextUtils.join(",", config.subtypes)), e);
-                result = null;
+                        TextUtils.join(",", taskArgs.config.subtypes)), e);
+                result = Pair.create(-1, new ArrayList<>());
             }
-            synchronized (lock) {
-                if (MdnsConfigs.useSessionIdToScheduleMdnsTask()) {
-                    // In case that the task is not canceled successfully, use session ID to check
-                    // if this task should continue to schedule more.
-                    if (sessionId != currentSessionId) {
-                        return;
-                    }
-                }
-
-                if (MdnsConfigs.shouldCancelScanTaskWhenFutureIsNull()) {
-                    if (nextQueryTaskFuture == null) {
-                        // If requestTaskFuture is set to null, the task is cancelled. We can't use
-                        // isCancelled() here because this QueryTask is different from the future
-                        // that is returned from executor.schedule(). See b/71646910.
-                        return;
-                    }
-                }
-                if ((result != null)) {
-                    for (int i = 0; i < listeners.size(); i++) {
-                        listeners.keyAt(i).onDiscoveryQuerySent(result.second, result.first);
-                    }
-                }
-                if (shouldRemoveServiceAfterTtlExpires()) {
-                    Iterator<MdnsResponse> iter = instanceNameToResponse.values().iterator();
-                    while (iter.hasNext()) {
-                        MdnsResponse existingResponse = iter.next();
-                        if (existingResponse.hasServiceRecord()
-                                && existingResponse
-                                .getServiceRecord()
-                                .getRemainingTTL(clock.elapsedRealtime())
-                                == 0) {
-                            iter.remove();
-                            for (int i = 0; i < listeners.size(); i++) {
-                                if (!responseMatchesOptions(existingResponse,
-                                        listeners.valueAt(i)))  {
-                                    continue;
-                                }
-                                final MdnsServiceBrowserListener listener = listeners.keyAt(i);
-                                if (existingResponse.getServiceInstanceName() != null) {
-                                    final MdnsServiceInfo serviceInfo =
-                                            buildMdnsServiceInfoFromResponse(
-                                                    existingResponse, serviceTypeLabels);
-                                    if (existingResponse.isComplete()) {
-                                        sharedLog.log("TTL expired. onServiceRemoved: "
-                                                + serviceInfo);
-                                        listener.onServiceRemoved(serviceInfo);
-                                    }
-                                    sharedLog.log("TTL expired. onServiceNameRemoved: "
-                                            + serviceInfo);
-                                    listener.onServiceNameRemoved(serviceInfo);
-                                }
-                            }
-                        }
-                    }
-                }
-                QueryTaskConfig nextRunConfig = this.config.getConfigForNextRun();
-                final long now = clock.elapsedRealtime();
-                lastSentTime = now;
-                final long minRemainingTtl = getMinRemainingTtlLocked(now);
-                final long timeToRun = calculateTimeToRun(this, nextRunConfig, now,
-                        minRemainingTtl, lastSentTime);
-                nextQueryTaskFuture = scheduleNextRunLocked(nextRunConfig,
-                        minRemainingTtl, now, timeToRun, lastScheduledTask.sessionId);
-            }
+            dependencies.sendMessage(
+                    handler, handler.obtainMessage(EVENT_QUERY_RESULT,
+                            new QuerySentResult(result.first, result.second, taskArgs)));
         }
     }
 
-    private static long calculateTimeToRun(@NonNull QueryTask lastScheduledTask,
+    private static long calculateTimeToRun(@NonNull ScheduledQueryTaskArgs taskArgs,
             QueryTaskConfig queryTaskConfig, long now, long minRemainingTtl, long lastSentTime) {
         final long baseDelayInMs = queryTaskConfig.delayUntilNextTaskWithoutBackoffMs;
         if (!queryTaskConfig.shouldUseQueryBackoff()) {
@@ -749,18 +815,17 @@
         }
         // If the next TTL expiration time hasn't changed, then use previous calculated timeToRun.
         if (lastSentTime < now
-                && lastScheduledTask.minTtlExpirationTimeWhenScheduled == now + minRemainingTtl) {
+                && taskArgs.minTtlExpirationTimeWhenScheduled == now + minRemainingTtl) {
             // Use the original scheduling time if the TTL has not changed, to avoid continuously
             // rescheduling to 80% of the remaining TTL as time passes
-            return lastScheduledTask.timeToRun;
+            return taskArgs.timeToRun;
         }
         return Math.max(now + (long) (0.8 * minRemainingTtl), lastSentTime + baseDelayInMs);
     }
 
-    @GuardedBy("lock")
-    private long getMinRemainingTtlLocked(long now) {
+    private long getMinRemainingTtl(long now) {
         long minRemainingTtl = Long.MAX_VALUE;
-        for (MdnsResponse response : instanceNameToResponse.values()) {
+        for (MdnsResponse response : serviceCache.getCachedServices(serviceType, socketKey)) {
             if (!response.isComplete()) {
                 continue;
             }
@@ -777,19 +842,19 @@
         return minRemainingTtl == Long.MAX_VALUE ? 0 : minRemainingTtl;
     }
 
-    @GuardedBy("lock")
     @NonNull
-    private Future<?> scheduleNextRunLocked(@NonNull QueryTaskConfig nextRunConfig,
+    private void scheduleNextRun(@NonNull QueryTaskConfig nextRunConfig,
             long minRemainingTtl,
             long timeWhenScheduled, long timeToRun, long sessionId) {
-        lastScheduledTask = new QueryTask(nextRunConfig, timeToRun,
+        lastScheduledQueryTaskArgs = new ScheduledQueryTaskArgs(nextRunConfig, timeToRun,
                 minRemainingTtl + timeWhenScheduled, sessionId);
         // The timeWhenScheduled could be greater than the timeToRun if the Runnable is delayed.
         long timeToNextTasksWithBackoffInMs = Math.max(timeToRun - timeWhenScheduled, 0);
-        sharedLog.log(
-                String.format("Next run: sessionId: %d, in %d ms", lastScheduledTask.sessionId,
-                        timeToNextTasksWithBackoffInMs));
-        return executor.schedule(lastScheduledTask, timeToNextTasksWithBackoffInMs,
-                MILLISECONDS);
+        sharedLog.log(String.format("Next run: sessionId: %d, in %d ms",
+                lastScheduledQueryTaskArgs.sessionId, timeToNextTasksWithBackoffInMs));
+        dependencies.sendMessageDelayed(
+                handler,
+                handler.obtainMessage(EVENT_START_QUERYTASK, lastScheduledQueryTaskArgs),
+                timeToNextTasksWithBackoffInMs);
     }
 }
\ No newline at end of file
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index 368860e..d03cac6 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -495,8 +495,11 @@
         final AutomaticOnOffKeepalive autoKi;
         try {
             autoKi = target.withKeepaliveInfo(res.second);
-            // Close the duplicated fd.
-            target.close();
+            // Only automatic keepalives duplicate the fd.
+            if (target.mAutomaticOnOffState != STATE_ALWAYS_ON) {
+                // Close the duplicated fd.
+                target.close();
+            }
         } catch (InvalidSocketException e) {
             Log.wtf(TAG, "Fail to create AutomaticOnOffKeepalive", e);
             return;
diff --git a/service/src/com/android/server/connectivity/KeepaliveTracker.java b/service/src/com/android/server/connectivity/KeepaliveTracker.java
index 125c269..2ce186f 100644
--- a/service/src/com/android/server/connectivity/KeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveTracker.java
@@ -553,6 +553,8 @@
 
     private KeepaliveInfo handleUpdateKeepaliveForClat(KeepaliveInfo ki)
             throws InvalidSocketException, InvalidPacketException {
+        // Translation applies to only NAT-T keepalive
+        if (ki.mType != KeepaliveInfo.TYPE_NATT) return ki;
         // Only try to translate address if the packet source address is the clat's source address.
         if (!ki.mPacket.getSrcAddress().equals(ki.getNai().getClatv4SrcAddress())) return ki;
 
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/DumpOnFailureRule.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/DumpOnFailureRule.java
index 78ae7b8..07434b1 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/DumpOnFailureRule.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/DumpOnFailureRule.java
@@ -64,6 +64,8 @@
                     "dumpsys usagestats appstandby",
                     "dumpsys connectivity trafficcontroller",
                     "dumpsys netd trafficcontroller",
+                    "dumpsys platform_compat", // TODO (b/279829773): Remove this dump
+                    "dumpsys jobscheduler " + TEST_APP2_PKG, // TODO (b/288220398): Remove this dump
             }) {
                 dumpCommandOutput(out, cmd);
             }
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java b/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java
index a7d5590..d112425 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java
@@ -26,16 +26,12 @@
     private static final String TEST_CLASS = TEST_PKG + ".ConnOnActivityStartTest";
     @Before
     public void setUp() throws Exception {
-        super.setUp();
-
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
     @After
     public void tearDown() throws Exception {
-        super.tearDown();
-
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java
index 5d7ad62..d8e7a2c 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java
@@ -23,14 +23,12 @@
 
     @Before
     public void setUp() throws Exception {
-        super.setUp();
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
     @After
     public void tearDown() throws Exception {
-        super.tearDown();
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java
index 40f5f59..3ddb88b 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java
@@ -23,14 +23,12 @@
 public class HostsideNetworkPolicyManagerTests extends HostsideNetworkTestCase {
     @Before
     public void setUp() throws Exception {
-        super.setUp();
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
     @After
     public void tearDown() throws Exception {
-        super.tearDown();
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java
index c896168..566d9da 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java
@@ -16,16 +16,21 @@
 
 package com.android.cts.net;
 
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
 
 import com.android.ddmlib.Log;
 import com.android.modules.utils.build.testing.DeviceSdkLevel;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.invoker.TestInformation;
+import com.android.tradefed.targetprep.BuildError;
 import com.android.tradefed.targetprep.TargetSetupError;
+import com.android.tradefed.targetprep.suite.SuiteApkInstaller;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+import com.android.tradefed.testtype.junit4.AfterClassWithInfo;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
-import com.android.tradefed.testtype.junit4.DeviceTestRunOptions;
+import com.android.tradefed.testtype.junit4.BeforeClassWithInfo;
 import com.android.tradefed.util.RunUtil;
 
 import org.junit.runner.RunWith;
@@ -40,34 +45,61 @@
     protected static final String TEST_APP2_PKG = "com.android.cts.net.hostside.app2";
     protected static final String TEST_APP2_APK = "CtsHostsideNetworkTestsApp2.apk";
 
-    protected void setUp() throws Exception {
-        DeviceSdkLevel deviceSdkLevel = new DeviceSdkLevel(getDevice());
-        String testApk = deviceSdkLevel.isDeviceAtLeastT() ? TEST_APK_NEXT
-                : TEST_APK;
+    @BeforeClassWithInfo
+    public static void setUpOnce(TestInformation testInfo) throws Exception {
+        DeviceSdkLevel deviceSdkLevel = new DeviceSdkLevel(testInfo.getDevice());
+        String testApk = deviceSdkLevel.isDeviceAtLeastT() ? TEST_APK_NEXT : TEST_APK;
 
-        uninstallPackage(TEST_PKG, false);
-        installPackage(testApk);
+        uninstallPackage(testInfo, TEST_PKG, false);
+        installPackage(testInfo, testApk);
     }
 
-    protected void tearDown() throws Exception {
-        uninstallPackage(TEST_PKG, true);
+    @AfterClassWithInfo
+    public static void tearDownOnce(TestInformation testInfo) throws DeviceNotAvailableException {
+        uninstallPackage(testInfo, TEST_PKG, true);
+    }
+
+    // Custom static method to install the specified package, this is used to bypass auto-cleanup
+    // per test in BaseHostJUnit4.
+    protected static void installPackage(TestInformation testInfo, String apk)
+            throws DeviceNotAvailableException, TargetSetupError {
+        assertNotNull(testInfo);
+        final int userId = testInfo.getDevice().getCurrentUser();
+        final SuiteApkInstaller installer = new SuiteApkInstaller();
+        // Force the apk clean up
+        installer.setCleanApk(true);
+        installer.addTestFileName(apk);
+        installer.setUserId(userId);
+        installer.setShouldGrantPermission(true);
+        installer.addInstallArg("-t");
+        try {
+            installer.setUp(testInfo);
+        } catch (BuildError e) {
+            throw new TargetSetupError(
+                    e.getMessage(), e, testInfo.getDevice().getDeviceDescriptor(), e.getErrorId());
+        }
     }
 
     protected void installPackage(String apk) throws DeviceNotAvailableException, TargetSetupError {
-        final DeviceTestRunOptions installOptions = new DeviceTestRunOptions(
-                null /* packageName */);
-        final int userId = getDevice().getCurrentUser();
-        installPackageAsUser(apk, true /* grantPermission */, userId, "-t");
+        installPackage(getTestInformation(), apk);
     }
 
-    protected void uninstallPackage(String packageName, boolean shouldSucceed)
+    protected static void uninstallPackage(TestInformation testInfo, String packageName,
+            boolean shouldSucceed)
             throws DeviceNotAvailableException {
-        final String result = uninstallPackage(packageName);
+        assertNotNull(testInfo);
+        final String result = testInfo.getDevice().uninstallPackage(packageName);
         if (shouldSucceed) {
             assertNull("uninstallPackage(" + packageName + ") failed: " + result, result);
         }
     }
 
+    protected void uninstallPackage(String packageName,
+            boolean shouldSucceed)
+            throws DeviceNotAvailableException {
+        uninstallPackage(getTestInformation(), packageName, shouldSucceed);
+    }
+
     protected void assertPackageUninstalled(String packageName) throws DeviceNotAvailableException,
             InterruptedException {
         final String command = "cmd package list packages " + packageName;
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java b/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java
index 0977deb..57b26bd 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java
@@ -32,16 +32,12 @@
 
     @Before
     public void setUp() throws Exception {
-        super.setUp();
-
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
     @After
     public void tearDown() throws Exception {
-        super.tearDown();
-
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideVpnTests.java b/tests/cts/hostside/src/com/android/cts/net/HostsideVpnTests.java
index 242fd5d..691ac90 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideVpnTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideVpnTests.java
@@ -26,16 +26,12 @@
 
     @Before
     public void setUp() throws Exception {
-        super.setUp();
-
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
     @After
     public void tearDown() throws Exception {
-        super.tearDown();
-
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
diff --git a/tests/cts/net/src/android/net/cts/NetworkStatsManagerTest.java b/tests/cts/net/src/android/net/cts/NetworkStatsManagerTest.java
index 83b9b81..7bccbde 100644
--- a/tests/cts/net/src/android/net/cts/NetworkStatsManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/NetworkStatsManagerTest.java
@@ -82,9 +82,9 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.BufferedInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.InputStreamReader;
 import java.net.HttpURLConnection;
 import java.net.URL;
 import java.net.UnknownHostException;
@@ -220,7 +220,7 @@
         } else {
             Log.w(LOG_TAG, "Network: " + networkInfo.toString());
         }
-        InputStreamReader in = null;
+        BufferedInputStream in = null;
         HttpURLConnection urlc = null;
         String originalKeepAlive = System.getProperty("http.keepAlive");
         System.setProperty("http.keepAlive", "false");
@@ -236,10 +236,10 @@
             urlc.connect();
             boolean ping = urlc.getResponseCode() == 200;
             if (ping) {
-                in = new InputStreamReader((InputStream) urlc.getContent());
-                // Since the test doesn't really care about the precise amount of data, instead
-                // of reading all contents, just read few bytes at the beginning.
-                in.read();
+                in = new BufferedInputStream((InputStream) urlc.getContent());
+                while (in.read() != -1) {
+                    // Comments to suppress lint error.
+                }
             }
         } catch (Exception e) {
             Log.i(LOG_TAG, "Badness during exercising remote server: " + e);
@@ -377,9 +377,14 @@
                 .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
                 .build(), callback);
         synchronized (this) {
-            try {
-                wait((int) (TIMEOUT_MILLIS * 2.4));
-            } catch (InterruptedException e) {
+            long now = System.currentTimeMillis();
+            final long deadline = (long) (now + TIMEOUT_MILLIS * 2.4);
+            while (!callback.success && now < deadline) {
+                try {
+                    wait(deadline - now);
+                } catch (InterruptedException e) {
+                }
+                now = System.currentTimeMillis();
             }
         }
         if (callback.success) {
diff --git a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
index 6c411cf..49620b0 100644
--- a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
@@ -80,8 +80,6 @@
 import com.android.modules.utils.build.SdkLevel.isAtLeastU
 import com.android.net.module.util.ArrayTrackRecord
 import com.android.net.module.util.TrackRecord
-import com.android.networkstack.apishim.NsdShimImpl
-import com.android.networkstack.apishim.common.NsdShim
 import com.android.testutils.ConnectivityModuleTest
 import com.android.testutils.DevSdkIgnoreRule
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
@@ -133,8 +131,6 @@
 private const val DBG = false
 private const val TEST_PORT = 12345
 
-private val nsdShim = NsdShimImpl.newInstance()
-
 @AppModeFull(reason = "Socket cannot bind in instant app mode")
 @RunWith(DevSdkIgnoreRunner::class)
 @SmallTest
@@ -293,7 +289,7 @@
             val serviceFound = expectCallbackEventually<ServiceFound> {
                 it.serviceInfo.serviceName == serviceName &&
                         (expectedNetwork == null ||
-                                expectedNetwork == nsdShim.getNetwork(it.serviceInfo))
+                                expectedNetwork == it.serviceInfo.network)
             }.serviceInfo
             // Discovered service types have a dot at the end
             assertEquals("$serviceType.", serviceFound.serviceType)
@@ -331,7 +327,7 @@
         }
     }
 
-    private class NsdServiceInfoCallbackRecord : NsdShim.ServiceInfoCallbackShim,
+    private class NsdServiceInfoCallbackRecord : NsdManager.ServiceInfoCallback,
             NsdRecord<NsdServiceInfoCallbackRecord.ServiceInfoCallbackEvent>() {
         sealed class ServiceInfoCallbackEvent : NsdEvent {
             data class RegisterCallbackFailed(val errorCode: Int) : ServiceInfoCallbackEvent()
@@ -361,11 +357,9 @@
     fun setUp() {
         handlerThread.start()
 
-        if (TestUtils.shouldTestTApis()) {
-            runAsShell(MANAGE_TEST_NETWORKS) {
-                testNetwork1 = createTestNetwork()
-                testNetwork2 = createTestNetwork()
-            }
+        runAsShell(MANAGE_TEST_NETWORKS) {
+            testNetwork1 = createTestNetwork()
+            testNetwork2 = createTestNetwork()
         }
     }
 
@@ -450,12 +444,10 @@
 
     @After
     fun tearDown() {
-        if (TestUtils.shouldTestTApis()) {
-            runAsShell(MANAGE_TEST_NETWORKS) {
-                // Avoid throwing here if initializing failed in setUp
-                if (this::testNetwork1.isInitialized) testNetwork1.close(cm)
-                if (this::testNetwork2.isInitialized) testNetwork2.close(cm)
-            }
+        runAsShell(MANAGE_TEST_NETWORKS) {
+            // Avoid throwing here if initializing failed in setUp
+            if (this::testNetwork1.isInitialized) testNetwork1.close(cm)
+            if (this::testNetwork2.isInitialized) testNetwork2.close(cm)
         }
         handlerThread.waitForIdle(TIMEOUT_MS)
         handlerThread.quitSafely()
@@ -601,9 +593,6 @@
 
     @Test
     fun testNsdManager_DiscoverOnNetwork() {
-        // This test requires shims supporting T+ APIs (discovering on specific network)
-        assumeTrue(TestUtils.shouldTestTApis())
-
         val si = NsdServiceInfo()
         si.serviceType = serviceType
         si.serviceName = this.serviceName
@@ -614,19 +603,19 @@
 
         tryTest {
             val discoveryRecord = NsdDiscoveryRecord()
-            nsdShim.discoverServices(nsdManager, serviceType, NsdManager.PROTOCOL_DNS_SD,
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
                     testNetwork1.network, Executor { it.run() }, discoveryRecord)
 
             val foundInfo = discoveryRecord.waitForServiceDiscovered(
                     serviceName, serviceType, testNetwork1.network)
-            assertEquals(testNetwork1.network, nsdShim.getNetwork(foundInfo))
+            assertEquals(testNetwork1.network, foundInfo.network)
 
             // Rewind to ensure the service is not found on the other interface
             discoveryRecord.nextEvents.rewind(0)
             assertNull(discoveryRecord.nextEvents.poll(timeoutMs = 100L) {
                 it is ServiceFound &&
                         it.serviceInfo.serviceName == registeredInfo.serviceName &&
-                        nsdShim.getNetwork(it.serviceInfo) != testNetwork1.network
+                        it.serviceInfo.network != testNetwork1.network
             }, "The service should not be found on this network")
         } cleanup {
             nsdManager.unregisterService(registrationRecord)
@@ -635,9 +624,6 @@
 
     @Test
     fun testNsdManager_DiscoverWithNetworkRequest() {
-        // This test requires shims supporting T+ APIs (discovering on network request)
-        assumeTrue(TestUtils.shouldTestTApis())
-
         val si = NsdServiceInfo()
         si.serviceType = serviceType
         si.serviceName = this.serviceName
@@ -652,7 +638,7 @@
 
         tryTest {
             val specifier = TestNetworkSpecifier(testNetwork1.iface.interfaceName)
-            nsdShim.discoverServices(nsdManager, serviceType, NsdManager.PROTOCOL_DNS_SD,
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
                     NetworkRequest.Builder()
                             .removeCapability(NET_CAPABILITY_TRUSTED)
                             .addTransportType(TRANSPORT_TEST)
@@ -667,27 +653,27 @@
             assertEquals(registeredInfo1.serviceName, serviceDiscovered.serviceInfo.serviceName)
             // Discovered service types have a dot at the end
             assertEquals("$serviceType.", serviceDiscovered.serviceInfo.serviceType)
-            assertEquals(testNetwork1.network, nsdShim.getNetwork(serviceDiscovered.serviceInfo))
+            assertEquals(testNetwork1.network, serviceDiscovered.serviceInfo.network)
 
             // Unregister, then register the service back: it should be lost and found again
             nsdManager.unregisterService(registrationRecord)
             val serviceLost1 = discoveryRecord.expectCallback<ServiceLost>()
             assertEquals(registeredInfo1.serviceName, serviceLost1.serviceInfo.serviceName)
-            assertEquals(testNetwork1.network, nsdShim.getNetwork(serviceLost1.serviceInfo))
+            assertEquals(testNetwork1.network, serviceLost1.serviceInfo.network)
 
             registrationRecord.expectCallback<ServiceUnregistered>()
             val registeredInfo2 = registerService(registrationRecord, si, executor)
             val serviceDiscovered2 = discoveryRecord.expectCallback<ServiceFound>()
             assertEquals(registeredInfo2.serviceName, serviceDiscovered2.serviceInfo.serviceName)
             assertEquals("$serviceType.", serviceDiscovered2.serviceInfo.serviceType)
-            assertEquals(testNetwork1.network, nsdShim.getNetwork(serviceDiscovered2.serviceInfo))
+            assertEquals(testNetwork1.network, serviceDiscovered2.serviceInfo.network)
 
             // Teardown, then bring back up a network on the test interface: the service should
             // go away, then come back
             testNetwork1.agent.unregister()
             val serviceLost = discoveryRecord.expectCallback<ServiceLost>()
             assertEquals(registeredInfo2.serviceName, serviceLost.serviceInfo.serviceName)
-            assertEquals(testNetwork1.network, nsdShim.getNetwork(serviceLost.serviceInfo))
+            assertEquals(testNetwork1.network, serviceLost.serviceInfo.network)
 
             val newAgent = runAsShell(MANAGE_TEST_NETWORKS) {
                 registerTestNetworkAgent(testNetwork1.iface.interfaceName)
@@ -696,7 +682,7 @@
             val serviceDiscovered3 = discoveryRecord.expectCallback<ServiceFound>()
             assertEquals(registeredInfo2.serviceName, serviceDiscovered3.serviceInfo.serviceName)
             assertEquals("$serviceType.", serviceDiscovered3.serviceInfo.serviceType)
-            assertEquals(newNetwork, nsdShim.getNetwork(serviceDiscovered3.serviceInfo))
+            assertEquals(newNetwork, serviceDiscovered3.serviceInfo.network)
         } cleanupStep {
             nsdManager.stopServiceDiscovery(discoveryRecord)
             discoveryRecord.expectCallback<DiscoveryStopped>()
@@ -707,9 +693,6 @@
 
     @Test
     fun testNsdManager_DiscoverWithNetworkRequest_NoMatchingNetwork() {
-        // This test requires shims supporting T+ APIs (discovering on network request)
-        assumeTrue(TestUtils.shouldTestTApis())
-
         val si = NsdServiceInfo()
         si.serviceType = serviceType
         si.serviceName = this.serviceName
@@ -722,7 +705,7 @@
         val specifier = TestNetworkSpecifier(testNetwork1.iface.interfaceName)
 
         tryTest {
-            nsdShim.discoverServices(nsdManager, serviceType, NsdManager.PROTOCOL_DNS_SD,
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
                     NetworkRequest.Builder()
                             .removeCapability(NET_CAPABILITY_TRUSTED)
                             .addTransportType(TRANSPORT_TEST)
@@ -754,9 +737,6 @@
 
     @Test
     fun testNsdManager_ResolveOnNetwork() {
-        // This test requires shims supporting T+ APIs (NsdServiceInfo.network)
-        assumeTrue(TestUtils.shouldTestTApis())
-
         val si = NsdServiceInfo()
         si.serviceType = serviceType
         si.serviceName = this.serviceName
@@ -772,21 +752,21 @@
 
             val foundInfo1 = discoveryRecord.waitForServiceDiscovered(
                     serviceName, serviceType, testNetwork1.network)
-            assertEquals(testNetwork1.network, nsdShim.getNetwork(foundInfo1))
+            assertEquals(testNetwork1.network, foundInfo1.network)
             // Rewind as the service could be found on each interface in any order
             discoveryRecord.nextEvents.rewind(0)
             val foundInfo2 = discoveryRecord.waitForServiceDiscovered(
                     serviceName, serviceType, testNetwork2.network)
-            assertEquals(testNetwork2.network, nsdShim.getNetwork(foundInfo2))
+            assertEquals(testNetwork2.network, foundInfo2.network)
 
-            nsdShim.resolveService(nsdManager, foundInfo1, Executor { it.run() }, resolveRecord)
+            nsdManager.resolveService(foundInfo1, Executor { it.run() }, resolveRecord)
             val cb = resolveRecord.expectCallback<ServiceResolved>()
             cb.serviceInfo.let {
                 // Resolved service type has leading dot
                 assertEquals(".$serviceType", it.serviceType)
                 assertEquals(registeredInfo.serviceName, it.serviceName)
                 assertEquals(si.port, it.port)
-                assertEquals(testNetwork1.network, nsdShim.getNetwork(it))
+                assertEquals(testNetwork1.network, it.network)
                 checkAddressScopeId(testNetwork1.iface, it.hostAddresses)
             }
             // TODO: check that MDNS packets are sent only on testNetwork1.
@@ -799,9 +779,6 @@
 
     @Test
     fun testNsdManager_RegisterOnNetwork() {
-        // This test requires shims supporting T+ APIs (NsdServiceInfo.network)
-        assumeTrue(TestUtils.shouldTestTApis())
-
         val si = NsdServiceInfo()
         si.serviceType = serviceType
         si.serviceName = this.serviceName
@@ -817,27 +794,27 @@
 
         tryTest {
             // Discover service on testNetwork1.
-            nsdShim.discoverServices(nsdManager, serviceType, NsdManager.PROTOCOL_DNS_SD,
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
                 testNetwork1.network, Executor { it.run() }, discoveryRecord)
             // Expect that service is found on testNetwork1
             val foundInfo = discoveryRecord.waitForServiceDiscovered(
                 serviceName, serviceType, testNetwork1.network)
-            assertEquals(testNetwork1.network, nsdShim.getNetwork(foundInfo))
+            assertEquals(testNetwork1.network, foundInfo.network)
 
             // Discover service on testNetwork2.
-            nsdShim.discoverServices(nsdManager, serviceType, NsdManager.PROTOCOL_DNS_SD,
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
                 testNetwork2.network, Executor { it.run() }, discoveryRecord2)
             // Expect that discovery is started then no other callbacks.
             discoveryRecord2.expectCallback<DiscoveryStarted>()
             discoveryRecord2.assertNoCallback()
 
             // Discover service on all networks (not specify any network).
-            nsdShim.discoverServices(nsdManager, serviceType, NsdManager.PROTOCOL_DNS_SD,
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
                 null as Network? /* network */, Executor { it.run() }, discoveryRecord3)
             // Expect that service is found on testNetwork1
             val foundInfo3 = discoveryRecord3.waitForServiceDiscovered(
                     serviceName, serviceType, testNetwork1.network)
-            assertEquals(testNetwork1.network, nsdShim.getNetwork(foundInfo3))
+            assertEquals(testNetwork1.network, foundInfo3.network)
         } cleanupStep {
             nsdManager.stopServiceDiscovery(discoveryRecord2)
             discoveryRecord2.expectCallback<DiscoveryStopped>()
@@ -970,9 +947,6 @@
 
     @Test
     fun testStopServiceResolution() {
-        // This test requires shims supporting U+ APIs (NsdManager.stopServiceResolution)
-        assumeTrue(TestUtils.shouldTestUApis())
-
         val si = NsdServiceInfo()
         si.serviceType = this@NsdManagerTest.serviceType
         si.serviceName = this@NsdManagerTest.serviceName
@@ -981,8 +955,8 @@
         val resolveRecord = NsdResolveRecord()
         // Try to resolve an unknown service then stop it immediately.
         // Expected ResolutionStopped callback.
-        nsdShim.resolveService(nsdManager, si, { it.run() }, resolveRecord)
-        nsdShim.stopServiceResolution(nsdManager, resolveRecord)
+        nsdManager.resolveService(si, { it.run() }, resolveRecord)
+        nsdManager.stopServiceResolution(resolveRecord)
         val stoppedCb = resolveRecord.expectCallback<ResolutionStopped>()
         assertEquals(si.serviceName, stoppedCb.serviceInfo.serviceName)
         assertEquals(si.serviceType, stoppedCb.serviceInfo.serviceType)
@@ -990,9 +964,6 @@
 
     @Test
     fun testRegisterServiceInfoCallback() {
-        // This test requires shims supporting U+ APIs (NsdManager.registerServiceInfoCallback)
-        assumeTrue(TestUtils.shouldTestUApis())
-
         val lp = cm.getLinkProperties(testNetwork1.network)
         assertNotNull(lp)
         val addresses = lp.addresses
@@ -1013,13 +984,13 @@
         val cbRecord = NsdServiceInfoCallbackRecord()
         tryTest {
             // Discover service on the network.
-            nsdShim.discoverServices(nsdManager, serviceType, NsdManager.PROTOCOL_DNS_SD,
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
                     testNetwork1.network, Executor { it.run() }, discoveryRecord)
             val foundInfo = discoveryRecord.waitForServiceDiscovered(
                     serviceName, serviceType, testNetwork1.network)
 
             // Register service callback and check the addresses are the same as network addresses
-            nsdShim.registerServiceInfoCallback(nsdManager, foundInfo, { it.run() }, cbRecord)
+            nsdManager.registerServiceInfoCallback(foundInfo, { it.run() }, cbRecord)
             val serviceInfoCb = cbRecord.expectCallback<ServiceUpdated>()
             assertEquals(foundInfo.serviceName, serviceInfoCb.serviceInfo.serviceName)
             val hostAddresses = serviceInfoCb.serviceInfo.hostAddresses
@@ -1035,7 +1006,7 @@
             cbRecord.expectCallback<ServiceUpdatedLost>()
         } cleanupStep {
             // Cancel subscription and check stop callback received.
-            nsdShim.unregisterServiceInfoCallback(nsdManager, cbRecord)
+            nsdManager.unregisterServiceInfoCallback(cbRecord)
             cbRecord.expectCallback<UnregisterCallbackSucceeded>()
         } cleanup {
             nsdManager.stopServiceDiscovery(discoveryRecord)
@@ -1045,9 +1016,6 @@
 
     @Test
     fun testStopServiceResolutionFailedCallback() {
-        // This test requires shims supporting U+ APIs (NsdManager.stopServiceResolution)
-        assumeTrue(TestUtils.shouldTestUApis())
-
         // It's not possible to make ResolutionListener#onStopResolutionFailed callback sending
         // because it is only sent in very edge-case scenarios when the legacy implementation is
         // used, and the legacy implementation is never used in the current AOSP builds. Considering
@@ -1115,7 +1083,7 @@
         si: NsdServiceInfo,
         executor: Executor = Executor { it.run() }
     ): NsdServiceInfo {
-        nsdShim.registerService(nsdManager, si, NsdManager.PROTOCOL_DNS_SD, executor, record)
+        nsdManager.registerService(si, NsdManager.PROTOCOL_DNS_SD, executor, record)
         // We may not always get the name that we tried to register;
         // This events tells us the name that was registered.
         val cb = record.expectCallback<ServiceRegistered>(REGISTRATION_TIMEOUT_MS)
@@ -1124,7 +1092,7 @@
 
     private fun resolveService(discoveredInfo: NsdServiceInfo): NsdServiceInfo {
         val record = NsdResolveRecord()
-        nsdShim.resolveService(nsdManager, discoveredInfo, Executor { it.run() }, record)
+        nsdManager.resolveService(discoveredInfo, Executor { it.run() }, record)
         val resolvedCb = record.expectCallback<ServiceResolved>()
         assertEquals(discoveredInfo.serviceName, resolvedCb.serviceInfo.serviceName)
 
diff --git a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
index eeffbe1..f4d3915 100644
--- a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
@@ -408,22 +408,22 @@
     @Test
     public void testIsAnyTcpSocketConnected_withTargetNetId() throws Exception {
         setupResponseWithSocketExisting();
-        mTestHandler.post(
-                () -> assertTrue(mAOOKeepaliveTracker.isAnyTcpSocketConnected(TEST_NETID)));
+        assertTrue(visibleOnHandlerThread(mTestHandler,
+                () -> mAOOKeepaliveTracker.isAnyTcpSocketConnected(TEST_NETID)));
     }
 
     @Test
     public void testIsAnyTcpSocketConnected_withIncorrectNetId() throws Exception {
         setupResponseWithSocketExisting();
-        mTestHandler.post(
-                () -> assertFalse(mAOOKeepaliveTracker.isAnyTcpSocketConnected(OTHER_NETID)));
+        assertFalse(visibleOnHandlerThread(mTestHandler,
+                () -> mAOOKeepaliveTracker.isAnyTcpSocketConnected(OTHER_NETID)));
     }
 
     @Test
     public void testIsAnyTcpSocketConnected_noSocketExists() throws Exception {
         setupResponseWithoutSocketExisting();
-        mTestHandler.post(
-                () -> assertFalse(mAOOKeepaliveTracker.isAnyTcpSocketConnected(TEST_NETID)));
+        assertFalse(visibleOnHandlerThread(mTestHandler,
+                () -> mAOOKeepaliveTracker.isAnyTcpSocketConnected(TEST_NETID)));
     }
 
     private void triggerEventKeepalive(int slot, int reason) {
@@ -500,9 +500,7 @@
         final AlarmManager.OnAlarmListener listener = listenerCaptor.getValue();
 
         // For realism, the listener should be posted on the handler
-        mTestHandler.post(() -> listener.onAlarm());
-        // Wait for the listener to be called. The listener enqueues a message to the handler.
-        HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
+        visibleOnHandlerThread(mTestHandler, () -> listener.onAlarm());
         // Wait for the message posted by the listener to be processed.
         HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
 
@@ -525,8 +523,7 @@
 
         doReturn(METRICS_COLLECTION_DURATION_MS).when(mDependencies).getElapsedRealtime();
         // For realism, the listener should be posted on the handler
-        mTestHandler.post(() -> listener.onAlarm());
-        HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
+        visibleOnHandlerThread(mTestHandler, () -> listener.onAlarm());
 
         verify(mKeepaliveStatsTracker).writeAndResetMetrics();
         // Alarm is rescheduled.
@@ -612,17 +609,60 @@
         verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
     }
 
-    @Test
-    public void testStartNattKeepalive_addressTranslationOnClat() throws Exception {
-        final InetAddress v6AddrSrc = InetAddresses.parseNumericAddress("2001:db8::1");
-        final InetAddress v6AddrDst = InetAddresses.parseNumericAddress("2001:db8::2");
-        doReturn(v6AddrDst).when(mNai).translateV4toClatV6(any());
-        doReturn(v6AddrSrc).when(mNai).getClatv6SrcAddress();
+    private void setupTestNaiForClat(InetAddress v6Src, InetAddress v6Dst) throws Exception {
+        doReturn(v6Dst).when(mNai).translateV4toClatV6(any());
+        doReturn(v6Src).when(mNai).getClatv6SrcAddress();
         doReturn(InetAddress.getByAddress(V4_SRC_ADDR)).when(mNai).getClatv4SrcAddress();
         // Setup nai to add clat address
         final LinkProperties stacked = new LinkProperties();
         stacked.setInterfaceName(TEST_V4_IFACE);
+        final InetAddress srcAddress = InetAddress.getByAddress(
+                new byte[] { (byte) 192, 0, 0, (byte) 129 });
+        mNai.linkProperties.addLinkAddress(new LinkAddress(srcAddress, 24));
         mNai.linkProperties.addStackedLink(stacked);
+    }
+
+    private TestKeepaliveInfo doStartTcpKeepalive(InetAddress srcAddr) throws Exception {
+        final KeepalivePacketData kpd = new TcpKeepalivePacketData(
+                srcAddr,
+                12345 /* srcPort */,
+                InetAddress.getByAddress(new byte[] { 8, 8, 8, 8}) /* dstAddr */,
+                12345 /* dstPort */, new byte[] {1},  111 /* tcpSeq */,
+                222 /* tcpAck */, 800 /* tcpWindow */, 2 /* tcpWindowScale */,
+                4 /* ipTos */, 64 /* ipTtl */);
+        final TestKeepaliveInfo testInfo = new TestKeepaliveInfo(kpd);
+
+        final KeepaliveInfo ki = mKeepaliveTracker.new KeepaliveInfo(
+                testInfo.socketKeepaliveCallback, mNai, kpd,
+                TEST_KEEPALIVE_INTERVAL_SEC, KeepaliveInfo.TYPE_TCP, testInfo.fd);
+        mKeepaliveTracker.setReturnedKeepaliveInfo(ki);
+
+        // Setup TCP keepalive.
+        mAOOKeepaliveTracker.startTcpKeepalive(mNai, testInfo.fd, TEST_KEEPALIVE_INTERVAL_SEC,
+                testInfo.socketKeepaliveCallback);
+        HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
+        return testInfo;
+    }
+    @Test
+    public void testStartTcpKeepalive_addressTranslationOnClat() throws Exception {
+        setupTestNaiForClat(InetAddresses.parseNumericAddress("2001:db8::1") /* v6Src */,
+                InetAddresses.parseNumericAddress("2001:db8::2") /* v6Dst */);
+        final InetAddress srcAddr = InetAddress.getByAddress(V4_SRC_ADDR);
+        doStartTcpKeepalive(srcAddr);
+        final ArgumentCaptor<TcpKeepalivePacketData> tpdCaptor =
+                ArgumentCaptor.forClass(TcpKeepalivePacketData.class);
+        verify(mNai).onStartTcpSocketKeepalive(
+                eq(TEST_SLOT), eq(TEST_KEEPALIVE_INTERVAL_SEC), tpdCaptor.capture());
+        final TcpKeepalivePacketData tpd = tpdCaptor.getValue();
+        // Verify the addresses still be the same address when clat is started.
+        assertEquals(srcAddr, tpd.getSrcAddress());
+    }
+
+    @Test
+    public void testStartNattKeepalive_addressTranslationOnClat() throws Exception {
+        final InetAddress v6AddrSrc = InetAddresses.parseNumericAddress("2001:db8::1");
+        final InetAddress v6AddrDst = InetAddresses.parseNumericAddress("2001:db8::2");
+        setupTestNaiForClat(v6AddrSrc, v6AddrDst);
 
         final TestKeepaliveInfo testInfo = doStartNattKeepalive();
         final ArgumentCaptor<NattKeepalivePacketData> kpdCaptor =
@@ -899,24 +939,8 @@
                 new byte[] { (byte) 192, 0, 0, (byte) 129 });
         mNai.linkProperties.addLinkAddress(new LinkAddress(srcAddress, 24));
 
-        final KeepalivePacketData kpd = new TcpKeepalivePacketData(
-                InetAddress.getByAddress(new byte[] { (byte) 192, 0, 0, (byte) 129 }) /* srcAddr */,
-                12345 /* srcPort */,
-                InetAddress.getByAddress(new byte[] { 8, 8, 8, 8}) /* dstAddr */,
-                12345 /* dstPort */, new byte[] {1},  111 /* tcpSeq */,
-                222 /* tcpAck */, 800 /* tcpWindow */, 2 /* tcpWindowScale */,
-                4 /* ipTos */, 64 /* ipTtl */);
-        final TestKeepaliveInfo testInfo = new TestKeepaliveInfo(kpd);
-
-        final KeepaliveInfo ki = mKeepaliveTracker.new KeepaliveInfo(
-                testInfo.socketKeepaliveCallback, mNai, kpd,
-                TEST_KEEPALIVE_INTERVAL_SEC, KeepaliveInfo.TYPE_TCP, testInfo.fd);
-        mKeepaliveTracker.setReturnedKeepaliveInfo(ki);
-
-        // Setup TCP keepalive.
-        mAOOKeepaliveTracker.startTcpKeepalive(mNai, testInfo.fd, TEST_KEEPALIVE_INTERVAL_SEC,
-                testInfo.socketKeepaliveCallback);
-        HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
+        final TestKeepaliveInfo testInfo =
+                doStartTcpKeepalive(InetAddress.getByAddress(V4_SRC_ADDR));
 
         // A closed socket will result in EVENT_HANGUP and trigger error to
         // FileDescriptorEventListener.
@@ -924,6 +948,6 @@
         HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
 
         // The keepalive should be removed in AutomaticOnOffKeepaliveTracker.
-        getAutoKiForBinder(testInfo.binder);
+        assertNull(getAutoKiForBinder(testInfo.binder));
     }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index 7829cb6..9ae727d 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -1965,7 +1965,16 @@
 
         vpn.startVpnProfile(TEST_VPN_PKG);
         final NetworkCallback nwCb = triggerOnAvailableAndGetCallback(underlyingNetworkCaps);
-        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
+        // There are 4 interactions with the executor.
+        // - Network available
+        // - LP change
+        // - NC change
+        // - schedule() calls in scheduleStartIkeSession()
+        // The first 3 calls are triggered from Executor.execute(). The execute() will also call to
+        // schedule() with 0 delay. Verify the exact interaction here so that it won't cause flakes
+        // in the follow-up flow.
+        verify(mExecutor, timeout(TEST_TIMEOUT_MS).times(4))
+                .schedule(any(Runnable.class), anyLong(), any());
         reset(mExecutor);
 
         // Mock the setup procedure by firing callbacks
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
index cf6275f..11c9653 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -16,6 +16,7 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsServiceTypeClient.EVENT_START_QUERYTASK;
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.junit.Assert.assertArrayEquals;
@@ -26,8 +27,10 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.never;
@@ -43,6 +46,7 @@
 import android.net.Network;
 import android.os.Handler;
 import android.os.HandlerThread;
+import android.os.Message;
 import android.text.TextUtils;
 
 import com.android.net.module.util.CollectionUtils;
@@ -112,6 +116,8 @@
     private MdnsResponseDecoder.Clock mockDecoderClock;
     @Mock
     private SharedLog mockSharedLog;
+    @Mock
+    private MdnsServiceTypeClient.Dependencies mockDeps;
     @Captor
     private ArgumentCaptor<MdnsServiceInfo> serviceInfoCaptor;
 
@@ -119,13 +125,16 @@
 
     private DatagramPacket[] expectedIPv4Packets;
     private DatagramPacket[] expectedIPv6Packets;
-    private ScheduledFuture<?>[] expectedSendFutures;
     private FakeExecutor currentThreadExecutor = new FakeExecutor();
 
     private MdnsServiceTypeClient client;
     private SocketKey socketKey;
     private HandlerThread thread;
     private Handler handler;
+    private MdnsServiceCache serviceCache;
+    private long latestDelayMs = 0;
+    private Message delayMessage = null;
+    private Handler realHandler = null;
 
     @Before
     @SuppressWarnings("DoNotMock")
@@ -135,15 +144,13 @@
 
         expectedIPv4Packets = new DatagramPacket[16];
         expectedIPv6Packets = new DatagramPacket[16];
-        expectedSendFutures = new ScheduledFuture<?>[16];
         socketKey = new SocketKey(mockNetwork, INTERFACE_INDEX);
 
-        for (int i = 0; i < expectedSendFutures.length; ++i) {
+        for (int i = 0; i < expectedIPv4Packets.length; ++i) {
             expectedIPv4Packets[i] = new DatagramPacket(buf, 0 /* offset */, 5 /* length */,
                     MdnsConstants.getMdnsIPv4Address(), MdnsConstants.MDNS_PORT);
             expectedIPv6Packets[i] = new DatagramPacket(buf, 0 /* offset */, 5 /* length */,
                     MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT);
-            expectedSendFutures[i] = Mockito.mock(ScheduledFuture.class);
         }
         when(mockPacketWriter.getPacket(IPV4_ADDRESS))
                 .thenReturn(expectedIPv4Packets[0])
@@ -184,9 +191,32 @@
         thread = new HandlerThread("MdnsServiceTypeClientTests");
         thread.start();
         handler = new Handler(thread.getLooper());
+        serviceCache = new MdnsServiceCache(thread.getLooper());
+
+        doAnswer(inv -> {
+            latestDelayMs = 0;
+            delayMessage = null;
+            return true;
+        }).when(mockDeps).removeMessages(any(Handler.class), eq(EVENT_START_QUERYTASK));
+
+        doAnswer(inv -> {
+            realHandler = (Handler) inv.getArguments()[0];
+            delayMessage = (Message) inv.getArguments()[1];
+            latestDelayMs = (long) inv.getArguments()[2];
+            return true;
+        }).when(mockDeps).sendMessageDelayed(any(Handler.class), any(Message.class), anyLong());
+
+        doAnswer(inv -> {
+            final Handler handler = (Handler) inv.getArguments()[0];
+            final Message message = (Message) inv.getArguments()[1];
+            runOnHandler(() -> handler.dispatchMessage(message));
+            return true;
+        }).when(mockDeps).sendMessage(any(Handler.class), any(Message.class));
+
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper()) {
+                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                        serviceCache) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -223,11 +253,18 @@
         runOnHandler(() -> client.notifySocketDestroyed());
     }
 
+    private void dispatchMessage() {
+        runOnHandler(() -> realHandler.dispatchMessage(delayMessage));
+        delayMessage = null;
+    }
+
     @Test
     public void sendQueries_activeScanMode() {
         MdnsSearchOptions searchOptions =
                 MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(false).build();
         startSendAndReceive(mockListenerOne, searchOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // First burst, 3 queries.
         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
@@ -265,10 +302,12 @@
                 13, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
         verifyAndSendQuery(
                 14, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        // Verify that Task is not removed before stopSendAndReceive was called.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // Stop sending packets.
         stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[15]).cancel(true);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
@@ -276,6 +315,8 @@
         MdnsSearchOptions searchOptions =
                 MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(false).build();
         startSendAndReceive(mockListenerOne, searchOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // First burst, first query is sent.
         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
@@ -289,7 +330,7 @@
                         .build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // The previous scheduled task should be canceled.
-        verify(expectedSendFutures[1]).cancel(true);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // Queries should continue to be sent.
         verifyAndSendQuery(1, 0, /* expectsUnicastResponse= */ true);
@@ -300,7 +341,7 @@
 
         // Stop sending packets.
         stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[5]).cancel(true);
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
@@ -308,6 +349,8 @@
         MdnsSearchOptions searchOptions =
                 MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(true).build();
         startSendAndReceive(mockListenerOne, searchOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // First burst, 3 query.
         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
@@ -324,7 +367,7 @@
 
         // Stop sending packets.
         stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[5]).cancel(true);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
@@ -333,6 +376,8 @@
                 MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(
                         false).setNumOfQueriesBeforeBackoff(11).build();
         startSendAndReceive(mockListenerOne, searchOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // First burst, 3 queries.
         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
@@ -367,16 +412,21 @@
         // 0.8 * smallestRemainingTtl is larger than time to next run.
         long currentTime = TEST_TTL / 2 + TEST_ELAPSED_REALTIME;
         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
+        doReturn(true).when(mockDeps).hasMessages(any(), eq(EVENT_START_QUERYTASK));
         processResponse(createResponse(
                 "service-instance-1", "192.0.2.123", 5353,
                 SERVICE_TYPE_LABELS,
                 Collections.emptyMap(), TEST_TTL), socketKey);
-        verifyAndSendQuery(12, (long) (TEST_TTL / 2 * 0.8), /* expectsUnicastResponse= */
-                false);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+        assertNotNull(delayMessage);
+        verifyAndSendQuery(12 /* index */, (long) (TEST_TTL / 2 * 0.8) /* timeInMs */,
+                false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                14 /* scheduledCount */);
         currentTime += (long) (TEST_TTL / 2 * 0.8);
         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
-        verifyAndSendQuery(
-                13, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        verifyAndSendQuery(13 /* index */, MdnsConfigs.timeBetweenQueriesInBurstMs(),
+                false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                15 /* scheduledCount */);
     }
 
     @Test
@@ -385,29 +435,36 @@
                 MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(
                         true).setNumOfQueriesBeforeBackoff(3).build();
         startSendAndReceive(mockListenerOne, searchOptions);
-        verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
-        verifyAndSendQuery(
-                1, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
-        verifyAndSendQuery(
-                2, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
-        verifyAndSendQuery(3, MdnsConfigs.timeBetweenBurstsMs(), /* expectsUnicastResponse= */
-                false);
-        assertEquals(4, currentThreadExecutor.getNumOfScheduledFuture());
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+
+        verifyAndSendQuery(0 /* index */, 0 /* timeInMs */, true /* expectsUnicastResponse */,
+                true /* multipleSocketDiscovery */, 1 /* scheduledCount */);
+        verifyAndSendQuery(1 /* index */, MdnsConfigs.timeBetweenQueriesInBurstMs(),
+                false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                2 /* scheduledCount */);
+        verifyAndSendQuery(2 /* index */, MdnsConfigs.timeBetweenQueriesInBurstMs(),
+                false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                3 /* scheduledCount */);
+        verifyAndSendQuery(3 /* index */, MdnsConfigs.timeBetweenBurstsMs(),
+                false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                4 /* scheduledCount */);
 
         // In backoff mode, the current scheduled task will be canceled and reschedule if the
         // 0.8 * smallestRemainingTtl is larger than time to next run.
         doReturn(TEST_ELAPSED_REALTIME + 20000).when(mockDecoderClock).elapsedRealtime();
+        doReturn(true).when(mockDeps).hasMessages(any(), eq(EVENT_START_QUERYTASK));
         processResponse(createResponse(
                 "service-instance-1", "192.0.2.123", 5353,
                 SERVICE_TYPE_LABELS,
                 Collections.emptyMap(), TEST_TTL), socketKey);
-        verify(expectedSendFutures[4]).cancel(true);
-        assertEquals(5, currentThreadExecutor.getNumOfScheduledFuture());
-        verifyAndSendQuery(4, 80000 /* timeInMs */, false /* expectsUnicastResponse */);
-        assertEquals(6, currentThreadExecutor.getNumOfScheduledFuture());
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+        assertNotNull(delayMessage);
+        verifyAndSendQuery(4 /* index */, 80000 /* timeInMs */, false /* expectsUnicastResponse */,
+                true /* multipleSocketDiscovery */, 6 /* scheduledCount */);
         // Next run should also be scheduled in 0.8 * smallestRemainingTtl
-        verifyAndSendQuery(5, 80000 /* timeInMs */, false /* expectsUnicastResponse */);
-        assertEquals(7, currentThreadExecutor.getNumOfScheduledFuture());
+        verifyAndSendQuery(5 /* index */, 80000 /* timeInMs */, false /* expectsUnicastResponse */,
+                true /* multipleSocketDiscovery */, 7 /* scheduledCount */);
 
         // If the records is not refreshed, the current scheduled task will not be canceled.
         doReturn(TEST_ELAPSED_REALTIME + 20001).when(mockDecoderClock).elapsedRealtime();
@@ -416,7 +473,7 @@
                 SERVICE_TYPE_LABELS,
                 Collections.emptyMap(), TEST_TTL,
                 TEST_ELAPSED_REALTIME - 1), socketKey);
-        verify(expectedSendFutures[7], never()).cancel(true);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // In backoff mode, the current scheduled task will not be canceled if the
         // 0.8 * smallestRemainingTtl is smaller than time to next run.
@@ -425,10 +482,10 @@
                 "service-instance-1", "192.0.2.123", 5353,
                 SERVICE_TYPE_LABELS,
                 Collections.emptyMap(), TEST_TTL), socketKey);
-        verify(expectedSendFutures[7], never()).cancel(true);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[7]).cancel(true);
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
@@ -436,6 +493,8 @@
         MdnsSearchOptions searchOptions =
                 MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(true).build();
         startSendAndReceive(mockListenerOne, searchOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // First burst, first query is sent.
         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
@@ -449,7 +508,7 @@
                         .build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // The previous scheduled task should be canceled.
-        verify(expectedSendFutures[1]).cancel(true);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // Queries should continue to be sent.
         verifyAndSendQuery(1, 0, /* expectsUnicastResponse= */ true);
@@ -460,7 +519,7 @@
 
         // Stop sending packets.
         stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[5]).cancel(true);
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
@@ -596,10 +655,9 @@
 
         // This time no query is submitted, only scheduled
         assertNull(currentThreadExecutor.getAndClearSubmittedRunnable());
-        assertNotNull(currentThreadExecutor.getAndClearLastScheduledRunnable());
         // This just skips the first query of the first burst
-        assertEquals(MdnsConfigs.timeBetweenQueriesInBurstMs(),
-                currentThreadExecutor.getAndClearLastScheduledDelayInMs());
+        verify(mockDeps).sendMessageDelayed(
+                any(), any(), eq(MdnsConfigs.timeBetweenQueriesInBurstMs()));
     }
 
     private static void verifyServiceInfo(MdnsServiceInfo serviceInfo, String serviceName,
@@ -853,7 +911,8 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper()) {
+                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                        serviceCache) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -877,6 +936,7 @@
         // Simulate the case where the response is under TTL.
         doReturn(TEST_ELAPSED_REALTIME + TEST_TTL - 1L).when(mockDecoderClock).elapsedRealtime();
         firstMdnsTask.run();
+        verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
 
         // Verify removed callback was not called.
         verifyServiceRemovedNoCallback(mockListenerOne);
@@ -884,6 +944,7 @@
         // Simulate the case where the response is after TTL.
         doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
         firstMdnsTask.run();
+        verify(mockDeps, times(2)).sendMessage(any(), any(Message.class));
 
         // Verify removed callback was called.
         verifyServiceRemovedCallback(
@@ -896,7 +957,8 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper()) {
+                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                        serviceCache) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -929,7 +991,8 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper()) {
+                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                        serviceCache) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -1049,7 +1112,8 @@
     @Test
     public void testProcessResponse_Resolve() throws Exception {
         client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                socketKey, mockSharedLog, thread.getLooper());
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String instanceName = "service-instance";
         final String[] hostname = new String[] { "testhost "};
@@ -1070,6 +1134,8 @@
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
                 srvTxtQueryCaptor.capture(),
                 eq(socketKey), eq(false));
+        verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
+        assertNotNull(delayMessage);
 
         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
                 new MdnsPacketReader(srvTxtQueryCaptor.getValue()));
@@ -1095,6 +1161,7 @@
         processResponse(srvTxtResponse, socketKey);
 
         // Expect a query for A/AAAA
+        dispatchMessage();
         final ArgumentCaptor<DatagramPacket> addressQueryCaptor =
                 ArgumentCaptor.forClass(DatagramPacket.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
@@ -1139,7 +1206,8 @@
     @Test
     public void testRenewTxtSrvInResolve() throws Exception {
         client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper());
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String instanceName = "service-instance";
         final String[] hostname = new String[] { "testhost "};
@@ -1160,6 +1228,8 @@
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
                 srvTxtQueryCaptor.capture(),
                 eq(socketKey), eq(false));
+        verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
+        assertNotNull(delayMessage);
 
         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
                 new MdnsPacketReader(srvTxtQueryCaptor.getValue()));
@@ -1187,6 +1257,7 @@
                 Collections.emptyList() /* authorityRecords */,
                 Collections.emptyList() /* additionalRecords */);
         processResponse(srvTxtResponse, socketKey);
+        dispatchMessage();
         inOrder.verify(mockListenerOne).onServiceNameDiscovered(any());
         inOrder.verify(mockListenerOne).onServiceFound(any());
 
@@ -1197,6 +1268,9 @@
         // Advance time so 75% of TTL passes and re-execute
         doReturn(TEST_ELAPSED_REALTIME + (long) (TEST_TTL * 0.75))
                 .when(mockDecoderClock).elapsedRealtime();
+        verify(mockDeps, times(2)).sendMessage(any(), any(Message.class));
+        assertNotNull(delayMessage);
+        dispatchMessage();
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
 
         // Expect a renewal query
@@ -1206,6 +1280,8 @@
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
                 renewalQueryCaptor.capture(),
                 eq(socketKey), eq(false));
+        verify(mockDeps, times(3)).sendMessage(any(), any(Message.class));
+        assertNotNull(delayMessage);
         inOrder.verify(mockListenerOne).onDiscoveryQuerySent(any(), anyInt());
         final MdnsPacket renewalPacket = MdnsPacket.parse(
                 new MdnsPacketReader(renewalQueryCaptor.getValue()));
@@ -1232,6 +1308,7 @@
                 Collections.emptyList() /* authorityRecords */,
                 Collections.emptyList() /* additionalRecords */);
         processResponse(refreshedSrvTxtResponse, socketKey);
+        dispatchMessage();
 
         // Advance time to updatedReceiptTime + 1, expected no refresh query because the cache
         // should contain the record that have update last receipt time.
@@ -1243,7 +1320,8 @@
     @Test
     public void testProcessResponse_ResolveExcludesOtherServices() {
         client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                socketKey, mockSharedLog, thread.getLooper());
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String requestedInstance = "instance1";
         final String otherInstance = "instance2";
@@ -1307,7 +1385,8 @@
     @Test
     public void testProcessResponse_SubtypeDiscoveryLimitedToSubtype() {
         client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                socketKey, mockSharedLog, thread.getLooper());
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String matchingInstance = "instance1";
         final String subtype = "_subtype";
@@ -1388,7 +1467,8 @@
     @Test
     public void testNotifySocketDestroyed() throws Exception {
         client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                socketKey, mockSharedLog, thread.getLooper());
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String requestedInstance = "instance1";
         final String otherInstance = "instance2";
@@ -1399,6 +1479,8 @@
                 .setResolveInstanceName("instance1").build();
 
         startSendAndReceive(mockListenerOne, resolveOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
         // Ensure the first task is executed so it schedules a future task
         currentThreadExecutor.getAndClearSubmittedFuture().get(
                 TEST_TIMEOUT_MS, TimeUnit.MILLISECONDS);
@@ -1407,7 +1489,7 @@
                         Integer.MAX_VALUE).build());
 
         // Filing the second request cancels the first future
-        verify(expectedSendFutures[0]).cancel(true);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // Ensure it gets executed too
         currentThreadExecutor.getAndClearSubmittedFuture().get(
@@ -1425,9 +1507,8 @@
                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
                 socketKey);
 
-        verify(expectedSendFutures[1], never()).cancel(true);
         notifySocketDestroyed();
-        verify(expectedSendFutures[1]).cancel(true);
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // mockListenerOne gets notified for the requested instance
         final InOrder inOrder1 = inOrder(mockListenerOne);
@@ -1442,16 +1523,109 @@
         verify(mockListenerOne, never()).onServiceNameRemoved(matchServiceName(otherInstance));
 
         // mockListenerTwo gets notified for both though
-        final InOrder inOrder2 = inOrder(mockListenerTwo);
-        inOrder2.verify(mockListenerTwo).onServiceNameDiscovered(
+        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));
+        verify(mockListenerTwo).onServiceFound(matchServiceName(requestedInstance));
+        verify(mockListenerTwo).onServiceNameDiscovered(matchServiceName(otherInstance));
+        verify(mockListenerTwo).onServiceFound(matchServiceName(otherInstance));
+        verify(mockListenerTwo).onServiceRemoved(matchServiceName(otherInstance));
+        verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(otherInstance));
+        verify(mockListenerTwo).onServiceRemoved(matchServiceName(requestedInstance));
+        verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(requestedInstance));
+    }
+
+    @Test
+    public void testServicesAreCached() throws Exception {
+        final String serviceName = "service-instance";
+        final String ipV4Address = "192.0.2.0";
+        // Register a listener
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+        InOrder inOrder = inOrder(mockListenerOne);
+
+        // Process a response which has ip address to make response become complete.
+        final String subtype = "ABCDE";
+        processResponse(createResponse(
+                        serviceName, ipV4Address, 5353, subtype,
+                        Collections.emptyMap(), TEST_TTL),
+                socketKey);
+
+        // Verify that onServiceNameDiscovered is called.
+        inOrder.verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
+        verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
+                serviceName,
+                SERVICE_TYPE_LABELS,
+                List.of(ipV4Address) /* ipv4Address */,
+                List.of() /* ipv6Address */,
+                5353 /* port */,
+                Collections.singletonList(subtype) /* subTypes */,
+                Collections.singletonMap("key", null) /* attributes */,
+                socketKey);
+
+        // Verify that onServiceFound is called.
+        inOrder.verify(mockListenerOne).onServiceFound(serviceInfoCaptor.capture());
+        verifyServiceInfo(serviceInfoCaptor.getAllValues().get(1),
+                serviceName,
+                SERVICE_TYPE_LABELS,
+                List.of(ipV4Address) /* ipv4Address */,
+                List.of() /* ipv6Address */,
+                5353 /* port */,
+                Collections.singletonList(subtype) /* subTypes */,
+                Collections.singletonMap("key", null) /* attributes */,
+                socketKey);
+
+        // Unregister the listener
+        stopSendAndReceive(mockListenerOne);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+
+        // Register another listener.
+        startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+        InOrder inOrder2 = inOrder(mockListenerTwo);
+
+        // The services are cached in MdnsServiceCache, verify that onServiceNameDiscovered is
+        // called immediately.
+        inOrder2.verify(mockListenerTwo).onServiceNameDiscovered(serviceInfoCaptor.capture());
+        verifyServiceInfo(serviceInfoCaptor.getAllValues().get(2),
+                serviceName,
+                SERVICE_TYPE_LABELS,
+                List.of(ipV4Address) /* ipv4Address */,
+                List.of() /* ipv6Address */,
+                5353 /* port */,
+                Collections.singletonList(subtype) /* subTypes */,
+                Collections.singletonMap("key", null) /* attributes */,
+                socketKey);
+
+        // The services are cached in MdnsServiceCache, verify that onServiceFound is
+        // called immediately.
+        inOrder2.verify(mockListenerTwo).onServiceFound(serviceInfoCaptor.capture());
+        verifyServiceInfo(serviceInfoCaptor.getAllValues().get(3),
+                serviceName,
+                SERVICE_TYPE_LABELS,
+                List.of(ipV4Address) /* ipv4Address */,
+                List.of() /* ipv6Address */,
+                5353 /* port */,
+                Collections.singletonList(subtype) /* subTypes */,
+                Collections.singletonMap("key", null) /* attributes */,
+                socketKey);
+
+        // Process a response with a different ip address, port and updated text attributes.
+        final String ipV6Address = "2001:db8::";
+        processResponse(createResponse(
+                serviceName, ipV6Address, 5354, subtype,
+                Collections.singletonMap("key", "value"), TEST_TTL), socketKey);
+
+        // Verify the onServiceUpdated is called.
+        inOrder2.verify(mockListenerTwo).onServiceUpdated(serviceInfoCaptor.capture());
+        verifyServiceInfo(serviceInfoCaptor.getAllValues().get(4),
+                serviceName,
+                SERVICE_TYPE_LABELS,
+                List.of(ipV4Address) /* ipv4Address */,
+                List.of(ipV6Address) /* ipv6Address */,
+                5354 /* port */,
+                Collections.singletonList(subtype) /* subTypes */,
+                Collections.singletonMap("key", "value") /* attributes */,
+                socketKey);
     }
 
     private static MdnsServiceInfo matchServiceName(String name) {
@@ -1461,13 +1635,17 @@
     // verifies that the right query was enqueued with the right delay, and send query by executing
     // the runnable.
     private void verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse) {
-        verifyAndSendQuery(
-                index, timeInMs, expectsUnicastResponse, true /* multipleSocketDiscovery */);
+        verifyAndSendQuery(index, timeInMs, expectsUnicastResponse,
+                true /* multipleSocketDiscovery */, index + 1 /* scheduledCount */);
     }
 
     private void verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse,
-            boolean multipleSocketDiscovery) {
-        assertEquals(timeInMs, currentThreadExecutor.getAndClearLastScheduledDelayInMs());
+            boolean multipleSocketDiscovery, int scheduledCount) {
+        // Dispatch the message
+        if (delayMessage != null && realHandler != null) {
+            dispatchMessage();
+        }
+        assertEquals(timeInMs, latestDelayMs);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         if (expectsUnicastResponse) {
             verify(mockSocketClient).sendPacketRequestingUnicastResponse(
@@ -1484,6 +1662,11 @@
                         expectedIPv6Packets[index], socketKey, false);
             }
         }
+        verify(mockDeps, times(index + 1))
+                .sendMessage(any(Handler.class), any(Message.class));
+        // Verify the task has been scheduled.
+        verify(mockDeps, times(scheduledCount))
+                .sendMessageDelayed(any(Handler.class), any(Message.class), anyLong());
     }
 
     private static String[] getTestServiceName(String instanceName) {
@@ -1528,7 +1711,7 @@
         public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
             lastScheduledDelayInMs = delay;
             lastScheduledRunnable = command;
-            return expectedSendFutures[futureIndex++];
+            return Mockito.mock(ScheduledFuture.class);
         }
 
         // Returns the delay of the last scheduled task, and clear it.
@@ -1556,10 +1739,6 @@
             lastSubmittedFuture = null;
             return val;
         }
-
-        public int getNumOfScheduledFuture() {
-            return futureIndex - 1;
-        }
     }
 
     private MdnsPacket createResponse(