diff --git a/Sources/NIOPerformanceTester/TCPThroughputBenchmark.swift b/Sources/NIOPerformanceTester/TCPThroughputBenchmark.swift index 6318a40413..6886267df7 100644 --- 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 private var context: ChannelHandlerContext! + init(_ connectionEstablishedPromise: EventLoopPromise) { + 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 = self.group.next().makePromise() + let connectionEstablishedPromise: EventLoopPromise = 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() {