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 d7d3679..54c1ee3 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
+++ b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
@@ -21,6 +21,7 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.greaterThan;
+import static org.junit.Assert.assertSame;
 
 import android.content.Context;
 import android.net.http.HttpEngine;
@@ -91,4 +92,17 @@
         request.getStatus(statusListener);
         statusListener.expectStatus(Status.INVALID);
     }
+
+    @Test
+    public void testUrlRequestCancel_CancelCalled() throws Exception {
+        UrlRequest request = buildUrlRequest(mTestServer.getSuccessUrl());
+        mCallback.setAutoAdvance(false);
+
+        request.start();
+        mCallback.waitForNextStep();
+        assertSame(mCallback.mResponseStep, ResponseStep.ON_RESPONSE_STARTED);
+
+        request.cancel();
+        mCallback.expectCallback(ResponseStep.ON_CANCELED);
+    }
 }
diff --git a/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java b/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
index 44d3ffc..9f8d9b1 100644
--- a/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
+++ b/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
@@ -2243,5 +2243,13 @@
         return mTetherClients;
     }
 
+    // Return map of upstream interface IPv4 address to interface index.
+    // This is used for testing only.
+    @NonNull
+    @VisibleForTesting
+    final HashMap<Inet4Address, Integer> getIpv4UpstreamIndicesForTesting() {
+        return mIpv4UpstreamIndices;
+    }
+
     private static native String[] getBpfCounterNames();
 }
diff --git a/Tethering/src/com/android/networkstack/tethering/Tethering.java b/Tethering/src/com/android/networkstack/tethering/Tethering.java
index e48019c..2e71fda 100644
--- a/Tethering/src/com/android/networkstack/tethering/Tethering.java
+++ b/Tethering/src/com/android/networkstack/tethering/Tethering.java
@@ -88,7 +88,6 @@
 import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.Network;
-import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
 import android.net.TetherStatesParcel;
 import android.net.TetheredClient;
@@ -1856,8 +1855,11 @@
             final Network newUpstream = (ns != null) ? ns.network : null;
             if (mTetherUpstream != newUpstream) {
                 mTetherUpstream = newUpstream;
-                mUpstreamNetworkMonitor.setCurrentUpstream(mTetherUpstream);
-                reportUpstreamChanged(ns);
+                reportUpstreamChanged(mTetherUpstream);
+                // Need to notify capabilities change after upstream network changed because new
+                // network's capabilities should be checked every time.
+                mNotificationUpdater.onUpstreamCapabilitiesChanged(
+                        (ns != null) ? ns.networkCapabilities : null);
             }
         }
 
@@ -2085,6 +2087,7 @@
                 if (mTetherUpstream != null) {
                     mTetherUpstream = null;
                     reportUpstreamChanged(null);
+                    mNotificationUpdater.onUpstreamCapabilitiesChanged(null);
                 }
                 mBpfCoordinator.stopPolling();
             }
@@ -2439,10 +2442,8 @@
         }
     }
 
