Merge "[Thread] create standalone border router test module" into main
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 3d7ea69..4f18fa2 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -1252,9 +1252,13 @@
          * {@link ConnectivityManager#registerNetworkAgent}
          * @hide
          */
-        public static Network registerNetworkAgentResult(
+        public static NetworkAndAgentRegistryParcelable registerNetworkAgentResult(
                 @Nullable final Network network, @Nullable final INetworkAgentRegistry registry) {
-            return network;
+            final NetworkAndAgentRegistryParcelable result =
+                    new NetworkAndAgentRegistryParcelable();
+            result.network = network;
+            result.registry = registry;
+            return result;
         }
     }
 
@@ -3968,7 +3972,8 @@
     @RequiresPermission(anyOf = {
             NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
             android.Manifest.permission.NETWORK_FACTORY})
-    public Network registerNetworkAgent(@NonNull INetworkAgent na, @NonNull NetworkInfo ni,
+    public NetworkAndAgentRegistryParcelable registerNetworkAgent(
+            @NonNull INetworkAgent na, @NonNull NetworkInfo ni,
             @NonNull LinkProperties lp, @NonNull NetworkCapabilities nc,
             @NonNull NetworkScore score, @NonNull NetworkAgentConfig config, int providerId) {
         return registerNetworkAgent(na, ni, lp, nc, null /* localNetworkConfig */, score, config,
@@ -3983,7 +3988,8 @@
     @RequiresPermission(anyOf = {
             NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
             android.Manifest.permission.NETWORK_FACTORY})
-    public Network registerNetworkAgent(@NonNull INetworkAgent na, @NonNull NetworkInfo ni,
+    public NetworkAndAgentRegistryParcelable registerNetworkAgent(
+            @NonNull INetworkAgent na, @NonNull NetworkInfo ni,
             @NonNull LinkProperties lp, @NonNull NetworkCapabilities nc,
             @Nullable LocalNetworkConfig localNetworkConfig, @NonNull NetworkScore score,
             @NonNull NetworkAgentConfig config, int providerId) {
diff --git a/framework/src/android/net/IConnectivityManager.aidl b/framework/src/android/net/IConnectivityManager.aidl
index 47b3316..a270684 100644
--- a/framework/src/android/net/IConnectivityManager.aidl
+++ b/framework/src/android/net/IConnectivityManager.aidl
@@ -30,6 +30,7 @@
 import android.net.LocalNetworkConfig;
 import android.net.Network;
 import android.net.NetworkAgentConfig;
+import android.net.NetworkAndAgentRegistryParcelable;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
 import android.net.NetworkRequest;
@@ -146,7 +147,8 @@
 
     void declareNetworkRequestUnfulfillable(in NetworkRequest request);
 
-    Network registerNetworkAgent(in INetworkAgent na, in NetworkInfo ni, in LinkProperties lp,
+    NetworkAndAgentRegistryParcelable registerNetworkAgent(
+            in INetworkAgent na, in NetworkInfo ni, in LinkProperties lp,
             in NetworkCapabilities nc, in NetworkScore score,
             in LocalNetworkConfig localNetworkConfig, in NetworkAgentConfig config,
             in int factorySerialNumber);
diff --git a/framework/src/android/net/INetworkAgent.aidl b/framework/src/android/net/INetworkAgent.aidl
index fa5175c..c6beeca 100644
--- a/framework/src/android/net/INetworkAgent.aidl
+++ b/framework/src/android/net/INetworkAgent.aidl
@@ -26,7 +26,7 @@
  * @hide
  */
 oneway interface INetworkAgent {
-    void onRegistered(in INetworkAgentRegistry registry);
+    void onRegistered();
     void onDisconnected();
     void onBandwidthUpdateRequested();
     void onValidationStatusChanged(int validationStatus,
diff --git a/framework/src/android/net/INetworkAgentRegistry.aidl b/framework/src/android/net/INetworkAgentRegistry.aidl
index 61b27b5..afdd1ee 100644
--- a/framework/src/android/net/INetworkAgentRegistry.aidl
+++ b/framework/src/android/net/INetworkAgentRegistry.aidl
@@ -30,7 +30,7 @@
  * Interface for NetworkAgents to send network properties.
  * @hide
  */
-oneway interface INetworkAgentRegistry {
+interface INetworkAgentRegistry {
     void sendNetworkCapabilities(in NetworkCapabilities nc);
     void sendLinkProperties(in LinkProperties lp);
     // TODO: consider replacing this by "markConnected()" and removing
diff --git a/framework/src/android/net/NetworkAgent.java b/framework/src/android/net/NetworkAgent.java
index cefa1ea..08f5ecd 100644
--- a/framework/src/android/net/NetworkAgent.java
+++ b/framework/src/android/net/NetworkAgent.java
@@ -98,6 +98,7 @@
     @Nullable
     private volatile Network mNetwork;
 
+    // Null before the agent is registered
     @Nullable
     private volatile INetworkAgentRegistry mRegistry;
 
@@ -121,6 +122,8 @@
     private NetworkInfo mNetworkInfo;
     @NonNull
     private final Object mRegisterLock = new Object();
+    // TODO : move the preconnected queue to the system server and remove this
+    private boolean mConnected = false;
 
     /**
      * The ID of the {@link NetworkProvider} that created this object, or
@@ -606,16 +609,16 @@
         public void handleMessage(Message msg) {
             switch (msg.what) {
                 case EVENT_AGENT_CONNECTED: {
-                    if (mRegistry != null) {
-                        log("Received new connection while already connected!");
-                    } else {
-                        if (VDBG) log("NetworkAgent fully connected");
-                        synchronized (mPreConnectedQueue) {
-                            final INetworkAgentRegistry registry = (INetworkAgentRegistry) msg.obj;
-                            mRegistry = registry;
+                    // TODO : move the pre-connected queue to the system server, and remove
+                    // handling this EVENT_AGENT_CONNECTED message.
+                    synchronized (mPreConnectedQueue) {
+                        if (mConnected) {
+                            log("Received new connection while already connected!");
+                        } else {
+                            if (VDBG) log("NetworkAgent fully connected");
                             for (RegistryAction a : mPreConnectedQueue) {
                                 try {
-                                    a.execute(registry);
+                                    a.execute(mRegistry);
                                 } catch (RemoteException e) {
                                     Log.wtf(LOG_TAG, "Communication error with registry", e);
                                     // Fall through
@@ -623,6 +626,7 @@
                             }
                             mPreConnectedQueue.clear();
                         }
+                        mConnected = true;
                     }
                     break;
                 }
@@ -631,7 +635,7 @@
                     // let the client know CS is done with us.
                     onNetworkUnwanted();
                     synchronized (mPreConnectedQueue) {
-                        mRegistry = null;
+                        mConnected = false;
                     }
                     break;
                 }
@@ -758,20 +762,32 @@
             }
             final ConnectivityManager cm = (ConnectivityManager) mInitialConfiguration.context
                     .getSystemService(Context.CONNECTIVITY_SERVICE);
+            final NetworkAndAgentRegistryParcelable result;
             if (mInitialConfiguration.localNetworkConfig == null) {
                 // Call registerNetworkAgent without localNetworkConfig argument to pass
                 // android.net.cts.NetworkAgentTest#testAgentStartsInConnecting in old cts
-                mNetwork = cm.registerNetworkAgent(new NetworkAgentBinder(mHandler),
+                result = cm.registerNetworkAgent(new NetworkAgentBinder(mHandler),
                         new NetworkInfo(mInitialConfiguration.info),
                         mInitialConfiguration.properties, mInitialConfiguration.capabilities,
                         mInitialConfiguration.score, mInitialConfiguration.config, providerId);
             } else {
-                mNetwork = cm.registerNetworkAgent(new NetworkAgentBinder(mHandler),
+                result = cm.registerNetworkAgent(new NetworkAgentBinder(mHandler),
                         new NetworkInfo(mInitialConfiguration.info),
                         mInitialConfiguration.properties, mInitialConfiguration.capabilities,
                         mInitialConfiguration.localNetworkConfig, mInitialConfiguration.score,
                         mInitialConfiguration.config, providerId);
             }
+            if (null == result && Process.isApplicationUid(Process.myUid())) {
+                // Let it slide in tests to allow mocking, since NetworkAndAgentRegistryParcelable
+                // is not public and can't be instantiated by CTS. The danger here is that if
+                // this happens in production for some reason the code will crash later instead
+                // of here. If this is a system app, it will still crash as expected.
+                Log.e(LOG_TAG, "registerNetworkAgent returned null. This agent will not work. "
+                        + "Is ConnectivityManager a mock ?");
+            } else {
+                mNetwork = result.network;
+                mRegistry = result.registry;
+            }
             mInitialConfiguration = null; // All this memory can now be GC'd
         }
         return mNetwork;
@@ -787,8 +803,8 @@
         }
 
         @Override
-        public void onRegistered(@NonNull INetworkAgentRegistry registry) {
-            mHandler.sendMessage(mHandler.obtainMessage(EVENT_AGENT_CONNECTED, registry));
+        public void onRegistered() {
+            mHandler.sendMessage(mHandler.obtainMessage(EVENT_AGENT_CONNECTED));
         }
 
         @Override
@@ -913,11 +929,13 @@
      *
      * @hide
      */
-    public INetworkAgent registerForTest(final Network network) {
+    public INetworkAgent registerForTest(final Network network,
+            final INetworkAgentRegistry registry) {
         log("Registering NetworkAgent for test");
         synchronized (mRegisterLock) {
             mNetwork = network;
             mInitialConfiguration = null;
+            mRegistry = registry;
         }
         return new NetworkAgentBinder(mHandler);
     }
@@ -958,7 +976,7 @@
                         FrameworkConnectivityStatsLog.CORE_NETWORKING_TERRIBLE_ERROR_OCCURRED__ERROR_TYPE__TYPE_MESSAGE_QUEUED_BEFORE_CONNECT
                 );
             }
-            if (mRegistry != null) {
+            if (mConnected) {
                 try {
                     action.execute(mRegistry);
                 } catch (RemoteException e) {
diff --git a/framework/src/android/net/NetworkAndAgentRegistryParcelable.aidl b/framework/src/android/net/NetworkAndAgentRegistryParcelable.aidl
new file mode 100644
index 0000000..8c01bbc
--- /dev/null
+++ b/framework/src/android/net/NetworkAndAgentRegistryParcelable.aidl
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net;
+
+import android.net.INetworkAgentRegistry;
+import android.net.Network;
+
+/**
+ * A pair of Network and NetworkAgentRegistry.
+ *
+ * {@hide}
+ */
+@JavaDerive(toString=true)
+parcelable NetworkAndAgentRegistryParcelable {
+  Network network;
+  INetworkAgentRegistry registry;
+}
diff --git a/service-t/Android.bp b/service-t/Android.bp
index ab38c7a..81378f5 100644
--- a/service-t/Android.bp
+++ b/service-t/Android.bp
@@ -109,13 +109,15 @@
         ":service-mdns-droidstubs",
     ],
     exclude_srcs: [
+        "src/com/android/server/connectivity/mdns/internal/MdnsRealtimeScheduler.java",
         "src/com/android/server/connectivity/mdns/internal/SocketNetlinkMonitor.java",
-        "src/com/android/server/connectivity/mdns/SocketNetLinkMonitorFactory.java",
         "src/com/android/server/connectivity/mdns/MdnsAdvertiser.java",
         "src/com/android/server/connectivity/mdns/MdnsAnnouncer.java",
         "src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java",
         "src/com/android/server/connectivity/mdns/MdnsProber.java",
         "src/com/android/server/connectivity/mdns/MdnsRecordRepository.java",
+        "src/com/android/server/connectivity/mdns/SchedulerFactory.java",
+        "src/com/android/server/connectivity/mdns/SocketNetLinkMonitorFactory.java",
     ],
     static_libs: [
         "net-utils-device-common-mdns-standalone-build-test",
@@ -132,7 +134,10 @@
 
 droidstubs {
     name: "service-mdns-droidstubs",
-    srcs: ["src/com/android/server/connectivity/mdns/SocketNetLinkMonitorFactory.java"],
+    srcs: [
+        "src/com/android/server/connectivity/mdns/SchedulerFactory.java",
+        "src/com/android/server/connectivity/mdns/SocketNetLinkMonitorFactory.java",
+    ],
     libs: [
         "net-utils-device-common-mdns-standalone-build-test",
         "service-connectivity-tiramisu-pre-jarjar",
diff --git a/service-t/src/com/android/server/connectivity/mdns/DiscoveryExecutor.java b/service-t/src/com/android/server/connectivity/mdns/DiscoveryExecutor.java
index 21af1a1..99354f8 100644
--- a/service-t/src/com/android/server/connectivity/mdns/DiscoveryExecutor.java
+++ b/service-t/src/com/android/server/connectivity/mdns/DiscoveryExecutor.java
@@ -49,7 +49,13 @@
     @NonNull
     private final ArrayList<Pair<Runnable, Long>> mPendingTasks = new ArrayList<>();
 
-    DiscoveryExecutor(@Nullable Looper defaultLooper) {
+    @GuardedBy("mPendingTasks")
+    @Nullable
+    Scheduler mScheduler;
+    @NonNull private final MdnsFeatureFlags mMdnsFeatureFlags;
+
+    DiscoveryExecutor(@Nullable Looper defaultLooper, @NonNull MdnsFeatureFlags mdnsFeatureFlags) {
+        mMdnsFeatureFlags = mdnsFeatureFlags;
         if (defaultLooper != null) {
             this.mHandlerThread = null;
             synchronized (mPendingTasks) {
@@ -62,7 +68,7 @@
                     synchronized (mPendingTasks) {
                         mHandler = new Handler(getLooper());
                         for (Pair<Runnable, Long> pendingTask : mPendingTasks) {
-                            mHandler.postDelayed(pendingTask.first, pendingTask.second);
+                            executeDelayed(pendingTask.first, pendingTask.second);
                         }
                         mPendingTasks.clear();
                     }
@@ -95,15 +101,33 @@
     /** Execute the given function after the specified amount of time elapses. */
     public void executeDelayed(Runnable function, long delayMillis) {
         final Handler handler;
+        final Scheduler scheduler;
         synchronized (mPendingTasks) {
             if (this.mHandler == null) {
                 mPendingTasks.add(Pair.create(function, delayMillis));
                 return;
             } else {
                 handler = this.mHandler;
+                if (mMdnsFeatureFlags.mIsAccurateDelayCallbackEnabled
+                        && this.mScheduler == null) {
+                    this.mScheduler = SchedulerFactory.createScheduler(mHandler);
+                }
+                scheduler = this.mScheduler;
             }
         }
-        handler.postDelayed(function, delayMillis);
+        if (scheduler != null) {
+            if (delayMillis == 0L) {
+                handler.post(function);
+                return;
+            }
+            if (HandlerUtils.isRunningOnHandlerThread(handler)) {
+                scheduler.postDelayed(function, delayMillis);
+            } else {
+                handler.post(() -> scheduler.postDelayed(function, delayMillis));
+            }
+        } else {
+            handler.postDelayed(function, delayMillis);
+        }
     }
 
     /** Shutdown the thread if necessary. */
@@ -111,6 +135,11 @@
         if (this.mHandlerThread != null) {
             this.mHandlerThread.quitSafely();
         }
+        synchronized (mPendingTasks) {
+            if (mScheduler != null) {
+                mScheduler.close();
+            }
+        }
     }
 
     /**
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
index 33bcb70..8cd3662 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
@@ -125,7 +125,7 @@
         this.sharedLog = sharedLog;
         this.perSocketServiceTypeClients = new PerSocketServiceTypeClients();
         this.mdnsFeatureFlags = mdnsFeatureFlags;
-        this.discoveryExecutor = new DiscoveryExecutor(socketClient.getLooper());
+        this.discoveryExecutor = new DiscoveryExecutor(socketClient.getLooper(), mdnsFeatureFlags);
     }
 
     /**
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
index 56d4b9a..95f4fff 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -37,7 +37,6 @@
 
 import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.DnsUtils;
-import com.android.net.module.util.RealtimeScheduler;
 import com.android.net.module.util.SharedLog;
 import com.android.server.connectivity.mdns.util.MdnsUtils;
 
@@ -96,9 +95,9 @@
     private final boolean removeServiceAfterTtlExpires =
             MdnsConfigs.removeServiceAfterTtlExpires();
     private final Clock clock;
-    // Use RealtimeScheduler for query scheduling, which allows for more accurate sending of
+    // Use MdnsRealtimeScheduler for query scheduling, which allows for more accurate sending of
     // queries.
-    @Nullable private final RealtimeScheduler realtimeScheduler;
+    @Nullable private final Scheduler scheduler;
 
     @Nullable private MdnsSearchOptions searchOptions;
 
@@ -193,7 +192,7 @@
                     sharedLog.log(String.format("Query sent with transactionId: %d. "
                                     + "Next run: sessionId: %d, in %d ms",
                             sentResult.transactionId, args.sessionId, timeToNextTaskMs));
-                    if (realtimeScheduler != null) {
+                    if (scheduler != null) {
                         setDelayedTask(args, timeToNextTaskMs);
                     } else {
                         dependencies.sendMessageDelayed(
@@ -264,11 +263,11 @@
         }
 
         /**
-         * @see RealtimeScheduler
+         * @see Scheduler
          */
         @Nullable
-        public RealtimeScheduler createRealtimeScheduler(@NonNull Handler handler) {
-            return new RealtimeScheduler(handler);
+        public Scheduler createScheduler(@NonNull Handler handler) {
+            return SchedulerFactory.createScheduler(handler);
         }
     }
 
@@ -317,8 +316,8 @@
         this.mdnsQueryScheduler = new MdnsQueryScheduler();
         this.cacheKey = new MdnsServiceCache.CacheKey(serviceType, socketKey);
         this.featureFlags = featureFlags;
-        this.realtimeScheduler = featureFlags.isAccurateDelayCallbackEnabled()
-                ? dependencies.createRealtimeScheduler(handler) : null;
+        this.scheduler = featureFlags.isAccurateDelayCallbackEnabled()
+                ? dependencies.createScheduler(handler) : null;
     }
 
     /**
@@ -328,8 +327,8 @@
         removeScheduledTask();
         mdnsQueryScheduler.cancelScheduledRun();
         serviceCache.unregisterServiceExpiredCallback(cacheKey);
-        if (realtimeScheduler != null) {
-            realtimeScheduler.close();
+        if (scheduler != null) {
+            scheduler.close();
         }
     }
 
@@ -339,8 +338,8 @@
     }
 
     private void setDelayedTask(ScheduledQueryTaskArgs args, long timeToNextTaskMs) {
-        realtimeScheduler.removeDelayedMessage(EVENT_START_QUERYTASK);
-        realtimeScheduler.sendDelayedMessage(
+        scheduler.removeDelayedMessage(EVENT_START_QUERYTASK);
+        scheduler.sendDelayedMessage(
                 handler.obtainMessage(EVENT_START_QUERYTASK, args), timeToNextTaskMs);
     }
 
@@ -404,7 +403,7 @@
             final long timeToNextTaskMs = calculateTimeToNextTask(args, now);
             sharedLog.log(String.format("Schedule a query. Next run: sessionId: %d, in %d ms",
                     args.sessionId, timeToNextTaskMs));
-            if (realtimeScheduler != null) {
+            if (scheduler != null) {
                 setDelayedTask(args, timeToNextTaskMs);
             } else {
                 dependencies.sendMessageDelayed(
@@ -451,8 +450,8 @@
     }
 
     private void removeScheduledTask() {
-        if (realtimeScheduler != null) {
-            realtimeScheduler.removeDelayedMessage(EVENT_START_QUERYTASK);
+        if (scheduler != null) {
+            scheduler.removeDelayedMessage(EVENT_START_QUERYTASK);
         } else {
             dependencies.removeMessages(handler, EVENT_START_QUERYTASK);
         }
@@ -541,8 +540,8 @@
                 }
             }
         }
-        final boolean hasScheduledTask = realtimeScheduler != null
-                ? realtimeScheduler.hasDelayedMessage(EVENT_START_QUERYTASK)
+        final boolean hasScheduledTask = scheduler != null
+                ? scheduler.hasDelayedMessage(EVENT_START_QUERYTASK)
                 : dependencies.hasMessages(handler, EVENT_START_QUERYTASK);
         if (hasScheduledTask) {
             final long now = clock.elapsedRealtime();
@@ -556,7 +555,7 @@
                 final long timeToNextTaskMs = calculateTimeToNextTask(args, now);
                 sharedLog.log(String.format("Reschedule a query. Next run: sessionId: %d, in %d ms",
                         args.sessionId, timeToNextTaskMs));
-                if (realtimeScheduler != null) {
+                if (scheduler != null) {
                     setDelayedTask(args, timeToNextTaskMs);
                 } else {
                     dependencies.sendMessageDelayed(
diff --git a/service-t/src/com/android/server/connectivity/mdns/Scheduler.java b/service-t/src/com/android/server/connectivity/mdns/Scheduler.java
new file mode 100644
index 0000000..85a8e76
--- /dev/null
+++ b/service-t/src/com/android/server/connectivity/mdns/Scheduler.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.connectivity.mdns;
+
+import android.os.Message;
+
+import androidx.annotation.NonNull;
+
+/**
+ * The interface for scheduler.
+ */
+public interface Scheduler {
+    /**
+     * Set a message to be sent after a specified delay.
+     */
+    boolean sendDelayedMessage(@NonNull Message message, long delayMs);
+
+    /**
+     * Remove a scheduled message.
+     */
+    void removeDelayedMessage(int what);
+
+    /**
+     * Check if there is a scheduled message.
+     */
+    boolean hasDelayedMessage(int what);
+
+    /**
+     * Set a runnable to be executed after a specified delay.
+     */
+    boolean postDelayed(@NonNull Runnable runnable, long delayMs);
+
+    /**
+     * Close this object.
+     */
+    void close();
+}
diff --git a/service-t/src/com/android/server/connectivity/mdns/SchedulerFactory.java b/service-t/src/com/android/server/connectivity/mdns/SchedulerFactory.java
new file mode 100644
index 0000000..1cc9a6b
--- /dev/null
+++ b/service-t/src/com/android/server/connectivity/mdns/SchedulerFactory.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.connectivity.mdns;
+
+import android.annotation.NonNull;
+import android.os.Handler;
+
+import com.android.server.connectivity.mdns.internal.MdnsRealtimeScheduler;
+
+/**
+ * The factory class for creating a scheduler.
+ */
+public class SchedulerFactory {
+
+    /**
+     * Creates an realtime delay callback.
+     */
+    public static Scheduler createScheduler(@NonNull Handler handler) {
+        return new MdnsRealtimeScheduler(handler);
+    }
+
+    private SchedulerFactory() {
+    }
+}
diff --git a/service-t/src/com/android/server/connectivity/mdns/internal/MdnsRealtimeScheduler.java b/service-t/src/com/android/server/connectivity/mdns/internal/MdnsRealtimeScheduler.java
new file mode 100644
index 0000000..eff7085
--- /dev/null
+++ b/service-t/src/com/android/server/connectivity/mdns/internal/MdnsRealtimeScheduler.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.connectivity.mdns.internal;
+
+import android.os.Handler;
+import android.os.Message;
+
+import androidx.annotation.NonNull;
+
+import com.android.net.module.util.RealtimeScheduler;
+import com.android.server.connectivity.mdns.Scheduler;
+
+/**
+ * The delay callback for delivering scheduled tasks accurately.
+ */
+public class MdnsRealtimeScheduler extends RealtimeScheduler implements
+        Scheduler {
+    private static final String TAG = MdnsRealtimeScheduler.class.getSimpleName();
+
+    public MdnsRealtimeScheduler(@NonNull Handler handler) {
+        super(handler);
+    }
+
+    public boolean sendDelayedMessage(@NonNull Message message, long delayMs) {
+        return super.sendDelayedMessage(message, delayMs);
+    }
+
+    public void removeDelayedMessage(int what) {
+        super.removeDelayedMessage(what);
+    }
+
+    public boolean hasDelayedMessage(int what) {
+        return super.hasDelayedMessage(what);
+    }
+
+    public boolean postDelayed(@NonNull Runnable runnable, long delayMs) {
+        return super.postDelayed(runnable, delayMs);
+    }
+
+    public void close() {
+        super.close();
+    }
+}
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 3ce3f02..a95d95c 100644
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -228,6 +228,7 @@
 import android.net.Network;
 import android.net.NetworkAgent;
 import android.net.NetworkAgentConfig;
+import android.net.NetworkAndAgentRegistryParcelable;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
 import android.net.NetworkInfo.DetailedState;
@@ -526,6 +527,7 @@
     private final boolean mBackgroundFirewallChainEnabled;
 
     private final boolean mUseDeclaredMethodsForCallbacksEnabled;
+    private final boolean mQueueNetworkAgentEventsInSystemServer;
 
     // Flag to delay callbacks for frozen apps, suppressing duplicate and stale callbacks.
     private final boolean mQueueCallbacksForFrozenApps;
@@ -1928,6 +1930,9 @@
         mUseDeclaredMethodsForCallbacksEnabled =
                 mDeps.isFeatureNotChickenedOut(context,
                         ConnectivityFlags.USE_DECLARED_METHODS_FOR_CALLBACKS);
+        mQueueNetworkAgentEventsInSystemServer =
+                mDeps.isFeatureNotChickenedOut(context,
+                        ConnectivityFlags.QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER);
         // registerUidFrozenStateChangedCallback is only available on U+
         mQueueCallbacksForFrozenApps = mDeps.isAtLeastU()
                 && mDeps.isFeatureNotChickenedOut(context, QUEUE_CALLBACKS_FOR_FROZEN_APPS);
@@ -4688,18 +4693,30 @@
         private void maybeHandleNetworkAgentMessage(Message msg) {
             final Pair<NetworkAgentInfo, Object> arg = (Pair<NetworkAgentInfo, Object>) msg.obj;
             final NetworkAgentInfo nai = arg.first;
-            if (!mNetworkAgentInfos.contains(nai)) {
-                if (VDBG) {
-                    log(String.format("%s from unknown NetworkAgent", eventName(msg.what)));
-                }
-                return;
-            }
 
             // If the network has been destroyed, the only thing that it can do is disconnect.
             if (nai.isDestroyed() && !isDisconnectRequest(msg)) {
                 return;
             }
 
+            if (mQueueNetworkAgentEventsInSystemServer && nai.maybeEnqueueMessage(msg)) {
+                // If the message is enqueued, the NAI will replay it immediately
+                // when registration is complete. It does this by sending all the
+                // messages in the order received immediately after the
+                // EVENT_AGENT_REGISTERED message.
+                return;
+            }
+
+            // If the nai has been registered (and doesn't enqueue), it should now be
+            // in the list of NAIs.
+            if (!mNetworkAgentInfos.contains(nai)) {
+                // TODO : this is supposed to be impossible
+                if (VDBG) {
+                    log(String.format("%s from unknown NetworkAgent", eventName(msg.what)));
+                }
+                return;
+            }
+
             switch (msg.what) {
                 case NetworkAgent.EVENT_NETWORK_CAPABILITIES_CHANGED: {
                     final NetworkCapabilities proposed = (NetworkCapabilities) arg.second;
@@ -9328,7 +9345,7 @@
      * @param providerId the ID of the provider owning this NetworkAgent.
      * @return the network created for this agent.
      */
-    public Network registerNetworkAgent(INetworkAgent na,
+    public NetworkAndAgentRegistryParcelable registerNetworkAgent(INetworkAgent na,
             NetworkInfo networkInfo,
             LinkProperties linkProperties,
             NetworkCapabilities networkCapabilities,
@@ -9371,7 +9388,8 @@
         }
     }
 
-    private Network registerNetworkAgentInternal(INetworkAgent na, NetworkInfo networkInfo,
+    private NetworkAndAgentRegistryParcelable registerNetworkAgentInternal(
+            INetworkAgent na, NetworkInfo networkInfo,
             LinkProperties linkProperties, NetworkCapabilities networkCapabilities,
             NetworkScore currentScore, NetworkAgentConfig networkAgentConfig,
             @Nullable LocalNetworkConfig localNetworkConfig, int providerId,
@@ -9403,8 +9421,11 @@
         // NetworkAgentInfo registration will finish when the NetworkMonitor is created.
         // If the network disconnects or sends any other event before that, messages are deferred by
         // NetworkAgent until nai.connect(), which will be called when finalizing the
-        // registration.
-        return nai.network;
+        // registration. TODO : have NetworkAgentInfo defer them instead.
+        final NetworkAndAgentRegistryParcelable result = new NetworkAndAgentRegistryParcelable();
+        result.network = nai.network;
+        result.registry = nai.getRegistry();
+        return result;
     }
 
     private void handleRegisterNetworkAgent(NetworkAgentInfo nai, INetworkMonitor networkMonitor) {
@@ -9416,8 +9437,6 @@
                 nai.getDeclaredCapabilitiesSanitized(mCarrierPrivilegeAuthenticator)));
         processLinkPropertiesFromAgent(nai, nai.linkProperties);
 
-        nai.onNetworkMonitorCreated(networkMonitor);
-
         mNetworkAgentInfos.add(nai);
         synchronized (mNetworkForNetId) {
             mNetworkForNetId.put(nai.network.getNetId(), nai);
@@ -9432,7 +9451,7 @@
         if (nai.isLocalNetwork()) {
             handleUpdateLocalNetworkConfig(nai, null /* oldConfig */, nai.localNetworkConfig);
         }
-        nai.notifyRegistered();
+        nai.notifyRegistered(networkMonitor);
         NetworkInfo networkInfo = nai.networkInfo;
         updateNetworkInfo(nai, networkInfo);
         updateVpnUids(nai, null, nai.networkCapabilities);
diff --git a/service/src/com/android/server/connectivity/ConnectivityFlags.java b/service/src/com/android/server/connectivity/ConnectivityFlags.java
index 136ea81..74bd235 100644
--- a/service/src/com/android/server/connectivity/ConnectivityFlags.java
+++ b/service/src/com/android/server/connectivity/ConnectivityFlags.java
@@ -62,6 +62,9 @@
     public static final String QUEUE_CALLBACKS_FOR_FROZEN_APPS =
             "queue_callbacks_for_frozen_apps";
 
+    public static final String QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER =
+            "queue_network_agent_events_in_system_server";
+
     private boolean mNoRematchAllRequestsOnRegister;
 
     /**
diff --git a/service/src/com/android/server/connectivity/NetworkAgentInfo.java b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
index e762a8e..4540f02 100644
--- a/service/src/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
@@ -62,6 +62,7 @@
 import android.net.TcpKeepalivePacketData;
 import android.os.Handler;
 import android.os.IBinder;
+import android.os.Message;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
 import android.os.SystemClock;
@@ -630,6 +631,7 @@
     // Used by ConnectivityService to keep track of 464xlat.
     public final Nat464Xlat clatd;
 
+    private final ArrayList<Message> mMessagesPendingRegistration = new ArrayList<>();
     // Set after asynchronous creation of the NetworkMonitor.
     private volatile NetworkMonitorManager mNetworkMonitor;
 
@@ -639,6 +641,7 @@
     private final ConnectivityService.Dependencies mConnServiceDeps;
     private final Context mContext;
     private final Handler mHandler;
+    private final NetworkAgentMessageHandler mRegistry;
     private final QosCallbackTracker mQosCallbackTracker;
     private final INetd mNetd;
 
@@ -673,6 +676,7 @@
         mNetd = netd;
         mContext = context;
         mHandler = handler;
+        mRegistry = new NetworkAgentMessageHandler(mHandler);
         this.factorySerialNumber = factorySerialNumber;
         this.creatorUid = creatorUid;
         mLingerDurationMs = lingerDurationMs;
@@ -698,10 +702,12 @@
      * Must be called from the ConnectivityService handler thread. A NetworkAgent can only be
      * registered once.
      */
-    public void notifyRegistered() {
+    public void notifyRegistered(final INetworkMonitor nm) {
+        HandlerUtils.ensureRunningOnHandlerThread(mHandler);
+        mNetworkMonitor = new NetworkMonitorManager(nm);
         try {
             networkAgent.asBinder().linkToDeath(mDeathMonitor, 0);
-            networkAgent.onRegistered(new NetworkAgentMessageHandler(mHandler));
+            networkAgent.onRegistered();
         } catch (RemoteException e) {
             Log.e(TAG, "Error registering NetworkAgent", e);
             maybeUnlinkDeathMonitor();
@@ -711,6 +717,29 @@
         }
 
         mHandler.obtainMessage(EVENT_AGENT_REGISTERED, ARG_AGENT_SUCCESS, 0, this).sendToTarget();
+        for (final Message enqueued : mMessagesPendingRegistration) {
+            mHandler.sendMessage(enqueued);
+        }
+        mMessagesPendingRegistration.clear();
+    }
+
+    /**
+     * Enqueues a message if it needs to be enqueued, and returns whether it was enqueued.
+     *
+     * The message is enqueued iff it can't be sent just yet. If it can be sent
+     * immediately, this method returns false and doesn't enqueue.
+     *
+     * If it enqueues, this method will make a copy of the message for enqueuing since
+     * messages can't be reused or recycled before the end of their processing by the
+     * handler.
+     */
+    public boolean maybeEnqueueMessage(final Message msg) {
+        HandlerUtils.ensureRunningOnHandlerThread(mHandler);
+        if (null != mNetworkMonitor) return false;
+        final Message m = mHandler.obtainMessage();
+        m.copyFrom(msg);
+        mMessagesPendingRegistration.add(m);
+        return true;
     }
 
     /**
@@ -1036,13 +1065,6 @@
     }
 
     /**
-     * Inform NetworkAgentInfo that a new NetworkMonitor was created.
-     */
-    public void onNetworkMonitorCreated(INetworkMonitor networkMonitor) {
-        mNetworkMonitor = new NetworkMonitorManager(networkMonitor);
-    }
-
-    /**
      * Set the NetworkCapabilities on this NetworkAgentInfo. Also attempts to notify NetworkMonitor
      * of the new capabilities, if NetworkMonitor has been created.
      *
@@ -1117,6 +1139,13 @@
         return mNetworkMonitor;
     }
 
+    /**
+     * Get the registry in this NetworkAgentInfo.
+     */
+    public INetworkAgentRegistry getRegistry() {
+        return mRegistry;
+    }
+
     // Functions for manipulating the requests satisfied by this network.
     //
     // These functions must only called on ConnectivityService's main thread.
diff --git a/staticlibs/Android.bp b/staticlibs/Android.bp
index b064723..abfc447 100644
--- a/staticlibs/Android.bp
+++ b/staticlibs/Android.bp
@@ -427,11 +427,7 @@
     srcs: [
         "device/com/android/net/module/util/FdEventsReader.java",
         "device/com/android/net/module/util/HandlerUtils.java",
-        "device/com/android/net/module/util/JniUtil.java",
-        "device/com/android/net/module/util/RealtimeScheduler.java",
         "device/com/android/net/module/util/SharedLog.java",
-        "device/com/android/net/module/util/ServiceConnectivityJni.java",
-        "device/com/android/net/module/util/TimerFdUtils.java",
         "framework/com/android/net/module/util/ByteUtils.java",
         "framework/com/android/net/module/util/CollectionUtils.java",
         "framework/com/android/net/module/util/DnsUtils.java",
diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index 5e035a2..678e1ca 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -154,10 +154,10 @@
 import java.security.MessageDigest
 import java.time.Duration
 import java.util.Arrays
-import java.util.Random
 import java.util.UUID
 import java.util.concurrent.Executors
 import kotlin.collections.ArrayList
+import kotlin.random.Random
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
@@ -267,7 +267,7 @@
     private class FakeConnectivityService {
         val mockRegistry = mock(INetworkAgentRegistry::class.java)
         private var agentField: INetworkAgent? = null
-        private val registry = object : INetworkAgentRegistry.Stub(),
+        val registry: INetworkAgentRegistry = object : INetworkAgentRegistry.Stub(),
                 INetworkAgentRegistry by mockRegistry {
             // asBinder has implementations in both INetworkAgentRegistry.Stub and mockRegistry, so
             // it needs to be disambiguated. Just fail the test as it should be unused here.
@@ -284,7 +284,7 @@
 
         fun connect(agent: INetworkAgent) {
             this.agentField = agent
-            agent.onRegistered(registry)
+            agent.onRegistered()
         }
 
         fun disconnect() = agent.onDisconnected()
@@ -413,7 +413,8 @@
     }
 
     private fun createNetworkAgentWithFakeCS() = createNetworkAgent().also {
-        mFakeConnectivityService.connect(it.registerForTest(Network(FAKE_NET_ID)))
+        val binder = it.registerForTest(Network(FAKE_NET_ID), mFakeConnectivityService.registry)
+        mFakeConnectivityService.connect(binder)
     }
 
     private fun TestableNetworkAgent.expectPostConnectionCallbacks(
@@ -1628,7 +1629,7 @@
         val s = Os.socket(AF_INET6, SOCK_DGRAM, 0)
         net.bindSocket(s)
         val content = ByteArray(16)
-        Random().nextBytes(content)
+        Random.nextBytes(content)
         Os.sendto(s, ByteBuffer.wrap(content), 0, REMOTE_ADDRESS, 7 /* port */)
         val match = reader.poll(DEFAULT_TIMEOUT_MS) {
             val udpStart = IPV6_HEADER_LEN + UDP_HEADER_LEN
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index c28a0f8..3eefa0f 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -2200,6 +2200,7 @@
                 case ConnectivityFlags.DELAY_DESTROY_SOCKETS:
                 case ConnectivityFlags.USE_DECLARED_METHODS_FOR_CALLBACKS:
                 case ConnectivityFlags.QUEUE_CALLBACKS_FOR_FROZEN_APPS:
+                case ConnectivityFlags.QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER:
                     return true;
                 default:
                     throw new UnsupportedOperationException("Unknown flag " + name
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/DiscoveryExecutorTest.kt b/tests/unit/java/com/android/server/connectivity/mdns/DiscoveryExecutorTest.kt
index 51539a0..67fb428 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/DiscoveryExecutorTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/DiscoveryExecutorTest.kt
@@ -17,6 +17,7 @@
 package com.android.server.connectivity.mdns
 
 import android.os.Build
+import android.os.Handler
 import android.os.HandlerThread
 import android.testing.TestableLooper
 import com.android.testutils.DevSdkIgnoreRule
@@ -36,6 +37,8 @@
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
 class DiscoveryExecutorTest {
     private val thread = HandlerThread(DiscoveryExecutorTest::class.simpleName).apply { start() }
+    private val handler by lazy { Handler(thread.looper) }
+    private val testableLooper by lazy { TestableLooper(thread.looper) }
 
     @After
     fun tearDown() {
@@ -45,8 +48,10 @@
 
     @Test
     fun testCheckAndRunOnHandlerThread() {
-        val testableLooper = TestableLooper(thread.looper)
-        val executor = DiscoveryExecutor(testableLooper.looper)
+        val executor = DiscoveryExecutor(
+                testableLooper.looper,
+                MdnsFeatureFlags.newBuilder().build()
+        )
         try {
             val future = CompletableFuture<Boolean>()
             executor.checkAndRunOnHandlerThread { future.complete(true) }
@@ -58,17 +63,17 @@
 
         // Create a DiscoveryExecutor with the null defaultLooper and verify the task can execute
         // normally.
-        val executor2 = DiscoveryExecutor(null /* defaultLooper */)
+        val executor2 = DiscoveryExecutor(
+                null /* defaultLooper */,
+                MdnsFeatureFlags.newBuilder().build()
+        )
         val future2 = CompletableFuture<Boolean>()
         executor2.checkAndRunOnHandlerThread { future2.complete(true) }
         assertTrue(future2.get(DEFAULT_TIMEOUT, TimeUnit.MILLISECONDS))
         executor2.shutDown()
     }
 
-    @Test
-    fun testExecute() {
-        val testableLooper = TestableLooper(thread.looper)
-        val executor = DiscoveryExecutor(testableLooper.looper)
+    private fun verifyExecute(executor: DiscoveryExecutor) {
         try {
             val future = CompletableFuture<Boolean>()
             executor.execute { future.complete(true) }
@@ -81,9 +86,27 @@
     }
 
     @Test
+    fun testExecute() {
+        verifyExecute(DiscoveryExecutor(
+                testableLooper.looper,
+                MdnsFeatureFlags.newBuilder().build()
+        ))
+    }
+
+    @Test
+    fun testExecute_RealtimeScheduler() {
+        verifyExecute(DiscoveryExecutor(
+                testableLooper.looper,
+                MdnsFeatureFlags.newBuilder().setIsAccurateDelayCallbackEnabled(true).build()
+        ))
+    }
+
+    @Test
     fun testExecuteDelayed() {
-        val testableLooper = TestableLooper(thread.looper)
-        val executor = DiscoveryExecutor(testableLooper.looper)
+        val executor = DiscoveryExecutor(
+                testableLooper.looper,
+                MdnsFeatureFlags.newBuilder().build()
+        )
         try {
             // Verify the executeDelayed method
             val future = CompletableFuture<Boolean>()
@@ -107,4 +130,21 @@
             testableLooper.destroy()
         }
     }
+
+    @Test
+    fun testExecuteDelayed_RealtimeScheduler() {
+        val executor = DiscoveryExecutor(
+                thread.looper,
+                MdnsFeatureFlags.newBuilder().setIsAccurateDelayCallbackEnabled(true).build()
+        )
+        try {
+            // Verify the executeDelayed method
+            val future = CompletableFuture<Boolean>()
+            // Schedule a task with 50ms delay
+            executor.executeDelayed({ future.complete(true) }, 50L)
+            assertTrue(future.get(500L, TimeUnit.MILLISECONDS))
+        } finally {
+            testableLooper.destroy()
+        }
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
index dad03e0..b9c0d2f 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -59,7 +59,6 @@
 import android.text.TextUtils;
 
 import com.android.net.module.util.CollectionUtils;
-import com.android.net.module.util.RealtimeScheduler;
 import com.android.net.module.util.SharedLog;
 import com.android.server.connectivity.mdns.MdnsServiceInfo.TextEntry;
 import com.android.server.connectivity.mdns.util.MdnsUtils;
@@ -129,7 +128,7 @@
     @Mock
     private MdnsServiceTypeClient.Dependencies mockDeps;
     @Mock
-    private RealtimeScheduler mockRealtimeScheduler;
+    private Scheduler mockScheduler;
     @Captor
     private ArgumentCaptor<MdnsServiceInfo> serviceInfoCaptor;
 
@@ -250,14 +249,14 @@
 
         doAnswer(inv -> {
             realHandler = (Handler) inv.getArguments()[0];
-            return mockRealtimeScheduler;
-        }).when(mockDeps).createRealtimeScheduler(any(Handler.class));
+            return mockScheduler;
+        }).when(mockDeps).createScheduler(any(Handler.class));
 
         doAnswer(inv -> {
             message = (Message) inv.getArguments()[0];
             latestDelayMs = (long) inv.getArguments()[1];
             return null;
-        }).when(mockRealtimeScheduler).sendDelayedMessage(any(), anyLong());
+        }).when(mockScheduler).sendDelayedMessage(any(), anyLong());
 
         client = makeMdnsServiceTypeClient(featureFlags);
     }
@@ -2137,7 +2136,7 @@
                 .setNumOfQueriesBeforeBackoff(numOfQueriesBeforeBackoff)
                 .build();
         startSendAndReceive(mockListenerOne, searchOptions);
-        verify(mockRealtimeScheduler, times(1)).removeDelayedMessage(EVENT_START_QUERYTASK);
+        verify(mockScheduler, times(1)).removeDelayedMessage(EVENT_START_QUERYTASK);
 
         // Verify that the first query has been sent.
         verifyAndSendQuery(0 /* index */, 0 /* timeInMs */, true /* expectsUnicastResponse */,
@@ -2159,13 +2158,13 @@
         // 0.8 * smallestRemainingTtl is larger than time to next run.
         long currentTime = TEST_TTL / 2 + TEST_ELAPSED_REALTIME;
         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
-        doReturn(true).when(mockRealtimeScheduler).hasDelayedMessage(EVENT_START_QUERYTASK);
+        doReturn(true).when(mockScheduler).hasDelayedMessage(EVENT_START_QUERYTASK);
         processResponse(createResponse(
                 "service-instance-1", "192.0.2.123", 5353,
                 SERVICE_TYPE_LABELS,
                 Collections.emptyMap(), TEST_TTL), socketKey);
         // Verify that the message removal occurred.
-        verify(mockRealtimeScheduler, times(6)).removeDelayedMessage(EVENT_START_QUERYTASK);
+        verify(mockScheduler, times(6)).removeDelayedMessage(EVENT_START_QUERYTASK);
         assertNotNull(message);
         verifyAndSendQuery(3 /* index */, (long) (TEST_TTL / 2 * 0.8) /* timeInMs */,
                 true /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
@@ -2174,7 +2173,7 @@
 
         // Stop sending packets.
         stopSendAndReceive(mockListenerOne);
-        verify(mockRealtimeScheduler, times(8)).removeDelayedMessage(EVENT_START_QUERYTASK);
+        verify(mockScheduler, times(8)).removeDelayedMessage(EVENT_START_QUERYTASK);
     }
 
     @Test
@@ -2184,12 +2183,12 @@
 
         // Start query
         startSendAndReceive(mockListenerOne, MdnsSearchOptions.newBuilder().build());
-        verify(mockRealtimeScheduler, times(1)).removeDelayedMessage(EVENT_START_QUERYTASK);
+        verify(mockScheduler, times(1)).removeDelayedMessage(EVENT_START_QUERYTASK);
 
         // Stop query and verify the close() method has been called.
         stopSendAndReceive(mockListenerOne);
-        verify(mockRealtimeScheduler, times(2)).removeDelayedMessage(EVENT_START_QUERYTASK);
-        verify(mockRealtimeScheduler).close();
+        verify(mockScheduler, times(2)).removeDelayedMessage(EVENT_START_QUERYTASK);
+        verify(mockScheduler).close();
     }
 
     private static MdnsServiceInfo matchServiceName(String name) {
@@ -2247,8 +2246,7 @@
                 .sendMessage(any(Handler.class), any(Message.class));
         // Verify the task has been scheduled.
         if (useAccurateDelayCallback) {
-            verify(mockRealtimeScheduler, times(scheduledCount))
-                    .sendDelayedMessage(any(), anyLong());
+            verify(mockScheduler, times(scheduledCount)).sendDelayedMessage(any(), anyLong());
         } else {
             verify(mockDeps, times(scheduledCount))
                     .sendMessageDelayed(any(Handler.class), any(Message.class), anyLong());
diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
index d7e781e..48333c5 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
@@ -169,6 +169,7 @@
         it[ConnectivityFlags.DELAY_DESTROY_SOCKETS] = true
         it[ConnectivityFlags.USE_DECLARED_METHODS_FOR_CALLBACKS] = true
         it[ConnectivityFlags.QUEUE_CALLBACKS_FOR_FROZEN_APPS] = true
+        it[ConnectivityFlags.QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER] = true
     }
     fun setFeatureEnabled(flag: String, enabled: Boolean) = enabledFeatures.set(flag, enabled)