Fix race in TCPThroughputBenchmark (#2724)

apple · May 20, 2024 · 196ceab · 196ceab
1 parent 609469f
commit 196ceab
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/Sources/NIOPerformanceTester/TCPThroughputBenchmark.swift b/Sources/NIOPerformanceTester/TCPThroughputBenchmark.swift
@@ -38,10 +38,16 @@ final class TCPThroughputBenchmark: Benchmark {
         public typealias InboundIn = ByteBuffer
         public typealias OutboundOut = ByteBuffer
 
+        private let connectionEstablishedPromise: EventLoopPromise<EventLoop>
         private var context: ChannelHandlerContext!
 
+        init(_ connectionEstablishedPromise: EventLoopPromise<EventLoop>) {
+            self.connectionEstablishedPromise = connectionEstablishedPromise
+        }
+
         public func channelActive(context: ChannelHandlerContext) {
             self.context = context
+            connectionEstablishedPromise.succeed(context.eventLoop)
         }
 
         public func send(_ message: ByteBuffer, times count: Int) {
@@ -106,12 +112,11 @@ final class TCPThroughputBenchmark: Benchmark {
     func setUp() throws {
         self.group = MultiThreadedEventLoopGroup(numberOfThreads: 4)
 
-        let connectionEstablished: EventLoopPromise<EventLoop> = self.group.next().makePromise()
+        let connectionEstablishedPromise: EventLoopPromise<EventLoop> = self.group.next().makePromise()
 
         self.serverChannel = try ServerBootstrap(group: self.group)
             .childChannelInitializer { channel in
-                self.serverHandler = ServerHandler()
-                connectionEstablished.succeed(channel.eventLoop)
+                self.serverHandler = ServerHandler(connectionEstablishedPromise)
                 return channel.pipeline.addHandler(self.serverHandler)
             }
             .bind(host: "127.0.0.1", port: 0)
@@ -134,7 +139,7 @@ final class TCPThroughputBenchmark: Benchmark {
         }
         self.message = message
 
-        self.serverEventLoop = try connectionEstablished.futureResult.wait()
+        self.serverEventLoop = try connectionEstablishedPromise.futureResult.wait()
     }
 
     func tearDown() {