From 3d51a99dc8a5b6c29df406ac58f7de18db6fa0b7 Mon Sep 17 00:00:00 2001 From: Paul Hu Date: Mon, 16 Sep 2024 16:27:08 +0000 Subject: [PATCH] Use timerfd to improve the accuracy of query timing Timerfd can provide more accurate delayed callbacks than Handler#postDelayed because it accounts for the time the device spends in deep sleep. Using timerfd allows for more precise query timing, which reduces the possibility of TTL expiration caused by inaccurate delayed callbacks. Bug: 366373064 Test: atest FrameworksNetTests NsdManagerTest Change-Id: I980ec8fd43f63976ef5387bb57c7ed70a1b68d82 --- .../connectivity/mdns/MdnsFeatureFlags.java | 31 +++++ .../mdns/MdnsServiceTypeClient.java | 87 +++++++++---- staticlibs/Android.bp | 3 + .../net/module/util/TimerFileDescriptor.java | 7 ++ .../mdns/MdnsServiceTypeClientTests.java | 116 +++++++++++++++--- 5 files changed, 206 insertions(+), 38 deletions(-) diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java b/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java index 4e27fef60e..c4a91103e2 100644 --- a/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java +++ b/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java @@ -89,6 +89,11 @@ public class MdnsFeatureFlags { "nsd_cached_services_retention_time"; public static final int DEFAULT_CACHED_SERVICES_RETENTION_TIME_MILLISECONDS = 10000; + /** + * A feature flag to control whether the accurate delay callback should be enabled. + */ + public static final String NSD_ACCURATE_DELAY_CALLBACK = "nsd_accurate_delay_callback"; + // Flag for offload feature public final boolean mIsMdnsOffloadFeatureEnabled; @@ -122,6 +127,9 @@ public class MdnsFeatureFlags { // Retention Time for cached services public final long mCachedServicesRetentionTime; + // Flag for accurate delay callback + public final boolean mIsAccurateDelayCallbackEnabled; + @Nullable private final FlagOverrideProvider mOverrideProvider; @@ -217,6 +225,14 @@ public class MdnsFeatureFlags { NSD_CACHED_SERVICES_RETENTION_TIME, (int) mCachedServicesRetentionTime); } + /** + * Indicates whether {@link #NSD_ACCURATE_DELAY_CALLBACK} is enabled, including for testing. + */ + public boolean isAccurateDelayCallbackEnabled() { + return mIsAccurateDelayCallbackEnabled + || isForceEnabledForTest(NSD_ACCURATE_DELAY_CALLBACK); + } + /** * The constructor for {@link MdnsFeatureFlags}. */ @@ -231,6 +247,7 @@ public class MdnsFeatureFlags { boolean avoidAdvertisingEmptyTxtRecords, boolean isCachedServicesRemovalEnabled, long cachedServicesRetentionTime, + boolean isAccurateDelayCallbackEnabled, @Nullable FlagOverrideProvider overrideProvider) { mIsMdnsOffloadFeatureEnabled = isOffloadFeatureEnabled; mIncludeInetAddressRecordsInProbing = includeInetAddressRecordsInProbing; @@ -243,6 +260,7 @@ public class MdnsFeatureFlags { mAvoidAdvertisingEmptyTxtRecords = avoidAdvertisingEmptyTxtRecords; mIsCachedServicesRemovalEnabled = isCachedServicesRemovalEnabled; mCachedServicesRetentionTime = cachedServicesRetentionTime; + mIsAccurateDelayCallbackEnabled = isAccurateDelayCallbackEnabled; mOverrideProvider = overrideProvider; } @@ -266,6 +284,7 @@ public class MdnsFeatureFlags { private boolean mAvoidAdvertisingEmptyTxtRecords; private boolean mIsCachedServicesRemovalEnabled; private long mCachedServicesRetentionTime; + private boolean mIsAccurateDelayCallbackEnabled; private FlagOverrideProvider mOverrideProvider; /** @@ -283,6 +302,7 @@ public class MdnsFeatureFlags { mAvoidAdvertisingEmptyTxtRecords = true; // Default enabled. mIsCachedServicesRemovalEnabled = false; mCachedServicesRetentionTime = DEFAULT_CACHED_SERVICES_RETENTION_TIME_MILLISECONDS; + mIsAccurateDelayCallbackEnabled = false; mOverrideProvider = null; } @@ -408,6 +428,16 @@ public class MdnsFeatureFlags { return this; } + /** + * Set whether the accurate delay callback is enabled. + * + * @see #NSD_ACCURATE_DELAY_CALLBACK + */ + public Builder setIsAccurateDelayCallbackEnabled(boolean isAccurateDelayCallbackEnabled) { + mIsAccurateDelayCallbackEnabled = isAccurateDelayCallbackEnabled; + return this; + } + /** * Builds a {@link MdnsFeatureFlags} with the arguments supplied to this builder. */ @@ -423,6 +453,7 @@ public class MdnsFeatureFlags { mAvoidAdvertisingEmptyTxtRecords, mIsCachedServicesRemovalEnabled, mCachedServicesRetentionTime, + mIsAccurateDelayCallbackEnabled, mOverrideProvider); } } 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 a43486eed9..7a93fecec1 100644 --- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java +++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java @@ -20,6 +20,7 @@ import static com.android.net.module.util.HandlerUtils.ensureRunningOnHandlerThr import static com.android.server.connectivity.mdns.MdnsSearchOptions.AGGRESSIVE_QUERY_MODE; import static com.android.server.connectivity.mdns.MdnsServiceCache.ServiceExpiredCallback; import static com.android.server.connectivity.mdns.MdnsServiceCache.findMatchedResponse; +import static com.android.server.connectivity.mdns.MdnsQueryScheduler.ScheduledQueryTaskArgs; import static com.android.server.connectivity.mdns.util.MdnsUtils.Clock; import static com.android.server.connectivity.mdns.util.MdnsUtils.buildMdnsServiceInfoFromResponse; @@ -37,6 +38,7 @@ import androidx.annotation.VisibleForTesting; import com.android.net.module.util.CollectionUtils; import com.android.net.module.util.DnsUtils; import com.android.net.module.util.SharedLog; +import com.android.net.module.util.TimerFileDescriptor; import com.android.server.connectivity.mdns.util.MdnsUtils; import java.io.IOException; @@ -94,6 +96,9 @@ public class MdnsServiceTypeClient { private final boolean removeServiceAfterTtlExpires = MdnsConfigs.removeServiceAfterTtlExpires(); private final Clock clock; + // Use TimerFileDescriptor for query scheduling, which allows for more accurate sending of + // queries. + @NonNull private final TimerFileDescriptor timerFd; @Nullable private MdnsSearchOptions searchOptions; @@ -139,8 +144,7 @@ public class MdnsServiceTypeClient { public void handleMessage(Message msg) { switch (msg.what) { case EVENT_START_QUERYTASK: { - final MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs = - (MdnsQueryScheduler.ScheduledQueryTaskArgs) msg.obj; + final ScheduledQueryTaskArgs taskArgs = (ScheduledQueryTaskArgs) msg.obj; // QueryTask should be run immediately after being created (not be scheduled in // advance). Because the result of "makeResponsesForResolve" depends on answers // that were received before it is called, so to take into account all answers @@ -174,7 +178,7 @@ public class MdnsServiceTypeClient { final long now = clock.elapsedRealtime(); lastSentTime = now; final long minRemainingTtl = getMinRemainingTtl(now); - MdnsQueryScheduler.ScheduledQueryTaskArgs args = + final ScheduledQueryTaskArgs args = mdnsQueryScheduler.scheduleNextRun( sentResult.taskArgs.config, minRemainingTtl, @@ -189,10 +193,14 @@ public class MdnsServiceTypeClient { sharedLog.log(String.format("Query sent with transactionId: %d. " + "Next run: sessionId: %d, in %d ms", sentResult.transactionId, args.sessionId, timeToNextTaskMs)); - dependencies.sendMessageDelayed( - handler, - handler.obtainMessage(EVENT_START_QUERYTASK, args), - timeToNextTaskMs); + if (featureFlags.isAccurateDelayCallbackEnabled()) { + setDelayedTask(args, timeToNextTaskMs); + } else { + dependencies.sendMessageDelayed( + handler, + handler.obtainMessage(EVENT_START_QUERYTASK, args), + timeToNextTaskMs); + } break; } default: @@ -254,6 +262,14 @@ public class MdnsServiceTypeClient { return List.of(new DatagramPacket(queryBuffer, 0, queryBuffer.length, address)); } } + + /** + * @see TimerFileDescriptor + */ + @Nullable + public TimerFileDescriptor createTimerFd(@NonNull Handler handler) { + return new TimerFileDescriptor(handler); + } } /** @@ -301,6 +317,7 @@ public class MdnsServiceTypeClient { this.mdnsQueryScheduler = new MdnsQueryScheduler(); this.cacheKey = new MdnsServiceCache.CacheKey(serviceType, socketKey); this.featureFlags = featureFlags; + this.timerFd = dependencies.createTimerFd(handler); } /** @@ -317,6 +334,13 @@ public class MdnsServiceTypeClient { ? serviceCache.getCachedServices(cacheKey) : Collections.emptyList(); } + private void setDelayedTask(ScheduledQueryTaskArgs args, long timeToNextTaskMs) { + timerFd.cancelTask(); + timerFd.setDelayedTask(new TimerFileDescriptor.MessageTask( + handler.obtainMessage(EVENT_START_QUERYTASK, args)), + timeToNextTaskMs); + } + /** * Registers {@code listener} for receiving discovery event of mDNS service instances, and * starts @@ -363,7 +387,7 @@ public class MdnsServiceTypeClient { } final long minRemainingTtl = getMinRemainingTtl(now); if (hadReply) { - MdnsQueryScheduler.ScheduledQueryTaskArgs args = + final ScheduledQueryTaskArgs args = mdnsQueryScheduler.scheduleNextRun( taskConfig, minRemainingTtl, @@ -377,10 +401,14 @@ public class MdnsServiceTypeClient { final long timeToNextTaskMs = calculateTimeToNextTask(args, now); sharedLog.log(String.format("Schedule a query. Next run: sessionId: %d, in %d ms", args.sessionId, timeToNextTaskMs)); - dependencies.sendMessageDelayed( - handler, - handler.obtainMessage(EVENT_START_QUERYTASK, args), - timeToNextTaskMs); + if (featureFlags.isAccurateDelayCallbackEnabled()) { + setDelayedTask(args, timeToNextTaskMs); + } else { + dependencies.sendMessageDelayed( + handler, + handler.obtainMessage(EVENT_START_QUERYTASK, args), + timeToNextTaskMs); + } } else { final List servicesToResolve = makeResponsesForResolve(socketKey); final QueryTask queryTask = new QueryTask( @@ -420,7 +448,11 @@ public class MdnsServiceTypeClient { } private void removeScheduledTask() { - dependencies.removeMessages(handler, EVENT_START_QUERYTASK); + if (featureFlags.isAccurateDelayCallbackEnabled()) { + timerFd.cancelTask(); + } else { + dependencies.removeMessages(handler, EVENT_START_QUERYTASK); + } sharedLog.log("Remove EVENT_START_QUERYTASK" + ", current session: " + currentSessionId); ++currentSessionId; @@ -506,10 +538,13 @@ public class MdnsServiceTypeClient { } } } - if (dependencies.hasMessages(handler, EVENT_START_QUERYTASK)) { + final boolean hasScheduledTask = featureFlags.isAccurateDelayCallbackEnabled() + ? timerFd.hasDelayedTask() + : dependencies.hasMessages(handler, EVENT_START_QUERYTASK); + if (hasScheduledTask) { final long now = clock.elapsedRealtime(); final long minRemainingTtl = getMinRemainingTtl(now); - MdnsQueryScheduler.ScheduledQueryTaskArgs args = + final ScheduledQueryTaskArgs args = mdnsQueryScheduler.maybeRescheduleCurrentRun(now, minRemainingTtl, lastSentTime, currentSessionId + 1, searchOptions.numOfQueriesBeforeBackoff()); @@ -518,10 +553,14 @@ public class MdnsServiceTypeClient { final long timeToNextTaskMs = calculateTimeToNextTask(args, now); sharedLog.log(String.format("Reschedule a query. Next run: sessionId: %d, in %d ms", args.sessionId, timeToNextTaskMs)); - dependencies.sendMessageDelayed( - handler, - handler.obtainMessage(EVENT_START_QUERYTASK, args), - timeToNextTaskMs); + if (featureFlags.isAccurateDelayCallbackEnabled()) { + setDelayedTask(args, timeToNextTaskMs); + } else { + dependencies.sendMessageDelayed( + handler, + handler.obtainMessage(EVENT_START_QUERYTASK, args), + timeToNextTaskMs); + } } } } @@ -686,10 +725,10 @@ public class MdnsServiceTypeClient { private static class QuerySentArguments { private final int transactionId; private final List subTypes = new ArrayList<>(); - private final MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs; + private final ScheduledQueryTaskArgs taskArgs; QuerySentArguments(int transactionId, @NonNull List subTypes, - @NonNull MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs) { + @NonNull ScheduledQueryTaskArgs taskArgs) { this.transactionId = transactionId; this.subTypes.addAll(subTypes); this.taskArgs = taskArgs; @@ -698,14 +737,14 @@ public class MdnsServiceTypeClient { // A FutureTask that enqueues a single query, and schedule a new FutureTask for the next task. private class QueryTask implements Runnable { - private final MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs; + private final ScheduledQueryTaskArgs taskArgs; private final List servicesToResolve = new ArrayList<>(); private final List subtypes = new ArrayList<>(); private final boolean sendDiscoveryQueries; private final List existingServices = new ArrayList<>(); private final boolean onlyUseIpv6OnIpv6OnlyNetworks; private final SocketKey socketKey; - QueryTask(@NonNull MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs, + QueryTask(@NonNull ScheduledQueryTaskArgs taskArgs, @NonNull Collection servicesToResolve, @NonNull Collection subtypes, boolean sendDiscoveryQueries, @NonNull Collection existingServices, @@ -771,7 +810,7 @@ public class MdnsServiceTypeClient { return minRemainingTtl == Long.MAX_VALUE ? 0 : minRemainingTtl; } - private static long calculateTimeToNextTask(MdnsQueryScheduler.ScheduledQueryTaskArgs args, + private static long calculateTimeToNextTask(ScheduledQueryTaskArgs args, long now) { return Math.max(args.timeToRun - now, 0); } diff --git a/staticlibs/Android.bp b/staticlibs/Android.bp index b5ee7892b6..046b0355f9 100644 --- a/staticlibs/Android.bp +++ b/staticlibs/Android.bp @@ -438,7 +438,10 @@ java_library { 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/SharedLog.java", + "device/com/android/net/module/util/TimerFdUtils.java", + "device/com/android/net/module/util/TimerFileDescriptor.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/staticlibs/device/com/android/net/module/util/TimerFileDescriptor.java b/staticlibs/device/com/android/net/module/util/TimerFileDescriptor.java index dbbccc54a5..a8c0f171fa 100644 --- a/staticlibs/device/com/android/net/module/util/TimerFileDescriptor.java +++ b/staticlibs/device/com/android/net/module/util/TimerFileDescriptor.java @@ -103,6 +103,13 @@ public class TimerFileDescriptor { public void post(Handler handler) { handler.sendMessage(mMessage); } + + /** + * Get scheduled message + */ + public Message getMessage() { + return mMessage; + } } /** 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 da0bc88236..ed95e4bc2d 100644 --- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java +++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java @@ -60,6 +60,7 @@ import android.text.TextUtils; import com.android.net.module.util.CollectionUtils; import com.android.net.module.util.SharedLog; +import com.android.net.module.util.TimerFileDescriptor; import com.android.server.connectivity.mdns.MdnsServiceInfo.TextEntry; import com.android.server.connectivity.mdns.util.MdnsUtils; import com.android.testutils.DevSdkIgnoreRule; @@ -127,6 +128,8 @@ public class MdnsServiceTypeClientTests { private SharedLog mockSharedLog; @Mock private MdnsServiceTypeClient.Dependencies mockDeps; + @Mock + private TimerFileDescriptor mockTimerFd; @Captor private ArgumentCaptor serviceInfoCaptor; @@ -145,6 +148,7 @@ public class MdnsServiceTypeClientTests { private Message delayMessage = null; private Handler realHandler = null; private MdnsFeatureFlags featureFlags = MdnsFeatureFlags.newBuilder().build(); + private TimerFileDescriptor.MessageTask task = null; @Before @SuppressWarnings("DoNotMock") @@ -244,10 +248,21 @@ public class MdnsServiceTypeClientTests { return true; }).when(mockDeps).sendMessage(any(Handler.class), any(Message.class)); - client = makeMdnsServiceTypeClient(); + doAnswer(inv -> { + realHandler = (Handler) inv.getArguments()[0]; + return mockTimerFd; + }).when(mockDeps).createTimerFd(any(Handler.class)); + + doAnswer(inv -> { + task = (TimerFileDescriptor.MessageTask) inv.getArguments()[0]; + latestDelayMs = (long) inv.getArguments()[1]; + return null; + }).when(mockTimerFd).setDelayedTask(any(), anyLong()); + + client = makeMdnsServiceTypeClient(featureFlags); } - private MdnsServiceTypeClient makeMdnsServiceTypeClient() { + private MdnsServiceTypeClient makeMdnsServiceTypeClient(MdnsFeatureFlags featureFlags) { return new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor, mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps, serviceCache, featureFlags); @@ -1926,9 +1941,7 @@ public class MdnsServiceTypeClientTests { @Test public void testSendQueryWithKnownAnswers() throws Exception { - client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor, - mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps, - serviceCache, + client = makeMdnsServiceTypeClient( MdnsFeatureFlags.newBuilder().setIsQueryWithKnownAnswerEnabled(true).build()); doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket( @@ -1990,9 +2003,7 @@ public class MdnsServiceTypeClientTests { @Test public void testSendQueryWithSubTypeWithKnownAnswers() throws Exception { - client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor, - mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps, - serviceCache, + client = makeMdnsServiceTypeClient( MdnsFeatureFlags.newBuilder().setIsQueryWithKnownAnswerEnabled(true).build()); doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket( @@ -2114,6 +2125,66 @@ public class MdnsServiceTypeClientTests { assertEquals(9680L, latestDelayMs); } + @Test + public void sendQueries_AccurateDelayCallback() { + client = makeMdnsServiceTypeClient( + MdnsFeatureFlags.newBuilder().setIsAccurateDelayCallbackEnabled(true).build()); + + final int numOfQueriesBeforeBackoff = 2; + final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder() + .addSubtype(SUBTYPE) + .setQueryMode(AGGRESSIVE_QUERY_MODE) + .setNumOfQueriesBeforeBackoff(numOfQueriesBeforeBackoff) + .build(); + startSendAndReceive(mockListenerOne, searchOptions); + verify(mockTimerFd, times(1)).cancelTask(); + + // Verify that the first query has been sent. + verifyAndSendQuery(0 /* index */, 0 /* timeInMs */, true /* expectsUnicastResponse */, + true /* multipleSocketDiscovery */, 1 /* scheduledCount */, + 1 /* sendMessageCount */, true /* useAccurateDelayCallback */); + // Verify that the task cancellation occurred before scheduling another query. + verify(mockTimerFd, times(2)).cancelTask(); + + // Verify that the second query has been sent + verifyAndSendQuery(1 /* index */, 0 /* timeInMs */, false /* expectsUnicastResponse */, + true /* multipleSocketDiscovery */, 2 /* scheduledCount */, + 2 /* sendMessageCount */, true /* useAccurateDelayCallback */); + // Verify that the task cancellation occurred before scheduling another query. + verify(mockTimerFd, times(3)).cancelTask(); + + // Verify that the third query has been sent + verifyAndSendQuery(2 /* index */, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS, + false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */, + 3 /* scheduledCount */, 3 /* sendMessageCount */, + true /* useAccurateDelayCallback */); + // Verify that the task cancellation occurred before scheduling another query. + verify(mockTimerFd, times(4)).cancelTask(); + + // In backoff mode, the current scheduled task will be canceled and reschedule if the + // 0.8 * smallestRemainingTtl is larger than time to next run. + long currentTime = TEST_TTL / 2 + TEST_ELAPSED_REALTIME; + doReturn(currentTime).when(mockDecoderClock).elapsedRealtime(); + doReturn(true).when(mockTimerFd).hasDelayedTask(); + processResponse(createResponse( + "service-instance-1", "192.0.2.123", 5353, + SERVICE_TYPE_LABELS, + Collections.emptyMap(), TEST_TTL), socketKey); + // Verify that the task cancellation occurred twice. + verify(mockTimerFd, times(6)).cancelTask(); + assertNotNull(task); + verifyAndSendQuery(3 /* index */, (long) (TEST_TTL / 2 * 0.8) /* timeInMs */, + true /* expectsUnicastResponse */, true /* multipleSocketDiscovery */, + 5 /* scheduledCount */, 4 /* sendMessageCount */, + true /* useAccurateDelayCallback */); + // Verify that the task cancellation occurred before scheduling another query. + verify(mockTimerFd, times(7)).cancelTask(); + + // Stop sending packets. + stopSendAndReceive(mockListenerOne); + verify(mockTimerFd, times(8)).cancelTask(); + } + private static MdnsServiceInfo matchServiceName(String name) { return argThat(info -> info.getServiceInstanceName().equals(name)); } @@ -2127,9 +2198,22 @@ public class MdnsServiceTypeClientTests { private void verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse, boolean multipleSocketDiscovery, int scheduledCount) { - // Dispatch the message - if (delayMessage != null && realHandler != null) { - dispatchMessage(); + verifyAndSendQuery(index, timeInMs, expectsUnicastResponse, + multipleSocketDiscovery, scheduledCount, index + 1 /* sendMessageCount */, + false /* useAccurateDelayCallback */); + } + + private void verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse, + boolean multipleSocketDiscovery, int scheduledCount, int sendMessageCount, + boolean useAccurateDelayCallback) { + if (useAccurateDelayCallback && task != null && realHandler != null) { + runOnHandler(() -> realHandler.dispatchMessage(task.getMessage())); + task = null; + } else { + // Dispatch the message + if (delayMessage != null && realHandler != null) { + dispatchMessage(); + } } assertEquals(timeInMs, latestDelayMs); currentThreadExecutor.getAndClearLastScheduledRunnable().run(); @@ -2152,11 +2236,15 @@ public class MdnsServiceTypeClientTests { eq(socketKey), eq(false)); } } - verify(mockDeps, times(index + 1)) + verify(mockDeps, times(sendMessageCount)) .sendMessage(any(Handler.class), any(Message.class)); // Verify the task has been scheduled. - verify(mockDeps, times(scheduledCount)) - .sendMessageDelayed(any(Handler.class), any(Message.class), anyLong()); + if (useAccurateDelayCallback) { + verify(mockTimerFd, times(scheduledCount)).setDelayedTask(any(), anyLong()); + } else { + verify(mockDeps, times(scheduledCount)) + .sendMessageDelayed(any(Handler.class), any(Message.class), anyLong()); + } } private static String[] getTestServiceName(String instanceName) {