diff --git a/Cronet/OWNERS b/Cronet/OWNERS
index 62c5737..c24680e 100644
--- a/Cronet/OWNERS
+++ b/Cronet/OWNERS
@@ -1,2 +1,2 @@
 set noparent
-file:platform/packages/modules/Connectivity:master:/OWNERS_core_networking
+file:platform/packages/modules/Connectivity:main:/OWNERS_core_networking
diff --git a/Cronet/tests/OWNERS b/Cronet/tests/OWNERS
index acb6ee6..a35a789 100644
--- a/Cronet/tests/OWNERS
+++ b/Cronet/tests/OWNERS
@@ -1,7 +1,7 @@
 # Bug component: 31808
 
 set noparent
-file:platform/packages/modules/Connectivity:master:/OWNERS_core_networking_xts
+file:platform/packages/modules/Connectivity:main:/OWNERS_core_networking_xts
 
 # TODO: Temp ownership to develop cronet CTS
 colibie@google.com #{LAST_RESORT_SUGGESTION}
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 0760e68..3a77097 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt
+++ b/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt
@@ -27,12 +27,15 @@
 import androidx.test.core.app.ApplicationProvider
 import com.android.testutils.DevSdkIgnoreRule
 import com.android.testutils.DevSdkIgnoreRunner
+import com.google.common.truth.Truth.assertThat
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import org.hamcrest.MatcherAssert
 import org.hamcrest.Matchers
 import org.junit.After
+import org.junit.AssumptionViolatedException
 import org.junit.Before
+import org.junit.Ignore
 import org.junit.runner.RunWith
 
 private const val URL = "https://source.android.com"
@@ -57,11 +60,6 @@
     @After
     @Throws(Exception::class)
     fun tearDown() {
-        // cancel active requests to enable engine shutdown.
-        stream?.let {
-            it.cancel()
-            callback.blockForDone()
-        }
         httpEngine.shutdown()
     }
 
@@ -71,14 +69,133 @@
 
     @Test
     @Throws(Exception::class)
+    @Ignore("b/292298108 Re-enable and confirm non-flaky after SLO")
     fun testBidirectionalStream_GetStream_CompletesSuccessfully() {
         stream = createBidirectionalStreamBuilder(URL).setHttpMethod("GET").build()
         stream!!.start()
-        callback.assumeCallback(ResponseStep.ON_SUCCEEDED)
+        // We call to a real server and hence the server may not be reachable, cancel this stream
+        // and rethrow the exception before tearDown,
+        // otherwise shutdown would fail with active request error.
+        try {
+            callback.assumeCallback(ResponseStep.ON_SUCCEEDED)
+        } catch (e: AssumptionViolatedException) {
+            stream!!.cancel()
+            callback.blockForDone()
+            throw e
+        }
+
         val info = callback.mResponseInfo
         assumeOKStatusCode(info)
         MatcherAssert.assertThat(
             "Received byte count must be > 0", info.receivedByteCount, Matchers.greaterThan(0L))
         assertEquals("h2", info.negotiatedProtocol)
     }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getHttpMethod() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val method = "GET"
+
+        builder.setHttpMethod(method)
+        stream = builder.build()
+        assertThat(stream!!.getHttpMethod()).isEqualTo(method)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_hasTrafficStatsTag() {
+        val builder = createBidirectionalStreamBuilder(URL)
+
+        builder.setTrafficStatsTag(10)
+        stream = builder.build()
+        assertThat(stream!!.hasTrafficStatsTag()).isTrue()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getTrafficStatsTag() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val trafficStatsTag = 10
+
+        builder.setTrafficStatsTag(trafficStatsTag)
+        stream = builder.build()
+        assertThat(stream!!.getTrafficStatsTag()).isEqualTo(trafficStatsTag)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_hasTrafficStatsUid() {
+        val builder = createBidirectionalStreamBuilder(URL)
+
+        builder.setTrafficStatsUid(10)
+        stream = builder.build()
+        assertThat(stream!!.hasTrafficStatsUid()).isTrue()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getTrafficStatsUid() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val trafficStatsUid = 10
+
+        builder.setTrafficStatsUid(trafficStatsUid)
+        stream = builder.build()
+        assertThat(stream!!.getTrafficStatsUid()).isEqualTo(trafficStatsUid)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getHeaders_asList() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val expectedHeaders = mapOf(
+          "Authorization" to "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ==",
+          "Max-Forwards" to "10",
+          "X-Client-Data" to "random custom header content").entries.toList()
+
+        for (header in expectedHeaders) {
+            builder.addHeader(header.key, header.value)
+        }
+
+        stream = builder.build()
+        assertThat(stream!!.getHeaders().getAsList()).containsAtLeastElementsIn(expectedHeaders)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getHeaders_asMap() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val expectedHeaders = mapOf(
+          "Authorization" to listOf("Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="),
+          "Max-Forwards" to listOf("10"),
+          "X-Client-Data" to listOf("random custom header content"))
+
+        for (header in expectedHeaders) {
+            builder.addHeader(header.key, header.value.get(0))
+        }
+
+        stream = builder.build()
+        assertThat(stream!!.getHeaders().getAsMap()).containsAtLeastEntriesIn(expectedHeaders)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getPriority() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val priority = BidirectionalStream.STREAM_PRIORITY_LOW
+
+        builder.setPriority(priority)
+        stream = builder.build()
+        assertThat(stream!!.getPriority()).isEqualTo(priority)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_isDelayRequestHeadersUntilFirstFlushEnabled() {
+        val builder = createBidirectionalStreamBuilder(URL)
+
+        builder.setDelayRequestHeadersUntilFirstFlushEnabled(true)
+        stream = builder.build()
+        assertThat(stream!!.isDelayRequestHeadersUntilFirstFlushEnabled()).isTrue()
+    }
 }
diff --git a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
index 07e7d45..3c4d134 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
+++ b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
@@ -363,6 +363,116 @@
                 .containsAtLeastElementsIn(expectedHeaders);
     }
 
+    @Test
+    public void testUrlRequest_getHttpMethod() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+        final String method = "POST";
+
+        builder.setHttpMethod(method);
+        UrlRequest request = builder.build();
+        assertThat(request.getHttpMethod()).isEqualTo(method);
+    }
+
+    @Test
+    public void testUrlRequest_getHeaders_asList() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+        final List<Map.Entry<String, String>> expectedHeaders = Arrays.asList(
+                Map.entry("Authorization", "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="),
+                Map.entry("Max-Forwards", "10"),
+                Map.entry("X-Client-Data", "random custom header content"));
+
+        for (Map.Entry<String, String> header : expectedHeaders) {
+            builder.addHeader(header.getKey(), header.getValue());
+        }
+
+        UrlRequest request = builder.build();
+        assertThat(request.getHeaders().getAsList()).containsAtLeastElementsIn(expectedHeaders);
+    }
+
+    @Test
+    public void testUrlRequest_getHeaders_asMap() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+        final Map<String, List<String>> expectedHeaders = Map.of(
+                "Authorization", Arrays.asList("Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="),
+                "Max-Forwards", Arrays.asList("10"),
+                "X-Client-Data", Arrays.asList("random custom header content"));
+
+        for (Map.Entry<String, List<String>> header : expectedHeaders.entrySet()) {
+            builder.addHeader(header.getKey(), header.getValue().get(0));
+        }
+
+        UrlRequest request = builder.build();
+        assertThat(request.getHeaders().getAsMap()).containsAtLeastEntriesIn(expectedHeaders);
+    }
+
+    @Test
+    public void testUrlRequest_isCacheDisabled() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+        final boolean isCacheDisabled = true;
+
+        builder.setCacheDisabled(isCacheDisabled);
+        UrlRequest request = builder.build();
+        assertThat(request.isCacheDisabled()).isEqualTo(isCacheDisabled);
+    }
+
+    @Test
+    public void testUrlRequest_isDirectExecutorAllowed() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+        final boolean isDirectExecutorAllowed = true;
+
+        builder.setDirectExecutorAllowed(isDirectExecutorAllowed);
+        UrlRequest request = builder.build();
+        assertThat(request.isDirectExecutorAllowed()).isEqualTo(isDirectExecutorAllowed);
+    }
+
+    @Test
+    public void testUrlRequest_getPriority() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+        final int priority = UrlRequest.REQUEST_PRIORITY_LOW;
+
+        builder.setPriority(priority);
+        UrlRequest request = builder.build();
+        assertThat(request.getPriority()).isEqualTo(priority);
+    }
+
+    @Test
+    public void testUrlRequest_hasTrafficStatsTag() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+
+        builder.setTrafficStatsTag(10);
+        UrlRequest request = builder.build();
+        assertThat(request.hasTrafficStatsTag()).isEqualTo(true);
+    }
+
+    @Test
+    public void testUrlRequest_getTrafficStatsTag() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+        final int trafficStatsTag = 10;
+
+        builder.setTrafficStatsTag(trafficStatsTag);
+        UrlRequest request = builder.build();
+        assertThat(request.getTrafficStatsTag()).isEqualTo(trafficStatsTag);
+    }
+
+    @Test
+    public void testUrlRequest_hasTrafficStatsUid() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+
+        builder.setTrafficStatsUid(10);
+        UrlRequest request = builder.build();
+        assertThat(request.hasTrafficStatsUid()).isEqualTo(true);
+    }
+
+    @Test
+    public void testUrlRequest_getTrafficStatsUid() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getSuccessUrl());
+        final int trafficStatsUid = 10;
+
+        builder.setTrafficStatsUid(trafficStatsUid);
+        UrlRequest request = builder.build();
+        assertThat(request.getTrafficStatsUid()).isEqualTo(trafficStatsUid);
+    }
+
     private static List<Map.Entry<String, String>> extractEchoedHeaders(HeaderBlock headers) {
         return headers.getAsList()
                 .stream()
diff --git a/Cronet/tests/mts/Android.bp b/Cronet/tests/mts/Android.bp
index 4e4251c..63905c8 100644
--- a/Cronet/tests/mts/Android.bp
+++ b/Cronet/tests/mts/Android.bp
@@ -38,7 +38,11 @@
 // tests need to inherit the NetHttpTests manifest.
 android_library {
     name: "NetHttpTestsLibPreJarJar",
-    static_libs: ["cronet_java_tests"],
+    static_libs: [
+        "cronet_aml_api_java",
+        "cronet_aml_java__testing",
+        "cronet_java_tests",
+    ],
     sdk_version: "module_current",
     min_sdk_version: "30",
 }
@@ -51,7 +55,8 @@
      static_libs: ["NetHttpTestsLibPreJarJar"],
      jarjar_rules: ":net-http-test-jarjar-rules",
      jni_libs: [
-        "cronet_aml_components_cronet_android_cronet_tests__testing"
+        "cronet_aml_components_cronet_android_cronet__testing",
+        "cronet_aml_components_cronet_android_cronet_tests__testing",
      ],
      test_suites: [
          "general-tests",
diff --git a/OWNERS b/OWNERS
index 07a775e..649efda 100644
--- a/OWNERS
+++ b/OWNERS
@@ -1,4 +1,4 @@
 set noparent
-file:platform/packages/modules/Connectivity:master:/OWNERS_core_networking
+file:platform/packages/modules/Connectivity:main:/OWNERS_core_networking
 
-per-file **IpSec* = file:platform/frameworks/base:master:/services/core/java/com/android/server/vcn/OWNERS
\ No newline at end of file
+per-file **IpSec* = file:platform/frameworks/base:main:/services/core/java/com/android/server/vcn/OWNERS
diff --git a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
index a9add20..9dad301 100644
--- a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
+++ b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
@@ -119,9 +119,11 @@
             (Inet4Address) parseNumericAddress("8.8.8.8");
     protected static final Inet6Address REMOTE_IP6_ADDR =
             (Inet6Address) parseNumericAddress("2002:db8:1::515:ca");
+    // The IPv6 network address translation of REMOTE_IP4_ADDR if pref64::/n is 64:ff9b::/96.
+    // For more information, see TetheringTester#PREF64_IPV4ONLY_ADDR, which assumes a prefix
+    // of 64:ff9b::/96.
     protected static final Inet6Address REMOTE_NAT64_ADDR =
             (Inet6Address) parseNumericAddress("64:ff9b::808:808");
-    protected static final IpPrefix TEST_NAT64PREFIX = new IpPrefix("64:ff9b::/96");
 
     // LOCAL_PORT is used by public port and private port. Assume port 9876 has not been used yet
     // before the testing that public port and private port are the same in the testing. Note that
@@ -165,13 +167,18 @@
         mHandlerThread.start();
         mHandler = new Handler(mHandlerThread.getLooper());
 
-        mRunTests = runAsShell(NETWORK_SETTINGS, TETHER_PRIVILEGED, () ->
-                mTm.isTetheringSupported());
+        mRunTests = isEthernetTetheringSupported();
         assumeTrue(mRunTests);
 
         mTetheredInterfaceRequester = new TetheredInterfaceRequester(mHandler, mEm);
     }
 
+    private boolean isEthernetTetheringSupported() throws Exception {
+        if (mEm == null) return false;
+
+        return runAsShell(NETWORK_SETTINGS, TETHER_PRIVILEGED, () -> mTm.isTetheringSupported());
+    }
+
     protected void maybeStopTapPacketReader(final TapPacketReader tapPacketReader)
             throws Exception {
         if (tapPacketReader != null) {
@@ -630,7 +637,6 @@
         final LinkProperties lp = new LinkProperties();
         lp.setLinkAddresses(addresses);
         lp.setDnsServers(dnses);
-        lp.setNat64Prefix(TEST_NAT64PREFIX);
 
         return runAsShell(MANAGE_TEST_NETWORKS, () -> initTestNetwork(mContext, lp, TIMEOUT_MS));
     }
diff --git a/Tethering/tests/integration/base/android/net/TetheringTester.java b/Tethering/tests/integration/base/android/net/TetheringTester.java
index 33baf93..3f3768e 100644
--- a/Tethering/tests/integration/base/android/net/TetheringTester.java
+++ b/Tethering/tests/integration/base/android/net/TetheringTester.java
@@ -16,6 +16,8 @@
 
 package android.net;
 
+import static android.net.DnsResolver.CLASS_IN;
+import static android.net.DnsResolver.TYPE_AAAA;
 import static android.net.InetAddresses.parseNumericAddress;
 import static android.system.OsConstants.ICMP_ECHO;
 import static android.system.OsConstants.ICMP_ECHOREPLY;
@@ -28,6 +30,8 @@
 
 import static com.android.net.module.util.DnsPacket.ANSECTION;
 import static com.android.net.module.util.DnsPacket.ARSECTION;
+import static com.android.net.module.util.DnsPacket.DnsHeader;
+import static com.android.net.module.util.DnsPacket.DnsRecord;
 import static com.android.net.module.util.DnsPacket.NSSECTION;
 import static com.android.net.module.util.DnsPacket.QDSECTION;
 import static com.android.net.module.util.HexDump.dumpHexString;
@@ -130,6 +134,27 @@
     // ICMP definition.
     private static final short ICMPECHO_CODE = 0x0;
 
+    // Prefix64 discovery definition. See RFC 7050 section 8.
+    // Note that the AAAA response Pref64::WKAs consisting of Pref64::/n and WKA.
+    // Use 64:ff9b::/96 as Pref64::/n and WKA 192.0.0.17{0|1} here.
+    //
+    // Host                                          DNS64 server
+    //   |                                                |
+    //   |  "AAAA" query for "ipv4only.arpa."             |
+    //   |----------------------------------------------->|
+    //   |                                                |
+    //   |  "AAAA" response with:                         |
+    //   |  "64:ff9b::192.0.0.170"                        |
+    //   |<-----------------------------------------------|
+    //
+    private static final String PREF64_IPV4ONLY_HOSTNAME = "ipv4only.arpa";
+    private static final InetAddress PREF64_IPV4ONLY_ADDR = parseNumericAddress(
+            "64:ff9b::192.0.0.170");
+
+    // DNS header definition.
+    private static final short FLAG = (short) 0x8100;  // qr, ra
+    private static final short TTL = (short) 0;
+
     public static final String DHCP_HOSTNAME = "testhostname";
 
     private final ArrayMap<MacAddress, TetheredDevice> mTetheredDevices;
@@ -490,6 +515,11 @@
             super(data);
         }
 
+        TestDnsPacket(@NonNull DnsHeader header, @Nullable ArrayList<DnsRecord> qd,
+                @Nullable ArrayList<DnsRecord> an) {
+            super(header, qd, an);
+        }
+
         @Nullable
         public static TestDnsPacket getTestDnsPacket(final ByteBuffer buf) {
             try {
@@ -862,10 +892,85 @@
         return null;
     }
 
+    @NonNull
+    private ByteBuffer buildUdpDnsPrefix64ReplyPacket(int dnsId, @NonNull final Inet6Address srcIp,
+            @NonNull final Inet6Address dstIp, short srcPort, short dstPort) throws Exception {
+        // [1] Build prefix64 DNS message.
+        final ArrayList<DnsRecord> qlist = new ArrayList<>();
+        // Fill QD section.
+        qlist.add(DnsRecord.makeQuestion(PREF64_IPV4ONLY_HOSTNAME, TYPE_AAAA, CLASS_IN));
+        final ArrayList<DnsRecord> alist = new ArrayList<>();
+        // Fill AN sections.
+        alist.add(DnsRecord.makeAOrAAAARecord(ANSECTION, PREF64_IPV4ONLY_HOSTNAME, CLASS_IN, TTL,
+                PREF64_IPV4ONLY_ADDR));
+        final TestDnsPacket dns = new TestDnsPacket(
+                new DnsHeader(dnsId, FLAG, qlist.size(), alist.size()), qlist, alist);
+
+        // [2] Build IPv6 UDP DNS packet.
+        return buildUdpPacket(srcIp, dstIp, srcPort, dstPort, ByteBuffer.wrap(dns.getBytes()));
+    }
+
+    private void maybeReplyUdpDnsPrefix64Discovery(@NonNull byte[] packet) {
+        final ByteBuffer buf = ByteBuffer.wrap(packet);
+
+        // [1] Parse the prefix64 discovery DNS query for hostname ipv4only.arpa.
+        // Parse IPv6 and UDP header.
+        Ipv6Header ipv6Header = null;
+        try {
+            ipv6Header = Struct.parse(Ipv6Header.class, buf);
+            if (ipv6Header == null || ipv6Header.nextHeader != IPPROTO_UDP) return;
+        } catch (Exception e) {
+            // Parsing packet fail means it is not IPv6 UDP packet.
+            return;
+        }
+        final UdpHeader udpHeader = Struct.parse(UdpHeader.class, buf);
+
+        // Parse DNS message.
+        final TestDnsPacket pref64Query = TestDnsPacket.getTestDnsPacket(buf);
+        if (pref64Query == null) return;
+        if (pref64Query.getHeader().isResponse()) return;
+        if (pref64Query.getQDCount() != 1) return;
+        if (pref64Query.getANCount() != 0) return;
+        if (pref64Query.getNSCount() != 0) return;
+        if (pref64Query.getARCount() != 0) return;
+
+        final List<DnsRecord> qdRecordList = pref64Query.getRecordList(QDSECTION);
+        if (qdRecordList.size() != 1) return;
+        if (!qdRecordList.get(0).dName.equals(PREF64_IPV4ONLY_HOSTNAME)) return;
+
+        // [2] Build prefix64 DNS discovery reply from received query.
+        // DNS response transaction id must be copied from DNS query. Used by the requester
+        // to match up replies to outstanding queries. See RFC 1035 section 4.1.1. Also reverse
+        // the source/destination address/port of query packet for building reply packet.
+        final ByteBuffer replyPacket;
+        try {
+            replyPacket = buildUdpDnsPrefix64ReplyPacket(pref64Query.getHeader().getId(),
+                    ipv6Header.dstIp /* srcIp */, ipv6Header.srcIp /* dstIp */,
+                    (short) udpHeader.dstPort /* srcPort */,
+                    (short) udpHeader.srcPort /* dstPort */);
+        } catch (Exception e) {
+            fail("Failed to build prefix64 discovery reply for " + ipv6Header.srcIp + ": " + e);
+            return;
+        }
+
+        Log.d(TAG, "Sending prefix64 discovery reply");
+        try {
+            sendDownloadPacket(replyPacket);
+        } catch (Exception e) {
+            fail("Failed to reply prefix64 discovery for " + ipv6Header.srcIp + ": " + e);
+        }
+    }
+
     private byte[] getUploadPacket(Predicate<byte[]> filter) {
         assertNotNull("Can't deal with upstream interface in local only mode", mUpstreamReader);
 
-        return mUpstreamReader.poll(PACKET_READ_TIMEOUT_MS, filter);
+        byte[] packet;
+        while ((packet = mUpstreamReader.poll(PACKET_READ_TIMEOUT_MS)) != null) {
+            if (filter.test(packet)) return packet;
+
+            maybeReplyUdpDnsPrefix64Discovery(packet);
+        }
+        return null;
     }
 
     private @NonNull byte[] verifyPacketNotNull(String message, @Nullable byte[] packet) {
diff --git a/bpf_progs/clatd.c b/bpf_progs/clatd.c
index 905b8fa..a104084 100644
--- a/bpf_progs/clatd.c
+++ b/bpf_progs/clatd.c
@@ -138,10 +138,11 @@
     }
 
     switch (proto) {
-        case IPPROTO_TCP:  // For TCP & UDP the checksum neutrality of the chosen IPv6
-        case IPPROTO_UDP:  // address means there is no need to update their checksums.
-        case IPPROTO_GRE:  // We do not need to bother looking at GRE/ESP headers,
-        case IPPROTO_ESP:  // since there is never a checksum to update.
+        case IPPROTO_TCP:      // For TCP, UDP & UDPLITE the checksum neutrality of the chosen
+        case IPPROTO_UDP:      // IPv6 address means there is no need to update their checksums.
+        case IPPROTO_UDPLITE:  //
+        case IPPROTO_GRE:      // We do not need to bother looking at GRE/ESP headers,
+        case IPPROTO_ESP:      // since there is never a checksum to update.
             break;
 
         default:  // do not know how to handle anything else
@@ -328,12 +329,13 @@
     if (ip4->frag_off & ~htons(IP_DF)) return TC_ACT_PIPE;
 
     switch (ip4->protocol) {
-        case IPPROTO_TCP:  // For TCP & UDP the checksum neutrality of the chosen IPv6
-        case IPPROTO_GRE:  // address means there is no need to update their checksums.
-        case IPPROTO_ESP:  // We do not need to bother looking at GRE/ESP headers,
-            break;         // since there is never a checksum to update.
+        case IPPROTO_TCP:      // For TCP, UDP & UDPLITE the checksum neutrality of the chosen
+        case IPPROTO_UDPLITE:  // IPv6 address means there is no need to update their checksums.
+        case IPPROTO_GRE:      // We do not need to bother looking at GRE/ESP headers,
+        case IPPROTO_ESP:      // since there is never a checksum to update.
+            break;
 
-        case IPPROTO_UDP:  // See above comment, but must also have UDP header...
+        case IPPROTO_UDP:      // See above comment, but must also have UDP header...
             if (data + sizeof(*ip4) + sizeof(struct udphdr) > data_end) return TC_ACT_PIPE;
             const struct udphdr* uh = (const struct udphdr*)(ip4 + 1);
             // If IPv4/UDP checksum is 0 then fallback to clatd so it can calculate the
diff --git a/framework/Android.bp b/framework/Android.bp
index 123f02a..cab11ae 100644
--- a/framework/Android.bp
+++ b/framework/Android.bp
@@ -109,8 +109,8 @@
         "framework-connectivity-defaults",
     ],
     static_libs: [
-        "cronet_aml_api_java",
-        "cronet_aml_java",
+        "httpclient_api",
+        "httpclient_impl",
     ],
     libs: [
         // This cannot be in the defaults clause above because if it were, it would be used
@@ -125,12 +125,12 @@
 
 java_defaults {
     name: "CronetJavaDefaults",
-    srcs: [":cronet_aml_api_sources"],
+    srcs: [":httpclient_api_sources"],
     libs: [
         "androidx.annotation_annotation",
     ],
     impl_only_static_libs: [
-        "cronet_aml_java",
+        "httpclient_impl",
     ],
 }
 
diff --git a/framework/src/android/net/LinkAddress.java b/framework/src/android/net/LinkAddress.java
index 90f55b3..8376963 100644
--- a/framework/src/android/net/LinkAddress.java
+++ b/framework/src/android/net/LinkAddress.java
@@ -37,6 +37,8 @@
 import android.os.SystemClock;
 import android.util.Pair;
 
+import com.android.net.module.util.ConnectivityUtils;
+
 import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetAddress;
@@ -146,11 +148,7 @@
      * Per RFC 4193 section 8, fc00::/7 identifies these addresses.
      */
     private boolean isIpv6ULA() {
-        if (isIpv6()) {
-            byte[] bytes = address.getAddress();
-            return ((bytes[0] & (byte)0xfe) == (byte)0xfc);
-        }
-        return false;
+        return ConnectivityUtils.isIPv6ULA(address);
     }
 
     /**
diff --git a/framework/src/android/net/NetworkCapabilities.java b/framework/src/android/net/NetworkCapabilities.java
index 92e9599..8e219a6 100644
--- a/framework/src/android/net/NetworkCapabilities.java
+++ b/framework/src/android/net/NetworkCapabilities.java
@@ -690,17 +690,10 @@
      */
     public static final int NET_CAPABILITY_PRIORITIZE_BANDWIDTH = 35;
 
-    private static final int MIN_NET_CAPABILITY = NET_CAPABILITY_MMS;
     private static final int MAX_NET_CAPABILITY = NET_CAPABILITY_PRIORITIZE_BANDWIDTH;
 
-    private static final int ALL_VALID_CAPABILITIES;
-    static {
-        int caps = 0;
-        for (int i = MIN_NET_CAPABILITY; i <= MAX_NET_CAPABILITY; ++i) {
-            caps |= 1 << i;
-        }
-        ALL_VALID_CAPABILITIES = caps;
-    }
+    // Set all bits up to the MAX_NET_CAPABILITY-th bit
+    private static final long ALL_VALID_CAPABILITIES = (2L << MAX_NET_CAPABILITY) - 1;
 
     /**
      * Network capabilities that are expected to be mutable, i.e., can change while a particular
@@ -2519,7 +2512,7 @@
     }
 
     private static boolean isValidCapability(@NetworkCapabilities.NetCapability int capability) {
-        return capability >= MIN_NET_CAPABILITY && capability <= MAX_NET_CAPABILITY;
+        return capability >= 0 && capability <= MAX_NET_CAPABILITY;
     }
 
     private static void checkValidCapability(@NetworkCapabilities.NetCapability int capability) {
diff --git a/remote_auth/OWNERS b/remote_auth/OWNERS
new file mode 100644
index 0000000..25a32b9
--- /dev/null
+++ b/remote_auth/OWNERS
@@ -0,0 +1,14 @@
+# Bug component: 1145231
+# Bug template url: http://b/new?component=1145231&template=1715387
+billyhuang@google.com
+boetger@google.com
+casbor@google.com
+derekjedral@google.com
+dlm@google.com
+igorzas@google.com
+jacobhobbie@google.com
+jasonsun@google.com
+jianbing@google.com
+jinjiechen@google.com
+justinmcclain@google.com
+salilr@google.com
diff --git a/remote_auth/README.md b/remote_auth/README.md
new file mode 100644
index 0000000..384fcf7
--- /dev/null
+++ b/remote_auth/README.md
@@ -0,0 +1,47 @@
+# RemoteAuth Mainline Module
+
+This directory contains code for the RemoteAuth module.
+
+## Directory Structure
+
+`framework`
+ - Contains client side APIs and AIDL files.
+
+`jni`
+ - JNI wrapper for invoking Android APIs from native code.
+
+`native`
+ - Native code implementation for RemoteAuth module services.
+
+`service`
+ - Server side implementation for RemoteAuth module services.
+
+`tests`
+ - Unit/Multi devices tests for RemoteAuth module (both Java and native code).
+
+## IDE setup
+
+### AIDEGen
+
+AIDEGen is deprecated, prefer ASfP [go/asfp](http://go/asfp)
+```sh
+$ source build/envsetup.sh && lunch <TARGET>
+$ cd packages/modules/Connectivity
+$ aidegen .
+# This will launch Intellij project for RemoteAuth module.
+```
+
+### ASfP
+
+See full instructions for ASfP at [go/asfp-getting-started](http://go/asfp-getting-started)
+
+## Build and Install
+
+```sh
+$ source build/envsetup.sh && lunch <TARGET>
+$ m com.google.android.tethering deapexer
+$ $ANDROID_BUILD_TOP/out/host/linux-x86/bin/deapexer decompress --input \
+    ${ANDROID_PRODUCT_OUT}/system/apex/com.google.android.tethering.capex \
+    --output /tmp/tethering.apex
+$ adb install -r /tmp/tethering.apex
+```
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/metrics/NetworkNsdReportedMetrics.java b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
new file mode 100644
index 0000000..bcedbef
--- /dev/null
+++ b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
@@ -0,0 +1,125 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.metrics;
+
+import static com.android.metrics.NetworkNsdReported.Builder;
+
+import android.stats.connectivity.MdnsQueryResult;
+import android.stats.connectivity.NsdEventType;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.ConnectivityStatsLog;
+
+/**
+ * Class to record the NetworkNsdReported into statsd. Each client should create this class to
+ * report its data.
+ */
+public class NetworkNsdReportedMetrics {
+    // Whether this client is using legacy backend.
+    private final boolean mIsLegacy;
+    // The client id.
+    private final int mClientId;
+    private final Dependencies mDependencies;
+
+    public NetworkNsdReportedMetrics(boolean isLegacy, int clientId) {
+        this(isLegacy, clientId, new Dependencies());
+    }
+
+    @VisibleForTesting
+    NetworkNsdReportedMetrics(boolean isLegacy, int clientId, Dependencies dependencies) {
+        mIsLegacy = isLegacy;
+        mClientId = clientId;
+        mDependencies = dependencies;
+    }
+
+    /**
+     * Dependencies of NetworkNsdReportedMetrics, for injection in tests.
+     */
+    public static class Dependencies {
+
+        /**
+         * @see ConnectivityStatsLog
+         */
+        public void statsWrite(NetworkNsdReported event) {
+            ConnectivityStatsLog.write(ConnectivityStatsLog.NETWORK_NSD_REPORTED,
+                    event.getIsLegacy(),
+                    event.getClientId(),
+                    event.getTransactionId(),
+                    event.getIsKnownService(),
+                    event.getType().getNumber(),
+                    event.getEventDurationMillisec(),
+                    event.getQueryResult().getNumber(),
+                    event.getFoundServiceCount(),
+                    event.getFoundCallbackCount(),
+                    event.getLostCallbackCount(),
+                    event.getRepliedRequestsCount());
+        }
+    }
+
+    private Builder makeReportedBuilder() {
+        final Builder builder = NetworkNsdReported.newBuilder();
+        builder.setIsLegacy(mIsLegacy);
+        builder.setClientId(mClientId);
+        return builder;
+    }
+
+    /**
+     * Report service registration succeeded metric data.
+     *
+     * @param transactionId The transaction id of service registration.
+     * @param durationMs The duration of service registration success.
+     */
+    public void reportServiceRegistrationSucceeded(int transactionId, long durationMs) {
+        final Builder builder = makeReportedBuilder();
+        builder.setTransactionId(transactionId);
+        builder.setType(NsdEventType.NET_REGISTER);
+        builder.setQueryResult(MdnsQueryResult.MQR_SERVICE_REGISTERED);
+        builder.setEventDurationMillisec(durationMs);
+        mDependencies.statsWrite(builder.build());
+    }
+
+    /**
+     * Report service registration failed metric data.
+     *
+     * @param transactionId The transaction id of service registration.
+     * @param durationMs The duration of service registration failed.
+     */
+    public void reportServiceRegistrationFailed(int transactionId, long durationMs) {
+        final Builder builder = makeReportedBuilder();
+        builder.setTransactionId(transactionId);
+        builder.setType(NsdEventType.NET_REGISTER);
+        builder.setQueryResult(MdnsQueryResult.MQR_SERVICE_REGISTRATION_FAILED);
+        builder.setEventDurationMillisec(durationMs);
+        mDependencies.statsWrite(builder.build());
+    }
+
+    /**
+     * Report service unregistration success metric data.
+     *
+     * @param transactionId The transaction id of service registration.
+     * @param durationMs The duration of service stayed registered.
+     */
+    public void reportServiceUnregistration(int transactionId, long durationMs) {
+        final Builder builder = makeReportedBuilder();
+        builder.setTransactionId(transactionId);
+        builder.setType(NsdEventType.NET_REGISTER);
+        builder.setQueryResult(MdnsQueryResult.MQR_SERVICE_UNREGISTERED);
+        builder.setEventDurationMillisec(durationMs);
+        // TODO: Report repliedRequestsCount
+        mDependencies.statsWrite(builder.build());
+    }
+}
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 25aa693..745c5bc 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -26,6 +26,7 @@
 
 import static com.android.modules.utils.build.SdkLevel.isAtLeastU;
 import static com.android.server.connectivity.mdns.MdnsRecord.MAX_LABEL_LENGTH;
+import static com.android.server.connectivity.mdns.util.MdnsUtils.Clock;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -68,6 +69,7 @@
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.internal.util.State;
 import com.android.internal.util.StateMachine;
+import com.android.metrics.NetworkNsdReportedMetrics;
 import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.DeviceConfigUtils;
 import com.android.net.module.util.InetAddressUtils;
@@ -160,6 +162,7 @@
     public static final boolean DBG = Log.isLoggable(TAG, Log.DEBUG);
     private static final long CLEANUP_DELAY_MS = 10000;
     private static final int IFACE_IDX_ANY = 0;
+    private static final int NO_TRANSACTION = -1;
     private static final SharedLog LOGGER = new SharedLog("serviceDiscovery");
 
     private final Context mContext;
@@ -176,6 +179,8 @@
     private final MdnsSocketProvider mMdnsSocketProvider;
     @NonNull
     private final MdnsAdvertiser mAdvertiser;
+    @NonNull
+    private final Clock mClock;
     private final SharedLog mServiceLogs = LOGGER.forSubComponent(TAG);
     // WARNING : Accessing these values in any thread is not safe, it must only be changed in the
     // state machine thread. If change this outside state machine, it will need to introduce
@@ -188,8 +193,8 @@
      */
     private final HashMap<NsdServiceConnector, ClientInfo> mClients = new HashMap<>();
 
-    /* A map from unique id to client info */
-    private final SparseArray<ClientInfo> mIdToClientInfoMap= new SparseArray<>();
+    /* A map from transaction(unique) id to client info */
+    private final SparseArray<ClientInfo> mTransactionIdToClientInfoMap = new SparseArray<>();
 
     // Note this is not final to avoid depending on the Wi-Fi service starting before NsdService
     @Nullable
@@ -211,16 +216,16 @@
     private int mClientNumberId = 1;
 
     private static class MdnsListener implements MdnsServiceBrowserListener {
-        protected final int mClientId;
+        protected final int mClientRequestId;
         protected final int mTransactionId;
         @NonNull
         protected final NsdServiceInfo mReqServiceInfo;
         @NonNull
         protected final String mListenedServiceType;
 
-        MdnsListener(int clientId, int transactionId, @NonNull NsdServiceInfo reqServiceInfo,
+        MdnsListener(int clientRequestId, int transactionId, @NonNull NsdServiceInfo reqServiceInfo,
                 @NonNull String listenedServiceType) {
-            mClientId = clientId;
+            mClientRequestId = clientRequestId;
             mTransactionId = transactionId;
             mReqServiceInfo = reqServiceInfo;
             mListenedServiceType = listenedServiceType;
@@ -261,67 +266,67 @@
 
     private class DiscoveryListener extends MdnsListener {
 
-        DiscoveryListener(int clientId, int transactionId, @NonNull NsdServiceInfo reqServiceInfo,
-                @NonNull String listenServiceType) {
-            super(clientId, transactionId, reqServiceInfo, listenServiceType);
+        DiscoveryListener(int clientRequestId, int transactionId,
+                @NonNull NsdServiceInfo reqServiceInfo, @NonNull String listenServiceType) {
+            super(clientRequestId, transactionId, reqServiceInfo, listenServiceType);
         }
 
         @Override
         public void onServiceNameDiscovered(@NonNull MdnsServiceInfo serviceInfo) {
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     NsdManager.SERVICE_FOUND,
-                    new MdnsEvent(mClientId, serviceInfo));
+                    new MdnsEvent(mClientRequestId, serviceInfo));
         }
 
         @Override
         public void onServiceNameRemoved(@NonNull MdnsServiceInfo serviceInfo) {
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     NsdManager.SERVICE_LOST,
-                    new MdnsEvent(mClientId, serviceInfo));
+                    new MdnsEvent(mClientRequestId, serviceInfo));
         }
     }
 
     private class ResolutionListener extends MdnsListener {
 
-        ResolutionListener(int clientId, int transactionId, @NonNull NsdServiceInfo reqServiceInfo,
-                @NonNull String listenServiceType) {
-            super(clientId, transactionId, reqServiceInfo, listenServiceType);
+        ResolutionListener(int clientRequestId, int transactionId,
+                @NonNull NsdServiceInfo reqServiceInfo, @NonNull String listenServiceType) {
+            super(clientRequestId, transactionId, reqServiceInfo, listenServiceType);
         }
 
         @Override
         public void onServiceFound(MdnsServiceInfo serviceInfo) {
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     NsdManager.RESOLVE_SERVICE_SUCCEEDED,
-                    new MdnsEvent(mClientId, serviceInfo));
+                    new MdnsEvent(mClientRequestId, serviceInfo));
         }
     }
 
     private class ServiceInfoListener extends MdnsListener {
 
-        ServiceInfoListener(int clientId, int transactionId, @NonNull NsdServiceInfo reqServiceInfo,
-                @NonNull String listenServiceType) {
-            super(clientId, transactionId, reqServiceInfo, listenServiceType);
+        ServiceInfoListener(int clientRequestId, int transactionId,
+                @NonNull NsdServiceInfo reqServiceInfo, @NonNull String listenServiceType) {
+            super(clientRequestId, transactionId, reqServiceInfo, listenServiceType);
         }
 
         @Override
         public void onServiceFound(@NonNull MdnsServiceInfo serviceInfo) {
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     NsdManager.SERVICE_UPDATED,
-                    new MdnsEvent(mClientId, serviceInfo));
+                    new MdnsEvent(mClientRequestId, serviceInfo));
         }
 
         @Override
         public void onServiceUpdated(@NonNull MdnsServiceInfo serviceInfo) {
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     NsdManager.SERVICE_UPDATED,
-                    new MdnsEvent(mClientId, serviceInfo));
+                    new MdnsEvent(mClientRequestId, serviceInfo));
         }
 
         @Override
         public void onServiceRemoved(@NonNull MdnsServiceInfo serviceInfo) {
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     NsdManager.SERVICE_UPDATED_LOST,
-                    new MdnsEvent(mClientId, serviceInfo));
+                    new MdnsEvent(mClientRequestId, serviceInfo));
         }
     }
 
@@ -409,8 +414,8 @@
             // Return early if NSD is not active, or not on any relevant network
             return -1;
         }
-        for (int i = 0; i < mIdToClientInfoMap.size(); i++) {
-            final ClientInfo clientInfo = mIdToClientInfoMap.valueAt(i);
+        for (int i = 0; i < mTransactionIdToClientInfoMap.size(); i++) {
+            final ClientInfo clientInfo = mTransactionIdToClientInfoMap.valueAt(i);
             if (!mRunningAppActiveUids.contains(clientInfo.mUid)) {
                 // Ignore non-active UIDs
                 continue;
@@ -427,12 +432,12 @@
      * Data class of mdns service callback information.
      */
     private static class MdnsEvent {
-        final int mClientId;
+        final int mClientRequestId;
         @NonNull
         final MdnsServiceInfo mMdnsServiceInfo;
 
-        MdnsEvent(int clientId, @NonNull MdnsServiceInfo mdnsServiceInfo) {
-            mClientId = clientId;
+        MdnsEvent(int clientRequestId, @NonNull MdnsServiceInfo mdnsServiceInfo) {
+            mClientRequestId = clientRequestId;
             mMdnsServiceInfo = mdnsServiceInfo;
         }
     }
@@ -471,7 +476,7 @@
         }
 
         private boolean isAnyRequestActive() {
-            return mIdToClientInfoMap.size() != 0;
+            return mTransactionIdToClientInfoMap.size() != 0;
         }
 
         private void scheduleStop() {
@@ -520,7 +525,7 @@
             @Override
             public boolean processMessage(Message msg) {
                 final ClientInfo cInfo;
-                final int clientId = msg.arg2;
+                final int clientRequestId = msg.arg2;
                 switch (msg.what) {
                     case NsdManager.REGISTER_CLIENT:
                         final ConnectorArgs arg = (ConnectorArgs) msg.obj;
@@ -528,11 +533,15 @@
                         try {
                             cb.asBinder().linkToDeath(arg.connector, 0);
                             final String tag = "Client" + arg.uid + "-" + mClientNumberId++;
+                            final NetworkNsdReportedMetrics metrics =
+                                    mDeps.makeNetworkNsdReportedMetrics(
+                                            !arg.useJavaBackend, (int) mClock.elapsedRealtime());
                             cInfo = new ClientInfo(cb, arg.uid, arg.useJavaBackend,
-                                    mServiceLogs.forSubComponent(tag));
+                                    mServiceLogs.forSubComponent(tag), metrics);
                             mClients.put(arg.connector, cInfo);
                         } catch (RemoteException e) {
-                            Log.w(TAG, "Client " + clientId + " has already died");
+                            Log.w(TAG, "Client request id " + clientRequestId
+                                    + " has already died");
                         }
                         break;
                     case NsdManager.UNREGISTER_CLIENT:
@@ -551,49 +560,49 @@
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
                             cInfo.onDiscoverServicesFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
                        break;
                     case NsdManager.STOP_DISCOVERY:
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
                             cInfo.onStopDiscoveryFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
                         break;
                     case NsdManager.REGISTER_SERVICE:
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
-                            cInfo.onRegisterServiceFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                            cInfo.onRegisterServiceFailedImmediately(
+                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
                         break;
                     case NsdManager.UNREGISTER_SERVICE:
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
                             cInfo.onUnregisterServiceFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
                         break;
                     case NsdManager.RESOLVE_SERVICE:
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
                             cInfo.onResolveServiceFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
                         break;
                     case NsdManager.STOP_RESOLUTION:
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
                             cInfo.onStopResolutionFailed(
-                                    clientId, NsdManager.FAILURE_OPERATION_NOT_RUNNING);
+                                    clientRequestId, NsdManager.FAILURE_OPERATION_NOT_RUNNING);
                         }
                         break;
                     case NsdManager.REGISTER_SERVICE_CALLBACK:
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
                             cInfo.onServiceInfoCallbackRegistrationFailed(
-                                    clientId, NsdManager.FAILURE_BAD_PARAMETERS);
+                                    clientRequestId, NsdManager.FAILURE_BAD_PARAMETERS);
                         }
                         break;
                     case NsdManager.DAEMON_CLEANUP:
@@ -644,27 +653,29 @@
                 return false;
             }
 
-            private void storeLegacyRequestMap(int clientId, int globalId, ClientInfo clientInfo,
-                    int what) {
-                clientInfo.mClientRequests.put(clientId, new LegacyClientRequest(globalId, what));
-                mIdToClientInfoMap.put(globalId, clientInfo);
+            private void storeLegacyRequestMap(int clientRequestId, int transactionId,
+                    ClientInfo clientInfo, int what) {
+                clientInfo.mClientRequests.put(clientRequestId, new LegacyClientRequest(
+                        transactionId, what, mClock, mClock.elapsedRealtime()));
+                mTransactionIdToClientInfoMap.put(transactionId, clientInfo);
                 // Remove the cleanup event because here comes a new request.
                 cancelStop();
             }
 
-            private void storeAdvertiserRequestMap(int clientId, int globalId,
+            private void storeAdvertiserRequestMap(int clientRequestId, int transactionId,
                     ClientInfo clientInfo, @Nullable Network requestedNetwork) {
-                clientInfo.mClientRequests.put(clientId,
-                        new AdvertiserClientRequest(globalId, requestedNetwork));
-                mIdToClientInfoMap.put(globalId, clientInfo);
+                clientInfo.mClientRequests.put(clientRequestId, new AdvertiserClientRequest(
+                        transactionId, requestedNetwork, mClock, mClock.elapsedRealtime()));
+                mTransactionIdToClientInfoMap.put(transactionId, clientInfo);
                 updateMulticastLock();
             }
 
-            private void removeRequestMap(int clientId, int globalId, ClientInfo clientInfo) {
-                final ClientRequest existing = clientInfo.mClientRequests.get(clientId);
+            private void removeRequestMap(
+                    int clientRequestId, int transactionId, ClientInfo clientInfo) {
+                final ClientRequest existing = clientInfo.mClientRequests.get(clientRequestId);
                 if (existing == null) return;
-                clientInfo.mClientRequests.remove(clientId);
-                mIdToClientInfoMap.remove(globalId);
+                clientInfo.mClientRequests.remove(clientRequestId);
+                mTransactionIdToClientInfoMap.remove(transactionId);
 
                 if (existing instanceof LegacyClientRequest) {
                     maybeScheduleStop();
@@ -674,12 +685,13 @@
                 }
             }
 
-            private void storeDiscoveryManagerRequestMap(int clientId, int globalId,
+            private void storeDiscoveryManagerRequestMap(int clientRequestId, int transactionId,
                     MdnsListener listener, ClientInfo clientInfo,
                     @Nullable Network requestedNetwork) {
-                clientInfo.mClientRequests.put(clientId,
-                        new DiscoveryManagerRequest(globalId, listener, requestedNetwork));
-                mIdToClientInfoMap.put(globalId, clientInfo);
+                clientInfo.mClientRequests.put(clientRequestId, new DiscoveryManagerRequest(
+                        transactionId, listener, requestedNetwork, mClock,
+                        mClock.elapsedRealtime()));
+                mTransactionIdToClientInfoMap.put(transactionId, clientInfo);
                 updateMulticastLock();
             }
 
@@ -695,17 +707,17 @@
                 return MdnsUtils.truncateServiceName(originalName, MAX_LABEL_LENGTH);
             }
 
-            private void stopDiscoveryManagerRequest(ClientRequest request, int clientId, int id,
-                    ClientInfo clientInfo) {
+            private void stopDiscoveryManagerRequest(ClientRequest request, int clientRequestId,
+                    int transactionId, ClientInfo clientInfo) {
                 clientInfo.unregisterMdnsListenerFromRequest(request);
-                removeRequestMap(clientId, id, clientInfo);
+                removeRequestMap(clientRequestId, transactionId, clientInfo);
             }
 
             @Override
             public boolean processMessage(Message msg) {
                 final ClientInfo clientInfo;
-                final int id;
-                final int clientId = msg.arg2;
+                final int transactionId;
+                final int clientRequestId = msg.arg2;
                 final ListenerArgs args;
                 switch (msg.what) {
                     case NsdManager.DISCOVER_SERVICES: {
@@ -722,12 +734,12 @@
 
                         if (requestLimitReached(clientInfo)) {
                             clientInfo.onDiscoverServicesFailed(
-                                    clientId, NsdManager.FAILURE_MAX_LIMIT);
+                                    clientRequestId, NsdManager.FAILURE_MAX_LIMIT);
                             break;
                         }
 
                         final NsdServiceInfo info = args.serviceInfo;
-                        id = getUniqueId();
+                        transactionId = getUniqueId();
                         final Pair<String, String> typeAndSubtype =
                                 parseTypeAndSubtype(info.getServiceType());
                         final String serviceType = typeAndSubtype == null
@@ -736,15 +748,15 @@
                                 || mDeps.isMdnsDiscoveryManagerEnabled(mContext)
                                 || useDiscoveryManagerForType(serviceType)) {
                             if (serviceType == null) {
-                                clientInfo.onDiscoverServicesFailed(clientId,
+                                clientInfo.onDiscoverServicesFailed(clientRequestId,
                                         NsdManager.FAILURE_INTERNAL_ERROR);
                                 break;
                             }
 
                             final String listenServiceType = serviceType + ".local";
                             maybeStartMonitoringSockets();
-                            final MdnsListener listener =
-                                    new DiscoveryListener(clientId, id, info, listenServiceType);
+                            final MdnsListener listener = new DiscoveryListener(clientRequestId,
+                                    transactionId, info, listenServiceType);
                             final MdnsSearchOptions.Builder optionsBuilder =
                                     MdnsSearchOptions.newBuilder()
                                             .setNetwork(info.getNetwork())
@@ -757,23 +769,24 @@
                             }
                             mMdnsDiscoveryManager.registerListener(
                                     listenServiceType, listener, optionsBuilder.build());
-                            storeDiscoveryManagerRequestMap(clientId, id, listener, clientInfo,
-                                    info.getNetwork());
-                            clientInfo.onDiscoverServicesStarted(clientId, info);
-                            clientInfo.log("Register a DiscoveryListener " + id
+                            storeDiscoveryManagerRequestMap(clientRequestId, transactionId,
+                                    listener, clientInfo, info.getNetwork());
+                            clientInfo.onDiscoverServicesStarted(clientRequestId, info);
+                            clientInfo.log("Register a DiscoveryListener " + transactionId
                                     + " for service type:" + listenServiceType);
                         } else {
                             maybeStartDaemon();
-                            if (discoverServices(id, info)) {
+                            if (discoverServices(transactionId, info)) {
                                 if (DBG) {
-                                    Log.d(TAG, "Discover " + msg.arg2 + " " + id
+                                    Log.d(TAG, "Discover " + msg.arg2 + " " + transactionId
                                             + info.getServiceType());
                                 }
-                                storeLegacyRequestMap(clientId, id, clientInfo, msg.what);
-                                clientInfo.onDiscoverServicesStarted(clientId, info);
+                                storeLegacyRequestMap(
+                                        clientRequestId, transactionId, clientInfo, msg.what);
+                                clientInfo.onDiscoverServicesStarted(clientRequestId, info);
                             } else {
-                                stopServiceDiscovery(id);
-                                clientInfo.onDiscoverServicesFailed(clientId,
+                                stopServiceDiscovery(transactionId);
+                                clientInfo.onDiscoverServicesFailed(clientRequestId,
                                         NsdManager.FAILURE_INTERNAL_ERROR);
                             }
                         }
@@ -791,26 +804,28 @@
                             break;
                         }
 
-                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
                         if (request == null) {
                             Log.e(TAG, "Unknown client request in STOP_DISCOVERY");
                             break;
                         }
-                        id = request.mGlobalId;
+                        transactionId = request.mTransactionId;
                         // Note isMdnsDiscoveryManagerEnabled may have changed to false at this
                         // point, so this needs to check the type of the original request to
                         // unregister instead of looking at the flag value.
                         if (request instanceof DiscoveryManagerRequest) {
-                            stopDiscoveryManagerRequest(request, clientId, id, clientInfo);
-                            clientInfo.onStopDiscoverySucceeded(clientId);
-                            clientInfo.log("Unregister the DiscoveryListener " + id);
+                            stopDiscoveryManagerRequest(
+                                    request, clientRequestId, transactionId, clientInfo);
+                            clientInfo.onStopDiscoverySucceeded(clientRequestId);
+                            clientInfo.log("Unregister the DiscoveryListener " + transactionId);
                         } else {
-                            removeRequestMap(clientId, id, clientInfo);
-                            if (stopServiceDiscovery(id)) {
-                                clientInfo.onStopDiscoverySucceeded(clientId);
+                            removeRequestMap(clientRequestId, transactionId, clientInfo);
+                            if (stopServiceDiscovery(transactionId)) {
+                                clientInfo.onStopDiscoverySucceeded(clientRequestId);
                             } else {
                                 clientInfo.onStopDiscoveryFailed(
-                                        clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                        clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                             }
                         }
                         break;
@@ -828,12 +843,12 @@
                         }
 
                         if (requestLimitReached(clientInfo)) {
-                            clientInfo.onRegisterServiceFailed(
-                                    clientId, NsdManager.FAILURE_MAX_LIMIT);
+                            clientInfo.onRegisterServiceFailedImmediately(
+                                    clientRequestId, NsdManager.FAILURE_MAX_LIMIT);
                             break;
                         }
 
-                        id = getUniqueId();
+                        transactionId = getUniqueId();
                         final NsdServiceInfo serviceInfo = args.serviceInfo;
                         final String serviceType = serviceInfo.getServiceType();
                         final Pair<String, String> typeSubtype = parseTypeAndSubtype(serviceType);
@@ -844,8 +859,8 @@
                                 || useAdvertiserForType(registerServiceType)) {
                             if (registerServiceType == null) {
                                 Log.e(TAG, "Invalid service type: " + serviceType);
-                                clientInfo.onRegisterServiceFailed(clientId,
-                                        NsdManager.FAILURE_INTERNAL_ERROR);
+                                clientInfo.onRegisterServiceFailedImmediately(
+                                        clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                                 break;
                             }
                             serviceInfo.setServiceType(registerServiceType);
@@ -857,19 +872,23 @@
                             // service type would generate service instance names like
                             // Name._subtype._sub._type._tcp, which is incorrect
                             // (it should be Name._type._tcp).
-                            mAdvertiser.addService(id, serviceInfo, typeSubtype.second);
-                            storeAdvertiserRequestMap(clientId, id, clientInfo,
+                            mAdvertiser.addService(transactionId, serviceInfo, typeSubtype.second);
+                            storeAdvertiserRequestMap(clientRequestId, transactionId, clientInfo,
                                     serviceInfo.getNetwork());
                         } else {
                             maybeStartDaemon();
-                            if (registerService(id, serviceInfo)) {
-                                if (DBG) Log.d(TAG, "Register " + clientId + " " + id);
-                                storeLegacyRequestMap(clientId, id, clientInfo, msg.what);
+                            if (registerService(transactionId, serviceInfo)) {
+                                if (DBG) {
+                                    Log.d(TAG, "Register " + clientRequestId
+                                            + " " + transactionId);
+                                }
+                                storeLegacyRequestMap(
+                                        clientRequestId, transactionId, clientInfo, msg.what);
                                 // Return success after mDns reports success
                             } else {
-                                unregisterService(id);
-                                clientInfo.onRegisterServiceFailed(
-                                        clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                unregisterService(transactionId);
+                                clientInfo.onRegisterServiceFailedImmediately(
+                                        clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                             }
 
                         }
@@ -886,26 +905,29 @@
                             Log.e(TAG, "Unknown connector in unregistration");
                             break;
                         }
-                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
                         if (request == null) {
                             Log.e(TAG, "Unknown client request in UNREGISTER_SERVICE");
                             break;
                         }
-                        id = request.mGlobalId;
-                        removeRequestMap(clientId, id, clientInfo);
+                        transactionId = request.mTransactionId;
+                        removeRequestMap(clientRequestId, transactionId, clientInfo);
 
                         // Note isMdnsAdvertiserEnabled may have changed to false at this point,
                         // so this needs to check the type of the original request to unregister
                         // instead of looking at the flag value.
                         if (request instanceof AdvertiserClientRequest) {
-                            mAdvertiser.removeService(id);
-                            clientInfo.onUnregisterServiceSucceeded(clientId);
+                            mAdvertiser.removeService(transactionId);
+                            clientInfo.onUnregisterServiceSucceeded(clientRequestId, transactionId,
+                                    request.calculateRequestDurationMs());
                         } else {
-                            if (unregisterService(id)) {
-                                clientInfo.onUnregisterServiceSucceeded(clientId);
+                            if (unregisterService(transactionId)) {
+                                clientInfo.onUnregisterServiceSucceeded(clientRequestId,
+                                        transactionId, request.calculateRequestDurationMs());
                             } else {
                                 clientInfo.onUnregisterServiceFailed(
-                                        clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                        clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                             }
                         }
                         break;
@@ -923,7 +945,7 @@
                         }
 
                         final NsdServiceInfo info = args.serviceInfo;
-                        id = getUniqueId();
+                        transactionId = getUniqueId();
                         final Pair<String, String> typeSubtype =
                                 parseTypeAndSubtype(info.getServiceType());
                         final String serviceType = typeSubtype == null
@@ -932,15 +954,15 @@
                                 ||  mDeps.isMdnsDiscoveryManagerEnabled(mContext)
                                 || useDiscoveryManagerForType(serviceType)) {
                             if (serviceType == null) {
-                                clientInfo.onResolveServiceFailed(clientId,
+                                clientInfo.onResolveServiceFailed(clientRequestId,
                                         NsdManager.FAILURE_INTERNAL_ERROR);
                                 break;
                             }
                             final String resolveServiceType = serviceType + ".local";
 
                             maybeStartMonitoringSockets();
-                            final MdnsListener listener =
-                                    new ResolutionListener(clientId, id, info, resolveServiceType);
+                            final MdnsListener listener = new ResolutionListener(clientRequestId,
+                                    transactionId, info, resolveServiceType);
                             final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
                                     .setNetwork(info.getNetwork())
                                     .setIsPassiveMode(true)
@@ -949,24 +971,25 @@
                                     .build();
                             mMdnsDiscoveryManager.registerListener(
                                     resolveServiceType, listener, options);
-                            storeDiscoveryManagerRequestMap(clientId, id, listener, clientInfo,
-                                    info.getNetwork());
-                            clientInfo.log("Register a ResolutionListener " + id
+                            storeDiscoveryManagerRequestMap(clientRequestId, transactionId,
+                                    listener, clientInfo, info.getNetwork());
+                            clientInfo.log("Register a ResolutionListener " + transactionId
                                     + " for service type:" + resolveServiceType);
                         } else {
                             if (clientInfo.mResolvedService != null) {
                                 clientInfo.onResolveServiceFailed(
-                                        clientId, NsdManager.FAILURE_ALREADY_ACTIVE);
+                                        clientRequestId, NsdManager.FAILURE_ALREADY_ACTIVE);
                                 break;
                             }
 
                             maybeStartDaemon();
-                            if (resolveService(id, info)) {
+                            if (resolveService(transactionId, info)) {
                                 clientInfo.mResolvedService = new NsdServiceInfo();
-                                storeLegacyRequestMap(clientId, id, clientInfo, msg.what);
+                                storeLegacyRequestMap(
+                                        clientRequestId, transactionId, clientInfo, msg.what);
                             } else {
                                 clientInfo.onResolveServiceFailed(
-                                        clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                        clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                             }
                         }
                         break;
@@ -983,26 +1006,28 @@
                             break;
                         }
 
-                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
                         if (request == null) {
                             Log.e(TAG, "Unknown client request in STOP_RESOLUTION");
                             break;
                         }
-                        id = request.mGlobalId;
+                        transactionId = request.mTransactionId;
                         // Note isMdnsDiscoveryManagerEnabled may have changed to false at this
                         // point, so this needs to check the type of the original request to
                         // unregister instead of looking at the flag value.
                         if (request instanceof DiscoveryManagerRequest) {
-                            stopDiscoveryManagerRequest(request, clientId, id, clientInfo);
-                            clientInfo.onStopResolutionSucceeded(clientId);
-                            clientInfo.log("Unregister the ResolutionListener " + id);
+                            stopDiscoveryManagerRequest(
+                                    request, clientRequestId, transactionId, clientInfo);
+                            clientInfo.onStopResolutionSucceeded(clientRequestId);
+                            clientInfo.log("Unregister the ResolutionListener " + transactionId);
                         } else {
-                            removeRequestMap(clientId, id, clientInfo);
-                            if (stopResolveService(id)) {
-                                clientInfo.onStopResolutionSucceeded(clientId);
+                            removeRequestMap(clientRequestId, transactionId, clientInfo);
+                            if (stopResolveService(transactionId)) {
+                                clientInfo.onStopResolutionSucceeded(clientRequestId);
                             } else {
                                 clientInfo.onStopResolutionFailed(
-                                        clientId, NsdManager.FAILURE_OPERATION_NOT_RUNNING);
+                                        clientRequestId, NsdManager.FAILURE_OPERATION_NOT_RUNNING);
                             }
                             clientInfo.mResolvedService = null;
                         }
@@ -1021,21 +1046,21 @@
                         }
 
                         final NsdServiceInfo info = args.serviceInfo;
-                        id = getUniqueId();
+                        transactionId = getUniqueId();
                         final Pair<String, String> typeAndSubtype =
                                 parseTypeAndSubtype(info.getServiceType());
                         final String serviceType = typeAndSubtype == null
                                 ? null : typeAndSubtype.first;
                         if (serviceType == null) {
-                            clientInfo.onServiceInfoCallbackRegistrationFailed(clientId,
+                            clientInfo.onServiceInfoCallbackRegistrationFailed(clientRequestId,
                                     NsdManager.FAILURE_BAD_PARAMETERS);
                             break;
                         }
                         final String resolveServiceType = serviceType + ".local";
 
                         maybeStartMonitoringSockets();
-                        final MdnsListener listener =
-                                new ServiceInfoListener(clientId, id, info, resolveServiceType);
+                        final MdnsListener listener = new ServiceInfoListener(clientRequestId,
+                                transactionId, info, resolveServiceType);
                         final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
                                 .setNetwork(info.getNetwork())
                                 .setIsPassiveMode(true)
@@ -1044,9 +1069,9 @@
                                 .build();
                         mMdnsDiscoveryManager.registerListener(
                                 resolveServiceType, listener, options);
-                        storeDiscoveryManagerRequestMap(clientId, id, listener, clientInfo,
-                                info.getNetwork());
-                        clientInfo.log("Register a ServiceInfoListener " + id
+                        storeDiscoveryManagerRequestMap(clientRequestId, transactionId, listener,
+                                clientInfo, info.getNetwork());
+                        clientInfo.log("Register a ServiceInfoListener " + transactionId
                                 + " for service type:" + resolveServiceType);
                         break;
                     }
@@ -1062,16 +1087,18 @@
                             break;
                         }
 
-                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
                         if (request == null) {
                             Log.e(TAG, "Unknown client request in UNREGISTER_SERVICE_CALLBACK");
                             break;
                         }
-                        id = request.mGlobalId;
+                        transactionId = request.mTransactionId;
                         if (request instanceof DiscoveryManagerRequest) {
-                            stopDiscoveryManagerRequest(request, clientId, id, clientInfo);
-                            clientInfo.onServiceInfoCallbackUnregistered(clientId);
-                            clientInfo.log("Unregister the ServiceInfoListener " + id);
+                            stopDiscoveryManagerRequest(
+                                    request, clientRequestId, transactionId, clientInfo);
+                            clientInfo.onServiceInfoCallbackUnregistered(clientRequestId);
+                            clientInfo.log("Unregister the ServiceInfoListener " + transactionId);
                         } else {
                             loge("Unregister failed with non-DiscoveryManagerRequest.");
                         }
@@ -1093,26 +1120,28 @@
                 return HANDLED;
             }
 
-            private boolean handleMDnsServiceEvent(int code, int id, Object obj) {
+            private boolean handleMDnsServiceEvent(int code, int transactionId, Object obj) {
                 NsdServiceInfo servInfo;
-                ClientInfo clientInfo = mIdToClientInfoMap.get(id);
+                ClientInfo clientInfo = mTransactionIdToClientInfoMap.get(transactionId);
                 if (clientInfo == null) {
-                    Log.e(TAG, String.format("id %d for %d has no client mapping", id, code));
+                    Log.e(TAG, String.format(
+                            "transactionId %d for %d has no client mapping", transactionId, code));
                     return false;
                 }
 
                 /* This goes in response as msg.arg2 */
-                int clientId = clientInfo.getClientId(id);
-                if (clientId < 0) {
+                int clientRequestId = clientInfo.getClientRequestId(transactionId);
+                if (clientRequestId < 0) {
                     // This can happen because of race conditions. For example,
                     // SERVICE_FOUND may race with STOP_SERVICE_DISCOVERY,
                     // and we may get in this situation.
-                    Log.d(TAG, String.format("%d for listener id %d that is no longer active",
-                            code, id));
+                    Log.d(TAG, String.format("%d for transactionId %d that is no longer active",
+                            code, transactionId));
                     return false;
                 }
                 if (DBG) {
-                    Log.d(TAG, String.format("MDns service event code:%d id=%d", code, id));
+                    Log.d(TAG, String.format(
+                            "MDns service event code:%d transactionId=%d", code, transactionId));
                 }
                 switch (code) {
                     case IMDnsEventListener.SERVICE_FOUND: {
@@ -1134,7 +1163,7 @@
                             break;
                         }
                         setServiceNetworkForCallback(servInfo, info.netId, info.interfaceIdx);
-                        clientInfo.onServiceFound(clientId, servInfo);
+                        clientInfo.onServiceFound(clientRequestId, servInfo);
                         break;
                     }
                     case IMDnsEventListener.SERVICE_LOST: {
@@ -1148,23 +1177,29 @@
                         // TODO: avoid returning null in that case, possibly by remembering
                         // found services on the same interface index and their network at the time
                         setServiceNetworkForCallback(servInfo, lostNetId, info.interfaceIdx);
-                        clientInfo.onServiceLost(clientId, servInfo);
+                        clientInfo.onServiceLost(clientRequestId, servInfo);
                         break;
                     }
                     case IMDnsEventListener.SERVICE_DISCOVERY_FAILED:
                         clientInfo.onDiscoverServicesFailed(
-                                clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         break;
                     case IMDnsEventListener.SERVICE_REGISTERED: {
                         final RegistrationInfo info = (RegistrationInfo) obj;
                         final String name = info.serviceName;
                         servInfo = new NsdServiceInfo(name, null /* serviceType */);
-                        clientInfo.onRegisterServiceSucceeded(clientId, servInfo);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
+                        clientInfo.onRegisterServiceSucceeded(clientRequestId, servInfo,
+                                transactionId, request.calculateRequestDurationMs());
                         break;
                     }
                     case IMDnsEventListener.SERVICE_REGISTRATION_FAILED:
-                        clientInfo.onRegisterServiceFailed(
-                                clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
+                        clientInfo.onRegisterServiceFailed(clientRequestId,
+                                NsdManager.FAILURE_INTERNAL_ERROR, transactionId,
+                                request.calculateRequestDurationMs());
                         break;
                     case IMDnsEventListener.SERVICE_RESOLVED: {
                         final ResolutionInfo info = (ResolutionInfo) obj;
@@ -1192,34 +1227,34 @@
                         serviceInfo.setTxtRecords(info.txtRecord);
                         // Network will be added after SERVICE_GET_ADDR_SUCCESS
 
-                        stopResolveService(id);
-                        removeRequestMap(clientId, id, clientInfo);
+                        stopResolveService(transactionId);
+                        removeRequestMap(clientRequestId, transactionId, clientInfo);
 
-                        final int id2 = getUniqueId();
-                        if (getAddrInfo(id2, info.hostname, info.interfaceIdx)) {
-                            storeLegacyRequestMap(clientId, id2, clientInfo,
+                        final int transactionId2 = getUniqueId();
+                        if (getAddrInfo(transactionId2, info.hostname, info.interfaceIdx)) {
+                            storeLegacyRequestMap(clientRequestId, transactionId2, clientInfo,
                                     NsdManager.RESOLVE_SERVICE);
                         } else {
                             clientInfo.onResolveServiceFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                             clientInfo.mResolvedService = null;
                         }
                         break;
                     }
                     case IMDnsEventListener.SERVICE_RESOLUTION_FAILED:
                         /* NNN resolveId errorCode */
-                        stopResolveService(id);
-                        removeRequestMap(clientId, id, clientInfo);
+                        stopResolveService(transactionId);
+                        removeRequestMap(clientRequestId, transactionId, clientInfo);
                         clientInfo.onResolveServiceFailed(
-                                clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         clientInfo.mResolvedService = null;
                         break;
                     case IMDnsEventListener.SERVICE_GET_ADDR_FAILED:
                         /* NNN resolveId errorCode */
-                        stopGetAddrInfo(id);
-                        removeRequestMap(clientId, id, clientInfo);
+                        stopGetAddrInfo(transactionId);
+                        removeRequestMap(clientRequestId, transactionId, clientInfo);
                         clientInfo.onResolveServiceFailed(
-                                clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         clientInfo.mResolvedService = null;
                         break;
                     case IMDnsEventListener.SERVICE_GET_ADDR_SUCCESS: {
@@ -1242,13 +1277,13 @@
                             setServiceNetworkForCallback(clientInfo.mResolvedService,
                                     netId, info.interfaceIdx);
                             clientInfo.onResolveServiceSucceeded(
-                                    clientId, clientInfo.mResolvedService);
+                                    clientRequestId, clientInfo.mResolvedService);
                         } else {
                             clientInfo.onResolveServiceFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
-                        stopGetAddrInfo(id);
-                        removeRequestMap(clientId, id, clientInfo);
+                        stopGetAddrInfo(transactionId);
+                        removeRequestMap(clientRequestId, transactionId, clientInfo);
                         clientInfo.mResolvedService = null;
                         break;
                     }
@@ -1305,7 +1340,7 @@
 
             private boolean handleMdnsDiscoveryManagerEvent(
                     int transactionId, int code, Object obj) {
-                final ClientInfo clientInfo = mIdToClientInfoMap.get(transactionId);
+                final ClientInfo clientInfo = mTransactionIdToClientInfoMap.get(transactionId);
                 if (clientInfo == null) {
                     Log.e(TAG, String.format(
                             "id %d for %d has no client mapping", transactionId, code));
@@ -1313,23 +1348,23 @@
                 }
 
                 final MdnsEvent event = (MdnsEvent) obj;
-                final int clientId = event.mClientId;
+                final int clientRequestId = event.mClientRequestId;
                 final NsdServiceInfo info = buildNsdServiceInfoFromMdnsEvent(event, code);
                 // Errors are already logged if null
                 if (info == null) return false;
-                if (DBG) {
-                    Log.d(TAG, String.format("MdnsDiscoveryManager event code=%s transactionId=%d",
-                            NsdManager.nameOf(code), transactionId));
-                }
+                mServiceLogs.log(String.format(
+                        "MdnsDiscoveryManager event code=%s transactionId=%d",
+                        NsdManager.nameOf(code), transactionId));
                 switch (code) {
                     case NsdManager.SERVICE_FOUND:
-                        clientInfo.onServiceFound(clientId, info);
+                        clientInfo.onServiceFound(clientRequestId, info);
                         break;
                     case NsdManager.SERVICE_LOST:
-                        clientInfo.onServiceLost(clientId, info);
+                        clientInfo.onServiceLost(clientRequestId, info);
                         break;
                     case NsdManager.RESOLVE_SERVICE_SUCCEEDED: {
-                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        final ClientRequest request =
+                                clientInfo.mClientRequests.get(clientRequestId);
                         if (request == null) {
                             Log.e(TAG, "Unknown client request in RESOLVE_SERVICE_SUCCEEDED");
                             break;
@@ -1349,11 +1384,11 @@
                         final List<InetAddress> addresses = getInetAddresses(serviceInfo);
                         if (addresses.size() != 0) {
                             info.setHostAddresses(addresses);
-                            clientInfo.onResolveServiceSucceeded(clientId, info);
+                            clientInfo.onResolveServiceSucceeded(clientRequestId, info);
                         } else {
                             // No address. Notify resolution failure.
                             clientInfo.onResolveServiceFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
 
                         // Unregister the listener immediately like IMDnsEventListener design
@@ -1361,7 +1396,8 @@
                             Log.wtf(TAG, "non-DiscoveryManager request in DiscoveryManager event");
                             break;
                         }
-                        stopDiscoveryManagerRequest(request, clientId, transactionId, clientInfo);
+                        stopDiscoveryManagerRequest(
+                                request, clientRequestId, transactionId, clientInfo);
                         break;
                     }
                     case NsdManager.SERVICE_UPDATED: {
@@ -1380,11 +1416,11 @@
 
                         final List<InetAddress> addresses = getInetAddresses(serviceInfo);
                         info.setHostAddresses(addresses);
-                        clientInfo.onServiceUpdated(clientId, info);
+                        clientInfo.onServiceUpdated(clientRequestId, info);
                         break;
                     }
                     case NsdManager.SERVICE_UPDATED_LOST:
-                        clientInfo.onServiceUpdatedLost(clientId);
+                        clientInfo.onServiceUpdatedLost(clientRequestId);
                         break;
                     default:
                         return false;
@@ -1543,6 +1579,7 @@
         handler.post(() -> mMdnsSocketClient.setCallback(mMdnsDiscoveryManager));
         mAdvertiser = deps.makeMdnsAdvertiser(handler.getLooper(), mMdnsSocketProvider,
                 new AdvertiserCallback(), LOGGER.forSubComponent("MdnsAdvertiser"));
+        mClock = deps.makeClock();
     }
 
     /**
@@ -1632,6 +1669,21 @@
         public int getCallingUid() {
             return Binder.getCallingUid();
         }
+
+        /**
+         * @see NetworkNsdReportedMetrics
+         */
+        public NetworkNsdReportedMetrics makeNetworkNsdReportedMetrics(
+                boolean isLegacy, int clientId) {
+            return new NetworkNsdReportedMetrics(isLegacy, clientId);
+        }
+
+        /**
+         * @see MdnsUtils.Clock
+         */
+        public Clock makeClock() {
+            return new Clock();
+        }
     }
 
     /**
@@ -1721,44 +1773,49 @@
 
     private class AdvertiserCallback implements MdnsAdvertiser.AdvertiserCallback {
         @Override
-        public void onRegisterServiceSucceeded(int serviceId, NsdServiceInfo registeredInfo) {
-            final ClientInfo clientInfo = getClientInfoOrLog(serviceId);
+        public void onRegisterServiceSucceeded(int transactionId, NsdServiceInfo registeredInfo) {
+            mServiceLogs.log("onRegisterServiceSucceeded: transactionId " + transactionId);
+            final ClientInfo clientInfo = getClientInfoOrLog(transactionId);
             if (clientInfo == null) return;
 
-            final int clientId = getClientIdOrLog(clientInfo, serviceId);
-            if (clientId < 0) return;
+            final int clientRequestId = getClientRequestIdOrLog(clientInfo, transactionId);
+            if (clientRequestId < 0) return;
 
             // onRegisterServiceSucceeded only has the service name in its info. This aligns with
             // historical behavior.
             final NsdServiceInfo cbInfo = new NsdServiceInfo(registeredInfo.getServiceName(), null);
-            clientInfo.onRegisterServiceSucceeded(clientId, cbInfo);
+            final ClientRequest request = clientInfo.mClientRequests.get(clientRequestId);
+            clientInfo.onRegisterServiceSucceeded(
+                    clientRequestId, cbInfo, transactionId, request.calculateRequestDurationMs());
         }
 
         @Override
-        public void onRegisterServiceFailed(int serviceId, int errorCode) {
-            final ClientInfo clientInfo = getClientInfoOrLog(serviceId);
+        public void onRegisterServiceFailed(int transactionId, int errorCode) {
+            final ClientInfo clientInfo = getClientInfoOrLog(transactionId);
             if (clientInfo == null) return;
 
-            final int clientId = getClientIdOrLog(clientInfo, serviceId);
-            if (clientId < 0) return;
-
-            clientInfo.onRegisterServiceFailed(clientId, errorCode);
+            final int clientRequestId = getClientRequestIdOrLog(clientInfo, transactionId);
+            if (clientRequestId < 0) return;
+            final ClientRequest request = clientInfo.mClientRequests.get(clientRequestId);
+            clientInfo.onRegisterServiceFailed(clientRequestId, errorCode, transactionId,
+                    request.calculateRequestDurationMs());
         }
 
-        private ClientInfo getClientInfoOrLog(int serviceId) {
-            final ClientInfo clientInfo = mIdToClientInfoMap.get(serviceId);
+        private ClientInfo getClientInfoOrLog(int transactionId) {
+            final ClientInfo clientInfo = mTransactionIdToClientInfoMap.get(transactionId);
             if (clientInfo == null) {
-                Log.e(TAG, String.format("Callback for service %d has no client", serviceId));
+                Log.e(TAG, String.format("Callback for service %d has no client", transactionId));
             }
             return clientInfo;
         }
 
-        private int getClientIdOrLog(@NonNull ClientInfo info, int serviceId) {
-            final int clientId = info.getClientId(serviceId);
-            if (clientId < 0) {
-                Log.e(TAG, String.format("Client ID not found for service %d", serviceId));
+        private int getClientRequestIdOrLog(@NonNull ClientInfo info, int transactionId) {
+            final int clientRequestId = info.getClientRequestId(transactionId);
+            if (clientRequestId < 0) {
+                Log.e(TAG, String.format(
+                        "Client request ID not found for service %d", transactionId));
             }
-            return clientId;
+            return clientRequestId;
         }
     }
 
@@ -1879,9 +1936,9 @@
         return mUniqueId;
     }
 
-    private boolean registerService(int regId, NsdServiceInfo service) {
+    private boolean registerService(int transactionId, NsdServiceInfo service) {
         if (DBG) {
-            Log.d(TAG, "registerService: " + regId + " " + service);
+            Log.d(TAG, "registerService: " + transactionId + " " + service);
         }
         String name = service.getServiceName();
         String type = service.getServiceType();
@@ -1892,28 +1949,29 @@
             Log.e(TAG, "Interface to register service on not found");
             return false;
         }
-        return mMDnsManager.registerService(regId, name, type, port, textRecord, registerInterface);
+        return mMDnsManager.registerService(
+                transactionId, name, type, port, textRecord, registerInterface);
     }
 
-    private boolean unregisterService(int regId) {
-        return mMDnsManager.stopOperation(regId);
+    private boolean unregisterService(int transactionId) {
+        return mMDnsManager.stopOperation(transactionId);
     }
 
-    private boolean discoverServices(int discoveryId, NsdServiceInfo serviceInfo) {
+    private boolean discoverServices(int transactionId, NsdServiceInfo serviceInfo) {
         final String type = serviceInfo.getServiceType();
         final int discoverInterface = getNetworkInterfaceIndex(serviceInfo);
         if (serviceInfo.getNetwork() != null && discoverInterface == IFACE_IDX_ANY) {
             Log.e(TAG, "Interface to discover service on not found");
             return false;
         }
-        return mMDnsManager.discover(discoveryId, type, discoverInterface);
+        return mMDnsManager.discover(transactionId, type, discoverInterface);
     }
 
-    private boolean stopServiceDiscovery(int discoveryId) {
-        return mMDnsManager.stopOperation(discoveryId);
+    private boolean stopServiceDiscovery(int transactionId) {
+        return mMDnsManager.stopOperation(transactionId);
     }
 
-    private boolean resolveService(int resolveId, NsdServiceInfo service) {
+    private boolean resolveService(int transactionId, NsdServiceInfo service) {
         final String name = service.getServiceName();
         final String type = service.getServiceType();
         final int resolveInterface = getNetworkInterfaceIndex(service);
@@ -1921,7 +1979,7 @@
             Log.e(TAG, "Interface to resolve service on not found");
             return false;
         }
-        return mMDnsManager.resolve(resolveId, name, type, "local.", resolveInterface);
+        return mMDnsManager.resolve(transactionId, name, type, "local.", resolveInterface);
     }
 
     /**
@@ -1970,16 +2028,16 @@
         return iface.getIndex();
     }
 
-    private boolean stopResolveService(int resolveId) {
-        return mMDnsManager.stopOperation(resolveId);
+    private boolean stopResolveService(int transactionId) {
+        return mMDnsManager.stopOperation(transactionId);
     }
 
-    private boolean getAddrInfo(int resolveId, String hostname, int interfaceIdx) {
-        return mMDnsManager.getServiceAddress(resolveId, hostname, interfaceIdx);
+    private boolean getAddrInfo(int transactionId, String hostname, int interfaceIdx) {
+        return mMDnsManager.getServiceAddress(transactionId, hostname, interfaceIdx);
     }
 
-    private boolean stopGetAddrInfo(int resolveId) {
-        return mMDnsManager.stopOperation(resolveId);
+    private boolean stopGetAddrInfo(int transactionId) {
+        return mMDnsManager.stopOperation(transactionId);
     }
 
     @Override
@@ -1999,18 +2057,28 @@
     }
 
     private abstract static class ClientRequest {
-        private final int mGlobalId;
+        private final int mTransactionId;
+        private final Clock mClock;
+        private final long mStartTimeMs;
 
-        private ClientRequest(int globalId) {
-            mGlobalId = globalId;
+        private ClientRequest(int transactionId, @NonNull Clock clock, long startTimeMs) {
+            mTransactionId = transactionId;
+            mClock = clock;
+            mStartTimeMs = startTimeMs;
+        }
+
+        public long calculateRequestDurationMs() {
+            final long stopTimeMs = mClock.elapsedRealtime();
+            return stopTimeMs - mStartTimeMs;
         }
     }
 
     private static class LegacyClientRequest extends ClientRequest {
         private final int mRequestCode;
 
-        private LegacyClientRequest(int globalId, int requestCode) {
-            super(globalId);
+        private LegacyClientRequest(int transactionId, int requestCode, @NonNull Clock clock,
+                long startTimeMs) {
+            super(transactionId, clock, startTimeMs);
             mRequestCode = requestCode;
         }
     }
@@ -2019,8 +2087,9 @@
         @Nullable
         private final Network mRequestedNetwork;
 
-        private JavaBackendClientRequest(int globalId, @Nullable Network requestedNetwork) {
-            super(globalId);
+        private JavaBackendClientRequest(int transactionId, @Nullable Network requestedNetwork,
+                @NonNull Clock clock, long startTimeMs) {
+            super(transactionId, clock, startTimeMs);
             mRequestedNetwork = requestedNetwork;
         }
 
@@ -2031,8 +2100,9 @@
     }
 
     private static class AdvertiserClientRequest extends JavaBackendClientRequest {
-        private AdvertiserClientRequest(int globalId, @Nullable Network requestedNetwork) {
-            super(globalId, requestedNetwork);
+        private AdvertiserClientRequest(int transactionId, @Nullable Network requestedNetwork,
+                @NonNull Clock clock, long startTimeMs) {
+            super(transactionId, requestedNetwork, clock, startTimeMs);
         }
     }
 
@@ -2040,9 +2110,9 @@
         @NonNull
         private final MdnsListener mListener;
 
-        private DiscoveryManagerRequest(int globalId, @NonNull MdnsListener listener,
-                @Nullable Network requestedNetwork) {
-            super(globalId, requestedNetwork);
+        private DiscoveryManagerRequest(int transactionId, @NonNull MdnsListener listener,
+                @Nullable Network requestedNetwork, @NonNull Clock clock, long startTimeMs) {
+            super(transactionId, requestedNetwork, clock, startTimeMs);
             mListener = listener;
         }
     }
@@ -2055,7 +2125,7 @@
         /* Remembers a resolved service until getaddrinfo completes */
         private NsdServiceInfo mResolvedService;
 
-        /* A map from client-side ID (listenerKey) to the request */
+        /* A map from client request ID (listenerKey) to the request */
         private final SparseArray<ClientRequest> mClientRequests = new SparseArray<>();
 
         // The target SDK of this client < Build.VERSION_CODES.S
@@ -2065,14 +2135,17 @@
         private final boolean mUseJavaBackend;
         // Store client logs
         private final SharedLog mClientLogs;
+        // Report the nsd metrics data
+        private final NetworkNsdReportedMetrics mMetrics;
 
         private ClientInfo(INsdManagerCallback cb, int uid, boolean useJavaBackend,
-                SharedLog sharedLog) {
+                SharedLog sharedLog, NetworkNsdReportedMetrics metrics) {
             mCb = cb;
             mUid = uid;
             mUseJavaBackend = useJavaBackend;
             mClientLogs = sharedLog;
             mClientLogs.log("New client. useJavaBackend=" + useJavaBackend);
+            mMetrics = metrics;
         }
 
         @Override
@@ -2083,10 +2156,10 @@
             sb.append("mUseJavaBackend ").append(mUseJavaBackend).append("\n");
             sb.append("mUid ").append(mUid).append("\n");
             for (int i = 0; i < mClientRequests.size(); i++) {
-                int clientID = mClientRequests.keyAt(i);
-                sb.append("clientId ")
-                        .append(clientID)
-                        .append(" mDnsId ").append(mClientRequests.valueAt(i).mGlobalId)
+                int clientRequestId = mClientRequests.keyAt(i);
+                sb.append("clientRequestId ")
+                        .append(clientRequestId)
+                        .append(" transactionId ").append(mClientRequests.valueAt(i).mTransactionId)
                         .append(" type ").append(
                                 mClientRequests.valueAt(i).getClass().getSimpleName())
                         .append("\n");
@@ -2115,13 +2188,14 @@
             mClientLogs.log("Client unregistered. expungeAllRequests!");
             // TODO: to keep handler responsive, do not clean all requests for that client at once.
             for (int i = 0; i < mClientRequests.size(); i++) {
-                final int clientId = mClientRequests.keyAt(i);
+                final int clientRequestId = mClientRequests.keyAt(i);
                 final ClientRequest request = mClientRequests.valueAt(i);
-                final int globalId = request.mGlobalId;
-                mIdToClientInfoMap.remove(globalId);
+                final int transactionId = request.mTransactionId;
+                mTransactionIdToClientInfoMap.remove(transactionId);
                 if (DBG) {
-                    Log.d(TAG, "Terminating client-ID " + clientId
-                            + " global-ID " + globalId + " type " + mClientRequests.get(clientId));
+                    Log.d(TAG, "Terminating clientRequestId " + clientRequestId
+                            + " transactionId " + transactionId
+                            + " type " + mClientRequests.get(clientRequestId));
                 }
 
                 if (request instanceof DiscoveryManagerRequest) {
@@ -2130,7 +2204,9 @@
                 }
 
                 if (request instanceof AdvertiserClientRequest) {
-                    mAdvertiser.removeService(globalId);
+                    mAdvertiser.removeService(transactionId);
+                    mMetrics.reportServiceUnregistration(
+                            transactionId, request.calculateRequestDurationMs());
                     continue;
                 }
 
@@ -2140,13 +2216,15 @@
 
                 switch (((LegacyClientRequest) request).mRequestCode) {
                     case NsdManager.DISCOVER_SERVICES:
-                        stopServiceDiscovery(globalId);
+                        stopServiceDiscovery(transactionId);
                         break;
                     case NsdManager.RESOLVE_SERVICE:
-                        stopResolveService(globalId);
+                        stopResolveService(transactionId);
                         break;
                     case NsdManager.REGISTER_SERVICE:
-                        unregisterService(globalId);
+                        unregisterService(transactionId);
+                        mMetrics.reportServiceUnregistration(
+                                transactionId, request.calculateRequestDurationMs());
                         break;
                     default:
                         break;
@@ -2175,12 +2253,11 @@
             return false;
         }
 
-        // mClientRequests is a sparse array of listener id -> ClientRequest.  For a given
-        // mDnsClient id, return the corresponding listener id.  mDnsClient id is also called a
-        // global id.
-        private int getClientId(final int globalId) {
+        // mClientRequests is a sparse array of client request id -> ClientRequest.  For a given
+        // transaction id, return the corresponding client request id.
+        private int getClientRequestId(final int transactionId) {
             for (int i = 0; i < mClientRequests.size(); i++) {
-                if (mClientRequests.valueAt(i).mGlobalId == globalId) {
+                if (mClientRequests.valueAt(i).mTransactionId == transactionId) {
                     return mClientRequests.keyAt(i);
                 }
             }
@@ -2239,7 +2316,13 @@
             }
         }
 
-        void onRegisterServiceFailed(int listenerKey, int error) {
+        void onRegisterServiceFailedImmediately(int listenerKey, int error) {
+            onRegisterServiceFailed(listenerKey, error, NO_TRANSACTION, 0 /* durationMs */);
+        }
+
+        void onRegisterServiceFailed(int listenerKey, int error, int transactionId,
+                long durationMs) {
+            mMetrics.reportServiceRegistrationFailed(transactionId, durationMs);
             try {
                 mCb.onRegisterServiceFailed(listenerKey, error);
             } catch (RemoteException e) {
@@ -2247,7 +2330,9 @@
             }
         }
 
-        void onRegisterServiceSucceeded(int listenerKey, NsdServiceInfo info) {
+        void onRegisterServiceSucceeded(int listenerKey, NsdServiceInfo info, int transactionId,
+                long durationMs) {
+            mMetrics.reportServiceRegistrationSucceeded(transactionId, durationMs);
             try {
                 mCb.onRegisterServiceSucceeded(listenerKey, info);
             } catch (RemoteException e) {
@@ -2263,7 +2348,8 @@
             }
         }
 
-        void onUnregisterServiceSucceeded(int listenerKey) {
+        void onUnregisterServiceSucceeded(int listenerKey, int transactionId, long durationMs) {
+            mMetrics.reportServiceUnregistration(transactionId, durationMs);
             try {
                 mCb.onUnregisterServiceSucceeded(listenerKey);
             } catch (RemoteException e) {
diff --git a/service-t/src/com/android/server/connectivity/mdns/AbstractSocketNetlink.java b/service-t/src/com/android/server/connectivity/mdns/AbstractSocketNetlinkMonitor.java
similarity index 95%
rename from service-t/src/com/android/server/connectivity/mdns/AbstractSocketNetlink.java
rename to service-t/src/com/android/server/connectivity/mdns/AbstractSocketNetlinkMonitor.java
index b792e46..bba3338 100644
--- a/service-t/src/com/android/server/connectivity/mdns/AbstractSocketNetlink.java
+++ b/service-t/src/com/android/server/connectivity/mdns/AbstractSocketNetlinkMonitor.java
@@ -19,7 +19,7 @@
 /**
  * The interface for netlink monitor.
  */
-public interface AbstractSocketNetlink {
+public interface AbstractSocketNetlinkMonitor {
 
     /**
      * Returns if the netlink monitor is supported or not. By default, it is not supported.
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 2d5bb00..b7417ed 100644
--- a/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
+++ b/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
@@ -16,9 +16,9 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsServiceTypeClient.INVALID_TRANSACTION_ID;
+
 import android.annotation.NonNull;
-import android.annotation.Nullable;
-import android.net.Network;
 import android.text.TextUtils;
 import android.util.Log;
 import android.util.Pair;
@@ -70,13 +70,13 @@
     private final List<String> subtypes;
     private final boolean expectUnicastResponse;
     private final int transactionId;
-    @Nullable
-    private final Network network;
+    @NonNull
+    private final SocketKey socketKey;
     private final boolean sendDiscoveryQueries;
     @NonNull
     private final List<MdnsResponse> servicesToResolve;
     @NonNull
-    private final MdnsResponseDecoder.Clock clock;
+    private final MdnsUtils.Clock clock;
     private final boolean onlyUseIpv6OnIpv6OnlyNetworks;
 
     EnqueueMdnsQueryCallable(
@@ -86,33 +86,37 @@
             @NonNull Collection<String> subtypes,
             boolean expectUnicastResponse,
             int transactionId,
-            @Nullable Network network,
+            @NonNull SocketKey socketKey,
             boolean onlyUseIpv6OnIpv6OnlyNetworks,
             boolean sendDiscoveryQueries,
             @NonNull Collection<MdnsResponse> servicesToResolve,
-            @NonNull MdnsResponseDecoder.Clock clock) {
+            @NonNull MdnsUtils.Clock clock) {
         weakRequestSender = new WeakReference<>(requestSender);
         this.packetWriter = packetWriter;
         serviceTypeLabels = TextUtils.split(serviceType, "\\.");
         this.subtypes = new ArrayList<>(subtypes);
         this.expectUnicastResponse = expectUnicastResponse;
         this.transactionId = transactionId;
-        this.network = network;
+        this.socketKey = socketKey;
         this.onlyUseIpv6OnIpv6OnlyNetworks = onlyUseIpv6OnIpv6OnlyNetworks;
         this.sendDiscoveryQueries = sendDiscoveryQueries;
         this.servicesToResolve = new ArrayList<>(servicesToResolve);
         this.clock = clock;
     }
 
+    /**
+     * Call to execute the mdns query.
+     *
+     * @return The pair of transaction id and the subtypes for the query.
+     */
     // 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(INVALID_TRANSACTION_ID, new ArrayList<>());
             }
 
             int numQuestions = 0;
@@ -159,7 +163,7 @@
 
             if (numQuestions == 0) {
                 // No query to send
-                return null;
+                return Pair.create(INVALID_TRANSACTION_ID, new ArrayList<>());
             }
 
             // Header.
@@ -198,7 +202,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(INVALID_TRANSACTION_ID, new ArrayList<>());
         }
     }
 
@@ -216,7 +220,7 @@
         if (expectUnicastResponse) {
             if (requestSender instanceof MdnsMultinetworkSocketClient) {
                 ((MdnsMultinetworkSocketClient) requestSender).sendPacketRequestingUnicastResponse(
-                        packet, network, onlyUseIpv6OnIpv6OnlyNetworks);
+                        packet, socketKey, onlyUseIpv6OnIpv6OnlyNetworks);
             } else {
                 requestSender.sendPacketRequestingUnicastResponse(
                         packet, onlyUseIpv6OnIpv6OnlyNetworks);
@@ -225,7 +229,7 @@
             if (requestSender instanceof MdnsMultinetworkSocketClient) {
                 ((MdnsMultinetworkSocketClient) requestSender)
                         .sendPacketRequestingMulticastResponse(
-                                packet, network, onlyUseIpv6OnIpv6OnlyNetworks);
+                                packet, socketKey, onlyUseIpv6OnIpv6OnlyNetworks);
             } else {
                 requestSender.sendPacketRequestingMulticastResponse(
                         packet, onlyUseIpv6OnIpv6OnlyNetworks);
diff --git a/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java b/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
index 0eebc61..161669b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
+++ b/service-t/src/com/android/server/connectivity/mdns/ExecutorProvider.java
@@ -16,6 +16,7 @@
 
 package com.android.server.connectivity.mdns;
 
+import android.annotation.NonNull;
 import android.util.ArraySet;
 
 import java.util.Set;
@@ -47,5 +48,17 @@
             }
             executor.shutdownNow();
         }
+        serviceTypeClientSchedulerExecutors.clear();
+    }
+
+    /**
+     * Shutdown one executor service and remove the executor service from the set.
+     * @param executorService the executorService to be shutdown
+     */
+    public void shutdownExecutorService(@NonNull ScheduledExecutorService executorService) {
+        if (!executorService.isShutdown()) {
+            executorService.shutdownNow();
+        }
+        serviceTypeClientSchedulerExecutors.remove(executorService);
     }
 }
\ No newline at end of file
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..d4aeacf 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java
@@ -50,14 +50,6 @@
         return false;
     }
 
-    public static boolean useSessionIdToScheduleMdnsTask() {
-        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 f0e1717..0c32cf1 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsConstants.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsConstants.java
@@ -18,14 +18,11 @@
 
 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
 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 afad3b7..d55098c 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());
         }
     }
 
@@ -194,13 +197,14 @@
                     }
 
                     @Override
-                    public void onAllSocketsDestroyed(@NonNull SocketKey socketKey) {
+                    public void onSocketDestroyed(@NonNull SocketKey socketKey) {
                         ensureRunningOnHandlerThread(handler);
                         final MdnsServiceTypeClient serviceTypeClient =
                                 perSocketServiceTypeClients.get(serviceType, socketKey);
                         if (serviceTypeClient == null) return;
                         // Notify all listeners that all services are removed from this socket.
                         serviceTypeClient.notifySocketDestroyed();
+                        executorProvider.shutdownExecutorService(serviceTypeClient.getExecutor());
                         perSocketServiceTypeClients.remove(serviceTypeClient);
                     }
                 });
@@ -235,6 +239,7 @@
             if (serviceTypeClient.stopSendAndReceive(listener)) {
                 // No listener is registered for the service type anymore, remove it from the list
                 // of the service type clients.
+                executorProvider.shutdownExecutorService(serviceTypeClient.getExecutor());
                 perSocketServiceTypeClients.remove(serviceTypeClient);
             }
         }
@@ -254,8 +259,7 @@
     private void handleOnResponseReceived(@NonNull MdnsPacket packet,
             @NonNull SocketKey socketKey) {
         for (MdnsServiceTypeClient serviceTypeClient : getMdnsServiceTypeClient(socketKey)) {
-            serviceTypeClient.processResponse(
-                    packet, socketKey.getInterfaceIndex(), socketKey.getNetwork());
+            serviceTypeClient.processResponse(packet, socketKey);
         }
     }
 
@@ -285,9 +289,11 @@
     MdnsServiceTypeClient createServiceTypeClient(@NonNull String serviceType,
             @NonNull SocketKey socketKey) {
         sharedLog.log("createServiceTypeClient for type:" + serviceType + " " + socketKey);
+        final String tag = serviceType + "-" + socketKey.getNetwork()
+                + "/" + socketKey.getInterfaceIndex();
         return new MdnsServiceTypeClient(
                 serviceType, socketClient,
                 executorProvider.newServiceTypeClientSchedulerExecutor(), socketKey,
-                sharedLog.forSubComponent(serviceType + "-" + socketKey));
+                sharedLog.forSubComponent(tag), handler.getLooper(), serviceCache);
     }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
index 1253444..d1fa57c 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
@@ -64,7 +64,7 @@
         @NonNull
         private final SocketCreationCallback mSocketCreationCallback;
         @NonNull
-        private final ArrayMap<MdnsInterfaceSocket, SocketKey> mActiveNetworkSockets =
+        private final ArrayMap<MdnsInterfaceSocket, SocketKey> mActiveSockets =
                 new ArrayMap<>();
 
         InterfaceSocketCallback(SocketCreationCallback socketCreationCallback) {
@@ -83,7 +83,7 @@
                 mSocketPacketHandlers.put(socket, handler);
             }
             socket.addPacketHandler(handler);
-            mActiveNetworkSockets.put(socket, socketKey);
+            mActiveSockets.put(socket, socketKey);
             mSocketCreationCallback.onSocketCreated(socketKey);
         }
 
@@ -95,16 +95,16 @@
         }
 
         private void notifySocketDestroyed(@NonNull MdnsInterfaceSocket socket) {
-            final SocketKey socketKey = mActiveNetworkSockets.remove(socket);
-            if (!isAnySocketActive(socketKey)) {
-                mSocketCreationCallback.onAllSocketsDestroyed(socketKey);
+            final SocketKey socketKey = mActiveSockets.remove(socket);
+            if (!isSocketActive(socket)) {
+                mSocketCreationCallback.onSocketDestroyed(socketKey);
             }
         }
 
         void onNetworkUnrequested() {
-            for (int i = mActiveNetworkSockets.size() - 1; i >= 0; i--) {
+            for (int i = mActiveSockets.size() - 1; i >= 0; i--) {
                 // Iterate from the end so the socket can be removed
-                final MdnsInterfaceSocket socket = mActiveNetworkSockets.keyAt(i);
+                final MdnsInterfaceSocket socket = mActiveSockets.keyAt(i);
                 notifySocketDestroyed(socket);
                 maybeCleanupPacketHandler(socket);
             }
@@ -114,17 +114,7 @@
     private boolean isSocketActive(@NonNull MdnsInterfaceSocket socket) {
         for (int i = 0; i < mRequestedNetworks.size(); i++) {
             final InterfaceSocketCallback isc = mRequestedNetworks.valueAt(i);
-            if (isc.mActiveNetworkSockets.containsKey(socket)) {
-                return true;
-            }
-        }
-        return false;
-    }
-
-    private boolean isAnySocketActive(@NonNull SocketKey socketKey) {
-        for (int i = 0; i < mRequestedNetworks.size(); i++) {
-            final InterfaceSocketCallback isc = mRequestedNetworks.valueAt(i);
-            if (isc.mActiveNetworkSockets.containsValue(socketKey)) {
+            if (isc.mActiveSockets.containsKey(socket)) {
                 return true;
             }
         }
@@ -135,7 +125,7 @@
         final ArrayMap<MdnsInterfaceSocket, SocketKey> sockets = new ArrayMap<>();
         for (int i = 0; i < mRequestedNetworks.size(); i++) {
             final InterfaceSocketCallback isc = mRequestedNetworks.valueAt(i);
-            sockets.putAll(isc.mActiveNetworkSockets);
+            sockets.putAll(isc.mActiveSockets);
         }
         return sockets;
     }
@@ -213,25 +203,22 @@
         return true;
     }
 
-    private void sendMdnsPacket(@NonNull DatagramPacket packet, @Nullable Network targetNetwork,
+    private void sendMdnsPacket(@NonNull DatagramPacket packet, @NonNull SocketKey targetSocketKey,
             boolean onlyUseIpv6OnIpv6OnlyNetworks) {
         final boolean isIpv6 = ((InetSocketAddress) packet.getSocketAddress()).getAddress()
                 instanceof Inet6Address;
         final boolean isIpv4 = ((InetSocketAddress) packet.getSocketAddress()).getAddress()
                 instanceof Inet4Address;
         final ArrayMap<MdnsInterfaceSocket, SocketKey> activeSockets = getActiveSockets();
-        boolean shouldQueryIpv6 = !onlyUseIpv6OnIpv6OnlyNetworks || isIpv6OnlyNetworks(
-                activeSockets, targetNetwork);
+        boolean shouldQueryIpv6 = !onlyUseIpv6OnIpv6OnlyNetworks || isIpv6OnlySockets(
+                activeSockets, targetSocketKey);
         for (int i = 0; i < activeSockets.size(); i++) {
             final MdnsInterfaceSocket socket = activeSockets.keyAt(i);
-            final Network network = activeSockets.valueAt(i).getNetwork();
+            final SocketKey socketKey = activeSockets.valueAt(i);
             // Check ip capability and network before sending packet
             if (((isIpv6 && socket.hasJoinedIpv6() && shouldQueryIpv6)
                     || (isIpv4 && socket.hasJoinedIpv4()))
-                    // Contrary to MdnsUtils.isNetworkMatched, only send packets targeting
-                    // the null network to interfaces that have the null network (tethering
-                    // downstream interfaces).
-                    && Objects.equals(network, targetNetwork)) {
+                    && Objects.equals(socketKey, targetSocketKey)) {
                 try {
                     socket.send(packet);
                 } catch (IOException e) {
@@ -241,13 +228,13 @@
         }
     }
 
-    private boolean isIpv6OnlyNetworks(
+    private boolean isIpv6OnlySockets(
             @NonNull ArrayMap<MdnsInterfaceSocket, SocketKey> activeSockets,
-            @Nullable Network targetNetwork) {
+            @NonNull SocketKey targetSocketKey) {
         for (int i = 0; i < activeSockets.size(); i++) {
             final MdnsInterfaceSocket socket = activeSockets.keyAt(i);
-            final Network network = activeSockets.valueAt(i).getNetwork();
-            if (Objects.equals(network, targetNetwork) && socket.hasJoinedIpv4()) {
+            final SocketKey socketKey = activeSockets.valueAt(i);
+            if (Objects.equals(socketKey, targetSocketKey) && socket.hasJoinedIpv4()) {
                 return false;
             }
         }
@@ -276,38 +263,35 @@
     }
 
     /**
-     * Send a mDNS request packet via given network that asks for multicast response.
-     *
-     * <p>The socket client may use a null network to identify some or all interfaces, in which case
-     * passing null sends the packet to these.
+     * Send a mDNS request packet via given socket key that asks for multicast response.
      */
     public void sendPacketRequestingMulticastResponse(@NonNull DatagramPacket packet,
-            @Nullable Network network, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
-        mHandler.post(() -> sendMdnsPacket(packet, network, onlyUseIpv6OnIpv6OnlyNetworks));
+            @NonNull SocketKey socketKey, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
+        mHandler.post(() -> sendMdnsPacket(packet, socketKey, onlyUseIpv6OnIpv6OnlyNetworks));
     }
 
     @Override
     public void sendPacketRequestingMulticastResponse(
             @NonNull DatagramPacket packet, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
-        sendPacketRequestingMulticastResponse(
-                packet, null /* network */, onlyUseIpv6OnIpv6OnlyNetworks);
+        throw new UnsupportedOperationException("This socket client need to specify the socket to"
+                + "send packet");
     }
 
     /**
-     * Send a mDNS request packet via given network that asks for unicast response.
+     * Send a mDNS request packet via given socket key that asks for unicast response.
      *
      * <p>The socket client may use a null network to identify some or all interfaces, in which case
      * passing null sends the packet to these.
      */
     public void sendPacketRequestingUnicastResponse(@NonNull DatagramPacket packet,
-            @Nullable Network network, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
-        mHandler.post(() -> sendMdnsPacket(packet, network, onlyUseIpv6OnIpv6OnlyNetworks));
+            @NonNull SocketKey socketKey, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
+        mHandler.post(() -> sendMdnsPacket(packet, socketKey, onlyUseIpv6OnIpv6OnlyNetworks));
     }
 
     @Override
     public void sendPacketRequestingUnicastResponse(
             @NonNull DatagramPacket packet, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
-        sendPacketRequestingUnicastResponse(
-                packet, null /* network */, onlyUseIpv6OnIpv6OnlyNetworks);
+        throw new UnsupportedOperationException("This socket client need to specify the socket to"
+                + "send packet");
     }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsResponseDecoder.java b/service-t/src/com/android/server/connectivity/mdns/MdnsResponseDecoder.java
index eff1880..a0a538e 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsResponseDecoder.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsResponseDecoder.java
@@ -19,7 +19,6 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.net.Network;
-import android.os.SystemClock;
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Pair;
@@ -40,10 +39,10 @@
     private final boolean allowMultipleSrvRecordsPerHost =
             MdnsConfigs.allowMultipleSrvRecordsPerHost();
     @Nullable private final String[] serviceType;
-    private final Clock clock;
+    private final MdnsUtils.Clock clock;
 
     /** Constructs a new decoder that will extract responses for the given service type. */
-    public MdnsResponseDecoder(@NonNull Clock clock, @Nullable String[] serviceType) {
+    public MdnsResponseDecoder(@NonNull MdnsUtils.Clock clock, @Nullable String[] serviceType) {
         this.clock = clock;
         this.serviceType = serviceType;
     }
@@ -330,10 +329,4 @@
         }
         return result == null ? List.of() : result;
     }
-
-    public static class Clock {
-        public long elapsedRealtime() {
-            return SystemClock.elapsedRealtime();
-        }
-    }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java
index 98c80ee..f09596d 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java
@@ -50,7 +50,8 @@
                             source.readBoolean(),
                             source.readParcelable(null),
                             source.readString(),
-                            (source.dataAvail() > 0) ? source.readBoolean() : false);
+                            source.readBoolean(),
+                            source.readInt());
                 }
 
                 @Override
@@ -62,9 +63,9 @@
     private final List<String> subtypes;
     @Nullable
     private final String resolveInstanceName;
-
     private final boolean isPassiveMode;
     private final boolean onlyUseIpv6OnIpv6OnlyNetworks;
+    private final int numOfQueriesBeforeBackoff;
     private final boolean removeExpiredService;
     // The target network for searching. Null network means search on all possible interfaces.
     @Nullable private final Network mNetwork;
@@ -76,13 +77,15 @@
             boolean removeExpiredService,
             @Nullable Network network,
             @Nullable String resolveInstanceName,
-            boolean onlyUseIpv6OnIpv6OnlyNetworks) {
+            boolean onlyUseIpv6OnIpv6OnlyNetworks,
+            int numOfQueriesBeforeBackoff) {
         this.subtypes = new ArrayList<>();
         if (subtypes != null) {
             this.subtypes.addAll(subtypes);
         }
         this.isPassiveMode = isPassiveMode;
         this.onlyUseIpv6OnIpv6OnlyNetworks = onlyUseIpv6OnIpv6OnlyNetworks;
+        this.numOfQueriesBeforeBackoff = numOfQueriesBeforeBackoff;
         this.removeExpiredService = removeExpiredService;
         mNetwork = network;
         this.resolveInstanceName = resolveInstanceName;
@@ -122,6 +125,14 @@
         return onlyUseIpv6OnIpv6OnlyNetworks;
     }
 
+    /**
+     *  Returns number of queries should be executed before backoff mode is enabled.
+     *  The default number is 3 if it is not set.
+     */
+    public int numOfQueriesBeforeBackoff() {
+        return numOfQueriesBeforeBackoff;
+    }
+
     /** Returns {@code true} if service will be removed after its TTL expires. */
     public boolean removeExpiredService() {
         return removeExpiredService;
@@ -159,6 +170,7 @@
         out.writeParcelable(mNetwork, 0);
         out.writeString(resolveInstanceName);
         out.writeBoolean(onlyUseIpv6OnIpv6OnlyNetworks);
+        out.writeInt(numOfQueriesBeforeBackoff);
     }
 
     /** A builder to create {@link MdnsSearchOptions}. */
@@ -166,6 +178,7 @@
         private final Set<String> subtypes;
         private boolean isPassiveMode = true;
         private boolean onlyUseIpv6OnIpv6OnlyNetworks = false;
+        private int numOfQueriesBeforeBackoff = 3;
         private boolean removeExpiredService;
         private Network mNetwork;
         private String resolveInstanceName;
@@ -219,6 +232,14 @@
         }
 
         /**
+         * Sets if the query backoff mode should be turned on.
+         */
+        public Builder setNumOfQueriesBeforeBackoff(int numOfQueriesBeforeBackoff) {
+            this.numOfQueriesBeforeBackoff = numOfQueriesBeforeBackoff;
+            return this;
+        }
+
+        /**
          * Sets if the service should be removed after TTL.
          *
          * @param removeExpiredService If set to {@code true}, the service will be removed after TTL
@@ -258,7 +279,8 @@
                     removeExpiredService,
                     mNetwork,
                     resolveInstanceName,
-                    onlyUseIpv6OnIpv6OnlyNetworks);
+                    onlyUseIpv6OnIpv6OnlyNetworks,
+                    numOfQueriesBeforeBackoff);
         }
     }
 }
\ 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 cd0be67..ec6af9b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceCache.java
@@ -22,7 +22,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.net.Network;
 import android.os.Handler;
 import android.os.Looper;
 import android.util.ArrayMap;
@@ -45,15 +44,15 @@
 public class MdnsServiceCache {
     private static class CacheKey {
         @NonNull final String mLowercaseServiceType;
-        @Nullable final Network mNetwork;
+        @NonNull final SocketKey mSocketKey;
 
-        CacheKey(@NonNull String serviceType, @Nullable Network network) {
+        CacheKey(@NonNull String serviceType, @NonNull SocketKey socketKey) {
             mLowercaseServiceType = toDnsLowerCase(serviceType);
-            mNetwork = network;
+            mSocketKey = socketKey;
         }
 
         @Override public int hashCode() {
-            return Objects.hash(mLowercaseServiceType, mNetwork);
+            return Objects.hash(mLowercaseServiceType, mSocketKey);
         }
 
         @Override public boolean equals(Object other) {
@@ -64,11 +63,11 @@
                 return false;
             }
             return Objects.equals(mLowercaseServiceType, ((CacheKey) other).mLowercaseServiceType)
-                    && Objects.equals(mNetwork, ((CacheKey) other).mNetwork);
+                    && Objects.equals(mSocketKey, ((CacheKey) other).mSocketKey);
         }
     }
     /**
-     * A map of cached services. Key is composed of service name, type and network. Value is the
+     * A map of cached services. Key is composed of service name, type and socket. Value is the
      * service which use the service type to discover from each socket.
      */
     @NonNull
@@ -81,23 +80,30 @@
     }
 
     /**
-     * Get the cache services which are queried from given service type and network.
+     * Get the cache services which are queried from given service type and socket.
      *
      * @param serviceType the target service type.
-     * @param network the target network
+     * @param socketKey the target socket
      * @return the set of services which matches the given service type.
      */
     @NonNull
     public List<MdnsResponse> getCachedServices(@NonNull String serviceType,
-            @Nullable Network network) {
+            @NonNull SocketKey socketKey) {
         ensureRunningOnHandlerThread(mHandler);
-        final CacheKey key = new CacheKey(serviceType, network);
+        final CacheKey key = new CacheKey(serviceType, socketKey);
         return mCachedServices.containsKey(key)
                 ? Collections.unmodifiableList(new ArrayList<>(mCachedServices.get(key)))
                 : 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())) {
@@ -112,15 +118,15 @@
      *
      * @param serviceName the target service name.
      * @param serviceType the target service type.
-     * @param network the target network
+     * @param socketKey the target socket
      * @return the service which matches given conditions.
      */
     @Nullable
     public MdnsResponse getCachedService(@NonNull String serviceName,
-            @NonNull String serviceType, @Nullable Network network) {
+            @NonNull String serviceType, @NonNull SocketKey socketKey) {
         ensureRunningOnHandlerThread(mHandler);
         final List<MdnsResponse> responses =
-                mCachedServices.get(new CacheKey(serviceType, network));
+                mCachedServices.get(new CacheKey(serviceType, socketKey));
         if (responses == null) {
             return null;
         }
@@ -132,14 +138,14 @@
      * Add or update a service.
      *
      * @param serviceType the service type.
-     * @param network the target network
+     * @param socketKey the target socket
      * @param response the response of the discovered service.
      */
-    public void addOrUpdateService(@NonNull String serviceType, @Nullable Network network,
+    public void addOrUpdateService(@NonNull String serviceType, @NonNull SocketKey socketKey,
             @NonNull MdnsResponse response) {
         ensureRunningOnHandlerThread(mHandler);
         final List<MdnsResponse> responses = mCachedServices.computeIfAbsent(
-                new CacheKey(serviceType, network), key -> new ArrayList<>());
+                new CacheKey(serviceType, socketKey), key -> new ArrayList<>());
         // Remove existing service if present.
         final MdnsResponse existing =
                 findMatchedResponse(responses, response.getServiceInstanceName());
@@ -148,18 +154,18 @@
     }
 
     /**
-     * Remove a service which matches the given service name, type and network.
+     * Remove a service which matches the given service name, type and socket.
      *
      * @param serviceName the target service name.
      * @param serviceType the target service type.
-     * @param network the target network.
+     * @param socketKey the target socket.
      */
     @Nullable
     public MdnsResponse removeService(@NonNull String serviceName, @NonNull String serviceType,
-            @Nullable Network network) {
+            @NonNull SocketKey socketKey) {
         ensureRunningOnHandlerThread(mHandler);
         final List<MdnsResponse> responses =
-                mCachedServices.get(new CacheKey(serviceType, network));
+                mCachedServices.get(new CacheKey(serviceType, socketKey));
         if (responses == null) {
             return null;
         }
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 a36eb1b..b5fd8a0 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -16,17 +16,20 @@
 
 package com.android.server.connectivity.mdns;
 
-import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static com.android.server.connectivity.mdns.MdnsServiceCache.findMatchedResponse;
+import static com.android.server.connectivity.mdns.util.MdnsUtils.Clock;
+import static com.android.server.connectivity.mdns.util.MdnsUtils.ensureRunningOnHandlerThread;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.net.Network;
+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;
@@ -36,13 +39,9 @@
 import java.net.Inet6Address;
 import java.util.ArrayList;
 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;
 
 /**
@@ -51,7 +50,12 @@
  */
 public class MdnsServiceTypeClient {
 
+    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;
+    static final int INVALID_TRANSACTION_ID = -1;
 
     private final String serviceType;
     private final String[] serviceTypeLabels;
@@ -60,14 +64,17 @@
     private final ScheduledExecutorService executor;
     @NonNull private final SocketKey socketKey;
     @NonNull private final SharedLog sharedLog;
-    private final Object lock = new Object();
+    @NonNull private final Handler handler;
+    @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.
-    private final Map<String, MdnsResponse> instanceNameToResponse = new HashMap<>();
     private final boolean removeServiceAfterTtlExpires =
             MdnsConfigs.removeServiceAfterTtlExpires();
-    private final MdnsResponseDecoder.Clock clock;
+    private final Clock clock;
 
     @Nullable private MdnsSearchOptions searchOptions;
 
@@ -76,9 +83,101 @@
     // new subtypes. It stays the same between packets for same subtypes.
     private long currentSessionId = 0;
 
-    @GuardedBy("lock")
     @Nullable
-    private Future<?> requestTaskFuture;
+    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 QuerySentArguments sentResult = (QuerySentArguments) msg.obj;
+                    // If a task is cancelled while the Executor is running it, EVENT_QUERY_RESULT
+                    // will still be sent when it ends. So use session ID to check if this task
+                    // should continue to schedule more.
+                    if (sentResult.taskArgs.sessionId != currentSessionId) {
+                        break;
+                    }
+
+                    if ((sentResult.transactionId != INVALID_TRANSACTION_ID)) {
+                        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}.
@@ -91,9 +190,11 @@
             @NonNull MdnsSocketClientBase socketClient,
             @NonNull ScheduledExecutorService executor,
             @NonNull SocketKey socketKey,
-            @NonNull SharedLog sharedLog) {
-        this(serviceType, socketClient, executor, new MdnsResponseDecoder.Clock(), socketKey,
-                sharedLog);
+            @NonNull SharedLog sharedLog,
+            @NonNull Looper looper,
+            @NonNull MdnsServiceCache serviceCache) {
+        this(serviceType, socketClient, executor, new Clock(), socketKey, sharedLog, looper,
+                new Dependencies(), serviceCache);
     }
 
     @VisibleForTesting
@@ -101,9 +202,12 @@
             @NonNull String serviceType,
             @NonNull MdnsSocketClientBase socketClient,
             @NonNull ScheduledExecutorService executor,
-            @NonNull MdnsResponseDecoder.Clock clock,
+            @NonNull Clock clock,
             @NonNull SocketKey socketKey,
-            @NonNull SharedLog sharedLog) {
+            @NonNull SharedLog sharedLog,
+            @NonNull Looper looper,
+            @NonNull Dependencies dependencies,
+            @NonNull MdnsServiceCache serviceCache) {
         this.serviceType = serviceType;
         this.socketClient = socketClient;
         this.executor = executor;
@@ -112,6 +216,9 @@
         this.clock = clock;
         this.socketKey = socketKey;
         this.sharedLog = sharedLog;
+        this.handler = new QueryTaskHandler(looper);
+        this.dependencies = dependencies;
+        this.serviceCache = serviceCache;
     }
 
     private static MdnsServiceInfo buildMdnsServiceInfoFromResponse(
@@ -171,49 +278,70 @@
      * @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) {
-        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;
-                    }
+        ensureRunningOnHandlerThread(handler);
+        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 (requestTaskFuture != 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(),
-                    currentSessionId,
-                    socketKey);
-            if (hadReply) {
-                requestTaskFuture = scheduleNextRunLocked(taskConfig);
-            } else {
-                requestTaskFuture = executor.submit(new QueryTask(taskConfig));
-            }
+        }
+        // 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() {
-        requestTaskFuture.cancel(true);
+    /**
+     * Get the executor service.
+     */
+    public ScheduledExecutorService getExecutor() {
+        return executor;
+    }
+
+    private void removeScheduledTask() {
+        dependencies.removeMessages(handler, EVENT_START_QUERYTASK);
+        sharedLog.log("Remove EVENT_START_QUERYTASK"
+                + ", current session: " + currentSessionId);
         ++currentSessionId;
-        requestTaskFuture = null;
+        lastScheduledQueryTaskArgs = null;
     }
 
     private boolean responseMatchesOptions(@NonNull MdnsResponse response,
@@ -244,62 +372,73 @@
      * listener}. Otherwise returns {@code false}.
      */
     public boolean stopSendAndReceive(@NonNull MdnsServiceBrowserListener listener) {
-        synchronized (lock) {
-            if (listeners.remove(listener) == null) {
-                return listeners.isEmpty();
-            }
-            if (listeners.isEmpty() && requestTaskFuture != null) {
-                cancelRequestTaskLocked();
-            }
+        ensureRunningOnHandlerThread(handler);
+        if (listeners.remove(listener) == null) {
             return listeners.isEmpty();
         }
-    }
-
-    public String[] getServiceTypeLabels() {
-        return serviceTypeLabels;
+        if (listeners.isEmpty()) {
+            removeScheduledTask();
+        }
+        return listeners.isEmpty();
     }
 
     /**
      * Process an incoming response packet.
      */
-    public synchronized void processResponse(@NonNull MdnsPacket packet, int interfaceIndex,
-            Network network) {
-        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(interfaceIndex,
-                    network);
-            for (MdnsResponse additionalResponse : additionalResponses) {
-                if (!instanceNameToResponse.containsKey(
-                        additionalResponse.getServiceInstanceName())) {
-                    currentList.add(additionalResponse);
-                }
+    public synchronized void processResponse(@NonNull MdnsPacket packet,
+            @NonNull SocketKey socketKey) {
+        ensureRunningOnHandlerThread(handler);
+        // 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, interfaceIndex, network);
+        }
+        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()) {
-                        onGoodbyeReceived(response.getServiceInstanceName());
-                    } else {
-                        onResponseModified(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 (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);
             }
         }
     }
 
     public synchronized void onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode) {
+        ensureRunningOnHandlerThread(handler);
         for (int i = 0; i < listeners.size(); i++) {
             listeners.keyAt(i).onFailedToParseMdnsResponse(receivedPacketNumber, errorCode);
         }
@@ -307,45 +446,41 @@
 
     /** Notify all services are removed because the socket is destroyed. */
     public void notifySocketDestroyed() {
-        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);
+        ensureRunningOnHandlerThread(handler);
+        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 (requestTaskFuture != null) {
-                cancelRequestTaskLocked();
+                sharedLog.log("Socket destroyed. onServiceNameRemoved: " + name);
+                listener.onServiceNameRemoved(serviceInfo);
             }
         }
+        removeScheduledTask();
     }
 
     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;
         }
@@ -378,7 +513,8 @@
     }
 
     private void onGoodbyeReceived(@Nullable String serviceInstanceName) {
-        final MdnsResponse response = instanceNameToResponse.remove(serviceInstanceName);
+        final MdnsResponse response =
+                serviceCache.removeService(serviceInstanceName, serviceType, socketKey);
         if (response == null) {
             return;
         }
@@ -408,115 +544,15 @@
         return new MdnsPacketWriter(DEFAULT_MTU);
     }
 
-    // A configuration for the PeriodicalQueryTask that contains parameters to build a query packet.
-    // Call to getConfigForNextRun returns a config that can be used to build the next query task.
-    @VisibleForTesting
-    static class QueryTaskConfig {
-
-        private static final int INITIAL_TIME_BETWEEN_BURSTS_MS =
-                (int) MdnsConfigs.initialTimeBetweenBurstsMs();
-        private static final int TIME_BETWEEN_BURSTS_MS = (int) MdnsConfigs.timeBetweenBurstsMs();
-        private static final int QUERIES_PER_BURST = (int) MdnsConfigs.queriesPerBurst();
-        private static final int TIME_BETWEEN_QUERIES_IN_BURST_MS =
-                (int) MdnsConfigs.timeBetweenQueriesInBurstMs();
-        private static final int QUERIES_PER_BURST_PASSIVE_MODE =
-                (int) MdnsConfigs.queriesPerBurstPassive();
-        private static final int UNSIGNED_SHORT_MAX_VALUE = 65536;
-        // The following fields are used by QueryTask so we need to test them.
-        @VisibleForTesting
-        final List<String> subtypes;
-        private final boolean alwaysAskForUnicastResponse =
-                MdnsConfigs.alwaysAskForUnicastResponseInEachBurst();
-        private final boolean usePassiveMode;
-        private final boolean onlyUseIpv6OnIpv6OnlyNetworks;
-        private final long sessionId;
-        @VisibleForTesting
-        int transactionId;
-        @VisibleForTesting
-        boolean expectUnicastResponse;
-        private int queriesPerBurst;
-        private int timeBetweenBurstsInMs;
-        private int burstCounter;
-        private int timeToRunNextTaskInMs;
-        private boolean isFirstBurst;
-        @NonNull private final SocketKey socketKey;
-
-        QueryTaskConfig(@NonNull Collection<String> subtypes,
-                boolean usePassiveMode,
-                boolean onlyUseIpv6OnIpv6OnlyNetworks,
-                long sessionId,
-                @Nullable SocketKey socketKey) {
-            this.usePassiveMode = usePassiveMode;
-            this.onlyUseIpv6OnIpv6OnlyNetworks = onlyUseIpv6OnIpv6OnlyNetworks;
-            this.subtypes = new ArrayList<>(subtypes);
-            this.queriesPerBurst = QUERIES_PER_BURST;
-            this.burstCounter = 0;
-            this.transactionId = 1;
-            this.expectUnicastResponse = true;
-            this.isFirstBurst = true;
-            this.sessionId = sessionId;
-            // Config the scan frequency based on the scan mode.
-            if (this.usePassiveMode) {
-                // In passive scan mode, sends a single burst of QUERIES_PER_BURST queries, and then
-                // in each TIME_BETWEEN_BURSTS interval, sends QUERIES_PER_BURST_PASSIVE_MODE
-                // queries.
-                this.timeBetweenBurstsInMs = TIME_BETWEEN_BURSTS_MS;
-            } else {
-                // In active scan mode, sends a burst of QUERIES_PER_BURST queries,
-                // TIME_BETWEEN_QUERIES_IN_BURST_MS apart, then waits for the scan interval, and
-                // then repeats. The scan interval starts as INITIAL_TIME_BETWEEN_BURSTS_MS and
-                // doubles until it maxes out at TIME_BETWEEN_BURSTS_MS.
-                this.timeBetweenBurstsInMs = INITIAL_TIME_BETWEEN_BURSTS_MS;
-            }
-            this.socketKey = socketKey;
-        }
-
-        QueryTaskConfig getConfigForNextRun() {
-            if (++transactionId > UNSIGNED_SHORT_MAX_VALUE) {
-                transactionId = 1;
-            }
-            // Only the first query expects uni-cast response.
-            expectUnicastResponse = false;
-            if (++burstCounter == queriesPerBurst) {
-                burstCounter = 0;
-
-                if (alwaysAskForUnicastResponse) {
-                    expectUnicastResponse = true;
-                }
-                // In passive scan mode, sends a single burst of QUERIES_PER_BURST queries, and
-                // then in each TIME_BETWEEN_BURSTS interval, sends QUERIES_PER_BURST_PASSIVE_MODE
-                // queries.
-                if (isFirstBurst) {
-                    isFirstBurst = false;
-                    if (usePassiveMode) {
-                        queriesPerBurst = QUERIES_PER_BURST_PASSIVE_MODE;
-                    }
-                }
-                // In active scan mode, sends a burst of QUERIES_PER_BURST queries,
-                // TIME_BETWEEN_QUERIES_IN_BURST_MS apart, then waits for the scan interval, and
-                // then repeats. The scan interval starts as INITIAL_TIME_BETWEEN_BURSTS_MS and
-                // doubles until it maxes out at TIME_BETWEEN_BURSTS_MS.
-                timeToRunNextTaskInMs = timeBetweenBurstsInMs;
-                if (timeBetweenBurstsInMs < TIME_BETWEEN_BURSTS_MS) {
-                    timeBetweenBurstsInMs = Math.min(timeBetweenBurstsInMs * 2,
-                            TIME_BETWEEN_BURSTS_MS);
-                }
-            } else {
-                timeToRunNextTaskInMs = TIME_BETWEEN_QUERIES_IN_BURST_MS;
-            }
-            return this;
-        }
-    }
-
-    private List<MdnsResponse> makeResponsesForResolve(int interfaceIndex,
-            @NonNull Network network) {
+    private List<MdnsResponse> makeResponsesForResolve(@NonNull SocketKey socketKey) {
         final List<MdnsResponse> resolveResponses = new ArrayList<>();
         for (int i = 0; i < listeners.size(); i++) {
             final String resolveName = listeners.valueAt(i).getResolveInstanceName();
             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);
@@ -524,36 +560,89 @@
                 instanceFullName.addAll(Arrays.asList(serviceTypeLabels));
                 knownResponse = new MdnsResponse(
                         0L /* lastUpdateTime */, instanceFullName.toArray(new String[0]),
-                        interfaceIndex, network);
+                        socketKey.getInterfaceIndex(), socketKey.getNetwork());
             }
             resolveResponses.add(knownResponse);
         }
         return resolveResponses;
     }
 
+    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;
+
+        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 QuerySentArguments {
+        private final int transactionId;
+        private final List<String> subTypes = new ArrayList<>();
+        private final ScheduledQueryTaskArgs taskArgs;
+
+        QuerySentArguments(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 QueryTaskConfig config;
+        private final ScheduledQueryTaskArgs taskArgs;
+        private final List<MdnsResponse> servicesToResolve = new ArrayList<>();
+        private final boolean sendDiscoveryQueries;
 
-        QueryTask(@NonNull QueryTaskConfig config) {
-            this.config = config;
+        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.
-                // Only the names are used to know which queries to send, other parameters like
-                // interfaceIndex do not matter.
-                servicesToResolve = makeResponsesForResolve(
-                        0 /* interfaceIndex */, config.socketKey.getNetwork());
-                sendDiscoveryQueries = servicesToResolve.size() < listeners.size();
-            }
             Pair<Integer, List<String>> result;
             try {
                 result =
@@ -561,83 +650,79 @@
                                 socketClient,
                                 createMdnsPacketWriter(),
                                 serviceType,
-                                config.subtypes,
-                                config.expectUnicastResponse,
-                                config.transactionId,
-                                config.socketKey.getNetwork(),
-                                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(INVALID_TRANSACTION_ID, 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 (config.sessionId != currentSessionId) {
-                        return;
-                    }
-                }
-
-                if (MdnsConfigs.shouldCancelScanTaskWhenFutureIsNull()) {
-                    if (requestTaskFuture == 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);
-                                }
-                            }
-                        }
-                    }
-                }
-                requestTaskFuture = scheduleNextRunLocked(this.config);
-            }
+            dependencies.sendMessage(
+                    handler, handler.obtainMessage(EVENT_QUERY_RESULT,
+                            new QuerySentArguments(result.first, result.second, taskArgs)));
         }
     }
 
+    private static long calculateTimeToRun(@NonNull ScheduledQueryTaskArgs taskArgs,
+            QueryTaskConfig queryTaskConfig, long now, long minRemainingTtl, long lastSentTime) {
+        final long baseDelayInMs = queryTaskConfig.delayUntilNextTaskWithoutBackoffMs;
+        if (!queryTaskConfig.shouldUseQueryBackoff()) {
+            return lastSentTime + baseDelayInMs;
+        }
+        if (minRemainingTtl <= 0) {
+            // There's no service, or there is an expired service. In any case, schedule for the
+            // minimum time, which is the base delay.
+            return lastSentTime + baseDelayInMs;
+        }
+        // If the next TTL expiration time hasn't changed, then use previous calculated timeToRun.
+        if (lastSentTime < now
+                && 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 taskArgs.timeToRun;
+        }
+        return Math.max(now + (long) (0.8 * minRemainingTtl), lastSentTime + baseDelayInMs);
+    }
+
+    private long getMinRemainingTtl(long now) {
+        long minRemainingTtl = Long.MAX_VALUE;
+        for (MdnsResponse response : serviceCache.getCachedServices(serviceType, socketKey)) {
+            if (!response.isComplete()) {
+                continue;
+            }
+            long remainingTtl =
+                    response.getServiceRecord().getRemainingTTL(now);
+            // remainingTtl is <= 0 means the service expired.
+            if (remainingTtl <= 0) {
+                return 0;
+            }
+            if (remainingTtl < minRemainingTtl) {
+                minRemainingTtl = remainingTtl;
+            }
+        }
+        return minRemainingTtl == Long.MAX_VALUE ? 0 : minRemainingTtl;
+    }
+
     @NonNull
-    private Future<?> scheduleNextRunLocked(@NonNull QueryTaskConfig lastRunConfig) {
-        QueryTaskConfig config = lastRunConfig.getConfigForNextRun();
-        return executor.schedule(new QueryTask(config), config.timeToRunNextTaskInMs, MILLISECONDS);
+    private void scheduleNextRun(@NonNull QueryTaskConfig nextRunConfig,
+            long minRemainingTtl,
+            long timeWhenScheduled, long timeToRun, long sessionId) {
+        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",
+                lastScheduledQueryTaskArgs.sessionId, timeToNextTasksWithBackoffInMs));
+        dependencies.sendMessageDelayed(
+                handler,
+                handler.obtainMessage(EVENT_START_QUERYTASK, lastScheduledQueryTaskArgs),
+                timeToNextTasksWithBackoffInMs);
     }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
index 5e4a8b5..b6000f0 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
@@ -82,6 +82,6 @@
         void onSocketCreated(@NonNull SocketKey socketKey);
 
         /*** Notify requested socket is destroyed */
-        void onAllSocketsDestroyed(@NonNull SocketKey socketKey);
+        void onSocketDestroyed(@NonNull SocketKey socketKey);
     }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketProvider.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketProvider.java
index 3df6313..6925b49 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketProvider.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketProvider.java
@@ -82,7 +82,7 @@
     @NonNull private final Dependencies mDependencies;
     @NonNull private final NetworkCallback mNetworkCallback;
     @NonNull private final TetheringEventCallback mTetheringEventCallback;
-    @NonNull private final AbstractSocketNetlink mSocketNetlinkMonitor;
+    @NonNull private final AbstractSocketNetlinkMonitor mSocketNetlinkMonitor;
     @NonNull private final SharedLog mSharedLog;
     private final ArrayMap<Network, SocketInfo> mNetworkSockets = new ArrayMap<>();
     private final ArrayMap<String, SocketInfo> mTetherInterfaceSockets = new ArrayMap<>();
@@ -253,16 +253,12 @@
             return iface.getIndex();
         }
         /*** Creates a SocketNetlinkMonitor */
-        public AbstractSocketNetlink createSocketNetlinkMonitor(@NonNull final Handler handler,
+        public AbstractSocketNetlinkMonitor createSocketNetlinkMonitor(
+                @NonNull final Handler handler,
                 @NonNull final SharedLog log,
                 @NonNull final NetLinkMonitorCallBack cb) {
             return SocketNetLinkMonitorFactory.createNetLinkMonitor(handler, log, cb);
         }
-
-        /*** Get interface index by given socket */
-        public int getInterfaceIndex(@NonNull MdnsInterfaceSocket socket) {
-            return socket.getInterface().getIndex();
-        }
     }
     /**
      * The callback interface for the netlink monitor messages.
@@ -323,11 +319,14 @@
         final MdnsInterfaceSocket mSocket;
         final List<LinkAddress> mAddresses;
         final int[] mTransports;
+        @NonNull final SocketKey mSocketKey;
 
-        SocketInfo(MdnsInterfaceSocket socket, List<LinkAddress> addresses, int[] transports) {
+        SocketInfo(MdnsInterfaceSocket socket, List<LinkAddress> addresses, int[] transports,
+                @NonNull SocketKey socketKey) {
             mSocket = socket;
             mAddresses = new ArrayList<>(addresses);
             mTransports = transports;
+            mSocketKey = socketKey;
         }
     }
 
@@ -447,7 +446,7 @@
         // Try to join the group again.
         socketInfo.mSocket.joinGroup(addresses);
 
-        notifyAddressesChanged(network, socketInfo.mSocket, addresses);
+        notifyAddressesChanged(network, socketInfo, addresses);
     }
     private LinkProperties createLPForTetheredInterface(@NonNull final String interfaceName,
             int ifaceIndex) {
@@ -528,21 +527,22 @@
                     networkInterface.getNetworkInterface(), MdnsConstants.MDNS_PORT, mLooper,
                     mPacketReadBuffer);
             final List<LinkAddress> addresses = lp.getLinkAddresses();
+            final Network network =
+                    networkKey == LOCAL_NET ? null : ((NetworkAsKey) networkKey).mNetwork;
+            final SocketKey socketKey = new SocketKey(network, networkInterface.getIndex());
             // TODO: technically transport types are mutable, although generally not in ways that
             // would meaningfully impact the logic using it here. Consider updating logic to
             // support transports being added/removed.
-            final SocketInfo socketInfo = new SocketInfo(socket, addresses, transports);
+            final SocketInfo socketInfo = new SocketInfo(socket, addresses, transports, socketKey);
             if (networkKey == LOCAL_NET) {
                 mTetherInterfaceSockets.put(interfaceName, socketInfo);
             } else {
-                mNetworkSockets.put(((NetworkAsKey) networkKey).mNetwork, socketInfo);
+                mNetworkSockets.put(network, socketInfo);
             }
             // Try to join IPv4/IPv6 group.
             socket.joinGroup(addresses);
 
             // Notify the listeners which need this socket.
-            final Network network =
-                    networkKey == LOCAL_NET ? null : ((NetworkAsKey) networkKey).mNetwork;
             notifySocketCreated(network, socketInfo);
         } catch (IOException e) {
             mSharedLog.e("Create socket failed ifName:" + interfaceName, e);
@@ -584,7 +584,7 @@
         if (socketInfo == null) return;
 
         socketInfo.mSocket.destroy();
-        notifyInterfaceDestroyed(network, socketInfo.mSocket);
+        notifyInterfaceDestroyed(network, socketInfo);
         mSocketRequestMonitor.onSocketDestroyed(network, socketInfo.mSocket);
         mSharedLog.log("Remove socket on net:" + network);
     }
@@ -593,7 +593,7 @@
         final SocketInfo socketInfo = mTetherInterfaceSockets.remove(interfaceName);
         if (socketInfo == null) return;
         socketInfo.mSocket.destroy();
-        notifyInterfaceDestroyed(null /* network */, socketInfo.mSocket);
+        notifyInterfaceDestroyed(null /* network */, socketInfo);
         mSocketRequestMonitor.onSocketDestroyed(null /* network */, socketInfo.mSocket);
         mSharedLog.log("Remove socket on ifName:" + interfaceName);
     }
@@ -602,9 +602,7 @@
         for (int i = 0; i < mCallbacksToRequestedNetworks.size(); i++) {
             final Network requestedNetwork = mCallbacksToRequestedNetworks.valueAt(i);
             if (isNetworkMatched(requestedNetwork, network)) {
-                final int ifaceIndex = mDependencies.getInterfaceIndex(socketInfo.mSocket);
-                final SocketKey socketKey = new SocketKey(network, ifaceIndex);
-                mCallbacksToRequestedNetworks.keyAt(i).onSocketCreated(socketKey,
+                mCallbacksToRequestedNetworks.keyAt(i).onSocketCreated(socketInfo.mSocketKey,
                         socketInfo.mSocket, socketInfo.mAddresses);
                 mSocketRequestMonitor.onSocketRequestFulfilled(network, socketInfo.mSocket,
                         socketInfo.mTransports);
@@ -612,25 +610,23 @@
         }
     }
 
-    private void notifyInterfaceDestroyed(Network network, MdnsInterfaceSocket socket) {
+    private void notifyInterfaceDestroyed(Network network, SocketInfo socketInfo) {
         for (int i = 0; i < mCallbacksToRequestedNetworks.size(); i++) {
             final Network requestedNetwork = mCallbacksToRequestedNetworks.valueAt(i);
             if (isNetworkMatched(requestedNetwork, network)) {
-                final int ifaceIndex = mDependencies.getInterfaceIndex(socket);
                 mCallbacksToRequestedNetworks.keyAt(i)
-                        .onInterfaceDestroyed(new SocketKey(network, ifaceIndex), socket);
+                        .onInterfaceDestroyed(socketInfo.mSocketKey, socketInfo.mSocket);
             }
         }
     }
 
-    private void notifyAddressesChanged(Network network, MdnsInterfaceSocket socket,
+    private void notifyAddressesChanged(Network network, SocketInfo socketInfo,
             List<LinkAddress> addresses) {
         for (int i = 0; i < mCallbacksToRequestedNetworks.size(); i++) {
             final Network requestedNetwork = mCallbacksToRequestedNetworks.valueAt(i);
             if (isNetworkMatched(requestedNetwork, network)) {
-                final int ifaceIndex = mDependencies.getInterfaceIndex(socket);
                 mCallbacksToRequestedNetworks.keyAt(i)
-                        .onAddressesChanged(new SocketKey(network, ifaceIndex), socket, addresses);
+                        .onAddressesChanged(socketInfo.mSocketKey, socketInfo.mSocket, addresses);
             }
         }
     }
@@ -647,9 +643,7 @@
             createSocket(new NetworkAsKey(network), lp);
         } else {
             // Notify the socket for requested network.
-            final int ifaceIndex = mDependencies.getInterfaceIndex(socketInfo.mSocket);
-            final SocketKey socketKey = new SocketKey(network, ifaceIndex);
-            cb.onSocketCreated(socketKey, socketInfo.mSocket, socketInfo.mAddresses);
+            cb.onSocketCreated(socketInfo.mSocketKey, socketInfo.mSocket, socketInfo.mAddresses);
             mSocketRequestMonitor.onSocketRequestFulfilled(network, socketInfo.mSocket,
                     socketInfo.mTransports);
         }
@@ -664,9 +658,7 @@
                     createLPForTetheredInterface(interfaceName, ifaceIndex));
         } else {
             // Notify the socket for requested network.
-            final int ifaceIndex = mDependencies.getInterfaceIndex(socketInfo.mSocket);
-            final SocketKey socketKey = new SocketKey(ifaceIndex);
-            cb.onSocketCreated(socketKey, socketInfo.mSocket, socketInfo.mAddresses);
+            cb.onSocketCreated(socketInfo.mSocketKey, socketInfo.mSocket, socketInfo.mAddresses);
             mSocketRequestMonitor.onSocketRequestFulfilled(null /* socketNetwork */,
                     socketInfo.mSocket, socketInfo.mTransports);
         }
diff --git a/service-t/src/com/android/server/connectivity/mdns/NetworkInterfaceWrapper.java b/service-t/src/com/android/server/connectivity/mdns/NetworkInterfaceWrapper.java
index 0ecae48..48c396e 100644
--- a/service-t/src/com/android/server/connectivity/mdns/NetworkInterfaceWrapper.java
+++ b/service-t/src/com/android/server/connectivity/mdns/NetworkInterfaceWrapper.java
@@ -57,6 +57,10 @@
         return networkInterface.getInterfaceAddresses();
     }
 
+    public int getIndex() {
+        return networkInterface.getIndex();
+    }
+
     @Override
     public String toString() {
         return networkInterface.toString();
diff --git a/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java b/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
new file mode 100644
index 0000000..19282b0
--- /dev/null
+++ b/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
@@ -0,0 +1,172 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.connectivity.mdns;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+/**
+ * A configuration for the PeriodicalQueryTask that contains parameters to build a query packet.
+ * Call to getConfigForNextRun returns a config that can be used to build the next query task.
+ */
+public class QueryTaskConfig {
+
+    private static final int INITIAL_TIME_BETWEEN_BURSTS_MS =
+            (int) MdnsConfigs.initialTimeBetweenBurstsMs();
+    private static final int TIME_BETWEEN_BURSTS_MS = (int) MdnsConfigs.timeBetweenBurstsMs();
+    private static final int QUERIES_PER_BURST = (int) MdnsConfigs.queriesPerBurst();
+    private static final int TIME_BETWEEN_QUERIES_IN_BURST_MS =
+            (int) MdnsConfigs.timeBetweenQueriesInBurstMs();
+    private static final int QUERIES_PER_BURST_PASSIVE_MODE =
+            (int) MdnsConfigs.queriesPerBurstPassive();
+    private static final int UNSIGNED_SHORT_MAX_VALUE = 65536;
+    // The following fields are used by QueryTask so we need to test them.
+    @VisibleForTesting
+    final List<String> subtypes;
+    private final boolean alwaysAskForUnicastResponse =
+            MdnsConfigs.alwaysAskForUnicastResponseInEachBurst();
+    private final boolean usePassiveMode;
+    final boolean onlyUseIpv6OnIpv6OnlyNetworks;
+    private final int numOfQueriesBeforeBackoff;
+    @VisibleForTesting
+    final int transactionId;
+    @VisibleForTesting
+    final boolean expectUnicastResponse;
+    private final int queriesPerBurst;
+    private final int timeBetweenBurstsInMs;
+    private final int burstCounter;
+    final long delayUntilNextTaskWithoutBackoffMs;
+    private final boolean isFirstBurst;
+    private final long queryCount;
+    @NonNull
+    final SocketKey socketKey;
+
+    QueryTaskConfig(@NonNull QueryTaskConfig other, long queryCount, int transactionId,
+            boolean expectUnicastResponse, boolean isFirstBurst, int burstCounter,
+            int queriesPerBurst, int timeBetweenBurstsInMs,
+            long delayUntilNextTaskWithoutBackoffMs) {
+        this.subtypes = new ArrayList<>(other.subtypes);
+        this.usePassiveMode = other.usePassiveMode;
+        this.onlyUseIpv6OnIpv6OnlyNetworks = other.onlyUseIpv6OnIpv6OnlyNetworks;
+        this.numOfQueriesBeforeBackoff = other.numOfQueriesBeforeBackoff;
+        this.transactionId = transactionId;
+        this.expectUnicastResponse = expectUnicastResponse;
+        this.queriesPerBurst = queriesPerBurst;
+        this.timeBetweenBurstsInMs = timeBetweenBurstsInMs;
+        this.burstCounter = burstCounter;
+        this.delayUntilNextTaskWithoutBackoffMs = delayUntilNextTaskWithoutBackoffMs;
+        this.isFirstBurst = isFirstBurst;
+        this.queryCount = queryCount;
+        this.socketKey = other.socketKey;
+    }
+    QueryTaskConfig(@NonNull Collection<String> subtypes,
+            boolean usePassiveMode,
+            boolean onlyUseIpv6OnIpv6OnlyNetworks,
+            int numOfQueriesBeforeBackoff,
+            @Nullable SocketKey socketKey) {
+        this.usePassiveMode = usePassiveMode;
+        this.onlyUseIpv6OnIpv6OnlyNetworks = onlyUseIpv6OnIpv6OnlyNetworks;
+        this.numOfQueriesBeforeBackoff = numOfQueriesBeforeBackoff;
+        this.subtypes = new ArrayList<>(subtypes);
+        this.queriesPerBurst = QUERIES_PER_BURST;
+        this.burstCounter = 0;
+        this.transactionId = 1;
+        this.expectUnicastResponse = true;
+        this.isFirstBurst = true;
+        // Config the scan frequency based on the scan mode.
+        if (this.usePassiveMode) {
+            // In passive scan mode, sends a single burst of QUERIES_PER_BURST queries, and then
+            // in each TIME_BETWEEN_BURSTS interval, sends QUERIES_PER_BURST_PASSIVE_MODE
+            // queries.
+            this.timeBetweenBurstsInMs = TIME_BETWEEN_BURSTS_MS;
+        } else {
+            // In active scan mode, sends a burst of QUERIES_PER_BURST queries,
+            // TIME_BETWEEN_QUERIES_IN_BURST_MS apart, then waits for the scan interval, and
+            // then repeats. The scan interval starts as INITIAL_TIME_BETWEEN_BURSTS_MS and
+            // doubles until it maxes out at TIME_BETWEEN_BURSTS_MS.
+            this.timeBetweenBurstsInMs = INITIAL_TIME_BETWEEN_BURSTS_MS;
+        }
+        this.socketKey = socketKey;
+        this.queryCount = 0;
+        this.delayUntilNextTaskWithoutBackoffMs = TIME_BETWEEN_QUERIES_IN_BURST_MS;
+    }
+
+    /**
+     * Get new QueryTaskConfig for next run.
+     */
+    public QueryTaskConfig getConfigForNextRun() {
+        long newQueryCount = queryCount + 1;
+        int newTransactionId = transactionId + 1;
+        if (newTransactionId > UNSIGNED_SHORT_MAX_VALUE) {
+            newTransactionId = 1;
+        }
+        boolean newExpectUnicastResponse = false;
+        boolean newIsFirstBurst = isFirstBurst;
+        int newQueriesPerBurst = queriesPerBurst;
+        int newBurstCounter = burstCounter + 1;
+        long newDelayUntilNextTaskWithoutBackoffMs = delayUntilNextTaskWithoutBackoffMs;
+        int newTimeBetweenBurstsInMs = timeBetweenBurstsInMs;
+        // Only the first query expects uni-cast response.
+        if (newBurstCounter == queriesPerBurst) {
+            newBurstCounter = 0;
+
+            if (alwaysAskForUnicastResponse) {
+                newExpectUnicastResponse = true;
+            }
+            // In passive scan mode, sends a single burst of QUERIES_PER_BURST queries, and
+            // then in each TIME_BETWEEN_BURSTS interval, sends QUERIES_PER_BURST_PASSIVE_MODE
+            // queries.
+            if (isFirstBurst) {
+                newIsFirstBurst = false;
+                if (usePassiveMode) {
+                    newQueriesPerBurst = QUERIES_PER_BURST_PASSIVE_MODE;
+                }
+            }
+            // In active scan mode, sends a burst of QUERIES_PER_BURST queries,
+            // TIME_BETWEEN_QUERIES_IN_BURST_MS apart, then waits for the scan interval, and
+            // then repeats. The scan interval starts as INITIAL_TIME_BETWEEN_BURSTS_MS and
+            // doubles until it maxes out at TIME_BETWEEN_BURSTS_MS.
+            newDelayUntilNextTaskWithoutBackoffMs = timeBetweenBurstsInMs;
+            if (timeBetweenBurstsInMs < TIME_BETWEEN_BURSTS_MS) {
+                newTimeBetweenBurstsInMs = Math.min(timeBetweenBurstsInMs * 2,
+                        TIME_BETWEEN_BURSTS_MS);
+            }
+        } else {
+            newDelayUntilNextTaskWithoutBackoffMs = TIME_BETWEEN_QUERIES_IN_BURST_MS;
+        }
+        return new QueryTaskConfig(this, newQueryCount, newTransactionId,
+                newExpectUnicastResponse, newIsFirstBurst, newBurstCounter, newQueriesPerBurst,
+                newTimeBetweenBurstsInMs, newDelayUntilNextTaskWithoutBackoffMs);
+    }
+
+    /**
+     * Determine if the query backoff should be used.
+     */
+    public boolean shouldUseQueryBackoff() {
+        // Don't enable backoff mode during the burst or in the first burst
+        if (burstCounter != 0 || isFirstBurst) {
+            return false;
+        }
+        return queryCount > numOfQueriesBeforeBackoff;
+    }
+}
diff --git a/service-t/src/com/android/server/connectivity/mdns/SocketKey.java b/service-t/src/com/android/server/connectivity/mdns/SocketKey.java
index a893acb..f13d0e0 100644
--- a/service-t/src/com/android/server/connectivity/mdns/SocketKey.java
+++ b/service-t/src/com/android/server/connectivity/mdns/SocketKey.java
@@ -43,6 +43,7 @@
         mInterfaceIndex = interfaceIndex;
     }
 
+    @Nullable
     public Network getNetwork() {
         return mNetwork;
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/SocketNetLinkMonitorFactory.java b/service-t/src/com/android/server/connectivity/mdns/SocketNetLinkMonitorFactory.java
index 6bc7941..77c8f9c 100644
--- a/service-t/src/com/android/server/connectivity/mdns/SocketNetLinkMonitorFactory.java
+++ b/service-t/src/com/android/server/connectivity/mdns/SocketNetLinkMonitorFactory.java
@@ -30,7 +30,7 @@
     /**
      * Creates a new netlink monitor.
      */
-    public static AbstractSocketNetlink createNetLinkMonitor(@NonNull final Handler handler,
+    public static AbstractSocketNetlinkMonitor createNetLinkMonitor(@NonNull final Handler handler,
             @NonNull SharedLog log, @NonNull MdnsSocketProvider.NetLinkMonitorCallBack cb) {
         return new SocketNetlinkMonitor(handler, log, cb);
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/internal/SocketNetlinkMonitor.java b/service-t/src/com/android/server/connectivity/mdns/internal/SocketNetlinkMonitor.java
index 451909c..c21c903 100644
--- a/service-t/src/com/android/server/connectivity/mdns/internal/SocketNetlinkMonitor.java
+++ b/service-t/src/com/android/server/connectivity/mdns/internal/SocketNetlinkMonitor.java
@@ -28,13 +28,13 @@
 import com.android.net.module.util.netlink.NetlinkMessage;
 import com.android.net.module.util.netlink.RtNetlinkAddressMessage;
 import com.android.net.module.util.netlink.StructIfaddrMsg;
-import com.android.server.connectivity.mdns.AbstractSocketNetlink;
+import com.android.server.connectivity.mdns.AbstractSocketNetlinkMonitor;
 import com.android.server.connectivity.mdns.MdnsSocketProvider;
 
 /**
  * The netlink monitor for MdnsSocketProvider.
  */
-public class SocketNetlinkMonitor extends NetlinkMonitor implements AbstractSocketNetlink {
+public class SocketNetlinkMonitor extends NetlinkMonitor implements AbstractSocketNetlinkMonitor {
 
     public static final String TAG = SocketNetlinkMonitor.class.getSimpleName();
 
diff --git a/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java b/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
index 3180a6f..df3bde8 100644
--- a/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
+++ b/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
@@ -20,6 +20,7 @@
 import android.annotation.Nullable;
 import android.net.Network;
 import android.os.Handler;
+import android.os.SystemClock;
 import android.util.ArraySet;
 
 import com.android.server.connectivity.mdns.MdnsConstants;
@@ -173,4 +174,14 @@
         return mdnsRecord.getTtl() > 0
                 && mdnsRecord.getRemainingTTL(now) <= mdnsRecord.getTtl() / 2;
     }
+
+    /** A wrapper class of {@link SystemClock} to be mocked in unit tests. */
+    public static class Clock {
+        /**
+         * @see SystemClock#elapsedRealtime
+         */
+        public long elapsedRealtime() {
+            return SystemClock.elapsedRealtime();
+        }
+    }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/ethernet/EthernetTracker.java b/service-t/src/com/android/server/ethernet/EthernetTracker.java
index 1f22b02..8141350 100644
--- a/service-t/src/com/android/server/ethernet/EthernetTracker.java
+++ b/service-t/src/com/android/server/ethernet/EthernetTracker.java
@@ -431,7 +431,7 @@
             for (String iface : getClientModeInterfaces(canUseRestrictedNetworks)) {
                 unicastInterfaceStateChange(listener, iface);
             }
-            if (mTetheringInterfaceMode == INTERFACE_MODE_SERVER) {
+            if (mTetheringInterface != null && mTetheringInterfaceMode == INTERFACE_MODE_SERVER) {
                 unicastInterfaceStateChange(listener, mTetheringInterface);
             }
 
diff --git a/service/ServiceConnectivityResources/res/values/config.xml b/service/ServiceConnectivityResources/res/values/config.xml
index 22d9b01..f30abc6 100644
--- a/service/ServiceConnectivityResources/res/values/config.xml
+++ b/service/ServiceConnectivityResources/res/values/config.xml
@@ -135,10 +135,17 @@
     <!-- Whether to cancel network notifications automatically when tapped -->
     <bool name="config_autoCancelNetworkNotifications">true</bool>
 
-    <!-- When no internet or partial connectivity is detected on a network, and a high priority
-         (heads up) notification would be shown due to the network being explicitly selected,
-         directly show the dialog that would normally be shown when tapping the notification
-         instead of showing the notification. -->
+    <!-- Configuration to let OEMs customize what to do when :
+         • Partial connectivity is detected on the network
+         • No internet is detected on the network, and
+           - the network was explicitly selected
+           - the system is configured to actively prefer bad wifi (see config_activelyPreferBadWifi)
+         The default behavior (false) is to post a notification with a PendingIntent so
+         the user is informed and can act if they wish.
+         Making this true instead will have the system fire the intent immediately instead
+         of showing a notification. OEMs who do this should have some intent receiver
+         listening to the intent and take the action they prefer (e.g. show a dialog,
+         show a customized notification etc).  -->
     <bool name="config_notifyNoInternetAsDialogWhenHighPriority">false</bool>
 
     <!-- When showing notifications indicating partial connectivity, display the same notifications
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index ec168dd..7aff6a4 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -40,6 +40,7 @@
 import android.app.StatsManager;
 import android.content.Context;
 import android.net.INetd;
+import android.os.Build;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
 import android.provider.DeviceConfig;
@@ -51,6 +52,8 @@
 import android.util.Pair;
 import android.util.StatsEvent;
 
+import androidx.annotation.RequiresApi;
+
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.modules.utils.BackgroundThread;
 import com.android.modules.utils.build.SdkLevel;
@@ -1140,19 +1143,48 @@
         }
     }
 
+    @RequiresApi(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
     private static native void native_init(boolean startSkDestroyListener);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_addNaughtyApp(int uid);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_removeNaughtyApp(int uid);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_addNiceApp(int uid);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_removeNiceApp(int uid);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_setChildChain(int childChain, boolean enable);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_replaceUidChain(String name, boolean isAllowlist, int[] uids);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_setUidRule(int childChain, int uid, int firewallRule);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_addUidInterfaceRules(String ifName, int[] uids);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_removeUidInterfaceRules(int[] uids);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_updateUidLockdownRule(int uid, boolean add);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native int native_swapActiveStatsMap();
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private native void native_setPermissionForUids(int permissions, int[] uids);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static native void native_dump(FileDescriptor fd, boolean verbose);
+
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static native int native_synchronizeKernelRCU();
 }
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 53127d1..83afa83 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -318,6 +318,7 @@
 import java.io.InterruptedIOException;
 import java.io.PrintWriter;
 import java.io.Writer;
+import java.lang.IllegalArgumentException;
 import java.net.Inet4Address;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
@@ -389,7 +390,7 @@
     // Timeout in case the "actively prefer bad wifi" feature is on
     private static final int ACTIVELY_PREFER_BAD_WIFI_INITIAL_TIMEOUT_MS = 20 * 1000;
     // Timeout in case the "actively prefer bad wifi" feature is off
-    private static final int DONT_ACTIVELY_PREFER_BAD_WIFI_INITIAL_TIMEOUT_MS = 8 * 1000;
+    private static final int DEFAULT_EVALUATION_TIMEOUT_MS = 8 * 1000;
 
     // Default to 30s linger time-out, and 5s for nascent network. Modifiable only for testing.
     private static final String LINGER_DELAY_PROPERTY = "persist.netmon.linger";
@@ -4010,7 +4011,7 @@
                     // the destroyed flag is only just above the "current satisfier wins"
                     // tie-breaker. But technically anything that affects scoring should rematch.
                     rematchAllNetworksAndRequests();
-                    mHandler.postDelayed(() -> disconnectAndDestroyNetwork(nai), timeoutMs);
+                    mHandler.postDelayed(() -> nai.disconnect(), timeoutMs);
                     break;
                 }
             }
@@ -4547,9 +4548,11 @@
 
     @VisibleForTesting
     protected static boolean shouldCreateNetworksImmediately() {
-        // Before U, physical networks are only created when the agent advances to CONNECTED.
-        // In U and above, all networks are immediately created when the agent is registered.
-        return SdkLevel.isAtLeastU();
+        // The feature of creating the networks immediately was slated for U, but race conditions
+        // detected late required this was flagged off.
+        // TODO : enable this in a Mainline update or in V, and re-enable the test for this
+        // in NetworkAgentTest.
+        return false;
     }
 
     private static boolean shouldCreateNativeNetwork(@NonNull NetworkAgentInfo nai,
@@ -4609,9 +4612,6 @@
         if (DBG) {
             log(nai.toShortString() + " disconnected, was satisfying " + nai.numNetworkRequests());
         }
-
-        nai.disconnect();
-
         // Clear all notifications of this network.
         mNotifier.clearNotification(nai.network.getNetId());
         // A network agent has disconnected.
@@ -5897,7 +5897,7 @@
                     final NetworkAgentInfo nai = getNetworkAgentInfoForNetwork((Network) msg.obj);
                     if (nai == null) break;
                     nai.onPreventAutomaticReconnect();
-                    disconnectAndDestroyNetwork(nai);
+                    nai.disconnect();
                     break;
                 case EVENT_SET_VPN_NETWORK_PREFERENCE:
                     handleSetVpnNetworkPreference((VpnNetworkPreferenceInfo) msg.obj);
@@ -9042,7 +9042,7 @@
                 break;
             }
         }
-        disconnectAndDestroyNetwork(nai);
+        nai.disconnect();
     }
 
     private void handleLingerComplete(NetworkAgentInfo oldNetwork) {
@@ -9584,10 +9584,7 @@
         updateLegacyTypeTrackerAndVpnLockdownForRematch(changes, nais);
 
         // Tear down all unneeded networks.
-        // Iterate in reverse order because teardownUnneededNetwork removes the nai from
-        // mNetworkAgentInfos.
-        for (int i = mNetworkAgentInfos.size() - 1; i >= 0; i--) {
-            final NetworkAgentInfo nai = mNetworkAgentInfos.valueAt(i);
+        for (NetworkAgentInfo nai : mNetworkAgentInfos) {
             if (unneeded(nai, UnneededFor.TEARDOWN)) {
                 if (nai.getInactivityExpiry() > 0) {
                     // This network has active linger timers and no requests, but is not
@@ -9940,10 +9937,25 @@
                 networkAgent.networkMonitor().notifyNetworkConnected(params.linkProperties,
                         params.networkCapabilities);
             }
-            final long delay = !avoidBadWifi() && activelyPreferBadWifi()
-                    ? ACTIVELY_PREFER_BAD_WIFI_INITIAL_TIMEOUT_MS
-                    : DONT_ACTIVELY_PREFER_BAD_WIFI_INITIAL_TIMEOUT_MS;
-            scheduleEvaluationTimeout(networkAgent.network, delay);
+            final long evaluationDelay;
+            if (!networkAgent.networkCapabilities.hasSingleTransport(TRANSPORT_WIFI)) {
+                // If the network is anything other than pure wifi, use the default timeout.
+                evaluationDelay = DEFAULT_EVALUATION_TIMEOUT_MS;
+            } else if (networkAgent.networkAgentConfig.isExplicitlySelected()) {
+                // If the network is explicitly selected, use the default timeout because it's
+                // shorter and the user is likely staring at the screen expecting it to validate
+                // right away.
+                evaluationDelay = DEFAULT_EVALUATION_TIMEOUT_MS;
+            } else if (avoidBadWifi() || !activelyPreferBadWifi()) {
+                // If avoiding bad wifi, or if not avoiding but also not preferring bad wifi
+                evaluationDelay = DEFAULT_EVALUATION_TIMEOUT_MS;
+            } else {
+                // It's wifi, automatically connected, and bad wifi is preferred : use the
+                // longer timeout to avoid the device switching to captive portals with bad
+                // signal or very slow response.
+                evaluationDelay = ACTIVELY_PREFER_BAD_WIFI_INITIAL_TIMEOUT_MS;
+            }
+            scheduleEvaluationTimeout(networkAgent.network, evaluationDelay);
 
             // Whether a particular NetworkRequest listen should cause signal strength thresholds to
             // be communicated to a particular NetworkAgent depends only on the network's immutable,
@@ -9970,6 +9982,7 @@
             // This has to happen after matching the requests, because callbacks are just requests.
             notifyNetworkCallbacks(networkAgent, ConnectivityManager.CALLBACK_PRECHECK);
         } else if (state == NetworkInfo.State.DISCONNECTED) {
+            networkAgent.disconnect();
             if (networkAgent.isVPN()) {
                 updateVpnUids(networkAgent, networkAgent.networkCapabilities, null);
             }
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index 6ba2033..d03cac6 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -291,6 +291,18 @@
             }
         }
 
+        /**
+         * Construct a new AutomaticOnOffKeepalive from existing AutomaticOnOffKeepalive with a
+         * new KeepaliveInfo.
+         */
+        public AutomaticOnOffKeepalive withKeepaliveInfo(KeepaliveTracker.KeepaliveInfo ki)
+                throws InvalidSocketException {
+            return new AutomaticOnOffKeepalive(
+                    ki,
+                    mAutomaticOnOffState != STATE_ALWAYS_ON /* autoOnOff */,
+                    mUnderpinnedNetwork);
+        }
+
         @Override
         public String toString() {
             return "AutomaticOnOffKeepalive [ "
@@ -470,13 +482,29 @@
      * The message is expected to contain a KeepaliveTracker.KeepaliveInfo.
      */
     public void handleStartKeepalive(Message message) {
-        final AutomaticOnOffKeepalive autoKi = (AutomaticOnOffKeepalive) message.obj;
-        final int error = mKeepaliveTracker.handleStartKeepalive(autoKi.mKi);
+        final AutomaticOnOffKeepalive target = (AutomaticOnOffKeepalive) message.obj;
+        final Pair<Integer, KeepaliveTracker.KeepaliveInfo> res =
+                mKeepaliveTracker.handleStartKeepalive(target.mKi);
+        final int error = res.first;
         if (error != SUCCESS) {
-            mEventLog.log("Failed to start keepalive " + autoKi.mCallback + " on "
-                    + autoKi.getNetwork() + " with error " + error);
+            mEventLog.log("Failed to start keepalive " + target.mCallback + " on "
+                    + target.getNetwork() + " with error " + error);
             return;
         }
+        // Generate a new auto ki with the started keepalive info.
+        final AutomaticOnOffKeepalive autoKi;
+        try {
+            autoKi = target.withKeepaliveInfo(res.second);
+            // 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;
+        }
+
         mEventLog.log("Start keepalive " + autoKi.mCallback + " on " + autoKi.getNetwork());
         mKeepaliveStatsTracker.onStartKeepalive(
                 autoKi.getNetwork(),
@@ -506,14 +534,19 @@
      * @return SUCCESS if the keepalive is successfully starting and the error reason otherwise.
      */
     private int handleResumeKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
-        final int error = mKeepaliveTracker.handleStartKeepalive(ki);
+        final Pair<Integer, KeepaliveTracker.KeepaliveInfo> res =
+                mKeepaliveTracker.handleStartKeepalive(ki);
+        final KeepaliveTracker.KeepaliveInfo startedKi = res.second;
+        final int error = res.first;
         if (error != SUCCESS) {
-            mEventLog.log("Failed to resume keepalive " + ki.mCallback + " on " + ki.mNai
-                    + " with error " + error);
+            mEventLog.log("Failed to resume keepalive " + startedKi.mCallback + " on "
+                    + startedKi.mNai + " with error " + error);
             return error;
         }
-        mKeepaliveStatsTracker.onResumeKeepalive(ki.getNai().network(), ki.getSlot());
-        mEventLog.log("Resumed successfully keepalive " + ki.mCallback + " on " + ki.mNai);
+
+        mKeepaliveStatsTracker.onResumeKeepalive(startedKi.getNai().network(), startedKi.getSlot());
+        mEventLog.log("Resumed successfully keepalive " + startedKi.mCallback
+                + " on " + startedKi.mNai);
 
         return SUCCESS;
     }
diff --git a/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java b/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
index d59d526..414aca3 100644
--- a/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
@@ -45,6 +45,7 @@
 import com.android.metrics.KeepaliveLifetimeForCarrier;
 import com.android.metrics.KeepaliveLifetimePerCarrier;
 import com.android.modules.utils.BackgroundThread;
+import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.CollectionUtils;
 import com.android.server.ConnectivityStatsLog;
 
@@ -251,6 +252,22 @@
         public long getElapsedRealtime() {
             return SystemClock.elapsedRealtime();
         }
+
+        /**
+         * Writes a DAILY_KEEPALIVE_INFO_REPORTED to ConnectivityStatsLog.
+         *
+         * @param dailyKeepaliveInfoReported the proto to write to statsD.
+         */
+        public void writeStats(DailykeepaliveInfoReported dailyKeepaliveInfoReported) {
+            ConnectivityStatsLog.write(
+                    ConnectivityStatsLog.DAILY_KEEPALIVE_INFO_REPORTED,
+                    dailyKeepaliveInfoReported.getDurationPerNumOfKeepalive().toByteArray(),
+                    dailyKeepaliveInfoReported.getKeepaliveLifetimePerCarrier().toByteArray(),
+                    dailyKeepaliveInfoReported.getKeepaliveRequests(),
+                    dailyKeepaliveInfoReported.getAutomaticKeepaliveRequests(),
+                    dailyKeepaliveInfoReported.getDistinctUserCount(),
+                    CollectionUtils.toIntArray(dailyKeepaliveInfoReported.getUidList()));
+        }
     }
 
     public KeepaliveStatsTracker(@NonNull Context context, @NonNull Handler handler) {
@@ -637,15 +654,15 @@
     /** Writes the stored metrics to ConnectivityStatsLog and resets.  */
     public void writeAndResetMetrics() {
         ensureRunningOnHandlerThread();
+        // Keepalive stats use repeated atoms, which are only supported on T+. If written to statsd
+        // on S- they will bootloop the system, so they must not be sent on S-. See b/289471411.
+        if (!SdkLevel.isAtLeastT()) {
+            Log.d(TAG, "KeepaliveStatsTracker is disabled before T, skipping write");
+            return;
+        }
+
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported = buildAndResetMetrics();
-        ConnectivityStatsLog.write(
-                ConnectivityStatsLog.DAILY_KEEPALIVE_INFO_REPORTED,
-                dailyKeepaliveInfoReported.getDurationPerNumOfKeepalive().toByteArray(),
-                dailyKeepaliveInfoReported.getKeepaliveLifetimePerCarrier().toByteArray(),
-                dailyKeepaliveInfoReported.getKeepaliveRequests(),
-                dailyKeepaliveInfoReported.getAutomaticKeepaliveRequests(),
-                dailyKeepaliveInfoReported.getDistinctUserCount(),
-                CollectionUtils.toIntArray(dailyKeepaliveInfoReported.getUidList()));
+        mDependencies.writeStats(dailyKeepaliveInfoReported);
     }
 
     private void ensureRunningOnHandlerThread() {
diff --git a/service/src/com/android/server/connectivity/KeepaliveTracker.java b/service/src/com/android/server/connectivity/KeepaliveTracker.java
index 76e97e2..2ce186f 100644
--- a/service/src/com/android/server/connectivity/KeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveTracker.java
@@ -54,6 +54,7 @@
 import android.system.ErrnoException;
 import android.system.Os;
 import android.util.Log;
+import android.util.Pair;
 
 import com.android.connectivity.resources.R;
 import com.android.internal.annotations.VisibleForTesting;
@@ -62,6 +63,8 @@
 import com.android.net.module.util.IpUtils;
 
 import java.io.FileDescriptor;
+import java.net.Inet4Address;
+import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
@@ -292,11 +295,15 @@
 
         private int checkSourceAddress() {
             // Check that we have the source address.
-            for (InetAddress address : mNai.linkProperties.getAddresses()) {
+            for (InetAddress address : mNai.linkProperties.getAllAddresses()) {
                 if (address.equals(mPacket.getSrcAddress())) {
                     return SUCCESS;
                 }
             }
+            // Or the address is the clat source address.
+            if (mPacket.getSrcAddress().equals(mNai.getClatv6SrcAddress())) {
+                return SUCCESS;
+            }
             return ERROR_INVALID_IP_ADDRESS;
         }
 
@@ -479,6 +486,15 @@
             return new KeepaliveInfo(mCallback, mNai, mPacket, mPid, mUid, mInterval, mType,
                     fd, mSlot, true /* resumed */);
         }
+
+        /**
+         * Construct a new KeepaliveInfo from existing KeepaliveInfo with a new KeepalivePacketData.
+         */
+        public KeepaliveInfo withPacketData(@NonNull KeepalivePacketData packet)
+                throws InvalidSocketException {
+            return new KeepaliveInfo(mCallback, mNai, packet, mPid, mUid, mInterval, mType,
+                    mFd, mSlot, mResumed);
+        }
     }
 
     void notifyErrorCallback(ISocketKeepaliveCallback cb, int error) {
@@ -512,15 +528,49 @@
      * Handle start keepalives with the message.
      *
      * @param ki the keepalive to start.
-     * @return SUCCESS if the keepalive is successfully starting and the error reason otherwise.
+     * @return Pair of (SUCCESS if the keepalive is successfully starting and the error reason
+     *         otherwise, the started KeepaliveInfo object)
      */
-    public int handleStartKeepalive(KeepaliveInfo ki) {
-        NetworkAgentInfo nai = ki.getNai();
+    public Pair<Integer, KeepaliveInfo> handleStartKeepalive(KeepaliveInfo ki) {
+        final KeepaliveInfo newKi;
+        try {
+            newKi = handleUpdateKeepaliveForClat(ki);
+        } catch (InvalidSocketException | InvalidPacketException e) {
+            Log.e(TAG, "Fail to construct keepalive packet");
+            notifyErrorCallback(ki.mCallback, ERROR_INVALID_IP_ADDRESS);
+            // Fail to create new keepalive packet for clat. Return the original keepalive info.
+            return new Pair<>(ERROR_INVALID_IP_ADDRESS, ki);
+        }
+
+        final NetworkAgentInfo nai = newKi.getNai();
         // If this was a paused keepalive, then reuse the same slot that was kept for it. Otherwise,
         // use the first free slot for this network agent.
-        final int slot = NO_KEEPALIVE != ki.mSlot ? ki.mSlot : findFirstFreeSlot(nai);
-        mKeepalives.get(nai).put(slot, ki);
-        return ki.start(slot);
+        final int slot = NO_KEEPALIVE != newKi.mSlot ? newKi.mSlot : findFirstFreeSlot(nai);
+        mKeepalives.get(nai).put(slot, newKi);
+
+        return new Pair<>(newKi.start(slot), newKi);
+    }
+
+    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;
+
+        final InetAddress dstAddr = ki.mPacket.getDstAddress();
+        // Do not perform translation for a v6 dst address.
+        if (!(dstAddr instanceof Inet4Address)) return ki;
+
+        final Inet6Address address = ki.getNai().translateV4toClatV6((Inet4Address) dstAddr);
+
+        if (address == null) return ki;
+
+        final int srcPort = ki.mPacket.getSrcPort();
+        final KeepaliveInfo newInfo = ki.withPacketData(NattKeepalivePacketData.nattKeepalivePacket(
+                ki.getNai().getClatv6SrcAddress(), srcPort, address, NATT_PORT));
+        Log.d(TAG, "Src is clat v4 address. Convert from " + ki + " to " + newInfo);
+        return newInfo;
     }
 
     public void handleStopAllKeepalives(NetworkAgentInfo nai, int reason) {
diff --git a/service/src/com/android/server/connectivity/Nat464Xlat.java b/service/src/com/android/server/connectivity/Nat464Xlat.java
index b315235..f9e07fd 100644
--- a/service/src/com/android/server/connectivity/Nat464Xlat.java
+++ b/service/src/com/android/server/connectivity/Nat464Xlat.java
@@ -567,7 +567,7 @@
         try {
             return (Inet6Address) Inet6Address.getByAddress(v6Addr);
         } catch (UnknownHostException e) {
-            Log.e(TAG, "getByAddress should never throw for a numeric address");
+            Log.wtf(TAG, "getByAddress should never throw for a numeric address", e);
             return null;
         }
     }
@@ -583,6 +583,21 @@
         return mIPv6Address;
     }
 
+    /**
+     * Get the generated v4 address of clat.
+     */
+    @Nullable
+    public Inet4Address getClatv4SrcAddress() {
+        // Variables in Nat464Xlat should only be accessed from handler thread.
+        ensureRunningOnHandlerThread();
+        if (!isStarted()) return null;
+
+        final LinkAddress v4Addr = getLinkAddress(mIface);
+        if (v4Addr == null) return null;
+
+        return (Inet4Address) v4Addr.getAddress();
+    }
+
     private void ensureRunningOnHandlerThread() {
         if (mNetwork.handler().getLooper().getThread() != Thread.currentThread()) {
             throw new IllegalStateException(
diff --git a/service/src/com/android/server/connectivity/NetworkAgentInfo.java b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
index 08c1455..845c04c 100644
--- a/service/src/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
@@ -1043,6 +1043,14 @@
     }
 
     /**
+     * Get the generated v4 address of clat.
+     */
+    @Nullable
+    public Inet4Address getClatv4SrcAddress() {
+        return clatd.getClatv4SrcAddress();
+    }
+
+    /**
      * Translate the input v4 address to v6 clat address.
      */
     @Nullable
diff --git a/service/src/com/android/server/connectivity/NetworkNotificationManager.java b/service/src/com/android/server/connectivity/NetworkNotificationManager.java
index 8b0cb7c..bc13592 100644
--- a/service/src/com/android/server/connectivity/NetworkNotificationManager.java
+++ b/service/src/com/android/server/connectivity/NetworkNotificationManager.java
@@ -322,7 +322,8 @@
 
     private boolean maybeNotifyViaDialog(Resources res, NotificationType notifyType,
             PendingIntent intent) {
-        if (notifyType != NotificationType.NO_INTERNET
+        if (notifyType != NotificationType.LOST_INTERNET
+                && notifyType != NotificationType.NO_INTERNET
                 && notifyType != NotificationType.PARTIAL_CONNECTIVITY) {
             return false;
         }
@@ -432,7 +433,8 @@
      * A notification with a higher number will take priority over a notification with a lower
      * number.
      */
-    private static int priority(NotificationType t) {
+    @VisibleForTesting
+    public static int priority(NotificationType t) {
         if (t == null) {
             return 0;
         }
diff --git a/service/src/com/android/server/connectivity/ProxyTracker.java b/service/src/com/android/server/connectivity/ProxyTracker.java
index 6a0918b..4415007 100644
--- a/service/src/com/android/server/connectivity/ProxyTracker.java
+++ b/service/src/com/android/server/connectivity/ProxyTracker.java
@@ -404,7 +404,7 @@
                 // network, so discount this case.
                 if (null == mGlobalProxy && !lp.getHttpProxy().getPacFileUrl()
                         .equals(defaultProxy.getPacFileUrl())) {
-                    throw new IllegalStateException("Unexpected discrepancy between proxy in LP of "
+                    Log.wtf(TAG, "Unexpected discrepancy between proxy in LP of "
                             + "default network and default proxy. The former has a PAC URL of "
                             + lp.getHttpProxy().getPacFileUrl() + " while the latter has "
                             + defaultProxy.getPacFileUrl());
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
new file mode 100644
index 0000000..cb01732
--- /dev/null
+++ b/tests/benchmark/Android.bp
@@ -0,0 +1,41 @@
+//
+// Copyright (C) 2023 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package {
+    // See: http://go/android-license-faq
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_test {
+    name: "ConnectivityBenchmarkTests",
+    defaults: [
+        "framework-connectivity-internal-test-defaults",
+    ],
+    platform_apis: true,
+    srcs: [
+        "src/**/*.kt",
+        "src/**/*.aidl",
+    ],
+    static_libs: [
+        "androidx.test.rules",
+        "net-tests-utils",
+        "service-connectivity-pre-jarjar",
+        "service-connectivity-tiramisu-pre-jarjar",
+    ],
+    test_suites: ["device-tests"],
+    jarjar_rules: ":connectivity-jarjar-rules",
+}
+
diff --git a/tests/benchmark/AndroidManifest.xml b/tests/benchmark/AndroidManifest.xml
new file mode 100644
index 0000000..bd2fce5
--- /dev/null
+++ b/tests/benchmark/AndroidManifest.xml
@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2023 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+        package="com.android.server.connectivity.benchmarktests">
+    <application>
+        <uses-library android:name="android.test.runner" />
+    </application>
+
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+         android:targetPackage="com.android.server.connectivity.benchmarktests"
+         android:label="Connectivity Benchmark Tests" />
+</manifest>
diff --git a/tests/benchmark/OWNERS b/tests/benchmark/OWNERS
new file mode 100644
index 0000000..3101da5
--- /dev/null
+++ b/tests/benchmark/OWNERS
@@ -0,0 +1,2 @@
+# Bug template url: http://b/new?component=31808
+# TODO: move bug template config to common owners file once b/226427845 is resolved
\ No newline at end of file
diff --git a/tests/benchmark/res/raw/netstats-many-uids-zip b/tests/benchmark/res/raw/netstats-many-uids-zip
new file mode 100644
index 0000000..22e8254
--- /dev/null
+++ b/tests/benchmark/res/raw/netstats-many-uids-zip
Binary files differ
diff --git a/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt
new file mode 100644
index 0000000..177014f
--- /dev/null
+++ b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsCollectionTest.kt
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.net.benchmarktests
+
+import android.net.NetworkStatsCollection
+import androidx.test.InstrumentationRegistry
+import com.android.internal.util.FileRotator.Reader
+import com.android.server.connectivity.benchmarktests.R
+import java.io.BufferedInputStream
+import java.io.DataInputStream
+import java.io.File
+import java.io.FileInputStream
+import java.io.FileOutputStream
+import java.nio.file.Files
+import java.util.concurrent.TimeUnit
+import java.util.zip.ZipInputStream
+import kotlin.test.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class NetworkStatsCollectionTest {
+    private val DEFAULT_BUFFER_SIZE = 8192
+    private val UID_COLLECTION_BUCKET_DURATION_MS = TimeUnit.HOURS.toMillis(2)
+
+    private val uidTestFiles: List<File> by lazy {
+        // These file generated by using real user dataset which has many uid records and agreed to
+        // share the dataset for testing purpose. These dataset can be extracted from rooted
+        // devices by using "adb pull /data/misc/apexdata/com.android.tethering/netstats" command.
+        val zipInputStream = ZipInputStream(getInputStreamForResource(R.raw.netstats_many_uids_zip))
+        getSortedListForPrefix(unzipToTempDir(zipInputStream), "uid")
+    }
+
+    @Test
+    fun testReadCollection_manyUids() {
+        val collection = NetworkStatsCollection(UID_COLLECTION_BUCKET_DURATION_MS)
+        for (file in uidTestFiles) {
+            readFile(file, collection)
+        }
+    }
+
+    private fun getInputStreamForResource(resourceId: Int): DataInputStream {
+        return DataInputStream(
+            InstrumentationRegistry.getContext()
+                .getResources().openRawResource(resourceId)
+        )
+    }
+
+    private fun unzipToTempDir(zis: ZipInputStream): File {
+        val statsDir =
+            Files.createTempDirectory(NetworkStatsCollectionTest::class.simpleName).toFile()
+        while (true) {
+            val entryName = zis.nextEntry?.name ?: break
+            val file = File(statsDir, entryName)
+            FileOutputStream(file).use { zis.copyTo(it, DEFAULT_BUFFER_SIZE) }
+        }
+        return statsDir
+    }
+
+    // List [xt|uid|uid_tag].<start>-<end> files under the given directory.
+    private fun getSortedListForPrefix(statsDir: File, prefix: String): List<File> {
+        assertTrue(statsDir.exists())
+        return (statsDir.list() ?: arrayOf()).mapNotNull {
+            if (it.startsWith("$prefix.")) File(statsDir, it) else null
+        }.sorted()
+    }
+
+    private fun readFile(file: File, reader: Reader) =
+        BufferedInputStream(FileInputStream(file)).use {
+            reader.read(it)
+        }
+}
diff --git a/tests/common/java/android/net/NattKeepalivePacketDataTest.kt b/tests/common/java/android/net/NattKeepalivePacketDataTest.kt
index dde1d86..e5806a6 100644
--- a/tests/common/java/android/net/NattKeepalivePacketDataTest.kt
+++ b/tests/common/java/android/net/NattKeepalivePacketDataTest.kt
@@ -28,6 +28,7 @@
 import com.android.testutils.assertEqualBothWays
 import com.android.testutils.assertParcelingIsLossless
 import com.android.testutils.parcelingRoundTrip
+import java.net.Inet6Address
 import java.net.InetAddress
 import kotlin.test.assertFailsWith
 import org.junit.Assert.assertEquals
@@ -44,10 +45,33 @@
 
     private val TEST_PORT = 4243
     private val TEST_PORT2 = 4244
+    // ::FFFF:1.2.3.4
+    private val SRC_V4_MAPPED_V6_ADDRESS_BYTES = byteArrayOf(
+        0x00.toByte(),
+        0x00.toByte(),
+        0x00.toByte(),
+        0x00.toByte(),
+        0x00.toByte(),
+        0x00.toByte(),
+        0x00.toByte(),
+        0x00.toByte(),
+        0x00.toByte(),
+        0x00.toByte(),
+        0xff.toByte(),
+        0xff.toByte(),
+        0x01.toByte(),
+        0x02.toByte(),
+        0x03.toByte(),
+        0x04.toByte()
+    )
     private val TEST_SRC_ADDRV4 = "198.168.0.2".address()
     private val TEST_DST_ADDRV4 = "198.168.0.1".address()
     private val TEST_ADDRV6 = "2001:db8::1".address()
-    private val TEST_ADDRV4MAPPEDV6 = "::ffff:1.2.3.4".address()
+    // This constant requires to be an Inet6Address, but InetAddresses.parseNumericAddress() will
+    // convert v4 mapped v6 address into an Inet4Address. So use Inet6Address.getByAddress() to
+    // create the address.
+    private val TEST_ADDRV4MAPPEDV6 = Inet6Address.getByAddress(null /* host */,
+        SRC_V4_MAPPED_V6_ADDRESS_BYTES, -1 /* scope_id */)
     private val TEST_ADDRV4 = "1.2.3.4".address()
 
     private fun String.address() = InetAddresses.parseNumericAddress(this)
diff --git a/tests/unit/java/android/net/nsd/NsdServiceInfoTest.java b/tests/common/java/android/net/nsd/NsdServiceInfoTest.java
similarity index 88%
rename from tests/unit/java/android/net/nsd/NsdServiceInfoTest.java
rename to tests/common/java/android/net/nsd/NsdServiceInfoTest.java
index 9ce0693..ffe0e91 100644
--- a/tests/unit/java/android/net/nsd/NsdServiceInfoTest.java
+++ b/tests/common/java/android/net/nsd/NsdServiceInfoTest.java
@@ -26,10 +26,10 @@
 import android.os.Build;
 import android.os.Bundle;
 import android.os.Parcel;
-import android.os.StrictMode;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.testutils.ConnectivityModuleTest;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
 
@@ -37,7 +37,6 @@
 import org.junit.runner.RunWith;
 
 import java.net.InetAddress;
-import java.net.UnknownHostException;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -45,22 +44,11 @@
 @RunWith(DevSdkIgnoreRunner.class)
 @SmallTest
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
+@ConnectivityModuleTest
 public class NsdServiceInfoTest {
 
     private static final InetAddress IPV4_ADDRESS = InetAddresses.parseNumericAddress("192.0.2.1");
     private static final InetAddress IPV6_ADDRESS = InetAddresses.parseNumericAddress("2001:db8::");
-    public final static InetAddress LOCALHOST;
-    static {
-        // Because test.
-        StrictMode.ThreadPolicy policy = new StrictMode.ThreadPolicy.Builder().permitAll().build();
-        StrictMode.setThreadPolicy(policy);
-
-        InetAddress _host = null;
-        try {
-            _host = InetAddress.getLocalHost();
-        } catch (UnknownHostException e) { }
-        LOCALHOST = _host;
-    }
 
     @Test
     public void testLimits() throws Exception {
@@ -89,10 +77,10 @@
         // Single key + value length too long.
         exceptionThrown = false;
         try {
-            String longValue = "loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
-                    "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
-                    "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
-                    "ooooooooooooooooooooooooooooong";  // 248 characters.
+            String longValue = "loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo"
+                    + "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo"
+                    + "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo"
+                    + "ooooooooooooooooooooooooooooong";  // 248 characters.
             info.setAttribute("longcat", longValue);  // Key + value == 255 characters.
         } catch (IllegalArgumentException e) {
             exceptionThrown = true;
@@ -127,7 +115,6 @@
         fullInfo.setServiceName("kitten");
         fullInfo.setServiceType("_kitten._tcp");
         fullInfo.setPort(4242);
-        fullInfo.setHost(LOCALHOST);
         fullInfo.setHostAddresses(List.of(IPV4_ADDRESS));
         fullInfo.setNetwork(new Network(123));
         fullInfo.setInterfaceIndex(456);
@@ -143,8 +130,7 @@
         attributedInfo.setServiceName("kitten");
         attributedInfo.setServiceType("_kitten._tcp");
         attributedInfo.setPort(4242);
-        attributedInfo.setHost(LOCALHOST);
-        fullInfo.setHostAddresses(List.of(IPV6_ADDRESS, IPV4_ADDRESS));
+        attributedInfo.setHostAddresses(List.of(IPV6_ADDRESS, IPV4_ADDRESS));
         attributedInfo.setAttribute("color", "pink");
         attributedInfo.setAttribute("sound", (new String("にゃあ")).getBytes("UTF-8"));
         attributedInfo.setAttribute("adorable", (String) null);
diff --git a/tests/cts/OWNERS b/tests/cts/OWNERS
index 8c2408b..286f9c8 100644
--- a/tests/cts/OWNERS
+++ b/tests/cts/OWNERS
@@ -1,7 +1,7 @@
 # Bug template url: http://b/new?component=31808
 # TODO: move bug template config to common owners file once b/226427845 is resolved
 set noparent
-file:platform/packages/modules/Connectivity:master:/OWNERS_core_networking_xts
+file:platform/packages/modules/Connectivity:main:/OWNERS_core_networking_xts
 
 # IPsec
 per-file **IpSec* = benedictwong@google.com, nharold@google.com
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractExpeditedJobTest.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractExpeditedJobTest.java
index a850e3b..7cac2af 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractExpeditedJobTest.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractExpeditedJobTest.java
@@ -74,6 +74,7 @@
     @RequiredProperties({APP_STANDBY_MODE})
     public void testNetworkAccess_appIdleState() throws Exception {
         turnBatteryOn();
+        setAppIdle(false);
         assertBackgroundNetworkAccess(true);
         assertExpeditedJobHasNetworkAccess();
 
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/app/src/com/android/cts/net/hostside/RestrictedModeTest.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/RestrictedModeTest.java
index 4266aad..35f1f1c 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/RestrictedModeTest.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/RestrictedModeTest.java
@@ -57,14 +57,18 @@
     @Test
     public void testNetworkAccess_withBatterySaver() throws Exception {
         setBatterySaverMode(true);
-        addPowerSaveModeWhitelist(TEST_APP2_PKG);
-        assertBackgroundNetworkAccess(true);
+        try {
+            addPowerSaveModeWhitelist(TEST_APP2_PKG);
+            assertBackgroundNetworkAccess(true);
 
-        setRestrictedNetworkingMode(true);
-        // App would be denied network access since Restricted mode is on.
-        assertBackgroundNetworkAccess(false);
-        setRestrictedNetworkingMode(false);
-        // Given that Restricted mode is turned off, app should be able to access network again.
-        assertBackgroundNetworkAccess(true);
+            setRestrictedNetworkingMode(true);
+            // App would be denied network access since Restricted mode is on.
+            assertBackgroundNetworkAccess(false);
+            setRestrictedNetworkingMode(false);
+            // Given that Restricted mode is turned off, app should be able to access network again.
+            assertBackgroundNetworkAccess(true);
+        } finally {
+            setBatterySaverMode(false);
+        }
     }
 }
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
index cd3b650..454940f 100755
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
@@ -1726,10 +1726,21 @@
         assertEquals(VpnManager.TYPE_VPN_SERVICE, ((VpnTransportInfo) ti).getType());
     }
 
-    private void assertDefaultProxy(ProxyInfo expected) {
+    private void assertDefaultProxy(ProxyInfo expected) throws Exception {
         assertEquals("Incorrect proxy config.", expected, mCM.getDefaultProxy());
         String expectedHost = expected == null ? null : expected.getHost();
         String expectedPort = expected == null ? null : String.valueOf(expected.getPort());
+
+        // ActivityThread may not have time to set it in the properties yet which will cause flakes.
+        // Wait for some time to deflake the test.
+        int attempt = 0;
+        while (!(Objects.equals(expectedHost, System.getProperty("http.proxyHost"))
+                && Objects.equals(expectedPort, System.getProperty("http.proxyPort")))
+                && attempt < 300) {
+            attempt++;
+            Log.d(TAG, "Wait for proxy being updated, attempt=" + attempt);
+            Thread.sleep(100);
+        }
         assertEquals("Incorrect proxy host system property.", expectedHost,
             System.getProperty("http.proxyHost"));
         assertEquals("Incorrect proxy port system property.", expectedPort,
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 cfd3130..69eba41 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java
@@ -18,36 +18,43 @@
 
 import android.platform.test.annotations.FlakyTest;
 
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.invoker.TestInformation;
+import com.android.tradefed.testtype.junit4.AfterClassWithInfo;
+import com.android.tradefed.testtype.junit4.BeforeClassWithInfo;
+
+import org.junit.Test;
+
 public class HostsideConnOnActivityStartTest extends HostsideNetworkTestCase {
     private static final String TEST_CLASS = TEST_PKG + ".ConnOnActivityStartTest";
-    @Override
-    public void setUp() throws Exception {
-        super.setUp();
-
-        uninstallPackage(TEST_APP2_PKG, false);
-        installPackage(TEST_APP2_APK);
+    @BeforeClassWithInfo
+    public static void setUpOnce(TestInformation testInfo) throws Exception {
+        uninstallPackage(testInfo, TEST_APP2_PKG, false);
+        installPackage(testInfo, TEST_APP2_APK);
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
-
-        uninstallPackage(TEST_APP2_PKG, true);
+    @AfterClassWithInfo
+    public static void tearDownOnce(TestInformation testInfo) throws DeviceNotAvailableException {
+        uninstallPackage(testInfo, TEST_APP2_PKG, true);
     }
 
+    @Test
     public void testStartActivity_batterySaver() throws Exception {
         runDeviceTests(TEST_PKG, TEST_CLASS, "testStartActivity_batterySaver");
     }
 
+    @Test
     public void testStartActivity_dataSaver() throws Exception {
         runDeviceTests(TEST_PKG, TEST_CLASS, "testStartActivity_dataSaver");
     }
 
     @FlakyTest(bugId = 231440256)
+    @Test
     public void testStartActivity_doze() throws Exception {
         runDeviceTests(TEST_PKG, TEST_CLASS, "testStartActivity_doze");
     }
 
+    @Test
     public void testStartActivity_appStandby() throws Exception {
         runDeviceTests(TEST_PKG, TEST_CLASS, "testStartActivity_appStandby");
     }
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 1312085..d8e7a2c 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java
@@ -14,26 +14,31 @@
  * limitations under the License.
  */
 package com.android.cts.net;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
 public class HostsideNetworkCallbackTests extends HostsideNetworkTestCase {
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    @Before
+    public void setUp() throws Exception {
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
+    @After
+    public void tearDown() throws Exception {
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
+    @Test
     public void testOnBlockedStatusChanged_dataSaver() throws Exception {
         runDeviceTests(TEST_PKG,
                 TEST_PKG + ".NetworkCallbackTest", "testOnBlockedStatusChanged_dataSaver");
     }
 
+    @Test
     public void testOnBlockedStatusChanged_powerSaver() throws Exception {
         runDeviceTests(TEST_PKG,
                 TEST_PKG + ".NetworkCallbackTest", "testOnBlockedStatusChanged_powerSaver");
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 fdb8876..3ddb88b 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java
@@ -16,49 +16,57 @@
 
 package com.android.cts.net;
 
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
 public class HostsideNetworkPolicyManagerTests extends HostsideNetworkTestCase {
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    @Before
+    public void setUp() throws Exception {
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
+    @After
+    public void tearDown() throws Exception {
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
+    @Test
     public void testIsUidNetworkingBlocked_withUidNotBlocked() throws Exception {
         runDeviceTests(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_withUidNotBlocked");
     }
 
+    @Test
     public void testIsUidNetworkingBlocked_withSystemUid() throws Exception {
         runDeviceTests(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest", "testIsUidNetworkingBlocked_withSystemUid");
     }
 
+    @Test
     public void testIsUidNetworkingBlocked_withDataSaverMode() throws Exception {
         runDeviceTests(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_withDataSaverMode");
     }
 
+    @Test
     public void testIsUidNetworkingBlocked_withRestrictedNetworkingMode() throws Exception {
         runDeviceTests(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_withRestrictedNetworkingMode");
     }
 
+    @Test
     public void testIsUidNetworkingBlocked_withPowerSaverMode() throws Exception {
         runDeviceTests(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_withPowerSaverMode");
     }
 
+    @Test
     public void testIsUidRestrictedOnMeteredNetworks() throws Exception {
         runDeviceTests(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest", "testIsUidRestrictedOnMeteredNetworks");
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 2aa1032..b89ab1f 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java
@@ -16,28 +16,27 @@
 
 package com.android.cts.net;
 
-import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
+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.ddmlib.testrunner.RemoteAndroidTestRunner;
-import com.android.ddmlib.testrunner.TestResult.TestStatus;
 import com.android.modules.utils.build.testing.DeviceSdkLevel;
-import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.device.DeviceNotAvailableException;
-import com.android.tradefed.result.CollectingTestListener;
-import com.android.tradefed.result.TestDescription;
-import com.android.tradefed.result.TestResult;
-import com.android.tradefed.result.TestRunResult;
-import com.android.tradefed.testtype.DeviceTestCase;
-import com.android.tradefed.testtype.IAbi;
-import com.android.tradefed.testtype.IAbiReceiver;
-import com.android.tradefed.testtype.IBuildReceiver;
+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.BeforeClassWithInfo;
 import com.android.tradefed.util.RunUtil;
 
-import java.io.FileNotFoundException;
-import java.util.Map;
+import org.junit.runner.RunWith;
 
-abstract class HostsideNetworkTestCase extends DeviceTestCase implements IAbiReceiver,
-        IBuildReceiver {
+@RunWith(DeviceJUnit4ClassRunner.class)
+abstract class HostsideNetworkTestCase extends BaseHostJUnit4Test {
     protected static final boolean DEBUG = false;
     protected static final String TAG = "HostsideNetworkTests";
     protected static final String TEST_PKG = "com.android.cts.net.hostside";
@@ -46,56 +45,62 @@
     protected static final String TEST_APP2_PKG = "com.android.cts.net.hostside.app2";
     protected static final String TEST_APP2_APK = "CtsHostsideNetworkTestsApp2.apk";
 
-    private IAbi mAbi;
-    private IBuildInfo mCtsBuild;
+    @BeforeClassWithInfo
+    public static void setUpOnceBase(TestInformation testInfo) throws Exception {
+        DeviceSdkLevel deviceSdkLevel = new DeviceSdkLevel(testInfo.getDevice());
+        String testApk = deviceSdkLevel.isDeviceAtLeastT() ? TEST_APK_NEXT : TEST_APK;
 
-    @Override
-    public void setAbi(IAbi abi) {
-        mAbi = abi;
+        uninstallPackage(testInfo, TEST_PKG, false);
+        installPackage(testInfo, testApk);
     }
 
-    @Override
-    public void setBuild(IBuildInfo buildInfo) {
-        mCtsBuild = buildInfo;
-    }
-
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-
-        assertNotNull(mAbi);
-        assertNotNull(mCtsBuild);
-
-        DeviceSdkLevel deviceSdkLevel = new DeviceSdkLevel(getDevice());
-        String testApk = deviceSdkLevel.isDeviceAtLeastT() ? TEST_APK_NEXT
-                : TEST_APK;
-
-        uninstallPackage(TEST_PKG, false);
-        installPackage(testApk);
-    }
-
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
-
-        uninstallPackage(TEST_PKG, true);
-    }
-
-    protected void installPackage(String apk) throws FileNotFoundException,
-            DeviceNotAvailableException {
-        CompatibilityBuildHelper buildHelper = new CompatibilityBuildHelper(mCtsBuild);
-        assertNull(getDevice().installPackage(buildHelper.getTestFile(apk),
-                false /* reinstall */, true /* grantPermissions */, "-t"));
-    }
-
-    protected void uninstallPackage(String packageName, boolean shouldSucceed)
+    @AfterClassWithInfo
+    public static void tearDownOnceBase(TestInformation testInfo)
             throws DeviceNotAvailableException {
-        final String result = getDevice().uninstallPackage(packageName);
+        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 {
+        installPackage(getTestInformation(), apk);
+    }
+
+    protected static void uninstallPackage(TestInformation testInfo, String packageName,
+            boolean shouldSucceed)
+            throws DeviceNotAvailableException {
+        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;
@@ -126,50 +131,6 @@
         fail("Package '" + packageName + "' not uinstalled after " + max_tries + " seconds");
     }
 
-    protected void runDeviceTests(String packageName, String testClassName)
-            throws DeviceNotAvailableException {
-        runDeviceTests(packageName, testClassName, null);
-    }
-
-    protected void runDeviceTests(String packageName, String testClassName, String methodName)
-            throws DeviceNotAvailableException {
-        RemoteAndroidTestRunner testRunner = new RemoteAndroidTestRunner(packageName,
-                "androidx.test.runner.AndroidJUnitRunner", getDevice().getIDevice());
-
-        if (testClassName != null) {
-            if (methodName != null) {
-                testRunner.setMethodName(testClassName, methodName);
-            } else {
-                testRunner.setClassName(testClassName);
-            }
-        }
-
-        final CollectingTestListener listener = new CollectingTestListener();
-        getDevice().runInstrumentationTests(testRunner, listener);
-
-        final TestRunResult result = listener.getCurrentRunResults();
-        if (result.isRunFailure()) {
-            throw new AssertionError("Failed to successfully run device tests for "
-                    + result.getName() + ": " + result.getRunFailureMessage());
-        }
-
-        if (result.hasFailedTests()) {
-            // build a meaningful error message
-            StringBuilder errorBuilder = new StringBuilder("on-device tests failed:\n");
-            for (Map.Entry<TestDescription, TestResult> resultEntry :
-                    result.getTestResults().entrySet()) {
-                final TestStatus testStatus = resultEntry.getValue().getStatus();
-                if (!TestStatus.PASSED.equals(testStatus)
-                        && !TestStatus.ASSUMPTION_FAILURE.equals(testStatus)) {
-                    errorBuilder.append(resultEntry.getKey().toString());
-                    errorBuilder.append(":\n");
-                    errorBuilder.append(resultEntry.getValue().getStackTrace());
-                }
-            }
-            throw new AssertionError(errorBuilder.toString());
-        }
-    }
-
     protected int getUid(String packageName) throws DeviceNotAvailableException {
         final int currentUser = getDevice().getCurrentUser();
         final String uidLines = runCommand(
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 21c78b7..57b26bd 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java
@@ -16,30 +16,33 @@
 
 package com.android.cts.net;
 
+import static org.junit.Assert.fail;
+
 import android.platform.test.annotations.SecurityTest;
 
 import com.android.ddmlib.Log;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.util.RunUtil;
 
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
 public class HostsideRestrictBackgroundNetworkTests extends HostsideNetworkTestCase {
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-
+    @Before
+    public void setUp() throws Exception {
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
-
+    @After
+    public void tearDown() throws Exception {
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
     @SecurityTest
+    @Test
     public void testDataWarningReceiver() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DataWarningReceiverTest",
                 "testSnoozeWarningNotReceived");
@@ -49,26 +52,31 @@
      * Data Saver Mode tests. *
      **************************/
 
+    @Test
     public void testDataSaverMode_disabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_disabled");
     }
 
+    @Test
     public void testDataSaverMode_whitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_whitelisted");
     }
 
+    @Test
     public void testDataSaverMode_enabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_enabled");
     }
 
+    @Test
     public void testDataSaverMode_blacklisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_blacklisted");
     }
 
+    @Test
     public void testDataSaverMode_reinstall() throws Exception {
         final int oldUid = getUid(TEST_APP2_PKG);
 
@@ -85,11 +93,13 @@
         assertRestrictBackgroundWhitelist(newUid, false);
     }
 
+    @Test
     public void testDataSaverMode_requiredWhitelistedPackages() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_requiredWhitelistedPackages");
     }
 
+    @Test
     public void testDataSaverMode_broadcastNotSentOnUnsupportedDevices() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testBroadcastNotSentOnUnsupportedDevices");
@@ -99,21 +109,25 @@
      * Battery Saver Mode tests. *
      *****************************/
 
+    @Test
     public void testBatterySaverModeMetered_disabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
+    @Test
     public void testBatterySaverModeMetered_whitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
+    @Test
     public void testBatterySaverModeMetered_enabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
+    @Test
     public void testBatterySaverMode_reinstall() throws Exception {
         if (!isDozeModeEnabled()) {
             Log.w(TAG, "testBatterySaverMode_reinstall() skipped because device does not support "
@@ -131,16 +145,19 @@
         assertPowerSaveModeWhitelist(TEST_APP2_PKG, false);
     }
 
+    @Test
     public void testBatterySaverModeNonMetered_disabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
+    @Test
     public void testBatterySaverModeNonMetered_whitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
+    @Test
     public void testBatterySaverModeNonMetered_enabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
@@ -150,26 +167,31 @@
      * App idle tests. *
      *******************/
 
+    @Test
     public void testAppIdleMetered_disabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
+    @Test
     public void testAppIdleMetered_whitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
+    @Test
     public void testAppIdleMetered_tempWhitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testBackgroundNetworkAccess_tempWhitelisted");
     }
 
+    @Test
     public void testAppIdleMetered_enabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
+    @Test
     public void testAppIdleMetered_idleWhitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testAppIdleNetworkAccess_idleWhitelisted");
@@ -180,41 +202,50 @@
     //    public void testAppIdle_reinstall() throws Exception {
     //    }
 
+    @Test
     public void testAppIdleNonMetered_disabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
+
+    @Test
     public void testAppIdleNonMetered_whitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
+    @Test
     public void testAppIdleNonMetered_tempWhitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testBackgroundNetworkAccess_tempWhitelisted");
     }
 
+    @Test
     public void testAppIdleNonMetered_enabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
+    @Test
     public void testAppIdleNonMetered_idleWhitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testAppIdleNetworkAccess_idleWhitelisted");
     }
 
+    @Test
     public void testAppIdleNonMetered_whenCharging() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testAppIdleNetworkAccess_whenCharging");
     }
 
+    @Test
     public void testAppIdleMetered_whenCharging() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testAppIdleNetworkAccess_whenCharging");
     }
 
+    @Test
     public void testAppIdle_toast() throws Exception {
         // Check that showing a toast doesn't bring an app out of standby
         runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
@@ -225,21 +256,25 @@
      * Doze Mode tests. *
      ********************/
 
+    @Test
     public void testDozeModeMetered_disabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
+    @Test
     public void testDozeModeMetered_whitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
+    @Test
     public void testDozeModeMetered_enabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
+    @Test
     public void testDozeModeMetered_enabledButWhitelistedOnNotificationAction() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
                 "testBackgroundNetworkAccess_enabledButWhitelistedOnNotificationAction");
@@ -250,21 +285,25 @@
     //    public void testDozeMode_reinstall() throws Exception {
     //    }
 
+    @Test
     public void testDozeModeNonMetered_disabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
+    @Test
     public void testDozeModeNonMetered_whitelisted() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
+    @Test
     public void testDozeModeNonMetered_enabled() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
+    @Test
     public void testDozeModeNonMetered_enabledButWhitelistedOnNotificationAction()
             throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
@@ -275,46 +314,55 @@
      * Mixed modes tests. *
      **********************/
 
+    @Test
     public void testDataAndBatterySaverModes_meteredNetwork() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDataAndBatterySaverModes_meteredNetwork");
     }
 
+    @Test
     public void testDataAndBatterySaverModes_nonMeteredNetwork() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDataAndBatterySaverModes_nonMeteredNetwork");
     }
 
+    @Test
     public void testDozeAndBatterySaverMode_powerSaveWhitelists() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDozeAndBatterySaverMode_powerSaveWhitelists");
     }
 
+    @Test
     public void testDozeAndAppIdle_powerSaveWhitelists() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDozeAndAppIdle_powerSaveWhitelists");
     }
 
+    @Test
     public void testAppIdleAndDoze_tempPowerSaveWhitelists() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testAppIdleAndDoze_tempPowerSaveWhitelists");
     }
 
+    @Test
     public void testAppIdleAndBatterySaver_tempPowerSaveWhitelists() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testAppIdleAndBatterySaver_tempPowerSaveWhitelists");
     }
 
+    @Test
     public void testDozeAndAppIdle_appIdleWhitelist() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDozeAndAppIdle_appIdleWhitelist");
     }
 
+    @Test
     public void testAppIdleAndDoze_tempPowerSaveAndAppIdleWhitelists() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testAppIdleAndDoze_tempPowerSaveAndAppIdleWhitelists");
     }
 
+    @Test
     public void testAppIdleAndBatterySaver_tempPowerSaveAndAppIdleWhitelists() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testAppIdleAndBatterySaver_tempPowerSaveAndAppIdleWhitelists");
@@ -323,11 +371,14 @@
     /**************************
      * Restricted mode tests. *
      **************************/
+
+    @Test
     public void testNetworkAccess_restrictedMode() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".RestrictedModeTest",
                 "testNetworkAccess");
     }
 
+    @Test
     public void testNetworkAccess_restrictedMode_withBatterySaver() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".RestrictedModeTest",
                 "testNetworkAccess_withBatterySaver");
@@ -337,10 +388,12 @@
      * Expedited job tests. *
      ************************/
 
+    @Test
     public void testMeteredNetworkAccess_expeditedJob() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".ExpeditedJobMeteredTest");
     }
 
+    @Test
     public void testNonMeteredNetworkAccess_expeditedJob() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".ExpeditedJobNonMeteredTest");
     }
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideSelfDeclaredNetworkCapabilitiesCheckTest.java b/tests/cts/hostside/src/com/android/cts/net/HostsideSelfDeclaredNetworkCapabilitiesCheckTest.java
index 4c2985d..c3bdb6d 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideSelfDeclaredNetworkCapabilitiesCheckTest.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideSelfDeclaredNetworkCapabilitiesCheckTest.java
@@ -15,6 +15,8 @@
  */
 package com.android.cts.net;
 
+import org.junit.Test;
+
 public class HostsideSelfDeclaredNetworkCapabilitiesCheckTest extends HostsideNetworkTestCase {
 
     private static final String TEST_WITH_PROPERTY_IN_CURRENT_SDK_APK =
@@ -34,6 +36,7 @@
             "requestNetwork_withoutRequestCapabilities";
 
 
+    @Test
     public void testRequestNetworkInCurrentSdkWithProperty() throws Exception {
         uninstallPackage(TEST_APP_PKG, false);
         installPackage(TEST_WITH_PROPERTY_IN_CURRENT_SDK_APK);
@@ -48,6 +51,7 @@
         uninstallPackage(TEST_APP_PKG, true);
     }
 
+    @Test
     public void testRequestNetworkInCurrentSdkWithoutProperty() throws Exception {
         uninstallPackage(TEST_APP_PKG, false);
         installPackage(TEST_WITHOUT_PROPERTY_IN_CURRENT_SDK_APK);
@@ -62,6 +66,7 @@
         uninstallPackage(TEST_APP_PKG, true);
     }
 
+    @Test
     public void testRequestNetworkInSdk33() throws Exception {
         uninstallPackage(TEST_APP_PKG, false);
         installPackage(TEST_IN_SDK_33_APK);
@@ -75,6 +80,7 @@
         uninstallPackage(TEST_APP_PKG, true);
     }
 
+    @Test
     public void testReinstallPackageWillUpdateProperty() throws Exception {
         uninstallPackage(TEST_APP_PKG, false);
         installPackage(TEST_WITHOUT_PROPERTY_IN_CURRENT_SDK_APK);
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 3ca4775..691ac90 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideVpnTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideVpnTests.java
@@ -18,95 +18,112 @@
 
 import android.platform.test.annotations.RequiresDevice;
 
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
 public class HostsideVpnTests extends HostsideNetworkTestCase {
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-
+    @Before
+    public void setUp() throws Exception {
         uninstallPackage(TEST_APP2_PKG, false);
         installPackage(TEST_APP2_APK);
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
-
+    @After
+    public void tearDown() throws Exception {
         uninstallPackage(TEST_APP2_PKG, true);
     }
 
+    @Test
     public void testChangeUnderlyingNetworks() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testChangeUnderlyingNetworks");
     }
 
+    @Test
     public void testDefault() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testDefault");
     }
 
+    @Test
     public void testAppAllowed() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testAppAllowed");
     }
 
+    @Test
     public void testAppDisallowed() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testAppDisallowed");
     }
 
+    @Test
     public void testSocketClosed() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testSocketClosed");
     }
 
+    @Test
     public void testGetConnectionOwnerUidSecurity() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testGetConnectionOwnerUidSecurity");
     }
 
+    @Test
     public void testSetProxy() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testSetProxy");
     }
 
+    @Test
     public void testSetProxyDisallowedApps() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testSetProxyDisallowedApps");
     }
 
+    @Test
     public void testNoProxy() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testNoProxy");
     }
 
+    @Test
     public void testBindToNetworkWithProxy() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testBindToNetworkWithProxy");
     }
 
+    @Test
     public void testVpnMeterednessWithNoUnderlyingNetwork() throws Exception {
         runDeviceTests(
                 TEST_PKG, TEST_PKG + ".VpnTest", "testVpnMeterednessWithNoUnderlyingNetwork");
     }
 
+    @Test
     public void testVpnMeterednessWithNullUnderlyingNetwork() throws Exception {
         runDeviceTests(
                 TEST_PKG, TEST_PKG + ".VpnTest", "testVpnMeterednessWithNullUnderlyingNetwork");
     }
 
+    @Test
     public void testVpnMeterednessWithNonNullUnderlyingNetwork() throws Exception {
         runDeviceTests(
                 TEST_PKG, TEST_PKG + ".VpnTest", "testVpnMeterednessWithNonNullUnderlyingNetwork");
     }
 
+    @Test
     public void testAlwaysMeteredVpnWithNullUnderlyingNetwork() throws Exception {
         runDeviceTests(
                 TEST_PKG, TEST_PKG + ".VpnTest", "testAlwaysMeteredVpnWithNullUnderlyingNetwork");
     }
 
     @RequiresDevice // Keepalive is not supported on virtual hardware
+    @Test
     public void testAutomaticOnOffKeepaliveModeClose() throws Exception {
         runDeviceTests(
                 TEST_PKG, TEST_PKG + ".VpnTest", "testAutomaticOnOffKeepaliveModeClose");
     }
 
     @RequiresDevice // Keepalive is not supported on virtual hardware
+    @Test
     public void testAutomaticOnOffKeepaliveModeNoClose() throws Exception {
         runDeviceTests(
                 TEST_PKG, TEST_PKG + ".VpnTest", "testAutomaticOnOffKeepaliveModeNoClose");
     }
 
+    @Test
     public void testAlwaysMeteredVpnWithNonNullUnderlyingNetwork() throws Exception {
         runDeviceTests(
                 TEST_PKG,
@@ -114,31 +131,38 @@
                 "testAlwaysMeteredVpnWithNonNullUnderlyingNetwork");
     }
 
+    @Test
     public void testB141603906() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testB141603906");
     }
 
+    @Test
     public void testDownloadWithDownloadManagerDisallowed() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest",
                 "testDownloadWithDownloadManagerDisallowed");
     }
 
+    @Test
     public void testExcludedRoutes() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testExcludedRoutes");
     }
 
+    @Test
     public void testIncludedRoutes() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testIncludedRoutes");
     }
 
+    @Test
     public void testInterleavedRoutes() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testInterleavedRoutes");
     }
 
+    @Test
     public void testBlockIncomingPackets() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testBlockIncomingPackets");
     }
 
+    @Test
     public void testSetVpnDefaultForUids() throws Exception {
         runDeviceTests(TEST_PKG, TEST_PKG + ".VpnTest", "testSetVpnDefaultForUids");
     }
diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index 9f8a05d..98ea224 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -164,6 +164,12 @@
     it.obj = obj
 }
 
+// On T and below, the native network is only created when the agent connects.
+// Starting in U, the native network was to be created as soon as the agent is registered,
+// but this has been flagged off for now pending resolution of race conditions.
+// TODO : enable this in a Mainline update or in V.
+private const val SHOULD_CREATE_NETWORKS_IMMEDIATELY = false
+
 @RunWith(DevSdkIgnoreRunner::class)
 // NetworkAgent is not updatable in R-, so this test does not need to be compatible with older
 // versions. NetworkAgent was also based on AsyncChannel before S so cannot be tested the same way.
@@ -1247,15 +1253,15 @@
 
         // Connect a third network. Because network1 is awaiting replacement, network3 is preferred
         // as soon as it validates (until then, it is outscored by network1).
-        // The fact that the first event seen by matchAllCallback is the connection of network3
+        // The fact that the first events seen by matchAllCallback is the connection of network3
         // implicitly ensures that no callbacks are sent since network1 was lost.
         val (agent3, network3) = connectNetwork()
+        matchAllCallback.expectAvailableThenValidatedCallbacks(network3)
+        testCallback.expectAvailableDoubleValidatedCallbacks(network3)
+
         // As soon as the replacement arrives, network1 is disconnected.
         // Check that this happens before the replacement timeout (5 seconds) fires.
-        matchAllCallback.expectAvailableCallbacks(network3, validated = false)
         matchAllCallback.expect<Lost>(network1, 2_000 /* timeoutMs */)
-        matchAllCallback.expectCaps(network3) { it.hasCapability(NET_CAPABILITY_VALIDATED) }
-        testCallback.expectAvailableDoubleValidatedCallbacks(network3)
         agent1.expectCallback<OnNetworkUnwanted>()
 
         // Test lingering:
@@ -1301,8 +1307,8 @@
         val callback = TestableNetworkCallback()
         requestNetwork(makeTestNetworkRequest(specifier = specifier6), callback)
         val agent6 = createNetworkAgent(specifier = specifier6)
-        agent6.register()
-        if (SdkLevel.isAtLeastU()) {
+        val network6 = agent6.register()
+        if (SHOULD_CREATE_NETWORKS_IMMEDIATELY) {
             agent6.expectCallback<OnNetworkCreated>()
         } else {
             // No callbacks are sent, so check LinkProperties to wait for the network to be created.
@@ -1316,9 +1322,10 @@
         val timeoutMs = agent6.DEFAULT_TIMEOUT_MS.toInt() + 1_000
         agent6.unregisterAfterReplacement(timeoutMs)
         agent6.expectCallback<OnNetworkUnwanted>()
-        if (!SdkLevel.isAtLeastT() || SdkLevel.isAtLeastU()) {
+        if (!SdkLevel.isAtLeastT() || SHOULD_CREATE_NETWORKS_IMMEDIATELY) {
             // Before T, onNetworkDestroyed is called even if the network was never created.
-            // On U+, the network was created by register(). Destroying it sends onNetworkDestroyed.
+            // If immediate native network creation is supported, the network was created by
+            // register(). Destroying it sends onNetworkDestroyed.
             agent6.expectCallback<OnNetworkDestroyed>()
         }
         // Poll for LinkProperties becoming null, because when onNetworkUnwanted is called, the
@@ -1368,9 +1375,8 @@
 
         val (newWifiAgent, newWifiNetwork) = connectNetwork(TRANSPORT_WIFI)
         testCallback.expectAvailableCallbacks(newWifiNetwork, validated = true)
-        matchAllCallback.expectAvailableCallbacks(newWifiNetwork, validated = false)
+        matchAllCallback.expectAvailableThenValidatedCallbacks(newWifiNetwork)
         matchAllCallback.expect<Lost>(wifiNetwork)
-        matchAllCallback.expectCaps(newWifiNetwork) { it.hasCapability(NET_CAPABILITY_VALIDATED) }
         wifiAgent.expectCallback<OnNetworkUnwanted>()
     }
 
@@ -1478,10 +1484,9 @@
 
     @Test
     fun testNativeNetworkCreation_PhysicalNetwork() {
-        // On T and below, the native network is only created when the agent connects.
-        // Starting in U, the native network is created as soon as the agent is registered.
-        doTestNativeNetworkCreation(expectCreatedImmediately = SdkLevel.isAtLeastU(),
-            intArrayOf(TRANSPORT_CELLULAR))
+        doTestNativeNetworkCreation(
+                expectCreatedImmediately = SHOULD_CREATE_NETWORKS_IMMEDIATELY,
+                intArrayOf(TRANSPORT_CELLULAR))
     }
 
     @Test
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/cts/net/src/android/net/cts/RateLimitTest.java b/tests/cts/net/src/android/net/cts/RateLimitTest.java
index 36b98fc..5c93738 100644
--- a/tests/cts/net/src/android/net/cts/RateLimitTest.java
+++ b/tests/cts/net/src/android/net/cts/RateLimitTest.java
@@ -36,6 +36,7 @@
 import android.icu.text.MessageFormat;
 import android.net.ConnectivityManager;
 import android.net.ConnectivitySettingsManager;
+import android.net.ConnectivityThread;
 import android.net.InetAddresses;
 import android.net.IpPrefix;
 import android.net.LinkAddress;
@@ -189,7 +190,19 @@
             // whatever happens, don't leave the device in rate limited state.
             ConnectivitySettingsManager.setIngressRateLimitInBytesPerSecond(mContext, -1);
         }
-        if (mSocket != null) mSocket.close();
+        if (mSocket == null) {
+            // HACK(b/272147742): dump ConnectivityThread if test initialization failed.
+            final StackTraceElement[] elements = ConnectivityThread.get().getStackTrace();
+            final StringBuilder sb = new StringBuilder();
+            // Skip first element as it includes the invocation of getStackTrace()
+            for (int i = 1; i < elements.length; i++) {
+                sb.append(elements[i]);
+                sb.append("\n");
+            }
+            Log.e(TAG, sb.toString());
+        } else {
+            mSocket.close();
+        }
         if (mNetworkAgent != null) mNetworkAgent.unregister();
         if (mTunInterface != null) mTunInterface.getFileDescriptor().close();
         if (mCm != null) mCm.unregisterNetworkCallback(mNetworkCallback);
diff --git a/tests/native/connectivity_native_test/OWNERS b/tests/native/connectivity_native_test/OWNERS
index fbfcf92..c9bfc40 100644
--- a/tests/native/connectivity_native_test/OWNERS
+++ b/tests/native/connectivity_native_test/OWNERS
@@ -1,4 +1,4 @@
 # Bug template url: http://b/new?component=31808
 # TODO: move bug template config to common owners file once b/226427845 is resolved
 set noparent
-file:platform/packages/modules/Connectivity:master:/OWNERS_core_networking_xts
+file:platform/packages/modules/Connectivity:main:/OWNERS_core_networking_xts
diff --git a/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt b/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
new file mode 100644
index 0000000..961c422
--- /dev/null
+++ b/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License")
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.metrics
+
+import android.os.Build
+import android.stats.connectivity.MdnsQueryResult
+import android.stats.connectivity.NsdEventType
+import com.android.testutils.DevSdkIgnoreRule
+import com.android.testutils.DevSdkIgnoreRunner
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.verify
+
+@RunWith(DevSdkIgnoreRunner::class)
+@DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.TIRAMISU)
+class NetworkNsdReportedMetricsTest {
+    private val deps = mock(NetworkNsdReportedMetrics.Dependencies::class.java)
+
+    @Test
+    fun testReportServiceRegistrationSucceeded() {
+        val clientId = 99
+        val transactionId = 100
+        val durationMs = 10L
+        val metrics = NetworkNsdReportedMetrics(true /* isLegacy */, clientId, deps)
+        metrics.reportServiceRegistrationSucceeded(transactionId, durationMs)
+
+        val eventCaptor = ArgumentCaptor.forClass(NetworkNsdReported::class.java)
+        verify(deps).statsWrite(eventCaptor.capture())
+        eventCaptor.value.let {
+            assertTrue(it.isLegacy)
+            assertEquals(clientId, it.clientId)
+            assertEquals(transactionId, it.transactionId)
+            assertEquals(NsdEventType.NET_REGISTER, it.type)
+            assertEquals(MdnsQueryResult.MQR_SERVICE_REGISTERED, it.queryResult)
+            assertEquals(durationMs, it.eventDurationMillisec)
+        }
+    }
+
+    @Test
+    fun testReportServiceRegistrationFailed() {
+        val clientId = 99
+        val transactionId = 100
+        val durationMs = 10L
+        val metrics = NetworkNsdReportedMetrics(false /* isLegacy */, clientId, deps)
+        metrics.reportServiceRegistrationFailed(transactionId, durationMs)
+
+        val eventCaptor = ArgumentCaptor.forClass(NetworkNsdReported::class.java)
+        verify(deps).statsWrite(eventCaptor.capture())
+        eventCaptor.value.let {
+            assertFalse(it.isLegacy)
+            assertEquals(clientId, it.clientId)
+            assertEquals(transactionId, it.transactionId)
+            assertEquals(NsdEventType.NET_REGISTER, it.type)
+            assertEquals(MdnsQueryResult.MQR_SERVICE_REGISTRATION_FAILED, it.queryResult)
+            assertEquals(durationMs, it.eventDurationMillisec)
+        }
+    }
+
+    @Test
+    fun testReportServiceUnregistration() {
+        val clientId = 99
+        val transactionId = 100
+        val durationMs = 10L
+        val metrics = NetworkNsdReportedMetrics(true /* isLegacy */, clientId, deps)
+        metrics.reportServiceUnregistration(transactionId, durationMs)
+
+        val eventCaptor = ArgumentCaptor.forClass(NetworkNsdReported::class.java)
+        verify(deps).statsWrite(eventCaptor.capture())
+        eventCaptor.value.let {
+            assertTrue(it.isLegacy)
+            assertEquals(clientId, it.clientId)
+            assertEquals(transactionId, it.transactionId)
+            assertEquals(NsdEventType.NET_REGISTER, it.type)
+            assertEquals(MdnsQueryResult.MQR_SERVICE_UNREGISTERED, it.queryResult)
+            assertEquals(durationMs, it.eventDurationMillisec)
+        }
+    }
+}
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 8de6a31..bbde9b4 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -2274,12 +2274,12 @@
             mDestroySocketsWrapper.destroyLiveTcpSocketsByOwnerUids(ownerUids);
         }
 
-        final ArrayTrackRecord<Long>.ReadHead mScheduledEvaluationTimeouts =
-                new ArrayTrackRecord<Long>().newReadHead();
+        final ArrayTrackRecord<Pair<Integer, Long>>.ReadHead mScheduledEvaluationTimeouts =
+                new ArrayTrackRecord<Pair<Integer, Long>>().newReadHead();
         @Override
         public void scheduleEvaluationTimeout(@NonNull Handler handler,
                 @NonNull final Network network, final long delayMs) {
-            mScheduledEvaluationTimeouts.add(delayMs);
+            mScheduledEvaluationTimeouts.add(new Pair<>(network.netId, delayMs));
             super.scheduleEvaluationTimeout(handler, network, delayMs);
         }
     }
@@ -2973,24 +2973,22 @@
         if (expectLingering) {
             generalCb.expectLosing(net1);
         }
+        generalCb.expectCaps(net2, c -> c.hasCapability(NET_CAPABILITY_VALIDATED));
+        defaultCb.expectAvailableDoubleValidatedCallbacks(net2);
 
         // Make sure cell 1 is unwanted immediately if the radio can't time share, but only
         // after some delay if it can.
         if (expectLingering) {
-            generalCb.expectCaps(net2, c -> c.hasCapability(NET_CAPABILITY_VALIDATED));
-            defaultCb.expectAvailableDoubleValidatedCallbacks(net2);
             net1.assertNotDisconnected(TEST_CALLBACK_TIMEOUT_MS); // always incurs the timeout
             generalCb.assertNoCallback();
             // assertNotDisconnected waited for TEST_CALLBACK_TIMEOUT_MS, so waiting for the
             // linger period gives TEST_CALLBACK_TIMEOUT_MS time for the event to process.
             net1.expectDisconnected(UNREASONABLY_LONG_ALARM_WAIT_MS);
-            generalCb.expect(LOST, net1);
         } else {
             net1.expectDisconnected(TEST_CALLBACK_TIMEOUT_MS);
-            generalCb.expect(LOST, net1);
-            generalCb.expectCaps(net2, c -> c.hasCapability(NET_CAPABILITY_VALIDATED));
-            defaultCb.expectAvailableDoubleValidatedCallbacks(net2);
         }
+        net1.disconnect();
+        generalCb.expect(LOST, net1);
 
         // Remove primary from net 2
         net2.setScore(new NetworkScore.Builder().build());
@@ -6153,7 +6151,7 @@
     }
 
     public void doTestPreferBadWifi(final boolean avoidBadWifi,
-            final boolean preferBadWifi,
+            final boolean preferBadWifi, final boolean explicitlySelected,
             @NonNull Predicate<Long> checkUnvalidationTimeout) throws Exception {
         // Pretend we're on a carrier that restricts switching away from bad wifi, and
         // depending on the parameter one that may indeed prefer bad wifi.
@@ -6177,10 +6175,13 @@
         mDefaultNetworkCallback.expectAvailableThenValidatedCallbacks(mCellAgent);
 
         mWiFiAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiAgent.explicitlySelected(explicitlySelected, false /* acceptUnvalidated */);
         mWiFiAgent.connect(false);
         wifiCallback.expectAvailableCallbacksUnvalidated(mWiFiAgent);
 
-        mDeps.mScheduledEvaluationTimeouts.poll(TIMEOUT_MS, t -> checkUnvalidationTimeout.test(t));
+        assertNotNull(mDeps.mScheduledEvaluationTimeouts.poll(TIMEOUT_MS,
+                t -> t.first == mWiFiAgent.getNetwork().netId
+                        && checkUnvalidationTimeout.test(t.second)));
 
         if (!avoidBadWifi && preferBadWifi) {
             expectUnvalidationCheckWillNotify(mWiFiAgent, NotificationType.LOST_INTERNET);
@@ -6196,27 +6197,33 @@
         // Starting with U this mode is no longer supported and can't actually be tested
         assumeFalse(mDeps.isAtLeastU());
         doTestPreferBadWifi(false /* avoidBadWifi */, false /* preferBadWifi */,
-                timeout -> timeout < 14_000);
+                false /* explicitlySelected */, timeout -> timeout < 14_000);
     }
 
     @Test
-    public void testPreferBadWifi_doNotAvoid_doPrefer() throws Exception {
+    public void testPreferBadWifi_doNotAvoid_doPrefer_notExplicit() throws Exception {
         doTestPreferBadWifi(false /* avoidBadWifi */, true /* preferBadWifi */,
-                timeout -> timeout > 14_000);
+                false /* explicitlySelected */, timeout -> timeout > 14_000);
+    }
+
+    @Test
+    public void testPreferBadWifi_doNotAvoid_doPrefer_explicitlySelected() throws Exception {
+        doTestPreferBadWifi(false /* avoidBadWifi */, true /* preferBadWifi */,
+                true /* explicitlySelected */, timeout -> timeout < 14_000);
     }
 
     @Test
     public void testPreferBadWifi_doAvoid_doNotPrefer() throws Exception {
         // If avoidBadWifi=true, then preferBadWifi should be irrelevant. Test anyway.
         doTestPreferBadWifi(true /* avoidBadWifi */, false /* preferBadWifi */,
-                timeout -> timeout < 14_000);
+                false /* explicitlySelected */, timeout -> timeout < 14_000);
     }
 
     @Test
     public void testPreferBadWifi_doAvoid_doPrefer() throws Exception {
         // If avoidBadWifi=true, then preferBadWifi should be irrelevant. Test anyway.
         doTestPreferBadWifi(true /* avoidBadWifi */, true /* preferBadWifi */,
-                timeout -> timeout < 14_000);
+                false /* explicitlySelected */, timeout -> timeout < 14_000);
     }
 
     @Test
@@ -6850,17 +6857,19 @@
 
     @Test
     public void testPacketKeepalives() throws Exception {
-        InetAddress myIPv4 = InetAddress.getByName("192.0.2.129");
+        final LinkAddress v4Addr = new LinkAddress("192.0.2.129/24");
+        final InetAddress myIPv4 = v4Addr.getAddress();
         InetAddress notMyIPv4 = InetAddress.getByName("192.0.2.35");
         InetAddress myIPv6 = InetAddress.getByName("2001:db8::1");
         InetAddress dstIPv4 = InetAddress.getByName("8.8.8.8");
         InetAddress dstIPv6 = InetAddress.getByName("2001:4860:4860::8888");
-
+        doReturn(getClatInterfaceConfigParcel(v4Addr)).when(mMockNetd)
+                .interfaceGetCfg(CLAT_MOBILE_IFNAME);
         final int validKaInterval = 15;
         final int invalidKaInterval = 9;
 
         LinkProperties lp = new LinkProperties();
-        lp.setInterfaceName("wlan12");
+        lp.setInterfaceName(MOBILE_IFNAME);
         lp.addLinkAddress(new LinkAddress(myIPv6, 64));
         lp.addLinkAddress(new LinkAddress(myIPv4, 25));
         lp.addRoute(new RouteInfo(InetAddress.getByName("fe80::1234")));
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index f51b28d..55384b3 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -43,7 +43,9 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.any;
@@ -95,6 +97,7 @@
 import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 
+import com.android.metrics.NetworkNsdReportedMetrics;
 import com.android.server.NsdService.Dependencies;
 import com.android.server.connectivity.mdns.MdnsAdvertiser;
 import com.android.server.connectivity.mdns.MdnsDiscoveryManager;
@@ -104,6 +107,7 @@
 import com.android.server.connectivity.mdns.MdnsServiceInfo;
 import com.android.server.connectivity.mdns.MdnsSocketProvider;
 import com.android.server.connectivity.mdns.MdnsSocketProvider.SocketRequestMonitor;
+import com.android.server.connectivity.mdns.util.MdnsUtils;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
 import com.android.testutils.HandlerUtils;
@@ -138,6 +142,7 @@
     static final int PROTOCOL = NsdManager.PROTOCOL_DNS_SD;
     private static final long CLEANUP_DELAY_MS = 500;
     private static final long TIMEOUT_MS = 500;
+    private static final long TEST_TIME_MS = 123L;
     private static final String SERVICE_NAME = "a_name";
     private static final String SERVICE_TYPE = "_test._tcp";
     private static final String SERVICE_FULL_NAME = SERVICE_NAME + "." + SERVICE_TYPE;
@@ -164,6 +169,8 @@
     @Mock WifiManager mWifiManager;
     @Mock WifiManager.MulticastLock mMulticastLock;
     @Mock ActivityManager mActivityManager;
+    @Mock NetworkNsdReportedMetrics mMetrics;
+    @Mock MdnsUtils.Clock mClock;
     SocketRequestMonitor mSocketRequestMonitor;
     OnUidImportanceListener mUidImportanceListener;
     HandlerThread mThread;
@@ -210,6 +217,9 @@
         doReturn(DEFAULT_RUNNING_APP_ACTIVE_IMPORTANCE_CUTOFF).when(mDeps).getDeviceConfigInt(
                 eq(NsdService.MDNS_CONFIG_RUNNING_APP_ACTIVE_IMPORTANCE_CUTOFF), anyInt());
         doReturn(mAdvertiser).when(mDeps).makeMdnsAdvertiser(any(), any(), any(), any());
+        doReturn(mMetrics).when(mDeps).makeNetworkNsdReportedMetrics(anyBoolean(), anyInt());
+        doReturn(mClock).when(mDeps).makeClock();
+        doReturn(TEST_TIME_MS).when(mClock).elapsedRealtime();
         mService = makeService();
         final ArgumentCaptor<SocketRequestMonitor> cbMonitorCaptor =
                 ArgumentCaptor.forClass(SocketRequestMonitor.class);
@@ -512,14 +522,16 @@
                 eq(SERVICE_NAME), eq(SERVICE_TYPE), eq(PORT), any(), eq(IFACE_IDX_ANY));
 
         // Register service successfully.
+        final int regId = regIdCaptor.getValue();
         final RegistrationInfo registrationInfo = new RegistrationInfo(
-                regIdCaptor.getValue(),
+                regId,
                 IMDnsEventListener.SERVICE_REGISTERED,
                 SERVICE_NAME,
                 SERVICE_TYPE,
                 PORT,
                 new byte[0] /* txtRecord */,
                 IFACE_IDX_ANY);
+        doReturn(TEST_TIME_MS + 10L).when(mClock).elapsedRealtime();
         eventListener.onServiceRegistrationStatus(registrationInfo);
 
         final ArgumentCaptor<NsdServiceInfo> registeredInfoCaptor =
@@ -528,19 +540,22 @@
                 .onServiceRegistered(registeredInfoCaptor.capture());
         final NsdServiceInfo registeredInfo = registeredInfoCaptor.getValue();
         assertEquals(SERVICE_NAME, registeredInfo.getServiceName());
+        verify(mMetrics).reportServiceRegistrationSucceeded(regId, 10L /* durationMs */);
 
         // Fail to register service.
         final RegistrationInfo registrationFailedInfo = new RegistrationInfo(
-                regIdCaptor.getValue(),
+                regId,
                 IMDnsEventListener.SERVICE_REGISTRATION_FAILED,
                 null /* serviceName */,
                 null /* registrationType */,
                 0 /* port */,
                 new byte[0] /* txtRecord */,
                 IFACE_IDX_ANY);
+        doReturn(TEST_TIME_MS + 20L).when(mClock).elapsedRealtime();
         eventListener.onServiceRegistrationStatus(registrationFailedInfo);
         verify(regListener, timeout(TIMEOUT_MS))
                 .onRegistrationFailed(any(), eq(FAILURE_INTERNAL_ERROR));
+        verify(mMetrics).reportServiceRegistrationFailed(regId, 20L /* durationMs */);
     }
 
     @Test
@@ -1215,17 +1230,22 @@
 
         // Verify onServiceRegistered callback
         final MdnsAdvertiser.AdvertiserCallback cb = cbCaptor.getValue();
-        cb.onRegisterServiceSucceeded(idCaptor.getValue(), regInfo);
+        final int regId = idCaptor.getValue();
+        doReturn(TEST_TIME_MS + 10L).when(mClock).elapsedRealtime();
+        cb.onRegisterServiceSucceeded(regId, regInfo);
 
         verify(regListener, timeout(TIMEOUT_MS)).onServiceRegistered(argThat(info -> matches(info,
                 new NsdServiceInfo(regInfo.getServiceName(), null))));
+        verify(mMetrics).reportServiceRegistrationSucceeded(regId, 10L /* durationMs */);
 
+        doReturn(TEST_TIME_MS + 100L).when(mClock).elapsedRealtime();
         client.unregisterService(regListener);
         waitForIdle();
         verify(mAdvertiser).removeService(idCaptor.getValue());
         verify(regListener, timeout(TIMEOUT_MS)).onServiceUnregistered(
                 argThat(info -> matches(info, regInfo)));
         verify(mSocketProvider, timeout(TIMEOUT_MS)).requestStopWhenInactive();
+        verify(mMetrics).reportServiceUnregistration(regId, 100L /* durationMs */);
     }
 
     @Test
@@ -1251,6 +1271,7 @@
 
         verify(regListener, timeout(TIMEOUT_MS)).onRegistrationFailed(
                 argThat(info -> matches(info, regInfo)), eq(FAILURE_INTERNAL_ERROR));
+        verify(mMetrics).reportServiceRegistrationFailed(anyInt(), anyLong());
     }
 
     @Test
@@ -1280,10 +1301,13 @@
 
         // Verify onServiceRegistered callback
         final MdnsAdvertiser.AdvertiserCallback cb = cbCaptor.getValue();
-        cb.onRegisterServiceSucceeded(idCaptor.getValue(), regInfo);
+        final int regId = idCaptor.getValue();
+        doReturn(TEST_TIME_MS + 10L).when(mClock).elapsedRealtime();
+        cb.onRegisterServiceSucceeded(regId, regInfo);
 
         verify(regListener, timeout(TIMEOUT_MS)).onServiceRegistered(
                 argThat(info -> matches(info, new NsdServiceInfo(regInfo.getServiceName(), null))));
+        verify(mMetrics).reportServiceRegistrationSucceeded(regId, 10L /* durationMs */);
     }
 
     @Test
diff --git a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
index 9e604e3..f4d3915 100644
--- a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
@@ -52,6 +52,7 @@
 import android.content.res.Resources;
 import android.net.INetd;
 import android.net.ISocketKeepaliveCallback;
+import android.net.InetAddresses;
 import android.net.KeepalivePacketData;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
@@ -116,7 +117,8 @@
     private static final int MOCK_RESOURCE_ID = 5;
     private static final int TEST_KEEPALIVE_INTERVAL_SEC = 10;
     private static final int TEST_KEEPALIVE_INVALID_INTERVAL_SEC = 9;
-
+    private static final byte[] V4_SRC_ADDR = new byte[] { (byte) 192, 0, 0, (byte) 129 };
+    private static final String TEST_V4_IFACE = "v4-testIface";
     private AutomaticOnOffKeepaliveTracker mAOOKeepaliveTracker;
     private HandlerThread mHandlerThread;
 
@@ -327,6 +329,8 @@
                 NetworkInfo.DetailedState.CONNECTED, "test reason", "test extra info");
         doReturn(new Network(TEST_NETID)).when(mNai).network();
         mNai.linkProperties = new LinkProperties();
+        doReturn(null).when(mNai).translateV4toClatV6(any());
+        doReturn(null).when(mNai).getClatv6SrcAddress();
 
         doReturn(PERMISSION_GRANTED).when(mCtx).checkPermission(any() /* permission */,
                 anyInt() /* pid */, anyInt() /* uid */);
@@ -404,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) {
@@ -429,8 +433,7 @@
     }
 
     private TestKeepaliveInfo doStartNattKeepalive(int intervalSeconds) throws Exception {
-        final InetAddress srcAddress = InetAddress.getByAddress(
-                new byte[] { (byte) 192, 0, 0, (byte) 129 });
+        final InetAddress srcAddress = InetAddress.getByAddress(V4_SRC_ADDR);
         final int srcPort = 12345;
         final InetAddress dstAddress = InetAddress.getByAddress(new byte[] {8, 8, 8, 8});
         final int dstPort = 12345;
@@ -497,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);
 
@@ -522,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.
@@ -609,6 +609,85 @@
         verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
     }
 
+    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 =
+                ArgumentCaptor.forClass(NattKeepalivePacketData.class);
+        verify(mNai).onStartNattSocketKeepalive(
+                eq(TEST_SLOT), eq(TEST_KEEPALIVE_INTERVAL_SEC), kpdCaptor.capture());
+        final NattKeepalivePacketData kpd = kpdCaptor.getValue();
+        // Verify the addresses are updated to v6 when clat is started.
+        assertEquals(v6AddrSrc, kpd.getSrcAddress());
+        assertEquals(v6AddrDst, kpd.getDstAddress());
+
+        triggerEventKeepalive(TEST_SLOT, SocketKeepalive.SUCCESS);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        // Remove clat address should stop the keepalive.
+        doReturn(null).when(mNai).getClatv6SrcAddress();
+        visibleOnHandlerThread(
+                mTestHandler, () -> mAOOKeepaliveTracker.handleCheckKeepalivesStillValid(mNai));
+        checkAndProcessKeepaliveStop();
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_IP_ADDRESS);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
     @Test
     public void testHandleEventSocketKeepalive_startingFailureHardwareError() throws Exception {
         final TestKeepaliveInfo testInfo = doStartNattKeepalive();
@@ -860,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.
@@ -885,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/KeepaliveStatsTrackerTest.java b/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
index 0d2e540..0d1b548 100644
--- a/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
@@ -19,6 +19,8 @@
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 
+import static com.android.testutils.DevSdkIgnoreRule.IgnoreAfter;
+import static com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 import static com.android.testutils.HandlerUtils.visibleOnHandlerThread;
 
 import static org.junit.Assert.assertArrayEquals;
@@ -31,6 +33,7 @@
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 import android.content.BroadcastReceiver;
@@ -62,6 +65,7 @@
 import com.android.testutils.HandlerUtils;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -103,6 +107,8 @@
                 .build();
     }
 
+    @Rule public final DevSdkIgnoreRule ignoreRule = new DevSdkIgnoreRule();
+
     private HandlerThread mHandlerThread;
     private Handler mTestHandler;
 
@@ -1192,4 +1198,43 @@
                     expectKeepaliveCarrierStats1, expectKeepaliveCarrierStats2
                 });
     }
+
+    @Test
+    @IgnoreAfter(Build.VERSION_CODES.S_V2)
+    public void testWriteMetrics_doNothingBeforeT() {
+        // Keepalive stats use repeated atoms, which are only supported on T+. If written to statsd
+        // on S- they will bootloop the system, so they must not be sent on S-. See b/289471411.
+        final int writeTime = 1000;
+        setElapsedRealtime(writeTime);
+        visibleOnHandlerThread(mTestHandler, () -> mKeepaliveStatsTracker.writeAndResetMetrics());
+        verify(mDependencies, never()).writeStats(any());
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testWriteMetrics() {
+        final int writeTime = 1000;
+
+        final ArgumentCaptor<DailykeepaliveInfoReported> dailyKeepaliveInfoReportedCaptor =
+                ArgumentCaptor.forClass(DailykeepaliveInfoReported.class);
+
+        setElapsedRealtime(writeTime);
+        visibleOnHandlerThread(mTestHandler, () -> mKeepaliveStatsTracker.writeAndResetMetrics());
+        // Ensure writeStats is called with the correct DailykeepaliveInfoReported metrics.
+        verify(mDependencies).writeStats(dailyKeepaliveInfoReportedCaptor.capture());
+        final DailykeepaliveInfoReported dailyKeepaliveInfoReported =
+                dailyKeepaliveInfoReportedCaptor.getValue();
+
+        // Same as the no keepalive case
+        final int[] expectRegisteredDurations = new int[] {writeTime};
+        final int[] expectActiveDurations = new int[] {writeTime};
+        assertDailyKeepaliveInfoReported(
+                dailyKeepaliveInfoReported,
+                /* expectRequestsCount= */ 0,
+                /* expectAutoRequestsCount= */ 0,
+                /* expectAppUids= */ new int[0],
+                expectRegisteredDurations,
+                expectActiveDurations,
+                new KeepaliveCarrierStats[0]);
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/NetworkNotificationManagerTest.java b/tests/unit/java/com/android/server/connectivity/NetworkNotificationManagerTest.java
index a27a0bf..b319c30 100644
--- a/tests/unit/java/com/android/server/connectivity/NetworkNotificationManagerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/NetworkNotificationManagerTest.java
@@ -50,6 +50,7 @@
 import android.content.Intent;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
+import android.content.pm.ResolveInfo;
 import android.content.res.Resources;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
@@ -60,8 +61,10 @@
 import android.telephony.TelephonyManager;
 import android.testing.PollingCheck;
 import android.util.DisplayMetrics;
+import android.util.Log;
 import android.widget.TextView;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.StringRes;
 import androidx.test.filters.SmallTest;
@@ -386,14 +389,37 @@
     }
 
     @Test
-    public void testNotifyNoInternetAsDialogWhenHighPriority() throws Exception {
-        doReturn(true).when(mResources).getBoolean(
+    public void testNotifyNoInternet_asNotification() throws Exception {
+        doTestNotifyNotificationAsDialogWhenHighPriority(false, NO_INTERNET);
+    }
+    @Test
+        public void testNotifyNoInternet_asDialog() throws Exception {
+        doTestNotifyNotificationAsDialogWhenHighPriority(true, NO_INTERNET);
+    }
+
+    @Test
+    public void testNotifyLostInternet_asNotification() throws Exception {
+        doTestNotifyNotificationAsDialogWhenHighPriority(false, LOST_INTERNET);
+    }
+
+    @Test
+    public void testNotifyLostInternet_asDialog() throws Exception {
+        doTestNotifyNotificationAsDialogWhenHighPriority(true, LOST_INTERNET);
+    }
+
+    public void doTestNotifyNotificationAsDialogWhenHighPriority(final boolean configActive,
+            @NonNull final NotificationType notifType) throws Exception {
+        doReturn(configActive).when(mResources).getBoolean(
                 R.bool.config_notifyNoInternetAsDialogWhenHighPriority);
 
         final Instrumentation instr = InstrumentationRegistry.getInstrumentation();
         final UiDevice uiDevice =  UiDevice.getInstance(instr);
         final Context ctx = instr.getContext();
         final PowerManager pm = ctx.getSystemService(PowerManager.class);
+        // If the prio of this notif is < that of NETWORK_SWITCH, it's the lowest prio and
+        // therefore it can't be tested whether it cancels other lower-prio notifs.
+        final boolean isLowestPrioNotif = NetworkNotificationManager.priority(notifType)
+                < NetworkNotificationManager.priority(NETWORK_SWITCH);
 
         // Wake up the device (it has no effect if the device is already awake).
         uiDevice.executeShellCommand("input keyevent KEYCODE_WAKEUP");
@@ -404,14 +430,34 @@
 
         // UiDevice.getLauncherPackageName() requires the test manifest to have a <queries> tag for
         // the launcher intent.
+        // Attempted workaround for b/286550950 where Settings is reported as the launcher
+        PollingCheck.check(
+                "Launcher package name was still settings after " + TEST_TIMEOUT_MS + "ms",
+                TEST_TIMEOUT_MS,
+                () -> {
+                    if ("com.android.settings".equals(uiDevice.getLauncherPackageName())) {
+                        final Intent intent = new Intent(Intent.ACTION_MAIN);
+                        intent.addCategory(Intent.CATEGORY_HOME);
+                        final List<ResolveInfo> acts = ctx.getPackageManager()
+                                .queryIntentActivities(intent, PackageManager.MATCH_DEFAULT_ONLY);
+                        Log.e(NetworkNotificationManagerTest.class.getSimpleName(),
+                                "Got settings as launcher name; launcher activities: " + acts);
+                        return false;
+                    }
+                    return true;
+                });
         final String launcherPackageName = uiDevice.getLauncherPackageName();
         assertTrue(String.format("Launcher (%s) is not shown", launcherPackageName),
                 uiDevice.wait(Until.hasObject(By.pkg(launcherPackageName)),
                         UI_AUTOMATOR_WAIT_TIME_MILLIS));
 
-        mManager.showNotification(TEST_NOTIF_ID, NETWORK_SWITCH, mWifiNai, mCellNai, null, false);
-        // Non-"no internet" notifications are not affected
-        verify(mNotificationManager).notify(eq(TEST_NOTIF_TAG), eq(NETWORK_SWITCH.eventId), any());
+        if (!isLowestPrioNotif) {
+            mManager.showNotification(TEST_NOTIF_ID, NETWORK_SWITCH, mWifiNai, mCellNai,
+                    null, false);
+            // Non-"no internet" notifications are not affected
+            verify(mNotificationManager).notify(eq(TEST_NOTIF_TAG), eq(NETWORK_SWITCH.eventId),
+                    any());
+        }
 
         final String testAction = "com.android.connectivity.coverage.TEST_DIALOG";
         final Intent intent = new Intent(testAction)
@@ -420,22 +466,30 @@
         final PendingIntent pendingIntent = PendingIntent.getActivity(ctx, 0 /* requestCode */,
                 intent, PendingIntent.FLAG_CANCEL_CURRENT | PendingIntent.FLAG_IMMUTABLE);
 
-        mManager.showNotification(TEST_NOTIF_ID, NO_INTERNET, mWifiNai, null /* switchToNai */,
+        mManager.showNotification(TEST_NOTIF_ID, notifType, mWifiNai, null /* switchToNai */,
                 pendingIntent, true /* highPriority */);
 
-        // Previous notifications are still dismissed
-        verify(mNotificationManager).cancel(TEST_NOTIF_TAG, NETWORK_SWITCH.eventId);
+        if (!isLowestPrioNotif) {
+            // Previous notifications are still dismissed
+            verify(mNotificationManager).cancel(TEST_NOTIF_TAG, NETWORK_SWITCH.eventId);
+        }
 
-        // Verify that the activity is shown (the activity shows the action on screen)
-        final UiObject actionText = uiDevice.findObject(new UiSelector().text(testAction));
-        assertTrue("Activity not shown", actionText.waitForExists(TEST_TIMEOUT_MS));
+        if (configActive) {
+            // Verify that the activity is shown (the activity shows the action on screen)
+            final UiObject actionText = uiDevice.findObject(new UiSelector().text(testAction));
+            assertTrue("Activity not shown", actionText.waitForExists(TEST_TIMEOUT_MS));
 
-        // Tapping the text should dismiss the dialog
-        actionText.click();
-        assertTrue("Activity not dismissed", actionText.waitUntilGone(TEST_TIMEOUT_MS));
+            // Tapping the text should dismiss the dialog
+            actionText.click();
+            assertTrue("Activity not dismissed", actionText.waitUntilGone(TEST_TIMEOUT_MS));
 
-        // Verify no NO_INTERNET notification was posted
-        verify(mNotificationManager, never()).notify(any(), eq(NO_INTERNET.eventId), any());
+            // Verify that the notification was not posted
+            verify(mNotificationManager, never()).notify(any(), eq(notifType.eventId), any());
+        } else {
+            // Notification should have been posted, and will have overridden the previous
+            // one because it has the same id (hence no cancel).
+            verify(mNotificationManager).notify(eq(TEST_NOTIF_TAG), eq(notifType.eventId), any());
+        }
     }
 
     private void doNotificationTextTest(NotificationType type, @StringRes int expectedTitleRes,
diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index dc50773..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
@@ -2768,23 +2777,30 @@
                 new PersistableBundle());
     }
 
-    private void verifyMobikeTriggered(List<Network> expected) {
+    private void verifyMobikeTriggered(List<Network> expected, int retryIndex) {
+        // Verify retry is scheduled
+        final long expectedDelayMs = mTestDeps.getValidationFailRecoveryMs(retryIndex);
+        final ArgumentCaptor<Long> delayCaptor = ArgumentCaptor.forClass(Long.class);
+        verify(mExecutor, times(retryIndex + 1)).schedule(
+                any(Runnable.class), delayCaptor.capture(), eq(TimeUnit.MILLISECONDS));
+        final List<Long> delays = delayCaptor.getAllValues();
+        assertEquals(expectedDelayMs, (long) delays.get(delays.size() - 1));
+
         final ArgumentCaptor<Network> networkCaptor = ArgumentCaptor.forClass(Network.class);
-        verify(mIkeSessionWrapper).setNetwork(networkCaptor.capture(),
-                anyInt() /* ipVersion */, anyInt() /* encapType */, anyInt() /* keepaliveDelay */);
+        verify(mIkeSessionWrapper, timeout(TEST_TIMEOUT_MS + expectedDelayMs))
+                .setNetwork(networkCaptor.capture(), anyInt() /* ipVersion */,
+                        anyInt() /* encapType */, anyInt() /* keepaliveDelay */);
         assertEquals(expected, Collections.singletonList(networkCaptor.getValue()));
     }
 
     @Test
     public void testDataStallInIkev2VpnMobikeDisabled() throws Exception {
-        verifySetupPlatformVpn(
+        final PlatformVpnSnapshot vpnSnapShot = verifySetupPlatformVpn(
                 createIkeConfig(createIkeConnectInfo(), false /* isMobikeEnabled */));
 
         doReturn(TEST_NETWORK).when(mMockNetworkAgent).getNetwork();
-        final ConnectivityDiagnosticsCallback connectivityDiagCallback =
-                getConnectivityDiagCallback();
-        final DataStallReport report = createDataStallReport();
-        connectivityDiagCallback.onDataStallSuspected(report);
+        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
+                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
 
         // Should not trigger MOBIKE if MOBIKE is not enabled
         verify(mIkeSessionWrapper, never()).setNetwork(any() /* network */,
@@ -2797,19 +2813,11 @@
                 createIkeConfig(createIkeConnectInfo(), true /* isMobikeEnabled */));
 
         doReturn(TEST_NETWORK).when(mMockNetworkAgent).getNetwork();
-        final ConnectivityDiagnosticsCallback connectivityDiagCallback =
-                getConnectivityDiagCallback();
-        final DataStallReport report = createDataStallReport();
-        connectivityDiagCallback.onDataStallSuspected(report);
-
+        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
+                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
         // Verify MOBIKE is triggered
-        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks());
-
-        // Expect to skip other data stall event if MOBIKE was started.
-        reset(mIkeSessionWrapper);
-        connectivityDiagCallback.onDataStallSuspected(report);
-        verify(mIkeSessionWrapper, never()).setNetwork(any() /* network */,
-                anyInt() /* ipVersion */, anyInt() /* encapType */, anyInt() /* keepaliveDelay */);
+        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks(),
+                0 /* retryIndex */);
 
         reset(mIkev2SessionCreator);
 
@@ -2819,14 +2827,6 @@
                 NetworkAgent.VALIDATION_STATUS_VALID);
         verify(mIkev2SessionCreator, never()).createIkeSession(
                 any(), any(), any(), any(), any(), any());
-
-        // Send invalid result to verify no ike session reset since the data stall suspected
-        // variables(timer counter and boolean) was reset.
-        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
-                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
-        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
-        verify(mIkev2SessionCreator, never()).createIkeSession(
-                any(), any(), any(), any(), any(), any());
     }
 
     @Test
@@ -2834,31 +2834,46 @@
         final PlatformVpnSnapshot vpnSnapShot = verifySetupPlatformVpn(
                 createIkeConfig(createIkeConnectInfo(), true /* isMobikeEnabled */));
 
-        final ConnectivityDiagnosticsCallback connectivityDiagCallback =
-                getConnectivityDiagCallback();
-
+        int retry = 0;
         doReturn(TEST_NETWORK).when(mMockNetworkAgent).getNetwork();
-        final DataStallReport report = createDataStallReport();
-        connectivityDiagCallback.onDataStallSuspected(report);
-
-        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks());
+        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
+                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
+        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks(),
+                retry++);
 
         reset(mIkev2SessionCreator);
 
+        // Second validation status update.
+        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
+                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
+        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks(),
+                retry++);
+
+        // Use real delay to verify reset session will not be performed if there is an existing
+        // recovery for resetting the session.
+        mExecutor.delayMs = TestExecutor.REAL_DELAY;
+        mExecutor.executeDirect = true;
         // Send validation status update should result in ike session reset.
         ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
                 NetworkAgent.VALIDATION_STATUS_NOT_VALID);
 
-        // Verify reset is scheduled and run.
-        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
+        // Verify session reset is scheduled
+        long expectedDelay = mTestDeps.getValidationFailRecoveryMs(retry++);
+        final ArgumentCaptor<Long> delayCaptor = ArgumentCaptor.forClass(Long.class);
+        verify(mExecutor, times(retry)).schedule(any(Runnable.class), delayCaptor.capture(),
+                eq(TimeUnit.MILLISECONDS));
+        final List<Long> delays = delayCaptor.getAllValues();
+        assertEquals(expectedDelay, (long) delays.get(delays.size() - 1));
 
         // Another invalid status reported should not trigger other scheduled recovery.
-        reset(mExecutor);
+        expectedDelay = mTestDeps.getValidationFailRecoveryMs(retry++);
         ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
                 NetworkAgent.VALIDATION_STATUS_NOT_VALID);
-        verify(mExecutor, never()).schedule(any(Runnable.class), anyLong(), any());
+        verify(mExecutor, never()).schedule(
+                any(Runnable.class), eq(expectedDelay), eq(TimeUnit.MILLISECONDS));
 
-        verify(mIkev2SessionCreator, timeout(TEST_TIMEOUT_MS))
+        // Verify that session being reset
+        verify(mIkev2SessionCreator, timeout(TEST_TIMEOUT_MS + expectedDelay))
                 .createIkeSession(any(), any(), any(), any(), any(), any());
     }
 
@@ -3137,6 +3152,12 @@
         }
 
         @Override
+        public long getValidationFailRecoveryMs(int retryCount) {
+            // Simply return retryCount as the delay seconds for retrying.
+            return retryCount * 100L;
+        }
+
+        @Override
         public ScheduledThreadPoolExecutor newScheduledThreadPoolExecutor() {
             return mExecutor;
         }
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsAdvertiserTest.kt b/tests/unit/java/com/android/server/connectivity/mdns/MdnsAdvertiserTest.kt
index c467f45..6a0334f 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsAdvertiserTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsAdvertiserTest.kt
@@ -56,8 +56,8 @@
 private val TEST_ADDR = parseNumericAddress("2001:db8::123")
 private val TEST_LINKADDR = LinkAddress(TEST_ADDR, 64 /* prefixLength */)
 private val TEST_NETWORK_1 = mock(Network::class.java)
-private val TEST_SOCKETKEY_1 = mock(SocketKey::class.java)
-private val TEST_SOCKETKEY_2 = mock(SocketKey::class.java)
+private val TEST_SOCKETKEY_1 = SocketKey(1001 /* interfaceIndex */)
+private val TEST_SOCKETKEY_2 = SocketKey(1002 /* interfaceIndex */)
 private val TEST_HOSTNAME = arrayOf("Android_test", "local")
 private const val TEST_SUBTYPE = "_subtype"
 
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
index d2298fe..e869b91 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
@@ -19,7 +19,6 @@
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.never;
@@ -54,6 +53,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
 
 /** Tests for {@link MdnsDiscoveryManager}. */
 @RunWith(DevSdkIgnoreRunner.class)
@@ -81,6 +81,7 @@
     private static final Pair<String, SocketKey> PER_SOCKET_SERVICE_TYPE_2_NETWORK_2 =
             Pair.create(SERVICE_TYPE_2, SOCKET_KEY_NETWORK_2);
     @Mock private ExecutorProvider executorProvider;
+    @Mock private ScheduledExecutorService mockExecutorService;
     @Mock private MdnsSocketClientBase socketClient;
     @Mock private MdnsServiceTypeClient mockServiceTypeClientType1NullNetwork;
     @Mock private MdnsServiceTypeClient mockServiceTypeClientType1Network1;
@@ -129,6 +130,7 @@
                         return null;
                     }
                 };
+        doReturn(mockExecutorService).when(mockServiceTypeClientType1NullNetwork).getExecutor();
     }
 
     @After
@@ -166,11 +168,25 @@
         when(mockServiceTypeClientType1NullNetwork.stopSendAndReceive(mockListenerOne))
                 .thenReturn(true);
         runOnHandler(() -> discoveryManager.unregisterListener(SERVICE_TYPE_1, mockListenerOne));
+        verify(executorProvider).shutdownExecutorService(mockExecutorService);
         verify(mockServiceTypeClientType1NullNetwork).stopSendAndReceive(mockListenerOne);
         verify(socketClient).stopDiscovery();
     }
 
     @Test
+    public void onSocketDestroy_shutdownExecutorService() throws IOException {
+        final MdnsSearchOptions options =
+                MdnsSearchOptions.newBuilder().setNetwork(null /* network */).build();
+        final SocketCreationCallback callback = expectSocketCreationCallback(
+                SERVICE_TYPE_1, mockListenerOne, options);
+        runOnHandler(() -> callback.onSocketCreated(SOCKET_KEY_NULL_NETWORK));
+        verify(mockServiceTypeClientType1NullNetwork).startSendAndReceive(mockListenerOne, options);
+
+        runOnHandler(() -> callback.onSocketDestroyed(SOCKET_KEY_NULL_NETWORK));
+        verify(executorProvider).shutdownExecutorService(mockExecutorService);
+    }
+
+    @Test
     public void registerMultipleListeners() throws IOException {
         final MdnsSearchOptions options =
                 MdnsSearchOptions.newBuilder().setNetwork(null /* network */).build();
@@ -212,31 +228,31 @@
         runOnHandler(() -> discoveryManager.onResponseReceived(
                 responseForServiceTypeOne, SOCKET_KEY_NULL_NETWORK));
         // Packets for network null are only processed by the ServiceTypeClient for network null
-        verify(mockServiceTypeClientType1NullNetwork).processResponse(responseForServiceTypeOne,
-                SOCKET_KEY_NULL_NETWORK.getInterfaceIndex(), SOCKET_KEY_NULL_NETWORK.getNetwork());
-        verify(mockServiceTypeClientType1Network1, never()).processResponse(any(), anyInt(), any());
-        verify(mockServiceTypeClientType2Network2, never()).processResponse(any(), anyInt(), any());
+        verify(mockServiceTypeClientType1NullNetwork).processResponse(
+                responseForServiceTypeOne, SOCKET_KEY_NULL_NETWORK);
+        verify(mockServiceTypeClientType1Network1, never()).processResponse(any(), any());
+        verify(mockServiceTypeClientType2Network2, never()).processResponse(any(), any());
 
         final MdnsPacket responseForServiceTypeTwo = createMdnsPacket(SERVICE_TYPE_2);
         runOnHandler(() -> discoveryManager.onResponseReceived(
                 responseForServiceTypeTwo, SOCKET_KEY_NETWORK_1));
-        verify(mockServiceTypeClientType1NullNetwork, never()).processResponse(any(), anyInt(),
-                eq(SOCKET_KEY_NETWORK_1.getNetwork()));
-        verify(mockServiceTypeClientType1Network1).processResponse(responseForServiceTypeTwo,
-                SOCKET_KEY_NETWORK_1.getInterfaceIndex(), SOCKET_KEY_NETWORK_1.getNetwork());
-        verify(mockServiceTypeClientType2Network2, never()).processResponse(any(), anyInt(),
-                eq(SOCKET_KEY_NETWORK_1.getNetwork()));
+        verify(mockServiceTypeClientType1NullNetwork, never()).processResponse(any(),
+                eq(SOCKET_KEY_NETWORK_1));
+        verify(mockServiceTypeClientType1Network1).processResponse(
+                responseForServiceTypeTwo, SOCKET_KEY_NETWORK_1);
+        verify(mockServiceTypeClientType2Network2, never()).processResponse(any(),
+                eq(SOCKET_KEY_NETWORK_1));
 
         final MdnsPacket responseForSubtype =
                 createMdnsPacket("subtype._sub._googlecast._tcp.local");
         runOnHandler(() -> discoveryManager.onResponseReceived(
                 responseForSubtype, SOCKET_KEY_NETWORK_2));
-        verify(mockServiceTypeClientType1NullNetwork, never()).processResponse(any(), anyInt(),
-                eq(SOCKET_KEY_NETWORK_2.getNetwork()));
-        verify(mockServiceTypeClientType1Network1, never()).processResponse(any(), anyInt(),
-                eq(SOCKET_KEY_NETWORK_2.getNetwork()));
-        verify(mockServiceTypeClientType2Network2).processResponse(responseForSubtype,
-                SOCKET_KEY_NETWORK_2.getInterfaceIndex(), SOCKET_KEY_NETWORK_2.getNetwork());
+        verify(mockServiceTypeClientType1NullNetwork, never()).processResponse(any(),
+                eq(SOCKET_KEY_NETWORK_2));
+        verify(mockServiceTypeClientType1Network1, never()).processResponse(any(),
+                eq(SOCKET_KEY_NETWORK_2));
+        verify(mockServiceTypeClientType2Network2).processResponse(
+                responseForSubtype, SOCKET_KEY_NETWORK_2);
     }
 
     @Test
@@ -260,15 +276,13 @@
         // Receive a response, it should be processed on both clients.
         final MdnsPacket response = createMdnsPacket(SERVICE_TYPE_1);
         runOnHandler(() -> discoveryManager.onResponseReceived(response, SOCKET_KEY_NETWORK_1));
-        verify(mockServiceTypeClientType1Network1).processResponse(response,
-                SOCKET_KEY_NETWORK_1.getInterfaceIndex(), SOCKET_KEY_NETWORK_1.getNetwork());
-        verify(mockServiceTypeClientType2Network1).processResponse(response,
-                SOCKET_KEY_NETWORK_1.getInterfaceIndex(), SOCKET_KEY_NETWORK_1.getNetwork());
+        verify(mockServiceTypeClientType1Network1).processResponse(response, SOCKET_KEY_NETWORK_1);
+        verify(mockServiceTypeClientType2Network1).processResponse(response, SOCKET_KEY_NETWORK_1);
 
         // The first callback receives a notification that the network has been destroyed,
         // mockServiceTypeClientOne1 should send service removed notifications and remove from the
         // list of clients.
-        runOnHandler(() -> callback.onAllSocketsDestroyed(SOCKET_KEY_NETWORK_1));
+        runOnHandler(() -> callback.onSocketDestroyed(SOCKET_KEY_NETWORK_1));
         verify(mockServiceTypeClientType1Network1).notifySocketDestroyed();
 
         // Receive a response again, it should be processed only on
@@ -276,23 +290,23 @@
         // removed from the list of clients, it is no longer able to process responses.
         runOnHandler(() -> discoveryManager.onResponseReceived(response, SOCKET_KEY_NETWORK_1));
         // Still times(1) as a response was received once previously
-        verify(mockServiceTypeClientType1Network1, times(1)).processResponse(response,
-                SOCKET_KEY_NETWORK_1.getInterfaceIndex(), SOCKET_KEY_NETWORK_1.getNetwork());
-        verify(mockServiceTypeClientType2Network1, times(2)).processResponse(response,
-                SOCKET_KEY_NETWORK_1.getInterfaceIndex(), SOCKET_KEY_NETWORK_1.getNetwork());
+        verify(mockServiceTypeClientType1Network1, times(1)).processResponse(
+                response, SOCKET_KEY_NETWORK_1);
+        verify(mockServiceTypeClientType2Network1, times(2)).processResponse(
+                response, SOCKET_KEY_NETWORK_1);
 
         // The client for NETWORK_1 receives the callback that the NETWORK_2 has been destroyed,
         // mockServiceTypeClientTwo2 shouldn't send any notifications.
-        runOnHandler(() -> callback2.onAllSocketsDestroyed(SOCKET_KEY_NETWORK_2));
+        runOnHandler(() -> callback2.onSocketDestroyed(SOCKET_KEY_NETWORK_2));
         verify(mockServiceTypeClientType2Network1, never()).notifySocketDestroyed();
 
         // Receive a response again, mockServiceTypeClientType2Network1 is still in the list of
         // clients, it's still able to process responses.
         runOnHandler(() -> discoveryManager.onResponseReceived(response, SOCKET_KEY_NETWORK_1));
-        verify(mockServiceTypeClientType1Network1, times(1)).processResponse(response,
-                SOCKET_KEY_NETWORK_1.getInterfaceIndex(), SOCKET_KEY_NETWORK_1.getNetwork());
-        verify(mockServiceTypeClientType2Network1, times(3)).processResponse(response,
-                SOCKET_KEY_NETWORK_1.getInterfaceIndex(), SOCKET_KEY_NETWORK_1.getNetwork());
+        verify(mockServiceTypeClientType1Network1, times(1)).processResponse(
+                response, SOCKET_KEY_NETWORK_1);
+        verify(mockServiceTypeClientType2Network1, times(3)).processResponse(
+                response, SOCKET_KEY_NETWORK_1);
     }
 
     @Test
@@ -310,17 +324,17 @@
         final MdnsPacket response = createMdnsPacket(SERVICE_TYPE_1);
         final int ifIndex = 1;
         runOnHandler(() -> discoveryManager.onResponseReceived(response, SOCKET_KEY_NULL_NETWORK));
-        verify(mockServiceTypeClientType1NullNetwork).processResponse(response,
-                SOCKET_KEY_NULL_NETWORK.getInterfaceIndex(), SOCKET_KEY_NULL_NETWORK.getNetwork());
+        verify(mockServiceTypeClientType1NullNetwork).processResponse(
+                response, SOCKET_KEY_NULL_NETWORK);
 
-        runOnHandler(() -> callback.onAllSocketsDestroyed(SOCKET_KEY_NULL_NETWORK));
+        runOnHandler(() -> callback.onSocketDestroyed(SOCKET_KEY_NULL_NETWORK));
         verify(mockServiceTypeClientType1NullNetwork).notifySocketDestroyed();
 
         // Receive a response again, it should not be processed.
         runOnHandler(() -> discoveryManager.onResponseReceived(response, SOCKET_KEY_NULL_NETWORK));
         // Still times(1) as a response was received once previously
-        verify(mockServiceTypeClientType1NullNetwork, times(1)).processResponse(response,
-                SOCKET_KEY_NULL_NETWORK.getInterfaceIndex(), SOCKET_KEY_NULL_NETWORK.getNetwork());
+        verify(mockServiceTypeClientType1NullNetwork, times(1)).processResponse(
+                response, SOCKET_KEY_NULL_NETWORK);
 
         // Unregister the listener, notifyNetworkUnrequested should be called but other stop methods
         // won't be call because the service type client was unregistered and destroyed. But those
@@ -329,7 +343,7 @@
         verify(socketClient).notifyNetworkUnrequested(mockListenerOne);
         verify(mockServiceTypeClientType1NullNetwork, never()).stopSendAndReceive(any());
         // The stopDiscovery() is only used by MdnsSocketClient, which doesn't send
-        // onAllSocketsDestroyed(). So the socket clients that send onAllSocketsDestroyed() do not
+        // onSocketDestroyed(). So the socket clients that send onSocketDestroyed() do not
         // need to call stopDiscovery().
         verify(socketClient, never()).stopDiscovery();
     }
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
index b812fa6..29de272 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
@@ -28,7 +28,6 @@
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.net.InetAddresses;
 import android.net.Network;
@@ -67,18 +66,18 @@
     @Mock private MdnsServiceBrowserListener mListener;
     @Mock private MdnsSocketClientBase.Callback mCallback;
     @Mock private SocketCreationCallback mSocketCreationCallback;
-    @Mock private SocketKey mSocketKey;
     private MdnsMultinetworkSocketClient mSocketClient;
     private Handler mHandler;
+    private SocketKey mSocketKey;
 
     @Before
     public void setUp() throws SocketException {
         MockitoAnnotations.initMocks(this);
-        doReturn(mNetwork).when(mSocketKey).getNetwork();
 
         final HandlerThread thread = new HandlerThread("MdnsMultinetworkSocketClientTest");
         thread.start();
         mHandler = new Handler(thread.getLooper());
+        mSocketKey = new SocketKey(1000 /* interfaceIndex */);
         mSocketClient = new MdnsMultinetworkSocketClient(thread.getLooper(), mProvider);
         mHandler.post(() -> mSocketClient.setCallback(mCallback));
     }
@@ -125,10 +124,8 @@
             doReturn(createEmptyNetworkInterface()).when(socket).getInterface();
         }
 
-        final SocketKey tetherSocketKey1 = mock(SocketKey.class);
-        final SocketKey tetherSocketKey2 = mock(SocketKey.class);
-        doReturn(null).when(tetherSocketKey1).getNetwork();
-        doReturn(null).when(tetherSocketKey2).getNetwork();
+        final SocketKey tetherSocketKey1 = new SocketKey(1001 /* interfaceIndex */);
+        final SocketKey tetherSocketKey2 = new SocketKey(1002 /* interfaceIndex */);
         // Notify socket created
         callback.onSocketCreated(mSocketKey, mSocket, List.of());
         verify(mSocketCreationCallback).onSocketCreated(mSocketKey);
@@ -137,8 +134,8 @@
         callback.onSocketCreated(tetherSocketKey2, tetherIfaceSock2, List.of());
         verify(mSocketCreationCallback).onSocketCreated(tetherSocketKey2);
 
-        // Send packet to IPv4 with target network and verify sending has been called.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mNetwork,
+        // Send packet to IPv4 with mSocketKey and verify sending has been called.
+        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mSocketKey,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket).send(ipv4Packet);
@@ -146,30 +143,30 @@
         verify(tetherIfaceSock2, never()).send(any());
 
         // Send packet to IPv4 with onlyUseIpv6OnIpv6OnlyNetworks = true, the packet will be sent.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mNetwork,
+        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mSocketKey,
                 true /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket, times(2)).send(ipv4Packet);
         verify(tetherIfaceSock1, never()).send(any());
         verify(tetherIfaceSock2, never()).send(any());
 
-        // Send packet to IPv6 without target network and verify sending has been called.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv6Packet, null,
+        // Send packet to IPv6 with tetherSocketKey1 and verify sending has been called.
+        mSocketClient.sendPacketRequestingMulticastResponse(ipv6Packet, tetherSocketKey1,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket, never()).send(ipv6Packet);
         verify(tetherIfaceSock1).send(ipv6Packet);
-        verify(tetherIfaceSock2).send(ipv6Packet);
+        verify(tetherIfaceSock2, never()).send(ipv6Packet);
 
         // Send packet to IPv6 with onlyUseIpv6OnIpv6OnlyNetworks = true, the packet will not be
         // sent. Therefore, the tetherIfaceSock1.send() and tetherIfaceSock2.send() are still be
         // called once.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv6Packet, null,
+        mSocketClient.sendPacketRequestingMulticastResponse(ipv6Packet, tetherSocketKey1,
                 true /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket, never()).send(ipv6Packet);
         verify(tetherIfaceSock1, times(1)).send(ipv6Packet);
-        verify(tetherIfaceSock2, times(1)).send(ipv6Packet);
+        verify(tetherIfaceSock2, never()).send(ipv6Packet);
     }
 
     @Test
@@ -240,8 +237,8 @@
         doReturn(createEmptyNetworkInterface()).when(socket2).getInterface();
         doReturn(createEmptyNetworkInterface()).when(socket3).getInterface();
 
-        final SocketKey socketKey2 = mock(SocketKey.class);
-        final SocketKey socketKey3 = mock(SocketKey.class);
+        final SocketKey socketKey2 = new SocketKey(1001 /* interfaceIndex */);
+        final SocketKey socketKey3 = new SocketKey(1002 /* interfaceIndex */);
         callback.onSocketCreated(mSocketKey, mSocket, List.of());
         callback.onSocketCreated(socketKey2, socket2, List.of());
         callback.onSocketCreated(socketKey3, socket3, List.of());
@@ -249,8 +246,8 @@
         verify(mSocketCreationCallback).onSocketCreated(socketKey2);
         verify(mSocketCreationCallback).onSocketCreated(socketKey3);
 
-        // Send IPv4 packet on the non-null Network and verify sending has been called.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mNetwork,
+        // Send IPv4 packet on the mSocketKey and verify sending has been called.
+        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mSocketKey,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket).send(ipv4Packet);
@@ -278,43 +275,44 @@
         verify(socketCreationCb2).onSocketCreated(socketKey2);
         verify(socketCreationCb2).onSocketCreated(socketKey3);
 
-        // Send IPv4 packet to null network and verify sending to the 2 tethered interface sockets.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, null,
+        // Send IPv4 packet on socket2 and verify sending to the socket2 only.
+        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, socketKey2,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         // ipv4Packet still sent only once on mSocket: times(1) matches the packet sent earlier on
         // mNetwork
         verify(mSocket, times(1)).send(ipv4Packet);
         verify(socket2).send(ipv4Packet);
-        verify(socket3).send(ipv4Packet);
+        verify(socket3, never()).send(ipv4Packet);
 
         // Unregister the second request
         mHandler.post(() -> mSocketClient.notifyNetworkUnrequested(listener2));
         verify(mProvider, timeout(DEFAULT_TIMEOUT)).unrequestSocket(callback2);
 
         // Send IPv4 packet again and verify it's still sent a second time
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, null,
+        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, socketKey2,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(socket2, times(2)).send(ipv4Packet);
-        verify(socket3, times(2)).send(ipv4Packet);
+        verify(socket3, never()).send(ipv4Packet);
 
         // Unrequest remaining sockets
         mHandler.post(() -> mSocketClient.notifyNetworkUnrequested(mListener));
         verify(mProvider, timeout(DEFAULT_TIMEOUT)).unrequestSocket(callback);
 
         // Send IPv4 packet and verify no more sending.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, null,
+        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mSocketKey,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket, times(1)).send(ipv4Packet);
         verify(socket2, times(2)).send(ipv4Packet);
-        verify(socket3, times(2)).send(ipv4Packet);
+        verify(socket3, never()).send(ipv4Packet);
     }
 
     @Test
     public void testNotifyNetworkUnrequested_SocketsOnNullNetwork() {
         final MdnsInterfaceSocket otherSocket = mock(MdnsInterfaceSocket.class);
+        final SocketKey otherSocketKey = new SocketKey(1001 /* interfaceIndex */);
         final SocketCallback callback = expectSocketCallback(
                 mListener, null /* requestedNetwork */);
         doReturn(createEmptyNetworkInterface()).when(mSocket).getInterface();
@@ -322,33 +320,36 @@
 
         callback.onSocketCreated(mSocketKey, mSocket, List.of());
         verify(mSocketCreationCallback).onSocketCreated(mSocketKey);
-        callback.onSocketCreated(mSocketKey, otherSocket, List.of());
-        verify(mSocketCreationCallback, times(2)).onSocketCreated(mSocketKey);
+        callback.onSocketCreated(otherSocketKey, otherSocket, List.of());
+        verify(mSocketCreationCallback).onSocketCreated(otherSocketKey);
 
-        verify(mSocketCreationCallback, never()).onAllSocketsDestroyed(mSocketKey);
+        verify(mSocketCreationCallback, never()).onSocketDestroyed(mSocketKey);
+        verify(mSocketCreationCallback, never()).onSocketDestroyed(otherSocketKey);
         mHandler.post(() -> mSocketClient.notifyNetworkUnrequested(mListener));
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
 
         verify(mProvider).unrequestSocket(callback);
-        verify(mSocketCreationCallback).onAllSocketsDestroyed(mSocketKey);
+        verify(mSocketCreationCallback).onSocketDestroyed(mSocketKey);
+        verify(mSocketCreationCallback).onSocketDestroyed(otherSocketKey);
     }
 
     @Test
     public void testSocketCreatedAndDestroyed_NullNetwork() throws IOException {
         final MdnsInterfaceSocket otherSocket = mock(MdnsInterfaceSocket.class);
+        final SocketKey otherSocketKey = new SocketKey(1001 /* interfaceIndex */);
         final SocketCallback callback = expectSocketCallback(mListener, null /* network */);
         doReturn(createEmptyNetworkInterface()).when(mSocket).getInterface();
         doReturn(createEmptyNetworkInterface()).when(otherSocket).getInterface();
 
         callback.onSocketCreated(mSocketKey, mSocket, List.of());
         verify(mSocketCreationCallback).onSocketCreated(mSocketKey);
-        callback.onSocketCreated(mSocketKey, otherSocket, List.of());
-        verify(mSocketCreationCallback, times(2)).onSocketCreated(mSocketKey);
+        callback.onSocketCreated(otherSocketKey, otherSocket, List.of());
+        verify(mSocketCreationCallback).onSocketCreated(otherSocketKey);
 
         // Notify socket destroyed
         callback.onInterfaceDestroyed(mSocketKey, mSocket);
-        verifyNoMoreInteractions(mSocketCreationCallback);
-        callback.onInterfaceDestroyed(mSocketKey, otherSocket);
-        verify(mSocketCreationCallback).onAllSocketsDestroyed(mSocketKey);
+        verify(mSocketCreationCallback).onSocketDestroyed(mSocketKey);
+        callback.onInterfaceDestroyed(otherSocketKey, otherSocket);
+        verify(mSocketCreationCallback).onSocketDestroyed(otherSocketKey);
     }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsResponseDecoderTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsResponseDecoderTests.java
index 05eca84..d71bea4 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsResponseDecoderTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsResponseDecoderTests.java
@@ -18,7 +18,7 @@
 
 import static android.net.InetAddresses.parseNumericAddress;
 
-import static com.android.server.connectivity.mdns.MdnsResponseDecoder.Clock;
+import static com.android.server.connectivity.mdns.util.MdnsUtils.Clock;
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.junit.Assert.assertArrayEquals;
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceCacheTest.kt b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceCacheTest.kt
index f091eea..b43bcf7 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceCacheTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceCacheTest.kt
@@ -16,7 +16,6 @@
 
 package com.android.server.connectivity.mdns
 
-import android.net.Network
 import android.os.Build
 import android.os.Handler
 import android.os.HandlerThread
@@ -32,7 +31,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mockito.mock
 
 private const val SERVICE_NAME_1 = "service-instance-1"
 private const val SERVICE_NAME_2 = "service-instance-2"
@@ -44,7 +42,7 @@
 @RunWith(DevSdkIgnoreRunner::class)
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
 class MdnsServiceCacheTest {
-    private val network = mock(Network::class.java)
+    private val socketKey = SocketKey(null /* network */, INTERFACE_INDEX)
     private val thread = HandlerThread(MdnsServiceCacheTest::class.simpleName)
     private val handler by lazy {
         Handler(thread.looper)
@@ -71,39 +69,47 @@
         return future.get(DEFAULT_TIMEOUT_MS, TimeUnit.MILLISECONDS)
     }
 
-    private fun addOrUpdateService(serviceType: String, network: Network, service: MdnsResponse):
-            Unit = runningOnHandlerAndReturn {
-        serviceCache.addOrUpdateService(serviceType, network, service) }
+    private fun addOrUpdateService(
+            serviceType: String,
+            socketKey: SocketKey,
+            service: MdnsResponse
+    ): Unit = runningOnHandlerAndReturn {
+        serviceCache.addOrUpdateService(serviceType, socketKey, service)
+    }
 
-    private fun removeService(serviceName: String, serviceType: String, network: Network):
+    private fun removeService(serviceName: String, serviceType: String, socketKey: SocketKey):
             Unit = runningOnHandlerAndReturn {
-        serviceCache.removeService(serviceName, serviceType, network) }
+        serviceCache.removeService(serviceName, serviceType, socketKey) }
 
-    private fun getService(serviceName: String, serviceType: String, network: Network):
+    private fun getService(serviceName: String, serviceType: String, socketKey: SocketKey):
             MdnsResponse? = runningOnHandlerAndReturn {
-        serviceCache.getCachedService(serviceName, serviceType, network) }
+        serviceCache.getCachedService(serviceName, serviceType, socketKey) }
 
-    private fun getServices(serviceType: String, network: Network): List<MdnsResponse> =
-        runningOnHandlerAndReturn { serviceCache.getCachedServices(serviceType, network) }
+    private fun getServices(serviceType: String, socketKey: SocketKey): List<MdnsResponse> =
+        runningOnHandlerAndReturn { serviceCache.getCachedServices(serviceType, socketKey) }
 
     @Test
     fun testAddAndRemoveService() {
-        addOrUpdateService(SERVICE_TYPE_1, network, createResponse(SERVICE_NAME_1, SERVICE_TYPE_1))
-        var response = getService(SERVICE_NAME_1, SERVICE_TYPE_1, network)
+        addOrUpdateService(
+                SERVICE_TYPE_1, socketKey, createResponse(SERVICE_NAME_1, SERVICE_TYPE_1))
+        var response = getService(SERVICE_NAME_1, SERVICE_TYPE_1, socketKey)
         assertNotNull(response)
         assertEquals(SERVICE_NAME_1, response.serviceInstanceName)
-        removeService(SERVICE_NAME_1, SERVICE_TYPE_1, network)
-        response = getService(SERVICE_NAME_1, SERVICE_TYPE_1, network)
+        removeService(SERVICE_NAME_1, SERVICE_TYPE_1, socketKey)
+        response = getService(SERVICE_NAME_1, SERVICE_TYPE_1, socketKey)
         assertNull(response)
     }
 
     @Test
     fun testGetCachedServices_multipleServiceTypes() {
-        addOrUpdateService(SERVICE_TYPE_1, network, createResponse(SERVICE_NAME_1, SERVICE_TYPE_1))
-        addOrUpdateService(SERVICE_TYPE_1, network, createResponse(SERVICE_NAME_2, SERVICE_TYPE_1))
-        addOrUpdateService(SERVICE_TYPE_2, network, createResponse(SERVICE_NAME_2, SERVICE_TYPE_2))
+        addOrUpdateService(
+                SERVICE_TYPE_1, socketKey, createResponse(SERVICE_NAME_1, SERVICE_TYPE_1))
+        addOrUpdateService(
+                SERVICE_TYPE_1, socketKey, createResponse(SERVICE_NAME_2, SERVICE_TYPE_1))
+        addOrUpdateService(
+                SERVICE_TYPE_2, socketKey, createResponse(SERVICE_NAME_2, SERVICE_TYPE_2))
 
-        val responses1 = getServices(SERVICE_TYPE_1, network)
+        val responses1 = getServices(SERVICE_TYPE_1, socketKey)
         assertEquals(2, responses1.size)
         assertTrue(responses1.stream().anyMatch { response ->
             response.serviceInstanceName == SERVICE_NAME_1
@@ -111,19 +117,19 @@
         assertTrue(responses1.any { response ->
             response.serviceInstanceName == SERVICE_NAME_2
         })
-        val responses2 = getServices(SERVICE_TYPE_2, network)
+        val responses2 = getServices(SERVICE_TYPE_2, socketKey)
         assertEquals(1, responses2.size)
         assertTrue(responses2.any { response ->
             response.serviceInstanceName == SERVICE_NAME_2
         })
 
-        removeService(SERVICE_NAME_2, SERVICE_TYPE_1, network)
-        val responses3 = getServices(SERVICE_TYPE_1, network)
+        removeService(SERVICE_NAME_2, SERVICE_TYPE_1, socketKey)
+        val responses3 = getServices(SERVICE_TYPE_1, socketKey)
         assertEquals(1, responses3.size)
         assertTrue(responses3.any { response ->
             response.serviceInstanceName == SERVICE_NAME_1
         })
-        val responses4 = getServices(SERVICE_TYPE_2, network)
+        val responses4 = getServices(SERVICE_TYPE_2, socketKey)
         assertEquals(1, responses4.size)
         assertTrue(responses4.any { response ->
             response.serviceInstanceName == SERVICE_NAME_2
@@ -132,5 +138,5 @@
 
     private fun createResponse(serviceInstanceName: String, serviceType: String) = MdnsResponse(
         0 /* now */, "$serviceInstanceName.$serviceType".split(".").toTypedArray(),
-            INTERFACE_INDEX, network)
+            socketKey.interfaceIndex, socketKey.network)
 }
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 9892e9f..1fdfe09 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;
@@ -41,15 +44,20 @@
 import android.annotation.Nullable;
 import android.net.InetAddresses;
 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;
 import com.android.net.module.util.SharedLog;
 import com.android.server.connectivity.mdns.MdnsServiceInfo.TextEntry;
-import com.android.server.connectivity.mdns.MdnsServiceTypeClient.QueryTaskConfig;
+import com.android.server.connectivity.mdns.util.MdnsUtils;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
+import com.android.testutils.HandlerUtils;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
@@ -82,7 +90,9 @@
 @DevSdkIgnoreRule.IgnoreUpTo(SC_V2)
 public class MdnsServiceTypeClientTests {
     private static final int INTERFACE_INDEX = 999;
+    private static final long DEFAULT_TIMEOUT = 2000L;
     private static final String SERVICE_TYPE = "_googlecast._tcp.local";
+    private static final String SUBTYPE = "_subtype";
     private static final String[] SERVICE_TYPE_LABELS = TextUtils.split(SERVICE_TYPE, "\\.");
     private static final InetSocketAddress IPV4_ADDRESS = new InetSocketAddress(
             MdnsConstants.getMdnsIPv4Address(), MdnsConstants.MDNS_PORT);
@@ -104,9 +114,11 @@
     @Mock
     private Network mockNetwork;
     @Mock
-    private MdnsResponseDecoder.Clock mockDecoderClock;
+    private MdnsUtils.Clock mockDecoderClock;
     @Mock
     private SharedLog mockSharedLog;
+    @Mock
+    private MdnsServiceTypeClient.Dependencies mockDeps;
     @Captor
     private ArgumentCaptor<MdnsServiceInfo> serviceInfoCaptor;
 
@@ -114,11 +126,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")
@@ -128,15 +145,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])
@@ -174,9 +189,35 @@
                 .thenReturn(expectedIPv6Packets[14])
                 .thenReturn(expectedIPv6Packets[15]);
 
+        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) {
+                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                        serviceCache) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -184,11 +225,47 @@
                 };
     }
 
+    @After
+    public void tearDown() {
+        if (thread != null) {
+            thread.quitSafely();
+        }
+    }
+
+    private void runOnHandler(Runnable r) {
+        handler.post(r);
+        HandlerUtils.waitForIdle(handler, DEFAULT_TIMEOUT);
+    }
+
+    private void startSendAndReceive(MdnsServiceBrowserListener listener,
+            MdnsSearchOptions searchOptions) {
+        runOnHandler(() -> client.startSendAndReceive(listener, searchOptions));
+    }
+
+    private void processResponse(MdnsPacket packet, SocketKey socketKey) {
+        runOnHandler(() -> client.processResponse(packet, socketKey));
+    }
+
+    private void stopSendAndReceive(MdnsServiceBrowserListener listener) {
+        runOnHandler(() -> client.stopSendAndReceive(listener));
+    }
+
+    private void notifySocketDestroyed() {
+        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();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).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);
@@ -226,17 +303,21 @@
                 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.
-        client.stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[15]).cancel(true);
+        stopSendAndReceive(mockListenerOne);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
     public void sendQueries_reentry_activeScanMode() {
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(false).build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).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);
@@ -244,13 +325,13 @@
         // After the first query is sent, change the subtypes, and restart.
         searchOptions =
                 MdnsSearchOptions.newBuilder()
-                        .addSubtype("12345")
-                        .addSubtype("abcde")
+                        .addSubtype(SUBTYPE)
+                        .addSubtype("_subtype2")
                         .setIsPassiveMode(false)
                         .build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+        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);
@@ -260,15 +341,17 @@
                 3, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
 
         // Stop sending packets.
-        client.stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[5]).cancel(true);
+        stopSendAndReceive(mockListenerOne);
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
     public void sendQueries_passiveScanMode() {
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(true).build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).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);
@@ -284,15 +367,135 @@
                 false);
 
         // Stop sending packets.
-        client.stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[5]).cancel(true);
+        stopSendAndReceive(mockListenerOne);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+    }
+
+    @Test
+    public void sendQueries_activeScanWithQueryBackoff() {
+        MdnsSearchOptions searchOptions =
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).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);
+        verifyAndSendQuery(
+                1, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        verifyAndSendQuery(
+                2, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        // Second burst will be sent after initialTimeBetweenBurstsMs, 3 queries.
+        verifyAndSendQuery(
+                3, MdnsConfigs.initialTimeBetweenBurstsMs(), /* expectsUnicastResponse= */ false);
+        verifyAndSendQuery(
+                4, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        verifyAndSendQuery(
+                5, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        // Third burst will be sent after initialTimeBetweenBurstsMs * 2, 3 queries.
+        verifyAndSendQuery(
+                6, MdnsConfigs.initialTimeBetweenBurstsMs() * 2, /* expectsUnicastResponse= */
+                false);
+        verifyAndSendQuery(
+                7, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        verifyAndSendQuery(
+                8, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        // Forth burst will be sent after initialTimeBetweenBurstsMs * 4, 3 queries.
+        verifyAndSendQuery(
+                9, MdnsConfigs.initialTimeBetweenBurstsMs() * 4, /* expectsUnicastResponse= */
+                false);
+        verifyAndSendQuery(
+                10, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        verifyAndSendQuery(
+                11, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
+        // In backoff mode, the current scheduled task will be canceled and reschedule if the
+        // 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);
+        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 /* index */, MdnsConfigs.timeBetweenQueriesInBurstMs(),
+                false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                15 /* scheduledCount */);
+    }
+
+    @Test
+    public void sendQueries_passiveScanWithQueryBackoff() {
+        MdnsSearchOptions searchOptions =
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(
+                        true).setNumOfQueriesBeforeBackoff(3).build();
+        startSendAndReceive(mockListenerOne, searchOptions);
+        // 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(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 /* 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();
+        processResponse(createResponse(
+                "service-instance-1", "192.0.2.123", 5353,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), TEST_TTL,
+                TEST_ELAPSED_REALTIME - 1), socketKey);
+        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.
+        doReturn(TEST_ELAPSED_REALTIME).when(mockDecoderClock).elapsedRealtime();
+        processResponse(createResponse(
+                "service-instance-1", "192.0.2.123", 5353,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), TEST_TTL), socketKey);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+
+        stopSendAndReceive(mockListenerOne);
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
     public void sendQueries_reentry_passiveScanMode() {
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(true).build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).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);
@@ -300,13 +503,13 @@
         // After the first query is sent, change the subtypes, and restart.
         searchOptions =
                 MdnsSearchOptions.newBuilder()
-                        .addSubtype("12345")
-                        .addSubtype("abcde")
+                        .addSubtype(SUBTYPE)
+                        .addSubtype("_subtype2")
                         .setIsPassiveMode(true)
                         .build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+        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);
@@ -316,8 +519,8 @@
                 3, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
 
         // Stop sending packets.
-        client.stopSendAndReceive(mockListenerOne);
-        verify(expectedSendFutures[5]).cancel(true);
+        stopSendAndReceive(mockListenerOne);
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
     }
 
     @Test
@@ -325,10 +528,11 @@
     public void testQueryTaskConfig_alwaysAskForUnicastResponse() {
         //MdnsConfigsFlagsImpl.alwaysAskForUnicastResponseInEachBurst.override(true);
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(false).build();
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(false).build();
         QueryTaskConfig config = new QueryTaskConfig(
                 searchOptions.getSubtypes(), searchOptions.isPassiveMode(),
-                false /* onlyUseIpv6OnIpv6OnlyNetworks */, 1, socketKey);
+                false /* onlyUseIpv6OnIpv6OnlyNetworks */, 3 /* numOfQueriesBeforeBackoff */,
+                socketKey);
 
         // This is the first query. We will ask for unicast response.
         assertTrue(config.expectUnicastResponse);
@@ -355,10 +559,11 @@
     @Test
     public void testQueryTaskConfig_askForUnicastInFirstQuery() {
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(false).build();
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(false).build();
         QueryTaskConfig config = new QueryTaskConfig(
                 searchOptions.getSubtypes(), searchOptions.isPassiveMode(),
-                false /* onlyUseIpv6OnIpv6OnlyNetworks */, 1, socketKey);
+                false /* onlyUseIpv6OnIpv6OnlyNetworks */, 3 /* numOfQueriesBeforeBackoff */,
+                socketKey);
 
         // This is the first query. We will ask for unicast response.
         assertTrue(config.expectUnicastResponse);
@@ -385,18 +590,18 @@
     @Test
     public void testIfPreviousTaskIsCanceledWhenNewSessionStarts() {
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(true).build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(true).build();
+        startSendAndReceive(mockListenerOne, searchOptions);
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Change the sutypes and start a new session.
         searchOptions =
                 MdnsSearchOptions.newBuilder()
-                        .addSubtype("12345")
-                        .addSubtype("abcde")
+                        .addSubtype(SUBTYPE)
+                        .addSubtype("_subtype2")
                         .setIsPassiveMode(true)
                         .build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+        startSendAndReceive(mockListenerOne, searchOptions);
 
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
@@ -414,10 +619,10 @@
     public void testIfPreviousTaskIsCanceledWhenSessionStops() {
         //MdnsConfigsFlagsImpl.shouldCancelScanTaskWhenFutureIsNull.override(true);
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype("12345").setIsPassiveMode(true).build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(true).build();
+        startSendAndReceive(mockListenerOne, searchOptions);
         // Change the sutypes and start a new session.
-        client.stopSendAndReceive(mockListenerOne);
+        stopSendAndReceive(mockListenerOne);
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
 
@@ -432,35 +637,33 @@
     @Test
     public void testQueryScheduledWhenAnsweredFromCache() {
         final MdnsSearchOptions searchOptions = MdnsSearchOptions.getDefaultOptions();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+        startSendAndReceive(mockListenerOne, searchOptions);
         assertNotNull(currentThreadExecutor.getAndClearSubmittedRunnable());
 
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 "service-instance-1", "192.0.2.123", 5353,
                 SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), TEST_TTL), /* interfaceIndex= */ 20, mockNetwork);
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         verify(mockListenerOne).onServiceNameDiscovered(any());
         verify(mockListenerOne).onServiceFound(any());
 
         // File another identical query
-        client.startSendAndReceive(mockListenerTwo, searchOptions);
+        startSendAndReceive(mockListenerTwo, searchOptions);
 
         verify(mockListenerTwo).onServiceNameDiscovered(any());
         verify(mockListenerTwo).onServiceFound(any());
 
         // 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,
             String[] serviceType, List<String> ipv4Addresses, List<String> ipv6Addresses, int port,
-            List<String> subTypes, Map<String, String> attributes, int interfaceIndex,
-            Network network) {
+            List<String> subTypes, Map<String, String> attributes, SocketKey socketKey) {
         assertEquals(serviceName, serviceInfo.getServiceInstanceName());
         assertArrayEquals(serviceType, serviceInfo.getServiceType());
         assertEquals(ipv4Addresses, serviceInfo.getIpv4Addresses());
@@ -471,18 +674,18 @@
             assertTrue(attributes.containsKey(key));
             assertEquals(attributes.get(key), serviceInfo.getAttributeByKey(key));
         }
-        assertEquals(interfaceIndex, serviceInfo.getInterfaceIndex());
-        assertEquals(network, serviceInfo.getNetwork());
+        assertEquals(socketKey.getInterfaceIndex(), serviceInfo.getInterfaceIndex());
+        assertEquals(socketKey.getNetwork(), serviceInfo.getNetwork());
     }
 
     @Test
     public void processResponse_incompleteResponse() {
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
 
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 "service-instance-1", null /* host */, 0 /* port */,
                 SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), TEST_TTL), socketKey);
         verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
                 "service-instance-1",
@@ -492,8 +695,7 @@
                 /* port= */ 0,
                 /* subTypes= */ List.of(),
                 Collections.emptyMap(),
-                INTERFACE_INDEX,
-                mockNetwork);
+                socketKey);
 
         verify(mockListenerOne, never()).onServiceFound(any(MdnsServiceInfo.class));
         verify(mockListenerOne, never()).onServiceUpdated(any(MdnsServiceInfo.class));
@@ -502,20 +704,18 @@
     @Test
     public void processIPv4Response_completeResponseForNewServiceInstance() throws Exception {
         final String ipV4Address = "192.168.1.1";
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
 
         // Process the initial response.
-        client.processResponse(createResponse(
-                "service-instance-1", ipV4Address, 5353,
-                /* subtype= */ "ABCDE",
-                Collections.emptyMap(), TEST_TTL), /* interfaceIndex= */ 20, mockNetwork);
+        processResponse(createResponse(
+                "service-instance-1", ipV4Address, 5353, SUBTYPE,
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Process a second response with a different port and updated text attributes.
-        client.processResponse(createResponse(
-                        "service-instance-1", ipV4Address, 5354,
-                        /* subtype= */ "ABCDE",
+        processResponse(createResponse(
+                        "service-instance-1", ipV4Address, 5354, SUBTYPE,
                         Collections.singletonMap("key", "value"), TEST_TTL),
-                /* interfaceIndex= */ 20, mockNetwork);
+                socketKey);
 
         // Verify onServiceNameDiscovered was called once for the initial response.
         verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
@@ -525,10 +725,9 @@
                 List.of(ipV4Address) /* ipv4Address */,
                 List.of() /* ipv6Address */,
                 5353 /* port */,
-                Collections.singletonList("ABCDE") /* subTypes */,
+                Collections.singletonList(SUBTYPE) /* subTypes */,
                 Collections.singletonMap("key", null) /* attributes */,
-                20 /* interfaceIndex */,
-                mockNetwork);
+                socketKey);
 
         // Verify onServiceFound was called once for the initial response.
         verify(mockListenerOne).onServiceFound(serviceInfoCaptor.capture());
@@ -536,10 +735,10 @@
         assertEquals(initialServiceInfo.getServiceInstanceName(), "service-instance-1");
         assertEquals(initialServiceInfo.getIpv4Address(), ipV4Address);
         assertEquals(initialServiceInfo.getPort(), 5353);
-        assertEquals(initialServiceInfo.getSubtypes(), Collections.singletonList("ABCDE"));
+        assertEquals(initialServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
         assertNull(initialServiceInfo.getAttributeByKey("key"));
-        assertEquals(initialServiceInfo.getInterfaceIndex(), 20);
-        assertEquals(mockNetwork, initialServiceInfo.getNetwork());
+        assertEquals(socketKey.getInterfaceIndex(), initialServiceInfo.getInterfaceIndex());
+        assertEquals(socketKey.getNetwork(), initialServiceInfo.getNetwork());
 
         // Verify onServiceUpdated was called once for the second response.
         verify(mockListenerOne).onServiceUpdated(serviceInfoCaptor.capture());
@@ -548,29 +747,27 @@
         assertEquals(updatedServiceInfo.getIpv4Address(), ipV4Address);
         assertEquals(updatedServiceInfo.getPort(), 5354);
         assertTrue(updatedServiceInfo.hasSubtypes());
-        assertEquals(updatedServiceInfo.getSubtypes(), Collections.singletonList("ABCDE"));
+        assertEquals(updatedServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
         assertEquals(updatedServiceInfo.getAttributeByKey("key"), "value");
-        assertEquals(updatedServiceInfo.getInterfaceIndex(), 20);
-        assertEquals(mockNetwork, updatedServiceInfo.getNetwork());
+        assertEquals(socketKey.getInterfaceIndex(), updatedServiceInfo.getInterfaceIndex());
+        assertEquals(socketKey.getNetwork(), updatedServiceInfo.getNetwork());
     }
 
     @Test
     public void processIPv6Response_getCorrectServiceInfo() throws Exception {
         final String ipV6Address = "2000:3333::da6c:63ff:fe7c:7483";
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
 
         // Process the initial response.
-        client.processResponse(createResponse(
-                "service-instance-1", ipV6Address, 5353,
-                /* subtype= */ "ABCDE",
-                Collections.emptyMap(), TEST_TTL), /* interfaceIndex= */ 20, mockNetwork);
+        processResponse(createResponse(
+                "service-instance-1", ipV6Address, 5353, SUBTYPE,
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Process a second response with a different port and updated text attributes.
-        client.processResponse(createResponse(
-                        "service-instance-1", ipV6Address, 5354,
-                        /* subtype= */ "ABCDE",
+        processResponse(createResponse(
+                        "service-instance-1", ipV6Address, 5354, SUBTYPE,
                         Collections.singletonMap("key", "value"), TEST_TTL),
-                /* interfaceIndex= */ 20, mockNetwork);
+                socketKey);
 
         // Verify onServiceNameDiscovered was called once for the initial response.
         verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
@@ -580,10 +777,9 @@
                 List.of() /* ipv4Address */,
                 List.of(ipV6Address) /* ipv6Address */,
                 5353 /* port */,
-                Collections.singletonList("ABCDE") /* subTypes */,
+                Collections.singletonList(SUBTYPE) /* subTypes */,
                 Collections.singletonMap("key", null) /* attributes */,
-                20 /* interfaceIndex */,
-                mockNetwork);
+                socketKey);
 
         // Verify onServiceFound was called once for the initial response.
         verify(mockListenerOne).onServiceFound(serviceInfoCaptor.capture());
@@ -591,10 +787,10 @@
         assertEquals(initialServiceInfo.getServiceInstanceName(), "service-instance-1");
         assertEquals(initialServiceInfo.getIpv6Address(), ipV6Address);
         assertEquals(initialServiceInfo.getPort(), 5353);
-        assertEquals(initialServiceInfo.getSubtypes(), Collections.singletonList("ABCDE"));
+        assertEquals(initialServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
         assertNull(initialServiceInfo.getAttributeByKey("key"));
-        assertEquals(initialServiceInfo.getInterfaceIndex(), 20);
-        assertEquals(mockNetwork, initialServiceInfo.getNetwork());
+        assertEquals(socketKey.getInterfaceIndex(), initialServiceInfo.getInterfaceIndex());
+        assertEquals(socketKey.getNetwork(), initialServiceInfo.getNetwork());
 
         // Verify onServiceUpdated was called once for the second response.
         verify(mockListenerOne).onServiceUpdated(serviceInfoCaptor.capture());
@@ -603,10 +799,10 @@
         assertEquals(updatedServiceInfo.getIpv6Address(), ipV6Address);
         assertEquals(updatedServiceInfo.getPort(), 5354);
         assertTrue(updatedServiceInfo.hasSubtypes());
-        assertEquals(updatedServiceInfo.getSubtypes(), Collections.singletonList("ABCDE"));
+        assertEquals(updatedServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
         assertEquals(updatedServiceInfo.getAttributeByKey("key"), "value");
-        assertEquals(updatedServiceInfo.getInterfaceIndex(), 20);
-        assertEquals(mockNetwork, updatedServiceInfo.getNetwork());
+        assertEquals(socketKey.getInterfaceIndex(), updatedServiceInfo.getInterfaceIndex());
+        assertEquals(socketKey.getNetwork(), updatedServiceInfo.getNetwork());
     }
 
     private void verifyServiceRemovedNoCallback(MdnsServiceBrowserListener listener) {
@@ -615,61 +811,60 @@
     }
 
     private void verifyServiceRemovedCallback(MdnsServiceBrowserListener listener,
-            String serviceName, String[] serviceType, int interfaceIndex, Network network) {
+            String serviceName, String[] serviceType, SocketKey socketKey) {
         verify(listener).onServiceRemoved(argThat(
                 info -> serviceName.equals(info.getServiceInstanceName())
                         && Arrays.equals(serviceType, info.getServiceType())
-                        && info.getInterfaceIndex() == interfaceIndex
-                        && network.equals(info.getNetwork())));
+                        && info.getInterfaceIndex() == socketKey.getInterfaceIndex()
+                        && socketKey.getNetwork().equals(info.getNetwork())));
         verify(listener).onServiceNameRemoved(argThat(
                 info -> serviceName.equals(info.getServiceInstanceName())
                         && Arrays.equals(serviceType, info.getServiceType())
-                        && info.getInterfaceIndex() == interfaceIndex
-                        && network.equals(info.getNetwork())));
+                        && info.getInterfaceIndex() == socketKey.getInterfaceIndex()
+                        && socketKey.getNetwork().equals(info.getNetwork())));
     }
 
     @Test
     public void processResponse_goodBye() throws Exception {
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
-        client.startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
 
         final String serviceName = "service-instance-1";
         final String ipV6Address = "2000:3333::da6c:63ff:fe7c:7483";
         // Process the initial response.
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 serviceName, ipV6Address, 5353,
                 SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 "goodbye-service", ipV6Address, 5353,
                 SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), /* ptrTtlMillis= */ 0L), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), /* ptrTtlMillis= */ 0L), socketKey);
 
         // Verify removed callback won't be called if the service is not existed.
         verifyServiceRemovedNoCallback(mockListenerOne);
         verifyServiceRemovedNoCallback(mockListenerTwo);
 
         // Verify removed callback would be called.
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 serviceName, ipV6Address, 5353,
                 SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), 0L), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), 0L), socketKey);
         verifyServiceRemovedCallback(
-                mockListenerOne, serviceName, SERVICE_TYPE_LABELS, INTERFACE_INDEX, mockNetwork);
+                mockListenerOne, serviceName, SERVICE_TYPE_LABELS, socketKey);
         verifyServiceRemovedCallback(
-                mockListenerTwo, serviceName, SERVICE_TYPE_LABELS, INTERFACE_INDEX, mockNetwork);
+                mockListenerTwo, serviceName, SERVICE_TYPE_LABELS, socketKey);
     }
 
     @Test
     public void reportExistingServiceToNewlyRegisteredListeners() throws Exception {
         // Process the initial response.
-        client.processResponse(createResponse(
-                "service-instance-1", "192.168.1.1", 5353,
-                /* subtype= */ "ABCDE",
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+        processResponse(createResponse(
+                "service-instance-1", "192.168.1.1", 5353, SUBTYPE,
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
 
         // Verify onServiceNameDiscovered was called once for the existing response.
         verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
@@ -679,10 +874,9 @@
                 List.of("192.168.1.1") /* ipv4Address */,
                 List.of() /* ipv6Address */,
                 5353 /* port */,
-                Collections.singletonList("ABCDE") /* subTypes */,
+                Collections.singletonList(SUBTYPE) /* subTypes */,
                 Collections.singletonMap("key", null) /* attributes */,
-                INTERFACE_INDEX,
-                mockNetwork);
+                socketKey);
 
         // Verify onServiceFound was called once for the existing response.
         verify(mockListenerOne).onServiceFound(serviceInfoCaptor.capture());
@@ -690,16 +884,16 @@
         assertEquals(existingServiceInfo.getServiceInstanceName(), "service-instance-1");
         assertEquals(existingServiceInfo.getIpv4Address(), "192.168.1.1");
         assertEquals(existingServiceInfo.getPort(), 5353);
-        assertEquals(existingServiceInfo.getSubtypes(), Collections.singletonList("ABCDE"));
+        assertEquals(existingServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
         assertNull(existingServiceInfo.getAttributeByKey("key"));
 
         // Process a goodbye message for the existing response.
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 "service-instance-1", "192.168.1.1", 5353,
                 SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), /* ptrTtlMillis= */ 0L), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), /* ptrTtlMillis= */ 0L), socketKey);
 
-        client.startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
 
         // Verify onServiceFound was not called on the newly registered listener after the existing
         // response is gone.
@@ -713,21 +907,24 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog) {
+                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                        serviceCache) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
                     }
                 };
-        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder().setRemoveExpiredService(
-                true).build();
-        client.startSendAndReceive(mockListenerOne, searchOptions);
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .setRemoveExpiredService(true)
+                .setNumOfQueriesBeforeBackoff(Integer.MAX_VALUE)
+                .build();
+        startSendAndReceive(mockListenerOne, searchOptions);
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Process the initial response.
-        client.processResponse(createResponse(
-                serviceInstanceName, "192.168.1.1", 5353, /* subtype= */ "ABCDE",
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+        processResponse(createResponse(
+                serviceInstanceName, "192.168.1.1", 5353, SUBTYPE,
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
@@ -735,6 +932,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);
@@ -742,10 +940,11 @@
         // 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(mockListenerOne, serviceInstanceName, SERVICE_TYPE_LABELS,
-                INTERFACE_INDEX, mockNetwork);
+        verifyServiceRemovedCallback(
+                mockListenerOne, serviceInstanceName, SERVICE_TYPE_LABELS, socketKey);
     }
 
     @Test
@@ -754,19 +953,20 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog) {
+                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                        serviceCache) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
                     }
                 };
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Process the initial response.
-        client.processResponse(createResponse(
-                serviceInstanceName, "192.168.1.1", 5353, /* subtype= */ "ABCDE",
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+        processResponse(createResponse(
+                serviceInstanceName, "192.168.1.1", 5353, SUBTYPE,
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
@@ -787,19 +987,20 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog) {
+                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                        serviceCache) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
                     }
                 };
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Process the initial response.
-        client.processResponse(createResponse(
-                serviceInstanceName, "192.168.1.1", 5353, /* subtype= */ "ABCDE",
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+        processResponse(createResponse(
+                serviceInstanceName, "192.168.1.1", 5353, SUBTYPE,
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
@@ -809,8 +1010,8 @@
         firstMdnsTask.run();
 
         // Verify removed callback was called.
-        verifyServiceRemovedCallback(mockListenerOne, serviceInstanceName, SERVICE_TYPE_LABELS,
-                INTERFACE_INDEX, mockNetwork);
+        verifyServiceRemovedCallback(
+                mockListenerOne, serviceInstanceName, SERVICE_TYPE_LABELS, socketKey);
     }
 
     @Test
@@ -818,30 +1019,29 @@
         final String serviceName = "service-instance";
         final String ipV4Address = "192.0.2.0";
         final String ipV6Address = "2001:db8::";
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
         InOrder inOrder = inOrder(mockListenerOne);
 
         // Process the initial response which is incomplete.
-        final String subtype = "ABCDE";
-        client.processResponse(createResponse(
-                serviceName, null, 5353, subtype,
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+        processResponse(createResponse(
+                serviceName, null, 5353, SUBTYPE,
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Process a second response which has ip address to make response become complete.
-        client.processResponse(createResponse(
-                serviceName, ipV4Address, 5353, subtype,
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+        processResponse(createResponse(
+                serviceName, ipV4Address, 5353, SUBTYPE,
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Process a third response with a different ip address, port and updated text attributes.
-        client.processResponse(createResponse(
-                serviceName, ipV6Address, 5354, subtype,
-                Collections.singletonMap("key", "value"), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+        processResponse(createResponse(
+                serviceName, ipV6Address, 5354, SUBTYPE,
+                Collections.singletonMap("key", "value"), TEST_TTL), socketKey);
 
         // Process the last response which is goodbye message (with the main type, not subtype).
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                         serviceName, ipV6Address, 5354, SERVICE_TYPE_LABELS,
                         Collections.singletonMap("key", "value"), /* ptrTtlMillis= */ 0L),
-                INTERFACE_INDEX, mockNetwork);
+                socketKey);
 
         // Verify onServiceNameDiscovered was first called for the initial response.
         inOrder.verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
@@ -851,10 +1051,9 @@
                 List.of() /* ipv4Address */,
                 List.of() /* ipv6Address */,
                 5353 /* port */,
-                Collections.singletonList(subtype) /* subTypes */,
+                Collections.singletonList(SUBTYPE) /* subTypes */,
                 Collections.singletonMap("key", null) /* attributes */,
-                INTERFACE_INDEX,
-                mockNetwork);
+                socketKey);
 
         // Verify onServiceFound was second called for the second response.
         inOrder.verify(mockListenerOne).onServiceFound(serviceInfoCaptor.capture());
@@ -864,10 +1063,9 @@
                 List.of(ipV4Address) /* ipv4Address */,
                 List.of() /* ipv6Address */,
                 5353 /* port */,
-                Collections.singletonList(subtype) /* subTypes */,
+                Collections.singletonList(SUBTYPE) /* subTypes */,
                 Collections.singletonMap("key", null) /* attributes */,
-                INTERFACE_INDEX,
-                mockNetwork);
+                socketKey);
 
         // Verify onServiceUpdated was third called for the third response.
         inOrder.verify(mockListenerOne).onServiceUpdated(serviceInfoCaptor.capture());
@@ -877,10 +1075,9 @@
                 List.of(ipV4Address) /* ipv4Address */,
                 List.of(ipV6Address) /* ipv6Address */,
                 5354 /* port */,
-                Collections.singletonList(subtype) /* subTypes */,
+                Collections.singletonList(SUBTYPE) /* subTypes */,
                 Collections.singletonMap("key", "value") /* attributes */,
-                INTERFACE_INDEX,
-                mockNetwork);
+                socketKey);
 
         // Verify onServiceRemoved was called for the last response.
         inOrder.verify(mockListenerOne).onServiceRemoved(serviceInfoCaptor.capture());
@@ -890,10 +1087,9 @@
                 List.of(ipV4Address) /* ipv4Address */,
                 List.of(ipV6Address) /* ipv6Address */,
                 5354 /* port */,
-                Collections.singletonList("ABCDE") /* subTypes */,
+                Collections.singletonList(SUBTYPE) /* subTypes */,
                 Collections.singletonMap("key", "value") /* attributes */,
-                INTERFACE_INDEX,
-                mockNetwork);
+                socketKey);
 
         // Verify onServiceNameRemoved was called for the last response.
         inOrder.verify(mockListenerOne).onServiceNameRemoved(serviceInfoCaptor.capture());
@@ -903,16 +1099,16 @@
                 List.of(ipV4Address) /* ipv4Address */,
                 List.of(ipV6Address) /* ipv6Address */,
                 5354 /* port */,
-                Collections.singletonList("ABCDE") /* subTypes */,
+                Collections.singletonList(SUBTYPE) /* subTypes */,
                 Collections.singletonMap("key", "value") /* attributes */,
-                INTERFACE_INDEX,
-                mockNetwork);
+                socketKey);
     }
 
     @Test
     public void testProcessResponse_Resolve() throws Exception {
-        client = new MdnsServiceTypeClient(
-                SERVICE_TYPE, mockSocketClient, currentThreadExecutor, socketKey, mockSharedLog);
+        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String instanceName = "service-instance";
         final String[] hostname = new String[] { "testhost "};
@@ -922,7 +1118,7 @@
         final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
                 .setResolveInstanceName(instanceName).build();
 
-        client.startSendAndReceive(mockListenerOne, resolveOptions);
+        startSendAndReceive(mockListenerOne, resolveOptions);
         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
 
         // Verify a query for SRV/TXT was sent, but no PTR query
@@ -932,7 +1128,9 @@
         // Send twice for IPv4 and IPv6
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
                 srvTxtQueryCaptor.capture(),
-                eq(mockNetwork), eq(false));
+                eq(socketKey), eq(false));
+        verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
+        assertNotNull(delayMessage);
 
         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
                 new MdnsPacketReader(srvTxtQueryCaptor.getValue()));
@@ -955,15 +1153,16 @@
                 Collections.emptyList() /* authorityRecords */,
                 Collections.emptyList() /* additionalRecords */);
 
-        client.processResponse(srvTxtResponse, INTERFACE_INDEX, mockNetwork);
+        processResponse(srvTxtResponse, socketKey);
 
         // Expect a query for A/AAAA
+        dispatchMessage();
         final ArgumentCaptor<DatagramPacket> addressQueryCaptor =
                 ArgumentCaptor.forClass(DatagramPacket.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
                 addressQueryCaptor.capture(),
-                eq(mockNetwork), eq(false));
+                eq(socketKey), eq(false));
 
         final MdnsPacket addressQueryPacket = MdnsPacket.parse(
                 new MdnsPacketReader(addressQueryCaptor.getValue()));
@@ -985,7 +1184,7 @@
                 Collections.emptyList() /* additionalRecords */);
 
         inOrder.verify(mockListenerOne, never()).onServiceNameDiscovered(any());
-        client.processResponse(addressResponse, INTERFACE_INDEX, mockNetwork);
+        processResponse(addressResponse, socketKey);
 
         inOrder.verify(mockListenerOne).onServiceFound(serviceInfoCaptor.capture());
         verifyServiceInfo(serviceInfoCaptor.getValue(),
@@ -996,14 +1195,14 @@
                 1234 /* port */,
                 Collections.emptyList() /* subTypes */,
                 Collections.emptyMap() /* attributes */,
-                INTERFACE_INDEX,
-                mockNetwork);
+                socketKey);
     }
 
     @Test
     public void testRenewTxtSrvInResolve() throws Exception {
         client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                mockDecoderClock, socketKey, mockSharedLog);
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String instanceName = "service-instance";
         final String[] hostname = new String[] { "testhost "};
@@ -1013,7 +1212,7 @@
         final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
                 .setResolveInstanceName(instanceName).build();
 
-        client.startSendAndReceive(mockListenerOne, resolveOptions);
+        startSendAndReceive(mockListenerOne, resolveOptions);
         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
 
         // Get the query for SRV/TXT
@@ -1023,7 +1222,9 @@
         // Send twice for IPv4 and IPv6
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
                 srvTxtQueryCaptor.capture(),
-                eq(mockNetwork), eq(false));
+                eq(socketKey), eq(false));
+        verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
+        assertNotNull(delayMessage);
 
         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
                 new MdnsPacketReader(srvTxtQueryCaptor.getValue()));
@@ -1050,7 +1251,8 @@
                                 InetAddresses.parseNumericAddress(ipV6Address))),
                 Collections.emptyList() /* authorityRecords */,
                 Collections.emptyList() /* additionalRecords */);
-        client.processResponse(srvTxtResponse, INTERFACE_INDEX, mockNetwork);
+        processResponse(srvTxtResponse, socketKey);
+        dispatchMessage();
         inOrder.verify(mockListenerOne).onServiceNameDiscovered(any());
         inOrder.verify(mockListenerOne).onServiceFound(any());
 
@@ -1061,6 +1263,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
@@ -1069,7 +1274,9 @@
         // Second and later sends are sent as "expect multicast response" queries
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
                 renewalQueryCaptor.capture(),
-                eq(mockNetwork), eq(false));
+                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()));
@@ -1095,7 +1302,8 @@
                                 InetAddresses.parseNumericAddress(ipV6Address))),
                 Collections.emptyList() /* authorityRecords */,
                 Collections.emptyList() /* additionalRecords */);
-        client.processResponse(refreshedSrvTxtResponse, INTERFACE_INDEX, mockNetwork);
+        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.
@@ -1106,8 +1314,9 @@
 
     @Test
     public void testProcessResponse_ResolveExcludesOtherServices() {
-        client = new MdnsServiceTypeClient(
-                SERVICE_TYPE, mockSocketClient, currentThreadExecutor, socketKey, mockSharedLog);
+        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String requestedInstance = "instance1";
         final String otherInstance = "instance2";
@@ -1119,30 +1328,30 @@
                 // Use different case in the options
                 .setResolveInstanceName(capitalizedRequestInstance).build();
 
-        client.startSendAndReceive(mockListenerOne, resolveOptions);
-        client.startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, resolveOptions);
+        startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
 
         // Complete response from instanceName
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                         requestedInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
-                INTERFACE_INDEX, mockNetwork);
+                socketKey);
 
         // Complete response from otherInstanceName
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                         otherInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
-                INTERFACE_INDEX, mockNetwork);
+                socketKey);
 
         // Address update from otherInstanceName
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 otherInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Goodbye from otherInstanceName
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 otherInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), 0L /* ttl */), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), 0L /* ttl */), socketKey);
 
         // mockListenerOne gets notified for the requested instance
         verify(mockListenerOne).onServiceNameDiscovered(
@@ -1170,8 +1379,9 @@
 
     @Test
     public void testProcessResponse_SubtypeDiscoveryLimitedToSubtype() {
-        client = new MdnsServiceTypeClient(
-                SERVICE_TYPE, mockSocketClient, currentThreadExecutor, socketKey, mockSharedLog);
+        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String matchingInstance = "instance1";
         final String subtype = "_subtype";
@@ -1183,8 +1393,8 @@
                 // Search with different case. Note MdnsSearchOptions subtype doesn't start with "_"
                 .addSubtype("Subtype").build();
 
-        client.startSendAndReceive(mockListenerOne, options);
-        client.startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerOne, options);
+        startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
 
         // Complete response from instanceName
         final MdnsPacket packetWithoutSubtype = createResponse(
@@ -1207,23 +1417,23 @@
                 newAnswers,
                 packetWithoutSubtype.authorityRecords,
                 packetWithoutSubtype.additionalRecords);
-        client.processResponse(packetWithSubtype, INTERFACE_INDEX, mockNetwork);
+        processResponse(packetWithSubtype, socketKey);
 
         // Complete response from otherInstanceName, without subtype
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                         otherInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
-                INTERFACE_INDEX, mockNetwork);
+                socketKey);
 
         // Address update from otherInstanceName
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 otherInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), TEST_TTL), socketKey);
 
         // Goodbye from otherInstanceName
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                 otherInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
-                Collections.emptyMap(), 0L /* ttl */), INTERFACE_INDEX, mockNetwork);
+                Collections.emptyMap(), 0L /* ttl */), socketKey);
 
         // mockListenerOne gets notified for the requested instance
         final ArgumentMatcher<MdnsServiceInfo> subtypeInstanceMatcher = info ->
@@ -1251,44 +1461,49 @@
 
     @Test
     public void testNotifySocketDestroyed() throws Exception {
-        client = new MdnsServiceTypeClient(
-                SERVICE_TYPE, mockSocketClient, currentThreadExecutor, socketKey, mockSharedLog);
+        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
 
         final String requestedInstance = "instance1";
         final String otherInstance = "instance2";
         final String ipV4Address = "192.0.2.0";
 
         final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
+                .setNumOfQueriesBeforeBackoff(Integer.MAX_VALUE)
                 .setResolveInstanceName("instance1").build();
 
-        client.startSendAndReceive(mockListenerOne, resolveOptions);
+        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);
-        client.startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+        startSendAndReceive(mockListenerTwo,
+                MdnsSearchOptions.newBuilder().setNumOfQueriesBeforeBackoff(
+                        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(
                 TEST_TIMEOUT_MS, TimeUnit.MILLISECONDS);
 
         // Complete response from instanceName
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                         requestedInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
-                INTERFACE_INDEX, mockNetwork);
+                socketKey);
 
         // Complete response from otherInstanceName
-        client.processResponse(createResponse(
+        processResponse(createResponse(
                         otherInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
-                INTERFACE_INDEX, mockNetwork);
+                socketKey);
 
-        verify(expectedSendFutures[1], never()).cancel(true);
-        client.notifySocketDestroyed();
-        verify(expectedSendFutures[1]).cancel(true);
+        notifySocketDestroyed();
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
 
         // mockListenerOne gets notified for the requested instance
         final InOrder inOrder1 = inOrder(mockListenerOne);
@@ -1307,12 +1522,106 @@
         inOrder2.verify(mockListenerTwo).onServiceNameDiscovered(
                 matchServiceName(requestedInstance));
         inOrder2.verify(mockListenerTwo).onServiceFound(matchServiceName(requestedInstance));
-        inOrder2.verify(mockListenerTwo).onServiceNameDiscovered(matchServiceName(otherInstance));
-        inOrder2.verify(mockListenerTwo).onServiceFound(matchServiceName(otherInstance));
-        inOrder2.verify(mockListenerTwo).onServiceRemoved(matchServiceName(otherInstance));
-        inOrder2.verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(otherInstance));
         inOrder2.verify(mockListenerTwo).onServiceRemoved(matchServiceName(requestedInstance));
         inOrder2.verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(requestedInstance));
+        verify(mockListenerTwo).onServiceNameDiscovered(matchServiceName(otherInstance));
+        verify(mockListenerTwo).onServiceFound(matchServiceName(otherInstance));
+        verify(mockListenerTwo).onServiceRemoved(matchServiceName(otherInstance));
+        verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(otherInstance));
+    }
+
+    @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.
+
+        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) {
@@ -1322,29 +1631,38 @@
     // 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(currentThreadExecutor.getAndClearLastScheduledDelayInMs(), timeInMs);
+            boolean multipleSocketDiscovery, int scheduledCount) {
+        // Dispatch the message
+        if (delayMessage != null && realHandler != null) {
+            dispatchMessage();
+        }
+        assertEquals(timeInMs, latestDelayMs);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         if (expectsUnicastResponse) {
             verify(mockSocketClient).sendPacketRequestingUnicastResponse(
-                    expectedIPv4Packets[index], mockNetwork, false);
+                    expectedIPv4Packets[index], socketKey, false);
             if (multipleSocketDiscovery) {
                 verify(mockSocketClient).sendPacketRequestingUnicastResponse(
-                        expectedIPv6Packets[index], mockNetwork, false);
+                        expectedIPv6Packets[index], socketKey, false);
             }
         } else {
             verify(mockSocketClient).sendPacketRequestingMulticastResponse(
-                    expectedIPv4Packets[index], mockNetwork, false);
+                    expectedIPv4Packets[index], socketKey, false);
             if (multipleSocketDiscovery) {
                 verify(mockSocketClient).sendPacketRequestingMulticastResponse(
-                        expectedIPv6Packets[index], mockNetwork, false);
+                        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) {
@@ -1389,7 +1707,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.
@@ -1435,7 +1753,7 @@
                 textAttributes, ptrTtlMillis);
     }
 
-    // Creates a mDNS response.
+
     private MdnsPacket createResponse(
             @NonNull String serviceInstanceName,
             @Nullable String host,
@@ -1443,6 +1761,19 @@
             @NonNull String[] type,
             @NonNull Map<String, String> textAttributes,
             long ptrTtlMillis) {
+        return createResponse(serviceInstanceName, host, port, type, textAttributes, ptrTtlMillis,
+                TEST_ELAPSED_REALTIME);
+    }
+
+    // Creates a mDNS response.
+    private MdnsPacket createResponse(
+            @NonNull String serviceInstanceName,
+            @Nullable String host,
+            int port,
+            @NonNull String[] type,
+            @NonNull Map<String, String> textAttributes,
+            long ptrTtlMillis,
+            long receiptTimeMillis) {
 
         final ArrayList<MdnsRecord> answerRecords = new ArrayList<>();
 
@@ -1453,7 +1784,7 @@
         final String[] serviceName = serviceNameList.toArray(new String[0]);
         final MdnsPointerRecord pointerRecord = new MdnsPointerRecord(
                 type,
-                TEST_ELAPSED_REALTIME /* receiptTimeMillis */,
+                receiptTimeMillis,
                 false /* cacheFlush */,
                 ptrTtlMillis,
                 serviceName);
@@ -1462,7 +1793,7 @@
         // Set SRV record.
         final MdnsServiceRecord serviceRecord = new MdnsServiceRecord(
                 serviceName,
-                TEST_ELAPSED_REALTIME /* receiptTimeMillis */,
+                receiptTimeMillis,
                 false /* cacheFlush */,
                 TEST_TTL,
                 0 /* servicePriority */,
@@ -1476,7 +1807,7 @@
             final InetAddress addr = InetAddresses.parseNumericAddress(host);
             final MdnsInetAddressRecord inetAddressRecord = new MdnsInetAddressRecord(
                     new String[] {"hostname"} /* name */,
-                    TEST_ELAPSED_REALTIME /* receiptTimeMillis */,
+                    receiptTimeMillis,
                     false /* cacheFlush */,
                     TEST_TTL,
                     addr);
@@ -1490,7 +1821,7 @@
         }
         final MdnsTextRecord textRecord = new MdnsTextRecord(
                 serviceName,
-                TEST_ELAPSED_REALTIME /* receiptTimeMillis */,
+                receiptTimeMillis,
                 false /* cacheFlush */,
                 TEST_TTL,
                 textEntries);
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketProviderTest.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketProviderTest.java
index 0eac5ec..e971de7 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketProviderTest.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketProviderTest.java
@@ -157,7 +157,6 @@
                 TETHERED_IFACE_NAME);
         doReturn(789).when(mDeps).getNetworkInterfaceIndexByName(
                 WIFI_P2P_IFACE_NAME);
-        doReturn(TETHERED_IFACE_IDX).when(mDeps).getInterfaceIndex(any());
         final HandlerThread thread = new HandlerThread("MdnsSocketProviderTest");
         thread.start();
         mHandler = new Handler(thread.getLooper());