Skip to content

Commit e0cab21

Browse files
committed
Graceful connection close without sending a QUIT command first
1 parent 4ea66c4 commit e0cab21

File tree

2 files changed

+30
-52
lines changed

2 files changed

+30
-52
lines changed

Sources/RediStack/RedisConnection.swift

Lines changed: 14 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -184,12 +184,17 @@ public final class RedisConnection: RedisClient {
184184
self.channel.closeFuture.whenSuccess {
185185
// if our state is still open, that means we didn't cause the closeFuture to resolve.
186186
// update state, metrics, and logging
187-
guard self.state.isConnected else { return }
188-
187+
let oldState = self.state
189188
self.state = .closed
190-
self.defaultLogger.warning("connection was closed unexpectedly")
191189
RedisMetrics.activeConnectionCount.decrement()
192-
self.onUnexpectedClosure?()
190+
191+
switch oldState {
192+
case .shuttingDown, .closed:
193+
break
194+
case .open, .pubsub:
195+
self.defaultLogger.warning("connection was closed unexpectedly")
196+
self.onUnexpectedClosure?()
197+
}
193198
}
194199

195200
self.defaultLogger.trace("connection created")
@@ -300,65 +305,23 @@ extension RedisConnection {
300305

301306
// we're now in a shutdown state, starting with the command queue.
302307
self.state = .shuttingDown
303-
304-
let notification = self.sendQuitCommand(logger: logger) // send "QUIT" so that all the responses are written out
305-
.flatMap { self.closeChannel() } // close the channel from our end
306-
.hop(to: finalEventLoop)
308+
309+
// Inform ChannelHandler about close intent using "RedisGracefulConnectionCloseEvent"
310+
let promise = finalEventLoop.makePromise(of: Void.self)
311+
let notification = promise.futureResult
312+
self.channel.triggerUserOutboundEvent(RedisGracefulConnectionCloseEvent(), promise: promise)
307313

308314
notification.whenFailure {
309315
logger.warning("failed to close connection", metadata: [
310316
RedisLogging.MetadataKeys.error: "\($0)"
311317
])
312318
}
313319
notification.whenSuccess {
314-
self.state = .closed
315320
logger.trace("connection is now closed")
316-
RedisMetrics.activeConnectionCount.decrement()
317321
}
318322

319323
return notification
320324
}
321-
322-
/// Bypasses everything for a normal command and explicitly just sends a "QUIT" command to Redis.
323-
/// - Note: If the command fails, the `NIO.EventLoopFuture` will still succeed - as it's not critical for the command to succeed.
324-
private func sendQuitCommand(logger: Logger) -> EventLoopFuture<Void> {
325-
let payload: RedisCommandHandler.OutboundCommandPayload = (
326-
RedisCommand<Void>(keyword: "QUIT", arguments: []).serialized(),
327-
self.eventLoop.makePromise()
328-
)
329-
330-
logger.trace("sending QUIT command")
331-
332-
return self.channel
333-
.writeAndFlush(payload) // write the command
334-
.flatMap { payload.responsePromise.futureResult } // chain the callback to the response's
335-
.map { _ in logger.trace("sent QUIT command") } // ignore the result's value
336-
.recover { _ in logger.debug("recovered from error sending QUIT") } // if there's an error, just return to void
337-
}
338-
339-
/// Attempts to close the `NIO.Channel`.
340-
/// SwiftNIO throws a `NIO.EventLoopError.shutdown` if the channel is already closed,
341-
/// so that case is captured to let this method's `NIO.EventLoopFuture` still succeed.
342-
private func closeChannel() -> EventLoopFuture<Void> {
343-
let promise = self.channel.eventLoop.makePromise(of: Void.self)
344-
345-
self.channel.close(promise: promise)
346-
347-
// if we succeed, great, if not - check the error that happened
348-
return promise.futureResult
349-
.flatMapError { error in
350-
guard let e = error as? EventLoopError else {
351-
return self.eventLoop.makeFailedFuture(error)
352-
}
353-
354-
// if the error is that the channel is already closed, great - just succeed.
355-
// otherwise, fail the chain
356-
switch e {
357-
case .shutdown: return self.eventLoop.makeSucceededFuture(())
358-
default: return self.eventLoop.makeFailedFuture(e)
359-
}
360-
}
361-
}
362325
}
363326

364327
// MARK: Logging

Sources/RediStackTestUtils/EmbeddedMockRedisServer.swift

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ internal final class EmbeddedMockRedisServer {
6868
}
6969

7070
func createConnectedChannel() -> Channel {
71-
let channel = EmbeddedChannel(loop: self.loop)
71+
let channel = EmbeddedChannel(handler: GracefulShutdownToCloseHandler(), loop: self.loop)
7272
channel.closeFuture.whenComplete { _ in
7373
self.channels.removeAll(where: { $0 === channel })
7474
}
@@ -84,3 +84,18 @@ internal final class EmbeddedMockRedisServer {
8484
try self.loop.close()
8585
}
8686
}
87+
88+
/// A `ChannelHandler` that triggers a channel close once `RedisGracefulConnectionCloseEvent` is received
89+
private final class GracefulShutdownToCloseHandler: ChannelHandler, ChannelOutboundHandler {
90+
typealias OutboundIn = NIOAny
91+
92+
func triggerUserOutboundEvent(context: ChannelHandlerContext, event: Any, promise: EventLoopPromise<Void>?) {
93+
switch event {
94+
case is RedisGracefulConnectionCloseEvent:
95+
context.close(mode: .all, promise: promise)
96+
97+
default:
98+
context.triggerUserOutboundEvent(event, promise: promise)
99+
}
100+
}
101+
}

0 commit comments

Comments
 (0)