-    private void reportUpstreamChanged(UpstreamNetworkState ns) {
+    private void reportUpstreamChanged(final Network network) {
         final int length = mTetheringEventCallbacks.beginBroadcast();
-        final Network network = (ns != null) ? ns.network : null;
-        final NetworkCapabilities capabilities = (ns != null) ? ns.networkCapabilities : null;
         try {
             for (int i = 0; i < length; i++) {
                 try {
@@ -2454,9 +2455,6 @@
         } finally {
             mTetheringEventCallbacks.finishBroadcast();
         }
-        // Need to notify capabilities change after upstream network changed because new network's
-        // capabilities should be checked every time.
-        mNotificationUpdater.onUpstreamCapabilitiesChanged(capabilities);
     }
 
     private void reportConfigurationChanged(TetheringConfigurationParcel config) {
diff --git a/Tethering/src/com/android/networkstack/tethering/UpstreamNetworkMonitor.java b/Tethering/src/com/android/networkstack/tethering/UpstreamNetworkMonitor.java
index 16c031b..ac2aa7b 100644
--- a/Tethering/src/com/android/networkstack/tethering/UpstreamNetworkMonitor.java
+++ b/Tethering/src/com/android/networkstack/tethering/UpstreamNetworkMonitor.java
@@ -133,8 +133,6 @@
     private boolean mIsDefaultCellularUpstream;
     // The current system default network (not really used yet).
     private Network mDefaultInternetNetwork;
-    // The current upstream network used for tethering.
-    private Network mTetheringUpstreamNetwork;
     private boolean mPreferTestNetworks;
 
     public UpstreamNetworkMonitor(Context ctx, StateMachine tgt, SharedLog log, int what) {
@@ -191,7 +189,6 @@
         releaseCallback(mListenAllCallback);
         mListenAllCallback = null;
 
-        mTetheringUpstreamNetwork = null;
         mNetworkMap.clear();
     }
 
@@ -342,11 +339,6 @@
         return findFirstDunNetwork(mNetworkMap.values());
     }
 
-    /** Tell UpstreamNetworkMonitor which network is the current upstream of tethering. */
-    public void setCurrentUpstream(Network upstream) {
-        mTetheringUpstreamNetwork = upstream;
-    }
-
     /** Return local prefixes. */
     public Set<IpPrefix> getLocalPrefixes() {
         return (Set<IpPrefix>) mLocalPrefixes.clone();
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/BpfCoordinatorTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/BpfCoordinatorTest.java
index 1978e99..4f32f3c 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/BpfCoordinatorTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/BpfCoordinatorTest.java
@@ -16,6 +16,8 @@
 
 package com.android.networkstack.tethering;
 
+import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
+import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
 import static android.net.NetworkStats.METERED_NO;
 import static android.net.NetworkStats.ROAMING_NO;
@@ -77,6 +79,7 @@
 import android.app.usage.NetworkStatsManager;
 import android.net.INetd;
 import android.net.InetAddresses;
+import android.net.IpPrefix;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.MacAddress;
@@ -89,6 +92,7 @@
 import android.os.Build;
 import android.os.Handler;
 import android.os.test.TestLooper;
+import android.util.SparseArray;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -156,11 +160,13 @@
 
     private static final int INVALID_IFINDEX = 0;
     private static final int UPSTREAM_IFINDEX = 1001;
+    private static final int UPSTREAM_XLAT_IFINDEX = 1002;
     private static final int UPSTREAM_IFINDEX2 = 1003;
     private static final int DOWNSTREAM_IFINDEX = 2001;
     private static final int DOWNSTREAM_IFINDEX2 = 2002;
 
     private static final String UPSTREAM_IFACE = "rmnet0";
+    private static final String UPSTREAM_XLAT_IFACE = "v4-rmnet0";
     private static final String UPSTREAM_IFACE2 = "wlan0";
 
     private static final MacAddress DOWNSTREAM_MAC = MacAddress.fromString("12:34:56:78:90:ab");
@@ -183,6 +189,10 @@
     private static final Inet4Address PRIVATE_ADDR2 =
             (Inet4Address) InetAddresses.parseNumericAddress("192.168.90.12");
 
+    private static final Inet4Address XLAT_LOCAL_IPV4ADDR =
+            (Inet4Address) InetAddresses.parseNumericAddress("192.0.0.46");
+    private static final IpPrefix NAT64_IP_PREFIX = new IpPrefix("64:ff9b::/96");
+
     // Generally, public port and private port are the same in the NAT conntrack message.
     // TODO: consider using different private port and public port for testing.
     private static final short REMOTE_PORT = (short) 443;
@@ -194,6 +204,10 @@
     private static final InterfaceParams UPSTREAM_IFACE_PARAMS = new InterfaceParams(
             UPSTREAM_IFACE, UPSTREAM_IFINDEX, null /* macAddr, rawip */,
             NetworkStackConstants.ETHER_MTU);
+    private static final InterfaceParams UPSTREAM_XLAT_IFACE_PARAMS = new InterfaceParams(
+            UPSTREAM_XLAT_IFACE, UPSTREAM_XLAT_IFINDEX, null /* macAddr, rawip */,
+            NetworkStackConstants.ETHER_MTU - 28
+            /* mtu delta from external/android-clat/clatd.c */);
     private static final InterfaceParams UPSTREAM_IFACE_PARAMS2 = new InterfaceParams(
             UPSTREAM_IFACE2, UPSTREAM_IFINDEX2, MacAddress.fromString("44:55:66:00:00:0c"),
             NetworkStackConstants.ETHER_MTU);
@@ -2281,4 +2295,170 @@
         verifyAddTetherOffloadRule4Mtu(INVALID_MTU, false /* isKernelMtu */,
                 NetworkStackConstants.ETHER_MTU /* expectedMtu */);
     }
+
+    private static LinkProperties buildUpstreamLinkProperties(final String interfaceName,
+            boolean withIPv4, boolean withIPv6, boolean with464xlat) {
+        final LinkProperties prop = new LinkProperties();
+        prop.setInterfaceName(interfaceName);
+
+        if (withIPv4) {
+            // Assign the address no matter what the interface is. It is okay for now because
+            // only single upstream is available.
+            // TODO: consider to assign address by interface once we need to test two or more
+            // BPF supported upstreams or multi upstreams are supported.
+            prop.addLinkAddress(new LinkAddress(PUBLIC_ADDR, 24));
+        }
+
+        if (withIPv6) {
+            // TODO: make this to be constant. Currently, no test who uses this function cares what
+            // the upstream IPv6 address is.
+            prop.addLinkAddress(new LinkAddress("2001:db8::5175:15ca/64"));
+        }
+
+        if (with464xlat) {
+            final String clatInterface = "v4-" + interfaceName;
+            final LinkProperties stackedLink = new LinkProperties();
+            stackedLink.setInterfaceName(clatInterface);
+            stackedLink.addLinkAddress(new LinkAddress(XLAT_LOCAL_IPV4ADDR, 24));
+            prop.addStackedLink(stackedLink);
+            prop.setNat64Prefix(NAT64_IP_PREFIX);
+        }
+
+        return prop;
+    }
+
+    private void verifyIpv4Upstream(
+            @NonNull final HashMap<Inet4Address, Integer> ipv4UpstreamIndices,
+            @NonNull final SparseArray<String> interfaceNames) {
+        assertEquals(1, ipv4UpstreamIndices.size());
+        Integer upstreamIndex = ipv4UpstreamIndices.get(PUBLIC_ADDR);
+        assertNotNull(upstreamIndex);
+        assertEquals(UPSTREAM_IFINDEX, upstreamIndex.intValue());
+        assertEquals(1, interfaceNames.size());
+        assertTrue(interfaceNames.contains(UPSTREAM_IFINDEX));
+    }
+
+    private void verifyUpdateUpstreamNetworkState()
+            throws Exception {
+        final BpfCoordinator coordinator = makeBpfCoordinator();
+        final HashMap<Inet4Address, Integer> ipv4UpstreamIndices =
+                coordinator.getIpv4UpstreamIndicesForTesting();
+        assertTrue(ipv4UpstreamIndices.isEmpty());
+        final SparseArray<String> interfaceNames =
+                coordinator.getInterfaceNamesForTesting();
+        assertEquals(0, interfaceNames.size());
+
+        // Verify the following are added or removed after upstream changes.
+        // - BpfCoordinator#mIpv4UpstreamIndices (for building IPv4 offload rules)
+        // - BpfCoordinator#mInterfaceNames (for updating limit)
+        //
+        // +-------+-------+-----------------------+
+        // | Test  | Up    |       Protocol        |
+        // | Case# | stream+-------+-------+-------+
+        // |       |       | IPv4  | IPv6  | Xlat  |
+        // +-------+-------+-------+-------+-------+
+        // |   1   | Cell  |   O   |       |       |
+        // +-------+-------+-------+-------+-------+
+        // |   2   | Cell  |       |   O   |       |
+        // +-------+-------+-------+-------+-------+
+        // |   3   | Cell  |   O   |   O   |       |
+        // +-------+-------+-------+-------+-------+
+        // |   4   |   -   |       |       |       |
+        // +-------+-------+-------+-------+-------+
+        // |       | Cell  |   O   |       |       |
+        // |       +-------+-------+-------+-------+
+        // |   5   | Cell  |       |   O   |   O   | <-- doesn't support offload (xlat)
+        // |       +-------+-------+-------+-------+
+        // |       | Cell  |   O   |       |       |
+        // +-------+-------+-------+-------+-------+
+        // |   6   | Wifi  |   O   |   O   |       | <-- doesn't support offload (ether ip)
+        // +-------+-------+-------+-------+-------+
+
+        // [1] Mobile IPv4 only
+        coordinator.addUpstreamNameToLookupTable(UPSTREAM_IFINDEX, UPSTREAM_IFACE);
+        doReturn(UPSTREAM_IFACE_PARAMS).when(mDeps).getInterfaceParams(UPSTREAM_IFACE);
+        final UpstreamNetworkState mobileIPv4UpstreamState = new UpstreamNetworkState(
+                buildUpstreamLinkProperties(UPSTREAM_IFACE,
+                        true /* IPv4 */, false /* IPv6 */, false /* 464xlat */),
+                new NetworkCapabilities().addTransportType(TRANSPORT_CELLULAR),
+                new Network(TEST_NET_ID));
+        coordinator.updateUpstreamNetworkState(mobileIPv4UpstreamState);
+        verifyIpv4Upstream(ipv4UpstreamIndices, interfaceNames);
+
+        // [2] Mobile IPv6 only
+        final UpstreamNetworkState mobileIPv6UpstreamState = new UpstreamNetworkState(
+                buildUpstreamLinkProperties(UPSTREAM_IFACE,
+                        false /* IPv4 */, true /* IPv6 */, false /* 464xlat */),
+                new NetworkCapabilities().addTransportType(TRANSPORT_CELLULAR),
+                new Network(TEST_NET_ID));
+        coordinator.updateUpstreamNetworkState(mobileIPv6UpstreamState);
+        assertTrue(ipv4UpstreamIndices.isEmpty());
+        assertEquals(1, interfaceNames.size());
+        assertTrue(interfaceNames.contains(UPSTREAM_IFINDEX));
+
+        // [3] Mobile IPv4 and IPv6
+        final UpstreamNetworkState mobileDualStackUpstreamState = new UpstreamNetworkState(
+                buildUpstreamLinkProperties(UPSTREAM_IFACE,
+                        true /* IPv4 */, true /* IPv6 */, false /* 464xlat */),
+                new NetworkCapabilities().addTransportType(TRANSPORT_CELLULAR),
+                new Network(TEST_NET_ID));
+        coordinator.updateUpstreamNetworkState(mobileDualStackUpstreamState);
+        verifyIpv4Upstream(ipv4UpstreamIndices, interfaceNames);
+
+        // [4] Lost upstream
+        coordinator.updateUpstreamNetworkState(null);
+        assertTrue(ipv4UpstreamIndices.isEmpty());
+        assertEquals(1, interfaceNames.size());
+        assertTrue(interfaceNames.contains(UPSTREAM_IFINDEX));
+
+        // [5] verify xlat interface
+        // Expect that xlat interface information isn't added to mapping.
+        doReturn(UPSTREAM_XLAT_IFACE_PARAMS).when(mDeps).getInterfaceParams(
+                UPSTREAM_XLAT_IFACE);
+        final UpstreamNetworkState mobile464xlatUpstreamState = new UpstreamNetworkState(
+                buildUpstreamLinkProperties(UPSTREAM_IFACE,
+                        false /* IPv4 */, true /* IPv6 */, true /* 464xlat */),
+                new NetworkCapabilities().addTransportType(TRANSPORT_CELLULAR),
+                new Network(TEST_NET_ID));
+
+        // Need to add a valid IPv4 upstream to verify that xlat interface doesn't support.
+        // Mobile IPv4 only
+        coordinator.updateUpstreamNetworkState(mobileIPv4UpstreamState);
+        verifyIpv4Upstream(ipv4UpstreamIndices, interfaceNames);
+
+        // Mobile IPv6 and xlat
+        // IpServer doesn't add xlat interface mapping via #addUpstreamNameToLookupTable on
+        // S and T devices.
+        coordinator.updateUpstreamNetworkState(mobile464xlatUpstreamState);
+        // Upstream IPv4 address mapping is removed because xlat interface is not supported.
+        assertTrue(ipv4UpstreamIndices.isEmpty());
+        assertEquals(1, interfaceNames.size());
+        assertTrue(interfaceNames.contains(UPSTREAM_IFINDEX));
+
+        // Need to add a valid IPv4 upstream to verify that wifi interface doesn't support.
+        // Mobile IPv4 only
+        coordinator.updateUpstreamNetworkState(mobileIPv4UpstreamState);
+        verifyIpv4Upstream(ipv4UpstreamIndices, interfaceNames);
+
+        // [6] Wifi IPv4 and IPv6
+        // Expect that upstream index map is cleared because ether ip is not supported.
+        coordinator.addUpstreamNameToLookupTable(UPSTREAM_IFINDEX2, UPSTREAM_IFACE2);
+        doReturn(UPSTREAM_IFACE_PARAMS2).when(mDeps).getInterfaceParams(UPSTREAM_IFACE2);
+        final UpstreamNetworkState wifiDualStackUpstreamState = new UpstreamNetworkState(
+                buildUpstreamLinkProperties(UPSTREAM_IFACE2,
+                        true /* IPv4 */, true /* IPv6 */, false /* 464xlat */),
+                new NetworkCapabilities().addTransportType(TRANSPORT_WIFI),
+                new Network(TEST_NET_ID2));
+        coordinator.updateUpstreamNetworkState(wifiDualStackUpstreamState);
+        assertTrue(ipv4UpstreamIndices.isEmpty());
+        assertEquals(2, interfaceNames.size());
+        assertTrue(interfaceNames.contains(UPSTREAM_IFINDEX));
+        assertTrue(interfaceNames.contains(UPSTREAM_IFINDEX2));
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.R)
+    public void testUpdateUpstreamNetworkState() throws Exception {
+        verifyUpdateUpstreamNetworkState();
+    }
 }
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
index f90b3a4..98a3b1d 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
@@ -261,6 +261,8 @@
 
     private static final int DHCPSERVER_START_TIMEOUT_MS = 1000;
 
+    private static final Network[] NULL_NETWORK = new Network[] {null};
+
     @Mock private ApplicationInfo mApplicationInfo;
     @Mock private Context mContext;
     @Mock private NetworkStatsManager mStatsManager;
@@ -303,6 +305,7 @@
     private MockContentResolver mContentResolver;
     private BroadcastReceiver mBroadcastReceiver;
     private Tethering mTethering;
+    private TestTetheringEventCallback mTetheringEventCallback;
     private PhoneStateListener mPhoneStateListener;
     private InterfaceConfigurationParcel mInterfaceConfiguration;
     private TetheringConfiguration mConfig;
@@ -670,6 +673,7 @@
         verify(mStatsManager, times(1)).registerNetworkStatsProvider(anyString(), any());
         verify(mNetd).registerUnsolicitedEventListener(any());
         verifyDefaultNetworkRequestFiled();
+        mTetheringEventCallback = registerTetheringEventCallback();
 
         final ArgumentCaptor<PhoneStateListener> phoneListenerCaptor =
                 ArgumentCaptor.forClass(PhoneStateListener.class);
@@ -745,6 +749,16 @@
         return request;
     }
 
+    @NonNull
+    private TestTetheringEventCallback registerTetheringEventCallback() {
+        TestTetheringEventCallback callback = new TestTetheringEventCallback();
+        mTethering.registerTetheringEventCallback(callback);
+        mLooper.dispatchAll();
+        // Pull the first event which is filed immediately after the callback registration.
+        callback.expectUpstreamChanged(NULL_NETWORK);
+        return callback;
+    }
+
     @After
     public void tearDown() {
         mServiceContext.unregisterReceiver(mBroadcastReceiver);
@@ -924,9 +938,9 @@
         // tetherMatchingInterfaces() which starts by fetching all interfaces).
         verify(mNetd, times(1)).interfaceGetList();
 
-        // UpstreamNetworkMonitor should receive selected upstream
+        // Event callback should receive selected upstream
         verify(mUpstreamNetworkMonitor, times(1)).getCurrentPreferredUpstream();
-        verify(mUpstreamNetworkMonitor, times(1)).setCurrentUpstream(upstreamState.network);
+        mTetheringEventCallback.expectUpstreamChanged(upstreamState.network);
     }
 
     @Test
@@ -1181,7 +1195,7 @@
         verify(mUpstreamNetworkMonitor, times(1)).getCurrentPreferredUpstream();
         verify(mUpstreamNetworkMonitor, never()).selectPreferredUpstreamType(any());
 
-        verify(mUpstreamNetworkMonitor, times(1)).setCurrentUpstream(upstreamState.network);
+        mTetheringEventCallback.expectUpstreamChanged(upstreamState.network);
     }
 
     private void verifyDisableTryCellWhenTetheringStop(InOrder inOrder) {
@@ -1206,14 +1220,14 @@
         mobile.fakeConnect();
         mCm.makeDefaultNetwork(mobile, BROADCAST_FIRST);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(mobile.networkId);
 
         // Switch upstream to wifi.
         wifi.fakeConnect();
         mCm.makeDefaultNetwork(wifi, BROADCAST_FIRST);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(false);
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(wifi.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(wifi.networkId);
     }
 
     private void verifyAutomaticUpstreamSelection(boolean configAutomatic) throws Exception {
@@ -1230,30 +1244,30 @@
         // Switch upstreams a few times.
         mCm.makeDefaultNetwork(mobile, BROADCAST_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(mobile.networkId);
 
         mCm.makeDefaultNetwork(wifi, BROADCAST_FIRST, doDispatchAll);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(false);
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(wifi.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(wifi.networkId);
 
         mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(mobile.networkId);
 
         mCm.makeDefaultNetwork(wifi, CALLBACKS_FIRST);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(false);
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(wifi.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(wifi.networkId);
 
         mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(mobile.networkId);
 
         // Wifi disconnecting should not have any affect since it's not the current upstream.
         wifi.fakeDisconnect();
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(any());
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
 
         // Lose and regain upstream.
         assertTrue(mUpstreamNetworkMonitor.getCurrentPreferredUpstream().linkProperties
@@ -1263,13 +1277,13 @@
         mobile.fakeDisconnect();
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(true);
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(null);
+        mTetheringEventCallback.expectUpstreamChanged(NULL_NETWORK);
 
         mobile = new TestNetworkAgent(mCm, buildMobile464xlatUpstreamState());
         mobile.fakeConnect();
         mCm.makeDefaultNetwork(mobile, BROADCAST_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(mobile.networkId);
 
         // Check the IP addresses to ensure that the upstream is indeed not the same as the previous
         // mobile upstream, even though the netId is (unrealistically) the same.
@@ -1281,13 +1295,13 @@
         mobile.fakeDisconnect();
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(true);
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(null);
+        mTetheringEventCallback.expectUpstreamChanged(NULL_NETWORK);
 
         mobile = new TestNetworkAgent(mCm, buildMobileDualStackUpstreamState());
         mobile.fakeConnect();
         mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(mobile.networkId);
 
         assertTrue(mUpstreamNetworkMonitor.getCurrentPreferredUpstream().linkProperties
                 .hasIPv4Address());
@@ -1328,27 +1342,27 @@
         mLooper.dispatchAll();
         mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST, null);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(mobile.networkId);
 
         wifi.fakeConnect();
         mLooper.dispatchAll();
         mCm.makeDefaultNetwork(wifi, CALLBACKS_FIRST, null);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(false);
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(wifi.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(wifi.networkId);
 
         verifyDisableTryCellWhenTetheringStop(inOrder);
     }
 
     private void verifyWifiUpstreamAndUnregisterDunCallback(@NonNull final InOrder inOrder,
-            @NonNull final TestNetworkAgent wifi,
-            @NonNull final NetworkCallback currentDunCallack) throws Exception {
+            @NonNull final TestNetworkAgent wifi, @NonNull final NetworkCallback currentDunCallack)
+            throws Exception {
         assertNotNull(currentDunCallack);
 
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(false);
         inOrder.verify(mCm).unregisterNetworkCallback(eq(currentDunCallack));
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(wifi.networkId);
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(any());
+        mTetheringEventCallback.expectUpstreamChanged(wifi.networkId);
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
     }
 
     @Nullable
@@ -1363,11 +1377,11 @@
                     captor.capture());
             dunNetworkCallback = captor.getValue();
         }
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(null);
+        mTetheringEventCallback.expectUpstreamChanged(NULL_NETWORK);
         final Runnable doDispatchAll = () -> mLooper.dispatchAll();
         dun.fakeConnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(dun.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(dun.networkId);
 
         if (needToRequestNetwork) {
             assertNotNull(dunNetworkCallback);
@@ -1484,11 +1498,11 @@
         dun.fakeDisconnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(true);
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(null);
+        mTetheringEventCallback.expectUpstreamChanged(NULL_NETWORK);
         inOrder.verify(mCm, never()).unregisterNetworkCallback(any(NetworkCallback.class));
         dun.fakeConnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(dun.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(dun.networkId);
 
         verifyDisableTryCellWhenTetheringStop(inOrder);
     }
@@ -1543,8 +1557,8 @@
         // automatic mode would request dun again and choose it as upstream.
         mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST);
         mLooper.dispatchAll();
-        final NetworkCallback dunNetworkCallback2 =
-                verifyDunUpstream(inOrder, dun, true /* needToRequestNetwork */);
+        final NetworkCallback dunNetworkCallback2 = verifyDunUpstream(inOrder, dun,
+                true /* needToRequestNetwork */);
 
         // [3] When default network switch to wifi and mobile is still connected,
         // unregister dun request and choose wifi as upstream.
@@ -1556,7 +1570,7 @@
         final Runnable doDispatchAll = () -> mLooper.dispatchAll();
         mobile.fakeDisconnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(any());
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
 
         verifyDisableTryCellWhenTetheringStop(inOrder);
     }
@@ -1627,19 +1641,19 @@
         final Runnable doDispatchAll = () -> mLooper.dispatchAll();
         dun.fakeConnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(dun.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(dun.networkId);
 
         // [6] When mobile is connected and default network switch to mobile, keep dun as upstream.
         mobile.fakeConnect();
         mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(any());
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
 
         // [7] When mobile is disconnected, keep dun as upstream.
         mCm.makeDefaultNetwork(null, CALLBACKS_FIRST, doDispatchAll);
         mobile.fakeDisconnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(any());
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
 
         verifyDisableTryCellWhenTetheringStop(inOrder);
     }
@@ -1670,7 +1684,7 @@
         final Runnable doDispatchAll = () -> mLooper.dispatchAll();
         dun.fakeConnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(dun.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(dun.networkId);
 
         // [8] Lose and regain upstream again.
         dun.fakeDisconnect(CALLBACKS_FIRST, doDispatchAll);
@@ -1714,7 +1728,7 @@
         final Runnable doDispatchAll = () -> mLooper.dispatchAll();
         dun.fakeConnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(dun.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(dun.networkId);
 
         // [9] When wifi is connected and default network switch to wifi, unregister dun request
         // and choose wifi as upstream. When dun is disconnected, keep wifi as upstream.
@@ -1724,7 +1738,7 @@
         verifyWifiUpstreamAndUnregisterDunCallback(inOrder, wifi, dunNetworkCallback);
         dun.fakeDisconnect(CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(any());
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
 
         // [10] When mobile and mobile are connected and default network switch to mobile
         // (may have low signal), automatic mode would request dun again and choose it as
@@ -1752,7 +1766,7 @@
         mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(false);
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(any());
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
         // BUG: when wifi disconnect, the dun request would not be filed again because wifi is
         // no longer be default network which do not have CONNECTIVIY_ACTION broadcast.
         wifi.fakeDisconnect();
@@ -1799,20 +1813,21 @@
     }
 
     private void chooseDunUpstreamTestCommon(final boolean automatic, InOrder inOrder,
-            TestNetworkAgent mobile, TestNetworkAgent wifi, TestNetworkAgent dun) throws Exception {
+            TestNetworkAgent mobile, TestNetworkAgent wifi, TestNetworkAgent dun)
+            throws Exception {
         final NetworkCallback dunNetworkCallback = setupDunUpstreamTest(automatic, inOrder);
 
         // Pretend cellular connected and expect the upstream to be not set.
         mobile.fakeConnect();
         mCm.makeDefaultNetwork(mobile, BROADCAST_FIRST);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(mobile.networkId);
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
 
         // Pretend dun connected and expect choose dun as upstream.
         final Runnable doDispatchAll = () -> mLooper.dispatchAll();
         dun.fakeConnect(BROADCAST_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(dun.networkId);
+        mTetheringEventCallback.expectUpstreamChanged(dun.networkId);
 
         // When wifi connected, unregister dun request and choose wifi as upstream.
         wifi.fakeConnect();
@@ -1821,7 +1836,7 @@
         verifyWifiUpstreamAndUnregisterDunCallback(inOrder, wifi, dunNetworkCallback);
         dun.fakeDisconnect(BROADCAST_FIRST, doDispatchAll);
         mLooper.dispatchAll();
-        inOrder.verify(mUpstreamNetworkMonitor, never()).setCurrentUpstream(any());
+        mTetheringEventCallback.assertNoUpstreamChangeCallback();
     }
 
     private void runNcmTethering() {
@@ -2268,7 +2283,7 @@
         mTethering.registerTetheringEventCallback(callback);
         mLooper.dispatchAll();
         callback.expectTetheredClientChanged(Collections.emptyList());
-        callback.expectUpstreamChanged(new Network[] {null});
+        callback.expectUpstreamChanged(NULL_NETWORK);
         callback.expectConfigurationChanged(
                 mTethering.getTetheringConfiguration().toStableParcelable());
         TetherStatesParcel tetherState = callback.pollTetherStatesChanged();
@@ -2316,7 +2331,7 @@
         tetherState = callback2.pollTetherStatesChanged();
         assertArrayEquals(tetherState.availableList, new TetheringInterface[] {wifiIface});
         mLooper.dispatchAll();
-        callback2.expectUpstreamChanged(new Network[] {null});
+        callback2.expectUpstreamChanged(NULL_NETWORK);
         callback2.expectOffloadStatusChanged(TETHER_HARDWARE_OFFLOAD_STOPPED);
         callback.assertNoCallback();
     }
@@ -2663,12 +2678,19 @@
     public void testUpstreamNetworkChanged() {
         final Tethering.TetherMainSM stateMachine = (Tethering.TetherMainSM)
                 mTetheringDependencies.mUpstreamNetworkMonitorSM;
+        // Gain upstream.
         final UpstreamNetworkState upstreamState = buildMobileIPv4UpstreamState();
         initTetheringUpstream(upstreamState);
         stateMachine.chooseUpstreamType(true);
+        mTetheringEventCallback.expectUpstreamChanged(upstreamState.network);
+        verify(mNotificationUpdater)
+                .onUpstreamCapabilitiesChanged(upstreamState.networkCapabilities);
 
-        verify(mUpstreamNetworkMonitor, times(1)).setCurrentUpstream(eq(upstreamState.network));
-        verify(mNotificationUpdater, times(1)).onUpstreamCapabilitiesChanged(any());
+        // Lose upstream.
+        initTetheringUpstream(null);
+        stateMachine.chooseUpstreamType(true);
+        mTetheringEventCallback.expectUpstreamChanged(NULL_NETWORK);
+        verify(mNotificationUpdater).onUpstreamCapabilitiesChanged(null);
     }
 
     @Test
@@ -2682,7 +2704,8 @@
         stateMachine.handleUpstreamNetworkMonitorCallback(EVENT_ON_CAPABILITIES, upstreamState);
         // Should have two onUpstreamCapabilitiesChanged().
         // One is called by reportUpstreamChanged(). One is called by EVENT_ON_CAPABILITIES.
-        verify(mNotificationUpdater, times(2)).onUpstreamCapabilitiesChanged(any());
+        verify(mNotificationUpdater, times(2))
+                .onUpstreamCapabilitiesChanged(upstreamState.networkCapabilities);
         reset(mNotificationUpdater);
 
         // Verify that onUpstreamCapabilitiesChanged won't be called if not current upstream network
@@ -2695,6 +2718,27 @@
     }
 
     @Test
+    public void testUpstreamCapabilitiesChanged_startStopTethering() throws Exception {
+        final TestNetworkAgent wifi = new TestNetworkAgent(mCm, buildWifiUpstreamState());
+
+        // Start USB tethering with no current upstream.
+        prepareUsbTethering();
+        sendUsbBroadcast(true, true, TETHER_USB_RNDIS_FUNCTION);
+
+        // Pretend wifi connected and expect the upstream to be set.
+        wifi.fakeConnect();
+        mCm.makeDefaultNetwork(wifi, CALLBACKS_FIRST);
+        mLooper.dispatchAll();
+        verify(mNotificationUpdater).onUpstreamCapabilitiesChanged(
+                wifi.networkCapabilities);
+
+        // Stop tethering.
+        // Expect that TetherModeAliveState#exit sends capabilities change notification to null.
+        runStopUSBTethering();
+        verify(mNotificationUpdater).onUpstreamCapabilitiesChanged(null);
+    }
+
+    @Test
     public void testDumpTetheringLog() throws Exception {
         final FileDescriptor mockFd = mock(FileDescriptor.class);
         final PrintWriter mockPw = mock(PrintWriter.class);
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 40defd4..4224da9 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -2287,6 +2287,13 @@
                 mExecutor.execute(() -> {
                     try {
                         if (mSlot != null) {
+                            // TODO : this is incorrect, because in the presence of auto on/off
+                            // keepalive the slot associated with this keepalive can have
+                            // changed. Also, this actually causes another problem where some other
+                            // app might stop your keepalive if it just knows the network and
+                            // the slot and goes through the trouble of grabbing the aidl object.
+                            // This code should use the callback to identify what keepalive to
+                            // stop instead.
                             mService.stopKeepalive(mNetwork, mSlot);
                         }
                     } catch (RemoteException e) {
diff --git a/framework/src/android/net/NetworkAgent.java b/framework/src/android/net/NetworkAgent.java
index 732bd87..3ec00d9 100644
--- a/framework/src/android/net/NetworkAgent.java
+++ b/framework/src/android/net/NetworkAgent.java
@@ -281,7 +281,7 @@
      *
      *   arg1 = the hardware slot number of the keepalive to start
      *   arg2 = interval in seconds
-     *   obj = KeepalivePacketData object describing the data to be sent
+     *   obj = AutomaticKeepaliveInfo object
      *
      * Also used internally by ConnectivityService / KeepaliveTracker, with different semantics.
      * @hide
@@ -491,8 +491,7 @@
      * TCP sockets are open over a VPN. The system will check periodically for presence of
      * such open sockets, and this message is what triggers the re-evaluation.
      *
-     * arg1 = hardware slot number of the keepalive
-     * obj = {@link Network} that the keepalive is started on.
+     * obj = A Binder object associated with the keepalive.
      * @hide
      */
     public static final int CMD_MONITOR_AUTOMATIC_KEEPALIVE = BASE + 30;
diff --git a/framework/src/android/net/NetworkCapabilities.java b/framework/src/android/net/NetworkCapabilities.java
index e70d75d..2385f69 100644
--- a/framework/src/android/net/NetworkCapabilities.java
+++ b/framework/src/android/net/NetworkCapabilities.java
@@ -65,6 +65,24 @@
  * bandwidth. Similarly if an application needs an unmetered network for a bulk
  * transfer it can specify that rather than assuming all cellular based
  * connections are metered and all Wi-Fi based connections are not.
+ *
+ * <p> Starting from Android 14, if the developer wants the application to call
+ * {@link android.net.ConnectivityManager#requestNetwork} with some specific capabilities, the
+ * developer has to explicitly add the
+ * {@link android.content.pm.PackageManager#PROPERTY_SELF_CERTIFIED_NETWORK_CAPABILITIES}
+ * property in the AndroidManifest.xml, which points to a self_certified_network_capabilities.xml
+ * resource file. In self_certified_network_capabilities.xml, it declares what kind of
+ * network capabilities the application wants to have.
+ *
+ * Here is an example self_certified_network_capabilities.xml:
+ * <pre>
+ *  {@code
+ *  <network-capabilities-declaration xmlns:android="http://schemas.android.com/apk/res/android">
+ *     <uses-network-capability android:name="NET_CAPABILITY_PRIORITIZE_LATENCY"/>
+ *     <uses-network-capability android:name="NET_CAPABILITY_PRIORITIZE_BANDWIDTH"/>
+ * </network-capabilities-declaration>
+ *  }
+ *  </pre>
  */
 public final class NetworkCapabilities implements Parcelable {
     private static final String TAG = "NetworkCapabilities";
@@ -622,11 +640,17 @@
 
     /**
      * Indicates that this network should be able to prioritize latency for the internet.
+     *
+     * <p> Starting from Android 14, user must explicitly declare they want to use this
+     * capability in app. Please refer to {@link NetworkCapabilities} for more details.
      */
     public static final int NET_CAPABILITY_PRIORITIZE_LATENCY = 34;
 
     /**
      * Indicates that this network should be able to prioritize bandwidth for the internet.
+     *
+     * <p> Starting from Android 14, user must explicitly declare they want to use this
+     * capability in app. Please refer to {@link NetworkCapabilities} for more details.
      */
     public static final int NET_CAPABILITY_PRIORITIZE_BANDWIDTH = 35;
 
diff --git a/nearby/tests/integration/untrusted/src/androidx/test/uiautomator/LogcatWaitMixin.java b/nearby/tests/integration/untrusted/src/androidx/test/uiautomator/LogcatWaitMixin.java
index 86e39dc..8b5ac12 100644
--- a/nearby/tests/integration/untrusted/src/androidx/test/uiautomator/LogcatWaitMixin.java
+++ b/nearby/tests/integration/untrusted/src/androidx/test/uiautomator/LogcatWaitMixin.java
@@ -54,7 +54,7 @@
             @NonNull String specificLog, @NonNull Date startTime) {
         return new Condition<UiDevice, Boolean>() {
             @Override
-            Boolean apply(UiDevice device) {
+            public Boolean apply(UiDevice device) {
                 String logcatLogs;
                 try {
                     logcatLogs = device.executeShellCommand("logcat -v time -v year -d");
diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
index 4c37b8d..aeadb4a 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
@@ -50,6 +50,7 @@
 void NetworkTraceHandler::InitPerfettoTracing() {
   perfetto::TracingInitArgs args = {};
   args.backends |= perfetto::kSystemBackend;
+  args.enable_system_consumer = false;
   perfetto::Tracing::Initialize(args);
   NetworkTraceHandler::RegisterDataSource();
 }
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 5dcf860..4ad39e1 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -52,7 +52,6 @@
 import android.util.Log;
 import android.util.Pair;
 import android.util.SparseArray;
-import android.util.SparseIntArray;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.State;
@@ -119,13 +118,15 @@
     private final NsdStateMachine mNsdStateMachine;
     private final MDnsManager mMDnsManager;
     private final MDnsEventCallback mMDnsEventCallback;
-    @Nullable
+    @NonNull
+    private final Dependencies mDeps;
+    @NonNull
     private final MdnsMultinetworkSocketClient mMdnsSocketClient;
-    @Nullable
+    @NonNull
     private final MdnsDiscoveryManager mMdnsDiscoveryManager;
-    @Nullable
+    @NonNull
     private final MdnsSocketProvider mMdnsSocketProvider;
-    @Nullable
+    @NonNull
     private final MdnsAdvertiser mAdvertiser;
     // 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
@@ -312,21 +313,14 @@
             mIsMonitoringSocketsStarted = true;
         }
 
-        private void maybeStopMonitoringSockets() {
-            if (!mIsMonitoringSocketsStarted) {
-                if (DBG) Log.d(TAG, "Socket monitoring has not been started.");
-                return;
-            }
+        private void maybeStopMonitoringSocketsIfNoActiveRequest() {
+            if (!mIsMonitoringSocketsStarted) return;
+            if (isAnyRequestActive()) return;
+
             mMdnsSocketProvider.stopMonitoringSockets();
             mIsMonitoringSocketsStarted = false;
         }
 
-        private void maybeStopMonitoringSocketsIfNoActiveRequest() {
-            if (!isAnyRequestActive()) {
-                maybeStopMonitoringSockets();
-            }
-        }
-
         NsdStateMachine(String name, Handler handler) {
             super(name, handler);
             addState(mDefaultState);
@@ -358,17 +352,12 @@
                         final NsdServiceConnector connector = (NsdServiceConnector) msg.obj;
                         cInfo = mClients.remove(connector);
                         if (cInfo != null) {
-                            if (mMdnsDiscoveryManager != null) {
-                                cInfo.unregisterAllListeners();
-                            }
                             cInfo.expungeAllRequests();
-                            if (cInfo.isLegacy()) {
+                            if (cInfo.isPreSClient()) {
                                 mLegacyClientCount -= 1;
                             }
                         }
-                        if (mMdnsDiscoveryManager != null || mAdvertiser != null) {
-                            maybeStopMonitoringSocketsIfNoActiveRequest();
-                        }
+                        maybeStopMonitoringSocketsIfNoActiveRequest();
                         maybeScheduleStop();
                         break;
                     case NsdManager.DISCOVER_SERVICES:
@@ -429,7 +418,7 @@
                         cInfo = getClientInfoForReply(msg);
                         if (cInfo != null) {
                             cancelStop();
-                            cInfo.setLegacy();
+                            cInfo.setPreSClient();
                             mLegacyClientCount += 1;
                             maybeStartDaemon();
                         }
@@ -461,41 +450,45 @@
             }
 
             private boolean requestLimitReached(ClientInfo clientInfo) {
-                if (clientInfo.mClientIds.size() >= ClientInfo.MAX_LIMIT) {
+                if (clientInfo.mClientRequests.size() >= ClientInfo.MAX_LIMIT) {
                     if (DBG) Log.d(TAG, "Exceeded max outstanding requests " + clientInfo);
                     return true;
                 }
                 return false;
             }
 
-            private void storeRequestMap(int clientId, int globalId, ClientInfo clientInfo, int what) {
-                clientInfo.mClientIds.put(clientId, globalId);
-                clientInfo.mClientRequests.put(clientId, what);
+            private void storeLegacyRequestMap(int clientId, int globalId, ClientInfo clientInfo,
+                    int what) {
+                clientInfo.mClientRequests.put(clientId, new LegacyClientRequest(globalId, what));
                 mIdToClientInfoMap.put(globalId, clientInfo);
                 // Remove the cleanup event because here comes a new request.
                 cancelStop();
             }
 
-            private void removeRequestMap(int clientId, int globalId, ClientInfo clientInfo) {
-                clientInfo.mClientIds.delete(clientId);
-                clientInfo.mClientRequests.delete(clientId);
-                mIdToClientInfoMap.remove(globalId);
-                maybeScheduleStop();
-                maybeStopMonitoringSocketsIfNoActiveRequest();
-            }
-
-            private void storeListenerMap(int clientId, int transactionId, MdnsListener listener,
+            private void storeAdvertiserRequestMap(int clientId, int globalId,
                     ClientInfo clientInfo) {
-                clientInfo.mClientIds.put(clientId, transactionId);
-                clientInfo.mListeners.put(clientId, listener);
-                mIdToClientInfoMap.put(transactionId, clientInfo);
+                clientInfo.mClientRequests.put(clientId, new AdvertiserClientRequest(globalId));
+                mIdToClientInfoMap.put(globalId, clientInfo);
             }
 
-            private void removeListenerMap(int clientId, int transactionId, ClientInfo clientInfo) {
-                clientInfo.mClientIds.delete(clientId);
-                clientInfo.mListeners.delete(clientId);
-                mIdToClientInfoMap.remove(transactionId);
-                maybeStopMonitoringSocketsIfNoActiveRequest();
+            private void removeRequestMap(int clientId, int globalId, ClientInfo clientInfo) {
+                final ClientRequest existing = clientInfo.mClientRequests.get(clientId);
+                if (existing == null) return;
+                clientInfo.mClientRequests.remove(clientId);
+                mIdToClientInfoMap.remove(globalId);
+
+                if (existing instanceof LegacyClientRequest) {
+                    maybeScheduleStop();
+                } else {
+                    maybeStopMonitoringSocketsIfNoActiveRequest();
+                }
+            }
+
+            private void storeDiscoveryManagerRequestMap(int clientId, int globalId,
+                    MdnsListener listener, ClientInfo clientInfo) {
+                clientInfo.mClientRequests.put(clientId,
+                        new DiscoveryManagerRequest(globalId, listener));
+                mIdToClientInfoMap.put(globalId, clientInfo);
             }
 
             private void clearRegisteredServiceInfo(ClientInfo clientInfo) {
@@ -579,7 +572,7 @@
 
                         final NsdServiceInfo info = args.serviceInfo;
                         id = getUniqueId();
-                        if (mMdnsDiscoveryManager != null) {
+                        if (mDeps.isMdnsDiscoveryManagerEnabled(mContext)) {
                             final String serviceType = constructServiceType(info.getServiceType());
                             if (serviceType == null) {
                                 clientInfo.onDiscoverServicesFailed(clientId,
@@ -597,7 +590,7 @@
                                     .build();
                             mMdnsDiscoveryManager.registerListener(
                                     listenServiceType, listener, options);
-                            storeListenerMap(clientId, id, listener, clientInfo);
+                            storeDiscoveryManagerRequestMap(clientId, id, listener, clientInfo);
                             clientInfo.onDiscoverServicesStarted(clientId, info);
                         } else {
                             maybeStartDaemon();
@@ -606,7 +599,7 @@
                                     Log.d(TAG, "Discover " + msg.arg2 + " " + id
                                             + info.getServiceType());
                                 }
-                                storeRequestMap(clientId, id, clientInfo, msg.what);
+                                storeLegacyRequestMap(clientId, id, clientInfo, msg.what);
                                 clientInfo.onDiscoverServicesStarted(clientId, info);
                             } else {
                                 stopServiceDiscovery(id);
@@ -616,7 +609,7 @@
                         }
                         break;
                     }
-                    case NsdManager.STOP_DISCOVERY:
+                    case NsdManager.STOP_DISCOVERY: {
                         if (DBG) Log.d(TAG, "Stop service discovery");
                         args = (ListenerArgs) msg.obj;
                         clientInfo = mClients.get(args.connector);
@@ -628,23 +621,21 @@
                             break;
                         }
 
-                        try {
-                            id = clientInfo.mClientIds.get(clientId);
-                        } catch (NullPointerException e) {
-                            clientInfo.onStopDiscoveryFailed(
-                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        if (request == null) {
+                            Log.e(TAG, "Unknown client request in STOP_DISCOVERY");
                             break;
                         }
-                        if (mMdnsDiscoveryManager != null) {
-                            final MdnsListener listener = clientInfo.mListeners.get(clientId);
-                            if (listener == null) {
-                                clientInfo.onStopDiscoveryFailed(
-                                        clientId, NsdManager.FAILURE_INTERNAL_ERROR);
-                                break;
-                            }
+                        id = request.mGlobalId;
+                        // 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) {
+                            final MdnsListener listener =
+                                    ((DiscoveryManagerRequest) request).mListener;
                             mMdnsDiscoveryManager.unregisterListener(
                                     listener.getListenedServiceType(), listener);
-                            removeListenerMap(clientId, id, clientInfo);
+                            removeRequestMap(clientId, id, clientInfo);
                             clientInfo.onStopDiscoverySucceeded(clientId);
                         } else {
                             removeRequestMap(clientId, id, clientInfo);
@@ -656,7 +647,8 @@
                             }
                         }
                         break;
-                    case NsdManager.REGISTER_SERVICE:
+                    }
+                    case NsdManager.REGISTER_SERVICE: {
                         if (DBG) Log.d(TAG, "Register service");
                         args = (ListenerArgs) msg.obj;
                         clientInfo = mClients.get(args.connector);
@@ -675,7 +667,7 @@
                         }
 
                         id = getUniqueId();
-                        if (mAdvertiser != null) {
+                        if (mDeps.isMdnsAdvertiserEnabled(mContext)) {
                             final NsdServiceInfo serviceInfo = args.serviceInfo;
                             final String serviceType = serviceInfo.getServiceType();
                             final String registerServiceType = constructServiceType(serviceType);
@@ -691,12 +683,12 @@
 
                             maybeStartMonitoringSockets();
                             mAdvertiser.addService(id, serviceInfo);
-                            storeRequestMap(clientId, id, clientInfo, msg.what);
+                            storeAdvertiserRequestMap(clientId, id, clientInfo);
                         } else {
                             maybeStartDaemon();
                             if (registerService(id, args.serviceInfo)) {
                                 if (DBG) Log.d(TAG, "Register " + clientId + " " + id);
-                                storeRequestMap(clientId, id, clientInfo, msg.what);
+                                storeLegacyRequestMap(clientId, id, clientInfo, msg.what);
                                 // Return success after mDns reports success
                             } else {
                                 unregisterService(id);
@@ -706,7 +698,8 @@
 
                         }
                         break;
-                    case NsdManager.UNREGISTER_SERVICE:
+                    }
+                    case NsdManager.UNREGISTER_SERVICE: {
                         if (DBG) Log.d(TAG, "unregister service");
                         args = (ListenerArgs) msg.obj;
                         clientInfo = mClients.get(args.connector);
@@ -717,10 +710,18 @@
                             Log.e(TAG, "Unknown connector in unregistration");
                             break;
                         }
-                        id = clientInfo.mClientIds.get(clientId);
+                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        if (request == null) {
+                            Log.e(TAG, "Unknown client request in UNREGISTER_SERVICE");
+                            break;
+                        }
+                        id = request.mGlobalId;
                         removeRequestMap(clientId, id, clientInfo);
 
-                        if (mAdvertiser != null) {
+                        // 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);
                         } else {
@@ -732,6 +733,7 @@
                             }
                         }
                         break;
+                    }
                     case NsdManager.RESOLVE_SERVICE: {
                         if (DBG) Log.d(TAG, "Resolve service");
                         args = (ListenerArgs) msg.obj;
@@ -746,7 +748,7 @@
 
                         final NsdServiceInfo info = args.serviceInfo;
                         id = getUniqueId();
-                        if (mMdnsDiscoveryManager != null) {
+                        if (mDeps.isMdnsDiscoveryManagerEnabled(mContext)) {
                             final String serviceType = constructServiceType(info.getServiceType());
                             if (serviceType == null) {
                                 clientInfo.onResolveServiceFailed(clientId,
@@ -764,7 +766,7 @@
                                     .build();
                             mMdnsDiscoveryManager.registerListener(
                                     resolveServiceType, listener, options);
-                            storeListenerMap(clientId, id, listener, clientInfo);
+                            storeDiscoveryManagerRequestMap(clientId, id, listener, clientInfo);
                         } else {
                             if (clientInfo.mResolvedService != null) {
                                 clientInfo.onResolveServiceFailed(
@@ -775,7 +777,7 @@
                             maybeStartDaemon();
                             if (resolveService(id, args.serviceInfo)) {
                                 clientInfo.mResolvedService = new NsdServiceInfo();
-                                storeRequestMap(clientId, id, clientInfo, msg.what);
+                                storeLegacyRequestMap(clientId, id, clientInfo, msg.what);
                             } else {
                                 clientInfo.onResolveServiceFailed(
                                         clientId, NsdManager.FAILURE_INTERNAL_ERROR);
@@ -783,7 +785,7 @@
                         }
                         break;
                     }
-                    case NsdManager.STOP_RESOLUTION:
+                    case NsdManager.STOP_RESOLUTION: {
                         if (DBG) Log.d(TAG, "Stop service resolution");
                         args = (ListenerArgs) msg.obj;
                         clientInfo = mClients.get(args.connector);
@@ -795,7 +797,12 @@
                             break;
                         }
 
-                        id = clientInfo.mClientIds.get(clientId);
+                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        if (request == null) {
+                            Log.e(TAG, "Unknown client request in STOP_RESOLUTION");
+                            break;
+                        }
+                        id = request.mGlobalId;
                         removeRequestMap(clientId, id, clientInfo);
                         if (stopResolveService(id)) {
                             clientInfo.onStopResolutionSucceeded(clientId);
@@ -806,6 +813,7 @@
                         clientInfo.mResolvedService = null;
                         // TODO: Implement the stop resolution with MdnsDiscoveryManager.
                         break;
+                    }
                     case NsdManager.REGISTER_SERVICE_CALLBACK:
                         if (DBG) Log.d(TAG, "Register a service callback");
                         args = (ListenerArgs) msg.obj;
@@ -829,13 +837,13 @@
                         if (resolveService(id, args.serviceInfo)) {
                             clientInfo.mRegisteredService = new NsdServiceInfo();
                             clientInfo.mClientIdForServiceUpdates = clientId;
-                            storeRequestMap(clientId, id, clientInfo, msg.what);
+                            storeLegacyRequestMap(clientId, id, clientInfo, msg.what);
                         } else {
                             clientInfo.onServiceInfoCallbackRegistrationFailed(
                                     clientId, NsdManager.FAILURE_BAD_PARAMETERS);
                         }
                         break;
-                    case NsdManager.UNREGISTER_SERVICE_CALLBACK:
+                    case NsdManager.UNREGISTER_SERVICE_CALLBACK: {
                         if (DBG) Log.d(TAG, "Unregister a service callback");
                         args = (ListenerArgs) msg.obj;
                         clientInfo = mClients.get(args.connector);
@@ -847,7 +855,12 @@
                             break;
                         }
 
-                        id = clientInfo.mClientIds.get(clientId);
+                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        if (request == null) {
+                            Log.e(TAG, "Unknown client request in STOP_RESOLUTION");
+                            break;
+                        }
+                        id = request.mGlobalId;
                         removeRequestMap(clientId, id, clientInfo);
                         if (stopResolveService(id)) {
                             clientInfo.onServiceInfoCallbackUnregistered(clientId);
@@ -856,6 +869,7 @@
                         }
                         clearRegisteredServiceInfo(clientInfo);
                         break;
+                    }
                     case MDNS_SERVICE_EVENT:
                         if (!handleMDnsServiceEvent(msg.arg1, msg.arg2, msg.obj)) {
                             return NOT_HANDLED;
@@ -995,7 +1009,8 @@
 
                         final int id2 = getUniqueId();
                         if (getAddrInfo(id2, info.hostname, info.interfaceIdx)) {
-                            storeRequestMap(clientId, id2, clientInfo, NsdManager.RESOLVE_SERVICE);
+                            storeLegacyRequestMap(clientId, id2, clientInfo,
+                                    NsdManager.RESOLVE_SERVICE);
                         } else {
                             notifyResolveFailedResult(isListenedToUpdates, clientId, clientInfo,
                                     NsdManager.FAILURE_BAD_PARAMETERS);
@@ -1110,6 +1125,11 @@
                         clientInfo.onServiceLost(clientId, info);
                         break;
                     case NsdManager.RESOLVE_SERVICE_SUCCEEDED: {
+                        final ClientRequest request = clientInfo.mClientRequests.get(clientId);
+                        if (request == null) {
+                            Log.e(TAG, "Unknown client request in RESOLVE_SERVICE_SUCCEEDED");
+                            break;
+                        }
                         final MdnsServiceInfo serviceInfo = event.mMdnsServiceInfo;
                         // Add '.' in front of the service type that aligns with historical behavior
                         info.setServiceType("." + event.mRequestedServiceType);
@@ -1140,10 +1160,14 @@
                         }
 
                         // Unregister the listener immediately like IMDnsEventListener design
-                        final MdnsListener listener = clientInfo.mListeners.get(clientId);
+                        if (!(request instanceof DiscoveryManagerRequest)) {
+                            Log.wtf(TAG, "non-DiscoveryManager request in DiscoveryManager event");
+                            break;
+                        }
+                        final MdnsListener listener = ((DiscoveryManagerRequest) request).mListener;
                         mMdnsDiscoveryManager.unregisterListener(
                                 listener.getListenedServiceType(), listener);
-                        removeListenerMap(clientId, transactionId, clientInfo);
+                        removeRequestMap(clientId, transactionId, clientInfo);
                         break;
                     }
                     default:
@@ -1216,32 +1240,16 @@
         mNsdStateMachine.start();
         mMDnsManager = ctx.getSystemService(MDnsManager.class);
         mMDnsEventCallback = new MDnsEventCallback(mNsdStateMachine);
+        mDeps = deps;
 
-        final boolean discoveryManagerEnabled = deps.isMdnsDiscoveryManagerEnabled(ctx);
-        final boolean advertiserEnabled = deps.isMdnsAdvertiserEnabled(ctx);
-        if (discoveryManagerEnabled || advertiserEnabled) {
-            mMdnsSocketProvider = deps.makeMdnsSocketProvider(ctx, handler.getLooper());
-        } else {
-            mMdnsSocketProvider = null;
-        }
-
-        if (discoveryManagerEnabled) {
-            mMdnsSocketClient =
-                    new MdnsMultinetworkSocketClient(handler.getLooper(), mMdnsSocketProvider);
-            mMdnsDiscoveryManager =
-                    deps.makeMdnsDiscoveryManager(new ExecutorProvider(), mMdnsSocketClient);
-            handler.post(() -> mMdnsSocketClient.setCallback(mMdnsDiscoveryManager));
-        } else {
-            mMdnsSocketClient = null;
-            mMdnsDiscoveryManager = null;
-        }
-
-        if (advertiserEnabled) {
-            mAdvertiser = deps.makeMdnsAdvertiser(handler.getLooper(), mMdnsSocketProvider,
-                    new AdvertiserCallback());
-        } else {
-            mAdvertiser = null;
-        }
+        mMdnsSocketProvider = deps.makeMdnsSocketProvider(ctx, handler.getLooper());
+        mMdnsSocketClient =
+                new MdnsMultinetworkSocketClient(handler.getLooper(), mMdnsSocketProvider);
+        mMdnsDiscoveryManager =
+                deps.makeMdnsDiscoveryManager(new ExecutorProvider(), mMdnsSocketClient);
+        handler.post(() -> mMdnsSocketClient.setCallback(mMdnsDiscoveryManager));
+        mAdvertiser = deps.makeMdnsAdvertiser(handler.getLooper(), mMdnsSocketProvider,
+                new AdvertiserCallback());
     }
 
     /**
@@ -1604,6 +1612,39 @@
         mNsdStateMachine.dump(fd, pw, args);
     }
 
+    private abstract static class ClientRequest {
+        private final int mGlobalId;
+
+        private ClientRequest(int globalId) {
+            mGlobalId = globalId;
+        }
+    }
+
+    private static class LegacyClientRequest extends ClientRequest {
+        private final int mRequestCode;
+
+        private LegacyClientRequest(int globalId, int requestCode) {
+            super(globalId);
+            mRequestCode = requestCode;
+        }
+    }
+
+    private static class AdvertiserClientRequest extends ClientRequest {
+        private AdvertiserClientRequest(int globalId) {
+            super(globalId);
+        }
+    }
+
+    private static class DiscoveryManagerRequest extends ClientRequest {
+        @NonNull
+        private final MdnsListener mListener;
+
+        private DiscoveryManagerRequest(int globalId, @NonNull MdnsListener listener) {
+            super(globalId);
+            mListener = listener;
+        }
+    }
+
     /* Information tracked per client */
     private class ClientInfo {
 
@@ -1612,17 +1653,11 @@
         /* Remembers a resolved service until getaddrinfo completes */
         private NsdServiceInfo mResolvedService;
 
-        /* A map from client id to unique id sent to mDns */
-        private final SparseIntArray mClientIds = new SparseIntArray();
-
-        /* A map from client id to the type of the request we had received */
-        private final SparseIntArray mClientRequests = new SparseIntArray();
-
-        /* A map from client id to the MdnsListener */
-        private final SparseArray<MdnsListener> mListeners = new SparseArray<>();
+        /* A map from client-side ID (listenerKey) to the request */
+        private final SparseArray<ClientRequest> mClientRequests = new SparseArray<>();
 
         // The target SDK of this client < Build.VERSION_CODES.S
-        private boolean mIsLegacy = false;
+        private boolean mIsPreSClient = false;
 
         /*** The service that is registered to listen to its updates */
         private NsdServiceInfo mRegisteredService;
@@ -1638,38 +1673,59 @@
         public String toString() {
             StringBuilder sb = new StringBuilder();
             sb.append("mResolvedService ").append(mResolvedService).append("\n");
-            sb.append("mIsLegacy ").append(mIsLegacy).append("\n");
-            for(int i = 0; i< mClientIds.size(); i++) {
-                int clientID = mClientIds.keyAt(i);
-                sb.append("clientId ").append(clientID).
-                    append(" mDnsId ").append(mClientIds.valueAt(i)).
-                    append(" type ").append(mClientRequests.get(clientID)).append("\n");
+            sb.append("mIsLegacy ").append(mIsPreSClient).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)
+                        .append(" type ").append(
+                                mClientRequests.valueAt(i).getClass().getSimpleName())
+                        .append("\n");
             }
             return sb.toString();
         }
 
-        private boolean isLegacy() {
-            return mIsLegacy;
+        private boolean isPreSClient() {
+            return mIsPreSClient;
         }
 
-        private void setLegacy() {
-            mIsLegacy = true;
+        private void setPreSClient() {
+            mIsPreSClient = true;
         }
 
         // Remove any pending requests from the global map when we get rid of a client,
         // and send cancellations to the daemon.
         private void expungeAllRequests() {
-            int globalId, clientId, i;
             // TODO: to keep handler responsive, do not clean all requests for that client at once.
-            for (i = 0; i < mClientIds.size(); i++) {
-                clientId = mClientIds.keyAt(i);
-                globalId = mClientIds.valueAt(i);
+            for (int i = 0; i < mClientRequests.size(); i++) {
+                final int clientId = mClientRequests.keyAt(i);
+                final ClientRequest request = mClientRequests.valueAt(i);
+                final int globalId = request.mGlobalId;
                 mIdToClientInfoMap.remove(globalId);
                 if (DBG) {
                     Log.d(TAG, "Terminating client-ID " + clientId
                             + " global-ID " + globalId + " type " + mClientRequests.get(clientId));
                 }
-                switch (mClientRequests.get(clientId)) {
+
+                if (request instanceof DiscoveryManagerRequest) {
+                    final MdnsListener listener =
+                            ((DiscoveryManagerRequest) request).mListener;
+                    mMdnsDiscoveryManager.unregisterListener(
+                            listener.getListenedServiceType(), listener);
+                    continue;
+                }
+
+                if (request instanceof AdvertiserClientRequest) {
+                    mAdvertiser.removeService(globalId);
+                    continue;
+                }
+
+                if (!(request instanceof LegacyClientRequest)) {
+                    throw new IllegalStateException("Unknown request type: " + request.getClass());
+                }
+
+                switch (((LegacyClientRequest) request).mRequestCode) {
                     case NsdManager.DISCOVER_SERVICES:
                         stopServiceDiscovery(globalId);
                         break;
@@ -1677,37 +1733,25 @@
                         stopResolveService(globalId);
                         break;
                     case NsdManager.REGISTER_SERVICE:
-                        if (mAdvertiser != null) {
-                            mAdvertiser.removeService(globalId);
-                        } else {
-                            unregisterService(globalId);
-                        }
+                        unregisterService(globalId);
                         break;
                     default:
                         break;
                 }
             }
-            mClientIds.clear();
             mClientRequests.clear();
         }
 
-        void unregisterAllListeners() {
-            for (int i = 0; i < mListeners.size(); i++) {
-                final MdnsListener listener = mListeners.valueAt(i);
-                mMdnsDiscoveryManager.unregisterListener(
-                        listener.getListenedServiceType(), listener);
-            }
-            mListeners.clear();
-        }
-
-        // mClientIds is a sparse array of listener id -> mDnsClient id.  For a given mDnsClient id,
-        // return the corresponding listener id.  mDnsClient id is also called a global id.
+        // 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) {
-            int idx = mClientIds.indexOfValue(globalId);
-            if (idx < 0) {
-                return idx;
+            for (int i = 0; i < mClientRequests.size(); i++) {
+                if (mClientRequests.valueAt(i).mGlobalId == globalId) {
+                    return mClientRequests.keyAt(i);
+                }
             }
-            return mClientIds.keyAt(idx);
+            return -1;
         }
 
         private void maybeNotifyRegisteredServiceLost(@NonNull NsdServiceInfo info) {
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index eef4a0e..f1c68cb 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -270,6 +270,7 @@
 import com.android.networkstack.apishim.common.UnsupportedApiLevelException;
 import com.android.server.connectivity.AutodestructReference;
 import com.android.server.connectivity.AutomaticOnOffKeepaliveTracker;
+import com.android.server.connectivity.AutomaticOnOffKeepaliveTracker.AutomaticOnOffKeepalive;
 import com.android.server.connectivity.CarrierPrivilegeAuthenticator;
 import com.android.server.connectivity.ClatCoordinator;
 import com.android.server.connectivity.ConnectivityFlags;
@@ -5546,9 +5547,11 @@
                     break;
                 }
                 case NetworkAgent.CMD_MONITOR_AUTOMATIC_KEEPALIVE: {
-                    final Network network = (Network) msg.obj;
-                    final int slot = msg.arg1;
+                    final AutomaticOnOffKeepalive ki =
+                            mKeepaliveTracker.getKeepaliveForBinder((IBinder) msg.obj);
+                    if (null == ki) return; // The callback was unregistered before the alarm fired
 
+                    final Network network = ki.getNetwork();
                     boolean networkFound = false;
                     final ArrayList<NetworkAgentInfo> vpnsRunningOnThisNetwork = new ArrayList<>();
                     for (NetworkAgentInfo n : mNetworkAgentInfos) {
@@ -5563,18 +5566,22 @@
                     if (!networkFound) return;
 
                     if (!vpnsRunningOnThisNetwork.isEmpty()) {
-                        mKeepaliveTracker.handleMonitorAutomaticKeepalive(network, slot,
+                        mKeepaliveTracker.handleMonitorAutomaticKeepalive(ki,
                                 // TODO: check all the VPNs running on top of this network
                                 vpnsRunningOnThisNetwork.get(0).network.netId);
                     } else {
                         // If no VPN, then make sure the keepalive is running.
-                        mKeepaliveTracker.handleMaybeResumeKeepalive(network, slot);
+                        mKeepaliveTracker.handleMaybeResumeKeepalive(ki);
                     }
                     break;
                 }
                 // Sent by KeepaliveTracker to process an app request on the state machine thread.
                 case NetworkAgent.CMD_STOP_SOCKET_KEEPALIVE: {
                     NetworkAgentInfo nai = getNetworkAgentInfoForNetwork((Network) msg.obj);
+                    if (nai == null) {
+                        Log.e(TAG, "Attempt to stop keepalive on nonexistent network");
+                        return;
+                    }
                     int slot = msg.arg1;
                     int reason = msg.arg2;
                     mKeepaliveTracker.handleStopKeepalive(nai, slot, reason);
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index 27be545..18e2dd8 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -44,8 +44,8 @@
 import android.net.MarkMaskParcel;
 import android.net.Network;
 import android.net.NetworkAgent;
-import android.net.SocketKeepalive;
 import android.net.SocketKeepalive.InvalidSocketException;
+import android.os.Bundle;
 import android.os.FileUtils;
 import android.os.Handler;
 import android.os.IBinder;
@@ -61,6 +61,7 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.modules.utils.build.SdkLevel;
+import com.android.net.module.util.BinderUtils;
 import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.DeviceConfigUtils;
 import com.android.net.module.util.HexDump;
@@ -84,8 +85,7 @@
  * Manages automatic on/off socket keepalive requests.
  *
  * Provides methods to stop and start automatic keepalive requests, and keeps track of keepalives
- * across all networks. For non-automatic on/off keepalive request, this class just forwards the
- * requests to KeepaliveTracker. This class is tightly coupled to ConnectivityService. It is not
+ * across all networks. This class is tightly coupled to ConnectivityService. It is not
  * thread-safe and its handle* methods must be called only from the ConnectivityService handler
  * thread.
  */
@@ -94,15 +94,17 @@
     private static final int[] ADDRESS_FAMILIES = new int[] {AF_INET6, AF_INET};
     private static final String ACTION_TCP_POLLING_ALARM =
             "com.android.server.connectivity.KeepaliveTracker.TCP_POLLING_ALARM";
-    private static final String EXTRA_NETWORK = "network_id";
-    private static final String EXTRA_SLOT = "slot";
+    private static final String EXTRA_BINDER_TOKEN = "token";
     private static final long DEFAULT_TCP_POLLING_INTERVAL_MS = 120_000L;
     private static final String AUTOMATIC_ON_OFF_KEEPALIVE_VERSION =
             "automatic_on_off_keepalive_version";
     /**
      * States for {@code #AutomaticOnOffKeepalive}.
      *
-     * A new AutomaticOnOffKeepalive starts with STATE_ENABLED. The system will monitor
+     * If automatic mode is off for this keepalive, the state is STATE_ALWAYS_ON and it stays
+     * so for the entire lifetime of this object.
+     *
+     * If enabled, a new AutomaticOnOffKeepalive starts with STATE_ENABLED. The system will monitor
      * the TCP sockets on VPN networks running on top of the specified network, and turn off
      * keepalive if there is no TCP socket any of the VPN networks. Conversely, it will turn
      * keepalive back on if any TCP socket is open on any of the VPN networks.
@@ -118,10 +120,12 @@
      */
     private static final int STATE_ENABLED = 0;
     private static final int STATE_SUSPENDED = 1;
+    private static final int STATE_ALWAYS_ON = 2;
     @Retention(RetentionPolicy.SOURCE)
     @IntDef(prefix = { "STATE_" }, value = {
             STATE_ENABLED,
-            STATE_SUSPENDED
+            STATE_SUSPENDED,
+            STATE_ALWAYS_ON
     })
     private @interface AutomaticOnOffState {}
 
@@ -156,58 +160,82 @@
         public void onReceive(Context context, Intent intent) {
             if (ACTION_TCP_POLLING_ALARM.equals(intent.getAction())) {
                 Log.d(TAG, "Received TCP polling intent");
-                final Network network = intent.getParcelableExtra(EXTRA_NETWORK);
-                final int slot = intent.getIntExtra(EXTRA_SLOT, -1);
+                final IBinder token = intent.getBundleExtra(EXTRA_BINDER_TOKEN).getBinder(
+                        EXTRA_BINDER_TOKEN);
                 mConnectivityServiceHandler.obtainMessage(
-                        NetworkAgent.CMD_MONITOR_AUTOMATIC_KEEPALIVE,
-                        slot, 0 , network).sendToTarget();
+                        NetworkAgent.CMD_MONITOR_AUTOMATIC_KEEPALIVE, token).sendToTarget();
             }
         }
     };
 
-    private static class AutomaticOnOffKeepalive {
+    /**
+     * Information about a managed keepalive.
+     *
+     * The keepalive in mKi is managed by this object. This object can be in one of three states
+     * (in mAutomatiOnOffState) :
+     * • STATE_ALWAYS_ON : this keepalive is always on
+     * • STATE_ENABLED : this keepalive is currently on, and monitored for possibly being turned
+     *                   off if no TCP socket is open on the VPN.
+     * • STATE_SUSPENDED : this keepalive is currently off, and monitored for possibly being
+     *                     resumed if a TCP socket is open on the VPN.
+     * See the documentation for the states for more detail.
+     */
+    public class AutomaticOnOffKeepalive {
         @NonNull
         private final KeepaliveTracker.KeepaliveInfo mKi;
         @NonNull
+        private final ISocketKeepaliveCallback mCallback;
+        @Nullable
         private final FileDescriptor mFd;
-        @NonNull
+        @Nullable
         private final PendingIntent mTcpPollingAlarm;
-        private final int mSlot;
         @AutomaticOnOffState
-        private int mAutomaticOnOffState = STATE_ENABLED;
+        private int mAutomaticOnOffState;
 
-        AutomaticOnOffKeepalive(@NonNull KeepaliveTracker.KeepaliveInfo ki,
-                @NonNull Context context) throws InvalidSocketException {
+        AutomaticOnOffKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki,
+                final boolean autoOnOff, @NonNull Context context) throws InvalidSocketException {
             this.mKi = Objects.requireNonNull(ki);
-            // A null fd is acceptable in KeepaliveInfo for backward compatibility of
-            // PacketKeepalive API, but it should not happen here because legacy API cannot setup
-            // automatic keepalive.
-            Objects.requireNonNull(ki.mFd);
-
-            // Get the slot from keepalive because the slot information may be missing when the
-            // keepalive is stopped.
-            this.mSlot = ki.getSlot();
-            try {
-                this.mFd = Os.dup(ki.mFd);
-            } catch (ErrnoException e) {
-                Log.e(TAG, "Cannot dup fd: ", e);
-                throw new InvalidSocketException(ERROR_INVALID_SOCKET, e);
+            mCallback = ki.mCallback;
+            if (autoOnOff && mDependencies.isFeatureEnabled(AUTOMATIC_ON_OFF_KEEPALIVE_VERSION)) {
+                mAutomaticOnOffState = STATE_ENABLED;
+                if (null == ki.mFd) {
+                    throw new IllegalArgumentException("fd can't be null with automatic "
+                            + "on/off keepalives");
+                }
+                try {
+                    mFd = Os.dup(ki.mFd);
+                } catch (ErrnoException e) {
+                    Log.e(TAG, "Cannot dup fd: ", e);
+                    throw new InvalidSocketException(ERROR_INVALID_SOCKET, e);
+                }
+                mTcpPollingAlarm = createTcpPollingAlarmIntent(context, mCallback.asBinder());
+            } else {
+                mAutomaticOnOffState = STATE_ALWAYS_ON;
+                // A null fd is acceptable in KeepaliveInfo for backward compatibility of
+                // PacketKeepalive API, but it must never happen with automatic keepalives.
+                // TODO : remove mFd from KeepaliveInfo or from this class.
+                mFd = ki.mFd;
+                mTcpPollingAlarm = null;
             }
-            mTcpPollingAlarm = createTcpPollingAlarmIntent(
-                    context, ki.getNai().network(), ki.getSlot());
+        }
+
+        public Network getNetwork() {
+            return mKi.getNai().network;
         }
 
         public boolean match(Network network, int slot) {
-            return this.mKi.getNai().network().equals(network) && this.mSlot == slot;
+            return mKi.getNai().network().equals(network) && mKi.getSlot() == slot;
         }
 
-        private static PendingIntent createTcpPollingAlarmIntent(@NonNull Context context,
-                @NonNull Network network, int slot) {
+        private PendingIntent createTcpPollingAlarmIntent(@NonNull Context context,
+                @NonNull IBinder token) {
             final Intent intent = new Intent(ACTION_TCP_POLLING_ALARM);
-            intent.putExtra(EXTRA_NETWORK, network);
-            intent.putExtra(EXTRA_SLOT, slot);
-            return PendingIntent.getBroadcast(
-                    context, 0 /* requestCode */, intent, PendingIntent.FLAG_IMMUTABLE);
+            // Intent doesn't expose methods to put extra Binders, but Bundle does.
+            final Bundle b = new Bundle();
+            b.putBinder(EXTRA_BINDER_TOKEN, token);
+            intent.putExtra(EXTRA_BINDER_TOKEN, b);
+            return BinderUtils.withCleanCallingIdentity(() -> PendingIntent.getBroadcast(
+                    context, 0 /* requestCode */, intent, PendingIntent.FLAG_IMMUTABLE));
         }
     }
 
@@ -242,25 +270,32 @@
     /**
      * Determine if any state transition is needed for the specific automatic keepalive.
      */
-    public void handleMonitorAutomaticKeepalive(@NonNull Network network, int slot, int vpnNetId) {
-        final AutomaticOnOffKeepalive autoKi = findAutomaticOnOffKeepalive(network, slot);
-        // This may happen if the keepalive is removed by the app, and the alarm is fired at the
-        // same time.
-        if (autoKi == null) return;
+    public void handleMonitorAutomaticKeepalive(@NonNull final AutomaticOnOffKeepalive ki,
+            final int vpnNetId) {
+        // Might happen if the automatic keepalive was removed by the app just as the alarm fires.
+        if (!mAutomaticOnOffKeepalives.contains(ki)) return;
+        if (STATE_ALWAYS_ON == ki.mAutomaticOnOffState) {
+            throw new IllegalStateException("Should not monitor non-auto keepalive");
+        }
 
-        handleMonitorTcpConnections(autoKi, vpnNetId);
+        handleMonitorTcpConnections(ki, vpnNetId);
     }
 
     /**
      * Determine if disable or re-enable keepalive is needed or not based on TCP sockets status.
      */
     private void handleMonitorTcpConnections(@NonNull AutomaticOnOffKeepalive ki, int vpnNetId) {
+        // Might happen if the automatic keepalive was removed by the app just as the alarm fires.
+        if (!mAutomaticOnOffKeepalives.contains(ki)) return;
+        if (STATE_ALWAYS_ON == ki.mAutomaticOnOffState) {
+            throw new IllegalStateException("Should not monitor non-auto keepalive");
+        }
         if (!isAnyTcpSocketConnected(vpnNetId)) {
             // No TCP socket exists. Stop keepalive if ENABLED, and remain SUSPENDED if currently
             // SUSPENDED.
             if (ki.mAutomaticOnOffState == STATE_ENABLED) {
                 ki.mAutomaticOnOffState = STATE_SUSPENDED;
-                handleSuspendKeepalive(ki.mKi.mNai, ki.mSlot, SUCCESS);
+                handleSuspendKeepalive(ki.mKi);
             }
         } else {
             handleMaybeResumeKeepalive(ki);
@@ -270,17 +305,15 @@
     }
 
     /**
-     * Resume keepalive for this slot on this network, if it wasn't already resumed.
+     * Resume an auto on/off keepalive, unless it's already resumed
+     * @param autoKi the keepalive to resume
      */
-    public void handleMaybeResumeKeepalive(@NonNull final Network network, final int slot) {
-        final AutomaticOnOffKeepalive autoKi = findAutomaticOnOffKeepalive(network, slot);
-        // This may happen if the keepalive is removed by the app, and the alarm is fired at
-        // the same time.
-        if (autoKi == null) return;
-        handleMaybeResumeKeepalive(autoKi);
-    }
-
-    private void handleMaybeResumeKeepalive(@NonNull AutomaticOnOffKeepalive autoKi) {
+    public void handleMaybeResumeKeepalive(@NonNull AutomaticOnOffKeepalive autoKi) {
+        // Might happen if the automatic keepalive was removed by the app just as the alarm fires.
+        if (!mAutomaticOnOffKeepalives.contains(autoKi)) return;
+        if (STATE_ALWAYS_ON == autoKi.mAutomaticOnOffState) {
+            throw new IllegalStateException("Should not resume non-auto keepalive");
+        }
         if (autoKi.mAutomaticOnOffState == STATE_ENABLED) return;
         KeepaliveTracker.KeepaliveInfo newKi;
         try {
@@ -289,15 +322,14 @@
             newKi = autoKi.mKi.withFd(autoKi.mFd);
         } catch (InvalidSocketException | IllegalArgumentException | SecurityException e) {
             Log.e(TAG, "Fail to construct keepalive", e);
-            mKeepaliveTracker.notifyErrorCallback(autoKi.mKi.mCallback, ERROR_INVALID_SOCKET);
+            mKeepaliveTracker.notifyErrorCallback(autoKi.mCallback, ERROR_INVALID_SOCKET);
             return;
         }
         autoKi.mAutomaticOnOffState = STATE_ENABLED;
-        handleResumeKeepalive(mConnectivityServiceHandler.obtainMessage(
-                NetworkAgent.CMD_START_SOCKET_KEEPALIVE,
-                autoKi.mAutomaticOnOffState, 0, newKi));
+        handleResumeKeepalive(newKi);
     }
 
+    // TODO : this method should be removed ; the keepalives should always be indexed by callback
     private int findAutomaticOnOffKeepaliveIndex(@NonNull Network network, int slot) {
         ensureRunningOnHandlerThread();
 
@@ -311,6 +343,7 @@
         return -1;
     }
 
+    // TODO : this method should be removed ; the keepalives should always be indexed by callback
     @Nullable
     private AutomaticOnOffKeepalive findAutomaticOnOffKeepalive(@NonNull Network network,
             int slot) {
@@ -321,6 +354,18 @@
     }
 
     /**
+     * Find the AutomaticOnOffKeepalive associated with a given callback.
+     * @return the keepalive associated with this callback, or null if none
+     */
+    @Nullable
+    public AutomaticOnOffKeepalive getKeepaliveForBinder(@NonNull final IBinder token) {
+        ensureRunningOnHandlerThread();
+
+        return CollectionUtils.findFirst(mAutomaticOnOffKeepalives,
+                it -> it.mCallback.asBinder().equals(token));
+    }
+
+    /**
      * Handle keepalive events from lower layer.
      *
      * Forward to KeepaliveTracker.
@@ -347,61 +392,48 @@
      * The message is expected to contain a KeepaliveTracker.KeepaliveInfo.
      */
     public void handleStartKeepalive(Message message) {
-        mKeepaliveTracker.handleStartKeepalive(message);
+        final AutomaticOnOffKeepalive autoKi = (AutomaticOnOffKeepalive) message.obj;
+        mKeepaliveTracker.handleStartKeepalive(autoKi.mKi);
 
         // Add automatic on/off request into list to track its life cycle.
-        final boolean automaticOnOff = message.arg1 != 0
-                && mDependencies.isFeatureEnabled(AUTOMATIC_ON_OFF_KEEPALIVE_VERSION);
-        if (automaticOnOff) {
-            final KeepaliveTracker.KeepaliveInfo ki = (KeepaliveTracker.KeepaliveInfo) message.obj;
-            AutomaticOnOffKeepalive autoKi;
-            try {
-                // CAREFUL : mKeepaliveTracker.handleStartKeepalive will assign |ki.mSlot| after
-                // pulling |ki| from the message. The constructor below will read this member
-                // (through ki.getSlot()) and therefore actively relies on handleStartKeepalive
-                // having assigned this member before this is called.
-                // TODO : clean this up by assigning the slot at the start of this method instead
-                // and ideally removing the mSlot member from KeepaliveInfo.
-                autoKi = new AutomaticOnOffKeepalive(ki, mContext);
-            } catch (SocketKeepalive.InvalidSocketException | IllegalArgumentException e) {
-                Log.e(TAG, "Fail to construct keepalive", e);
-                mKeepaliveTracker.notifyErrorCallback(ki.mCallback, ERROR_INVALID_SOCKET);
-                return;
-            }
-            mAutomaticOnOffKeepalives.add(autoKi);
+        mAutomaticOnOffKeepalives.add(autoKi);
+        if (STATE_ALWAYS_ON != autoKi.mAutomaticOnOffState) {
             startTcpPollingAlarm(autoKi.mTcpPollingAlarm);
         }
     }
 
-    private void handleResumeKeepalive(Message message) {
-        mKeepaliveTracker.handleStartKeepalive(message);
+    private void handleResumeKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
+        mKeepaliveTracker.handleStartKeepalive(ki);
     }
 
-    private void handleSuspendKeepalive(NetworkAgentInfo nai, int slot, int reason) {
-        mKeepaliveTracker.handleStopKeepalive(nai, slot, reason);
+    private void handleSuspendKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
+        // TODO : mKT.handleStopKeepalive should take a KeepaliveInfo instead
+        // TODO : create a separate success code for suspend
+        mKeepaliveTracker.handleStopKeepalive(ki.getNai(), ki.getSlot(), SUCCESS);
     }
 
     /**
      * Handle stop keepalives on the specific network with given slot.
      */
-    public void handleStopKeepalive(NetworkAgentInfo nai, int slot, int reason) {
+    public void handleStopKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
         final AutomaticOnOffKeepalive autoKi = findAutomaticOnOffKeepalive(nai.network, slot);
-
-        // Let the original keepalive do the stop first, and then clean up the keepalive if it's an
-        // automatic keepalive.
-        if (autoKi == null || autoKi.mAutomaticOnOffState == STATE_ENABLED) {
-            mKeepaliveTracker.handleStopKeepalive(nai, slot, reason);
+        if (null == autoKi) {
+            Log.e(TAG, "Attempt to stop nonexistent keepalive " + slot + " on " + nai);
+            return;
         }
 
-        // Not an AutomaticOnOffKeepalive.
-        if (autoKi == null) return;
+        // Stop the keepalive unless it was suspended. This includes the case where it's managed
+        // but enabled, and the case where it's always on.
+        if (autoKi.mAutomaticOnOffState != STATE_SUSPENDED) {
+            mKeepaliveTracker.handleStopKeepalive(nai, slot, reason);
+        }
 
         cleanupAutoOnOffKeepalive(autoKi);
     }
 
     private void cleanupAutoOnOffKeepalive(@NonNull final AutomaticOnOffKeepalive autoKi) {
         ensureRunningOnHandlerThread();
-        mAlarmManager.cancel(autoKi.mTcpPollingAlarm);
+        if (null != autoKi.mTcpPollingAlarm) mAlarmManager.cancel(autoKi.mTcpPollingAlarm);
         // Close the duplicated fd that maintains the lifecycle of socket.
         FileUtils.closeQuietly(autoKi.mFd);
         mAutomaticOnOffKeepalives.remove(autoKi);
@@ -423,10 +455,15 @@
             int dstPort, boolean automaticOnOffKeepalives) {
         final KeepaliveTracker.KeepaliveInfo ki = mKeepaliveTracker.makeNattKeepaliveInfo(nai, fd,
                 intervalSeconds, cb, srcAddrString, srcPort, dstAddrString, dstPort);
-        if (null != ki) {
+        if (null == ki) return;
+        try {
+            final AutomaticOnOffKeepalive autoKi = new AutomaticOnOffKeepalive(ki,
+                    automaticOnOffKeepalives, mContext);
             mConnectivityServiceHandler.obtainMessage(NetworkAgent.CMD_START_SOCKET_KEEPALIVE,
                     // TODO : move ConnectivityService#encodeBool to a static lib.
-                    automaticOnOffKeepalives ? 1 : 0, 0, ki).sendToTarget();
+                    automaticOnOffKeepalives ? 1 : 0, 0, autoKi).sendToTarget();
+        } catch (InvalidSocketException e) {
+            mKeepaliveTracker.notifyErrorCallback(cb, e.error);
         }
     }
 
@@ -447,10 +484,15 @@
             boolean automaticOnOffKeepalives) {
         final KeepaliveTracker.KeepaliveInfo ki = mKeepaliveTracker.makeNattKeepaliveInfo(nai, fd,
                 resourceId, intervalSeconds, cb, srcAddrString, dstAddrString, dstPort);
-        if (null != ki) {
+        if (null == ki) return;
+        try {
+            final AutomaticOnOffKeepalive autoKi = new AutomaticOnOffKeepalive(ki,
+                    automaticOnOffKeepalives, mContext);
             mConnectivityServiceHandler.obtainMessage(NetworkAgent.CMD_START_SOCKET_KEEPALIVE,
                     // TODO : move ConnectivityService#encodeBool to a static lib.
-                    automaticOnOffKeepalives ? 1 : 0, 0, ki).sendToTarget();
+                    automaticOnOffKeepalives ? 1 : 0, 0, autoKi).sendToTarget();
+        } catch (InvalidSocketException e) {
+            mKeepaliveTracker.notifyErrorCallback(cb, e.error);
         }
     }
 
@@ -471,9 +513,14 @@
             @NonNull ISocketKeepaliveCallback cb) {
         final KeepaliveTracker.KeepaliveInfo ki = mKeepaliveTracker.makeTcpKeepaliveInfo(nai, fd,
                 intervalSeconds, cb);
-        if (null != ki) {
-            mConnectivityServiceHandler.obtainMessage(CMD_START_SOCKET_KEEPALIVE, ki)
+        if (null == ki) return;
+        try {
+            final AutomaticOnOffKeepalive autoKi = new AutomaticOnOffKeepalive(ki,
+                    false /* autoOnOff, tcp keepalives are never auto on/off */, mContext);
+            mConnectivityServiceHandler.obtainMessage(CMD_START_SOCKET_KEEPALIVE, autoKi)
                     .sendToTarget();
+        } catch (InvalidSocketException e) {
+            mKeepaliveTracker.notifyErrorCallback(cb, e.error);
         }
     }
 
diff --git a/service/src/com/android/server/connectivity/KeepaliveTracker.java b/service/src/com/android/server/connectivity/KeepaliveTracker.java
index 03f8f3e..7cb613b 100644
--- a/service/src/com/android/server/connectivity/KeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveTracker.java
@@ -49,7 +49,6 @@
 import android.os.Binder;
 import android.os.Handler;
 import android.os.IBinder;
-import android.os.Message;
 import android.os.Process;
 import android.os.RemoteException;
 import android.system.ErrnoException;
@@ -126,8 +125,9 @@
      * which is only returned when the hardware has successfully started the keepalive.
      */
     class KeepaliveInfo implements IBinder.DeathRecipient {
-        // Bookkeeping data.
+        // TODO : remove this member. Only AutoOnOffKeepalive should have a reference to this.
         public final ISocketKeepaliveCallback mCallback;
+        // Bookkeeping data.
         private final int mUid;
         private final int mPid;
         private final boolean mPrivileged;
@@ -456,8 +456,7 @@
     /**
      * Handle start keepalives with the message.
      */
-    public void handleStartKeepalive(Message message) {
-        KeepaliveInfo ki = (KeepaliveInfo) message.obj;
+    public void handleStartKeepalive(KeepaliveInfo ki) {
         NetworkAgentInfo nai = ki.getNai();
         int slot = findFirstFreeSlot(nai);
         mKeepalives.get(nai).put(slot, ki);
diff --git a/tests/common/AndroidTest_Coverage.xml b/tests/common/AndroidTest_Coverage.xml
index 48d26b8..c94ec27 100644
--- a/tests/common/AndroidTest_Coverage.xml
+++ b/tests/common/AndroidTest_Coverage.xml
@@ -13,7 +13,7 @@
      limitations under the License.
 -->
 <configuration description="Runs coverage tests for Connectivity">
-    <target_preparer class="com.android.tradefed.targetprep.TestAppInstallSetup">
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
       <option name="test-file-name" value="ConnectivityCoverageTests.apk" />
       <option name="install-arg" value="-t" />
     </target_preparer>
diff --git a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
index b7eb009..66e7713 100644
--- a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
@@ -44,6 +44,11 @@
 import android.net.cts.NsdManagerTest.NsdResolveRecord.ResolveEvent.ResolveStopped
 import android.net.cts.NsdManagerTest.NsdResolveRecord.ResolveEvent.ServiceResolved
 import android.net.cts.NsdManagerTest.NsdResolveRecord.ResolveEvent.StopResolutionFailed
+import android.net.cts.NsdManagerTest.NsdServiceInfoCallbackRecord.ServiceInfoCallbackEvent.RegisterCallbackFailed
+import android.net.cts.NsdManagerTest.NsdServiceInfoCallbackRecord.ServiceInfoCallbackEvent.ServiceUpdated
+import android.net.cts.NsdManagerTest.NsdServiceInfoCallbackRecord.ServiceInfoCallbackEvent.ServiceUpdatedLost
+import android.net.cts.NsdManagerTest.NsdServiceInfoCallbackRecord.ServiceInfoCallbackEvent.UnregisterCallbackSucceeded
+import android.net.cts.util.CtsNetUtils
 import android.net.nsd.NsdManager
 import android.net.nsd.NsdManager.DiscoveryListener
 import android.net.nsd.NsdManager.RegistrationListener
@@ -60,6 +65,7 @@
 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.TestableNetworkAgent
@@ -115,6 +121,7 @@
     private val serviceName = "NsdTest%09d".format(Random().nextInt(1_000_000_000))
     private val serviceType = "_nmt%09d._tcp".format(Random().nextInt(1_000_000_000))
     private val handlerThread = HandlerThread(NsdManagerTest::class.java.simpleName)
+    private val ctsNetUtils by lazy{ CtsNetUtils(context) }
 
     private lateinit var testNetwork1: TestTapNetwork
     private lateinit var testNetwork2: TestTapNetwork
@@ -157,7 +164,8 @@
 
         inline fun <reified V : NsdEvent> expectCallback(timeoutMs: Long = TIMEOUT_MS): V {
             val nextEvent = nextEvents.poll(timeoutMs)
-            assertNotNull(nextEvent, "No callback received after $timeoutMs ms")
+            assertNotNull(nextEvent, "No callback received after $timeoutMs ms, expected " +
+                    "${V::class.java.simpleName}")
             assertTrue(nextEvent is V, "Expected ${V::class.java.simpleName} but got " +
                     nextEvent.javaClass.simpleName)
             return nextEvent
@@ -287,6 +295,32 @@
         }
     }
 
+    private class NsdServiceInfoCallbackRecord : NsdShim.ServiceInfoCallbackShim,
+            NsdRecord<NsdServiceInfoCallbackRecord.ServiceInfoCallbackEvent>() {
+        sealed class ServiceInfoCallbackEvent : NsdEvent {
+            data class RegisterCallbackFailed(val errorCode: Int) : ServiceInfoCallbackEvent()
+            data class ServiceUpdated(val serviceInfo: NsdServiceInfo) : ServiceInfoCallbackEvent()
+            object ServiceUpdatedLost : ServiceInfoCallbackEvent()
+            object UnregisterCallbackSucceeded : ServiceInfoCallbackEvent()
+        }
+
+        override fun onServiceInfoCallbackRegistrationFailed(err: Int) {
+            add(RegisterCallbackFailed(err))
+        }
+
+        override fun onServiceUpdated(si: NsdServiceInfo) {
+            add(ServiceUpdated(si))
+        }
+
+        override fun onServiceLost() {
+            add(ServiceUpdatedLost)
+        }
+
+        override fun onServiceInfoCallbackUnregistered() {
+            add(UnregisterCallbackSucceeded)
+        }
+    }
+
     @Before
     fun setUp() {
         handlerThread.start()
@@ -772,6 +806,64 @@
         assertEquals(si.serviceType, stoppedCb.serviceInfo.serviceType)
     }
 
+    @Test
+    fun testRegisterServiceInfoCallback() {
+        // This test requires shims supporting U+ APIs (NsdManager.subscribeService)
+        assumeTrue(TestUtils.shouldTestUApis())
+
+        // Ensure Wi-Fi network connected and get addresses
+        val wifiNetwork = ctsNetUtils.ensureWifiConnected()
+        val lp = cm.getLinkProperties(wifiNetwork)
+        assertNotNull(lp)
+        val addresses = lp.addresses
+        assertFalse(addresses.isEmpty())
+
+        val si = NsdServiceInfo().apply {
+            serviceType = this@NsdManagerTest.serviceType
+            serviceName = this@NsdManagerTest.serviceName
+            network = wifiNetwork
+            port = 12345 // Test won't try to connect so port does not matter
+        }
+
+        // Register service on Wi-Fi network
+        val registrationRecord = NsdRegistrationRecord()
+        registerService(registrationRecord, si)
+
+        val discoveryRecord = NsdDiscoveryRecord()
+        val cbRecord = NsdServiceInfoCallbackRecord()
+        tryTest {
+            // Discover service on Wi-Fi network.
+            nsdShim.discoverServices(nsdManager, serviceType, NsdManager.PROTOCOL_DNS_SD,
+                    wifiNetwork, Executor { it.run() }, discoveryRecord)
+            val foundInfo = discoveryRecord.waitForServiceDiscovered(
+                    serviceName, wifiNetwork)
+
+            // Subscribe to service and check the addresses are the same as Wi-Fi addresses
+            nsdShim.registerServiceInfoCallback(nsdManager, foundInfo, { it.run() }, cbRecord)
+            for (i in addresses.indices) {
+                val subscribeCb = cbRecord.expectCallback<ServiceUpdated>()
+                assertEquals(foundInfo.serviceName, subscribeCb.serviceInfo.serviceName)
+                val hostAddresses = subscribeCb.serviceInfo.hostAddresses
+                assertEquals(i + 1, hostAddresses.size)
+                for (hostAddress in hostAddresses) {
+                    assertTrue(addresses.contains(hostAddress))
+                }
+            }
+        } cleanupStep {
+            nsdManager.unregisterService(registrationRecord)
+            registrationRecord.expectCallback<ServiceUnregistered>()
+            discoveryRecord.expectCallback<ServiceLost>()
+            cbRecord.expectCallback<ServiceUpdatedLost>()
+        } cleanupStep {
+            // Cancel subscription and check stop callback received.
+            nsdShim.unregisterServiceInfoCallback(nsdManager, cbRecord)
+            cbRecord.expectCallback<UnregisterCallbackSucceeded>()
+        } cleanup {
+            nsdManager.stopServiceDiscovery(discoveryRecord)
+            discoveryRecord.expectCallback<DiscoveryStopped>()
+        }
+    }
+
     /**
      * Register a service and return its registration record.
      */
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index 98a8ed2..a2c4b9b 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -45,6 +45,7 @@
 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 static org.mockito.Mockito.when;
 
 import android.compat.testing.PlatformCompatChangeRule;
@@ -170,6 +171,9 @@
         doReturn(true).when(mMockMDnsM).resolve(
                 anyInt(), anyString(), anyString(), anyString(), anyInt());
         doReturn(false).when(mDeps).isMdnsDiscoveryManagerEnabled(any(Context.class));
+        doReturn(mDiscoveryManager).when(mDeps).makeMdnsDiscoveryManager(any(), any());
+        doReturn(mSocketProvider).when(mDeps).makeMdnsSocketProvider(any(), any());
+        doReturn(mAdvertiser).when(mDeps).makeMdnsAdvertiser(any(), any(), any());
 
         mService = makeService();
     }
@@ -824,40 +828,50 @@
                 client.unregisterServiceInfoCallback(serviceInfoCallback));
     }
 
-    private void makeServiceWithMdnsDiscoveryManagerEnabled() {
+    private void setMdnsDiscoveryManagerEnabled() {
         doReturn(true).when(mDeps).isMdnsDiscoveryManagerEnabled(any(Context.class));
-        doReturn(mDiscoveryManager).when(mDeps).makeMdnsDiscoveryManager(any(), any());
-        doReturn(mSocketProvider).when(mDeps).makeMdnsSocketProvider(any(), any());
-
-        mService = makeService();
-        verify(mDeps).makeMdnsDiscoveryManager(any(), any());
-        verify(mDeps).makeMdnsSocketProvider(any(), any());
     }
 
-    private void makeServiceWithMdnsAdvertiserEnabled() {
+    private void setMdnsAdvertiserEnabled() {
         doReturn(true).when(mDeps).isMdnsAdvertiserEnabled(any(Context.class));
-        doReturn(mAdvertiser).when(mDeps).makeMdnsAdvertiser(any(), any(), any());
-        doReturn(mSocketProvider).when(mDeps).makeMdnsSocketProvider(any(), any());
-
-        mService = makeService();
-        verify(mDeps).makeMdnsAdvertiser(any(), any(), any());
-        verify(mDeps).makeMdnsSocketProvider(any(), any());
     }
 
     @Test
     public void testMdnsDiscoveryManagerFeature() {
         // Create NsdService w/o feature enabled.
-        connectClient(mService);
-        verify(mDeps, never()).makeMdnsDiscoveryManager(any(), any());
-        verify(mDeps, never()).makeMdnsSocketProvider(any(), any());
+        final NsdManager client = connectClient(mService);
+        final DiscoveryListener discListenerWithoutFeature = mock(DiscoveryListener.class);
+        client.discoverServices(SERVICE_TYPE, PROTOCOL, discListenerWithoutFeature);
+        waitForIdle();
 
-        // Create NsdService again w/ feature enabled.
-        makeServiceWithMdnsDiscoveryManagerEnabled();
+        final ArgumentCaptor<Integer> legacyIdCaptor = ArgumentCaptor.forClass(Integer.class);
+        verify(mMockMDnsM).discover(legacyIdCaptor.capture(), any(), anyInt());
+        verifyNoMoreInteractions(mDiscoveryManager);
+
+        setMdnsDiscoveryManagerEnabled();
+        final DiscoveryListener discListenerWithFeature = mock(DiscoveryListener.class);
+        client.discoverServices(SERVICE_TYPE, PROTOCOL, discListenerWithFeature);
+        waitForIdle();
+
+        final String serviceTypeWithLocalDomain = SERVICE_TYPE + ".local";
+        final ArgumentCaptor<MdnsServiceBrowserListener> listenerCaptor =
+                ArgumentCaptor.forClass(MdnsServiceBrowserListener.class);
+        verify(mDiscoveryManager).registerListener(eq(serviceTypeWithLocalDomain),
+                listenerCaptor.capture(), any());
+
+        client.stopServiceDiscovery(discListenerWithoutFeature);
+        waitForIdle();
+        verify(mMockMDnsM).stopOperation(legacyIdCaptor.getValue());
+
+        client.stopServiceDiscovery(discListenerWithFeature);
+        waitForIdle();
+        verify(mDiscoveryManager).unregisterListener(serviceTypeWithLocalDomain,
+                listenerCaptor.getValue());
     }
 
     @Test
     public void testDiscoveryWithMdnsDiscoveryManager() {
-        makeServiceWithMdnsDiscoveryManagerEnabled();
+        setMdnsDiscoveryManagerEnabled();
 
         final NsdManager client = connectClient(mService);
         final DiscoveryListener discListener = mock(DiscoveryListener.class);
@@ -922,7 +936,7 @@
 
     @Test
     public void testDiscoveryWithMdnsDiscoveryManager_FailedWithInvalidServiceType() {
-        makeServiceWithMdnsDiscoveryManagerEnabled();
+        setMdnsDiscoveryManagerEnabled();
 
         final NsdManager client = connectClient(mService);
         final DiscoveryListener discListener = mock(DiscoveryListener.class);
@@ -951,7 +965,7 @@
 
     @Test
     public void testResolutionWithMdnsDiscoveryManager() throws UnknownHostException {
-        makeServiceWithMdnsDiscoveryManagerEnabled();
+        setMdnsDiscoveryManagerEnabled();
 
         final NsdManager client = connectClient(mService);
         final ResolveListener resolveListener = mock(ResolveListener.class);
@@ -1005,8 +1019,43 @@
     }
 
     @Test
+    public void testMdnsAdvertiserFeatureFlagging() {
+        // Create NsdService w/o feature enabled.
+        final NsdManager client = connectClient(mService);
+        final NsdServiceInfo regInfo = new NsdServiceInfo(SERVICE_NAME, SERVICE_TYPE);
+        regInfo.setHost(parseNumericAddress("192.0.2.123"));
+        regInfo.setPort(12345);
+        final RegistrationListener regListenerWithoutFeature = mock(RegistrationListener.class);
+        client.registerService(regInfo, PROTOCOL, regListenerWithoutFeature);
+        waitForIdle();
+
+        final ArgumentCaptor<Integer> legacyIdCaptor = ArgumentCaptor.forClass(Integer.class);
+        verify(mMockMDnsM).registerService(legacyIdCaptor.capture(), any(), any(), anyInt(),
+                any(), anyInt());
+        verifyNoMoreInteractions(mAdvertiser);
+
+        setMdnsAdvertiserEnabled();
+        final RegistrationListener regListenerWithFeature = mock(RegistrationListener.class);
+        client.registerService(regInfo, PROTOCOL, regListenerWithFeature);
+        waitForIdle();
+
+        final ArgumentCaptor<Integer> serviceIdCaptor = ArgumentCaptor.forClass(Integer.class);
+        verify(mAdvertiser).addService(serviceIdCaptor.capture(),
+                argThat(info -> matches(info, regInfo)));
+
+        client.unregisterService(regListenerWithoutFeature);
+        waitForIdle();
+        verify(mMockMDnsM).stopOperation(legacyIdCaptor.getValue());
+        verify(mAdvertiser, never()).removeService(anyInt());
+
+        client.unregisterService(regListenerWithFeature);
+        waitForIdle();
+        verify(mAdvertiser).removeService(serviceIdCaptor.getValue());
+    }
+
+    @Test
     public void testAdvertiseWithMdnsAdvertiser() {
-        makeServiceWithMdnsAdvertiserEnabled();
+        setMdnsAdvertiserEnabled();
 
         final NsdManager client = connectClient(mService);
         final RegistrationListener regListener = mock(RegistrationListener.class);
@@ -1045,7 +1094,7 @@
 
     @Test
     public void testAdvertiseWithMdnsAdvertiser_FailedWithInvalidServiceType() {
-        makeServiceWithMdnsAdvertiserEnabled();
+        setMdnsAdvertiserEnabled();
 
         final NsdManager client = connectClient(mService);
         final RegistrationListener regListener = mock(RegistrationListener.class);
@@ -1070,7 +1119,7 @@
 
     @Test
     public void testAdvertiseWithMdnsAdvertiser_LongServiceName() {
-        makeServiceWithMdnsAdvertiserEnabled();
+        setMdnsAdvertiserEnabled();
 
         final NsdManager client = connectClient(mService);
         final RegistrationListener regListener = mock(RegistrationListener.class);