Skip to content

Commit

Permalink
fix connection leak (#17)
Browse files Browse the repository at this point in the history
motivation: recent PR to address NIO atomic changes has caused a conection leak

changes:
* implement a state machine to cache and manage connection state
* add connection count test
* refactor tests
* safer concurrency
  • Loading branch information
tomerd authored Mar 9, 2020
1 parent 9525518 commit 0f8c911
Show file tree
Hide file tree
Showing 3 changed files with 177 additions and 188 deletions.
43 changes: 41 additions & 2 deletions Sources/StatsdClient/StatsdClient.swift
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
//===----------------------------------------------------------------------===//

import CoreMetrics
import Dispatch
import NIO
import NIOConcurrencyHelpers

Expand Down Expand Up @@ -245,6 +246,15 @@ private final class Client {

private let isShutdown = NIOAtomic<Bool>.makeAtomic(value: false)

private var state = State.disconnected
private let lock = Lock()

private enum State {
case disconnected
case connecting(EventLoopFuture<Void>)
case connected(Channel)
}

init(eventLoopGroupProvider: StatsdClient.EventLoopGroupProvider, address: SocketAddress) {
self.eventLoopGroupProvider = eventLoopGroupProvider
switch self.eventLoopGroupProvider {
Expand Down Expand Up @@ -273,8 +283,37 @@ private final class Client {
}

func emit(_ metric: Metric) -> EventLoopFuture<Void> {
return self.connect().flatMap { channel in
channel.writeAndFlush(metric)
self.lock.lock()
switch self.state {
case .disconnected:
let promise = self.eventLoopGroup.next().makePromise(of: Void.self)
self.state = .connecting(promise.futureResult)
self.lock.unlock()
self.connect().flatMap { channel -> EventLoopFuture<Void> in
self.lock.withLock {
guard case .connecting = self.state else {
preconditionFailure("invalid state \(self.state)")
}
self.state = .connected(channel)
}
return self.emit(metric)
}.cascade(to: promise)
return promise.futureResult
case .connecting(let future):
let future = future.flatMap {
self.emit(metric)
}
self.state = .connecting(future)
self.lock.unlock()
return future
case .connected(let channel):
guard channel.isActive else {
self.state = .disconnected
self.lock.unlock()
return self.emit(metric)
}
self.lock.unlock()
return channel.writeAndFlush(metric)
}
}

Expand Down
1 change: 1 addition & 0 deletions Tests/StatsdClientTests/StatsdClientTests+XCTest.swift
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ extension StatsdClientTests {
("testRecorderInteger", testRecorderInteger),
("testRecorderDouble", testRecorderDouble),
("testCouncurrency", testCouncurrency),
("testNumberOfConnections", testNumberOfConnections),
]
}
}
Loading

0 comments on commit 0f8c911

Please sign in to comment.