From 1a3229b9fc43fb5a5b6ffcfd100fadc06f0df9dd Mon Sep 17 00:00:00 2001
From: Cory Benfield <lukasa@apple.com>
Date: Mon, 16 Dec 2024 17:53:38 +0000
Subject: [PATCH] Fix now-racy scheduled callback tests (#3031)

Motivation:

In PR 3026 we I fixed some warnings, but as part of doing so I
eliminated some complexity in the scheduled callback tests. Apparently I
did that a bit hard, as I lost a few synchronization edges that made the
tests lightly flaky.

Modifications:

- Re-add synchronization edges between the EL and the test when we
advance time
- Re-add synchronization edges after cancellation calls.

Result:

More reliable tests
---
 .../NIOScheduledCallbackTests.swift              | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/Tests/NIOPosixTests/NIOScheduledCallbackTests.swift b/Tests/NIOPosixTests/NIOScheduledCallbackTests.swift
index 9940b74e12..16dc984ac2 100644
--- a/Tests/NIOPosixTests/NIOScheduledCallbackTests.swift
+++ b/Tests/NIOPosixTests/NIOScheduledCallbackTests.swift
@@ -27,6 +27,8 @@ protocol ScheduledCallbackTestRequirements {
 
     // ELG-backed ELs need to be shutdown via the ELG.
     func shutdownEventLoop() async throws
+
+    func waitForLoopTick() async throws
 }
 
 final class MTELGScheduledCallbackTests: _BaseScheduledCallbackTests {
@@ -41,6 +43,10 @@ final class MTELGScheduledCallbackTests: _BaseScheduledCallbackTests {
         func shutdownEventLoop() async throws {
             try await self.group.shutdownGracefully()
         }
+
+        func waitForLoopTick() async throws {
+            try await self.loop.submit {}.get()
+        }
     }
 
     override func setUp() async throws {
@@ -60,6 +66,10 @@ final class NIOAsyncTestingEventLoopScheduledCallbackTests: _BaseScheduledCallba
         func shutdownEventLoop() async throws {
             await self._loop.shutdownGracefully()
         }
+
+        func waitForLoopTick() async throws {
+            try await self._loop.executeInContext {}
+        }
     }
 
     override func setUp() async throws {
@@ -83,6 +93,7 @@ extension _BaseScheduledCallbackTests {
 
     func advanceTime(by amount: TimeAmount) async throws {
         try await self.requirements.advanceTime(by: amount)
+        try await self.requirements.waitForLoopTick()
     }
 
     func shutdownEventLoop() async throws {
@@ -151,6 +162,7 @@ extension _BaseScheduledCallbackTests {
 
         let scheduledCallback = try self.loop.scheduleCallback(in: .milliseconds(1), handler: handler)
         scheduledCallback.cancel()
+        try await self.requirements.waitForLoopTick()
         handler.assert(callbackCount: 0, cancelCount: 1)
     }
 
@@ -161,6 +173,7 @@ extension _BaseScheduledCallbackTests {
         try await self.advanceTime(by: .milliseconds(1))
         try await handler.waitForCallback(timeout: .seconds(1))
         scheduledCallback.cancel()
+        try await self.requirements.waitForLoopTick()
         handler.assert(callbackCount: 1, cancelCount: 0)
     }
 
@@ -169,7 +182,9 @@ extension _BaseScheduledCallbackTests {
 
         let scheduledCallback = try self.loop.scheduleCallback(in: .milliseconds(1), handler: handler)
         scheduledCallback.cancel()
+        try await self.requirements.waitForLoopTick()
         scheduledCallback.cancel()
+        try await self.requirements.waitForLoopTick()
         handler.assert(callbackCount: 0, cancelCount: 1)
     }
 
@@ -197,6 +212,7 @@ extension _BaseScheduledCallbackTests {
 
         let handle = try self.loop.scheduleCallback(in: .milliseconds(1), handler: handler)
         handle.cancel()
+        try await self.requirements.waitForLoopTick()
         handler.assert(callbackCount: 0, cancelCount: 1)
 
         try await self.shutdownEventLoop()