@@ -201,16 +201,20 @@ public final class KafkaConsumer {
201
201
switch nextAction {
202
202
case . pollForAndYieldMessage( let client, let source) :
203
203
do {
204
- guard let message = try client. consumerPoll ( ) else {
205
- break
204
+ if let message = try client. consumerPoll ( ) {
205
+ // We do not support back pressure, we can ignore the yield result
206
+ _ = source. yield ( message)
206
207
}
207
- // We do not support back pressure, we can ignore the yield result
208
- _ = source. yield ( message)
209
208
} catch {
210
209
source. finish ( )
211
210
throw error
212
211
}
213
212
try await Task . sleep ( for: self . config. pollInterval)
213
+ case . pollUntilClosed( let client) :
214
+ // Ignore poll result, we are closing down and just polling to commit
215
+ // outstanding consumer state
216
+ _ = try client. consumerPoll ( )
217
+ try await Task . sleep ( for: self . config. pollInterval)
214
218
case . terminatePollLoop:
215
219
return
216
220
}
@@ -304,7 +308,9 @@ extension KafkaConsumer {
304
308
source: Producer . Source
305
309
)
306
310
/// The ``KafkaConsumer`` has been closed.
307
- case finished
311
+ ///
312
+ /// - Parameter client: Client used for handling the connection to the Kafka cluster.
313
+ case finished( client: KafkaClient )
308
314
}
309
315
310
316
/// The current state of the StateMachine.
@@ -335,6 +341,11 @@ extension KafkaConsumer {
335
341
client: KafkaClient ,
336
342
source: Producer . Source
337
343
)
344
+ /// The ``KafkaConsumer`` is in the process of closing down, but still needs to poll
345
+ /// to commit its state to the broker.
346
+ ///
347
+ /// - Parameter client: Client used for handling the connection to the Kafka cluster.
348
+ case pollUntilClosed( client: KafkaClient )
338
349
/// Terminate the poll loop.
339
350
case terminatePollLoop
340
351
}
@@ -351,8 +362,12 @@ extension KafkaConsumer {
351
362
fatalError ( " Subscribe to consumer group / assign to topic partition pair before reading messages " )
352
363
case . consuming( let client, let source) :
353
364
return . pollForAndYieldMessage( client: client, source: source)
354
- case . finished:
355
- return . terminatePollLoop
365
+ case . finished( let client) :
366
+ if client. isConsumerClosed {
367
+ return . terminatePollLoop
368
+ } else {
369
+ return . pollUntilClosed( client: client)
370
+ }
356
371
}
357
372
}
358
373
@@ -433,7 +448,7 @@ extension KafkaConsumer {
433
448
case . initializing:
434
449
fatalError ( " subscribe() / assign() should have been invoked before \( #function) " )
435
450
case . consuming( let client, let source) :
436
- self . state = . finished
451
+ self . state = . finished( client : client )
437
452
return . shutdownGracefullyAndFinishSource(
438
453
client: client,
439
454
source: source
0 commit comments