Implement RealtimeScheduler in DiscoveryExecutor

This provides more accurate delay callbacks, thereby enhancing
the executor's overall functionality.

Bug: 366373064
Test: atest FrameworksNetTests NsdManagerTest
Change-Id: I903851464a92427b6e256573d693bb89a96d013f
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..5b6ec9d 100644
--- a/service-t/src/com/android/server/connectivity/mdns/DiscoveryExecutor.java
+++ b/service-t/src/com/android/server/connectivity/mdns/DiscoveryExecutor.java
@@ -27,6 +27,7 @@
 import androidx.annotation.GuardedBy;
 
 import com.android.net.module.util.HandlerUtils;
+import com.android.net.module.util.RealtimeScheduler;
 
 import java.util.ArrayList;
 import java.util.concurrent.Executor;
@@ -49,7 +50,13 @@
     @NonNull
     private final ArrayList<Pair<Runnable, Long>> mPendingTasks = new ArrayList<>();
 
-    DiscoveryExecutor(@Nullable Looper defaultLooper) {
+    @GuardedBy("mPendingTasks")
+    @Nullable
+    RealtimeScheduler mRealtimeScheduler;
+    @NonNull private final MdnsFeatureFlags mMdnsFeatureFlags;
+
+    DiscoveryExecutor(@Nullable Looper defaultLooper, @NonNull MdnsFeatureFlags mdnsFeatureFlags) {
+        mMdnsFeatureFlags = mdnsFeatureFlags;
         if (defaultLooper != null) {
             this.mHandlerThread = null;
             synchronized (mPendingTasks) {
@@ -62,7 +69,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 +102,33 @@
     /** Execute the given function after the specified amount of time elapses. */
     public void executeDelayed(Runnable function, long delayMillis) {
         final Handler handler;
+        final RealtimeScheduler realtimeScheduler;
         synchronized (mPendingTasks) {
             if (this.mHandler == null) {
                 mPendingTasks.add(Pair.create(function, delayMillis));
                 return;
             } else {
                 handler = this.mHandler;
+                if (mMdnsFeatureFlags.mIsAccurateDelayCallbackEnabled
+                        && this.mRealtimeScheduler == null) {
+                    this.mRealtimeScheduler = new RealtimeScheduler(mHandler);
+                }
+                realtimeScheduler = this.mRealtimeScheduler;
             }
         }
-        handler.postDelayed(function, delayMillis);
+        if (realtimeScheduler != null) {
+            if (delayMillis == 0L) {
+                handler.post(function);
+                return;
+            }
+            if (HandlerUtils.isRunningOnHandlerThread(handler)) {
+                realtimeScheduler.postDelayed(function, delayMillis);
+            } else {
+                handler.post(() -> realtimeScheduler.postDelayed(function, delayMillis));
+            }
+        } else {
+            handler.postDelayed(function, delayMillis);
+        }
     }
 
     /** Shutdown the thread if necessary. */
@@ -111,6 +136,11 @@
         if (this.mHandlerThread != null) {
             this.mHandlerThread.quitSafely();
         }
+        synchronized (mPendingTasks) {
+            if (mRealtimeScheduler != null) {
+                mRealtimeScheduler.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/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()
+        }
+    }
 }