Skip to content

Commit

Permalink
Ensure kafka output flushes queue on shutdown (#679)
Browse files Browse the repository at this point in the history
* fix timeout settings in kafka output
* add tests for shutdown behavior
* Update CHANGELOG.md
  • Loading branch information
ekneg54 authored Oct 2, 2024
1 parent b0d8752 commit 78d4b24
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 8 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
* add manual how to use local images with minikube example setup to documentation
* move `Configuration` to top level of documentation
* add `CONTRIBUTING` file
* sets the default for `flush_timeout` and `send_timeout` in `kafka_output` connector to `0` seconds

### Bugfix

Expand Down
41 changes: 33 additions & 8 deletions logprep/connector/confluent_kafka/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,10 +148,22 @@ class Config(Output.Config):

topic: str = field(validator=validators.instance_of(str))
"""The topic into which the processed events should be written to."""
error_topic: str
"""The topic into which events should be written that couldn't be processed successfully."""
flush_timeout: float
send_timeout: int = field(validator=validators.instance_of(int), default=0)
error_topic: str = field(validator=validators.instance_of(str))
"""The topic into which the failed events should be written to."""
flush_timeout: float = field(
validator=validators.instance_of(float), converter=float, default=0
)
"""The maximum time in seconds to wait for the producer to flush the messages
to kafka broker. If the buffer is full, the producer will block until the buffer
is empty or the timeout is reached. This implies that the producer does not
wait for all messages to be send to the broker, if the timeout is reached
before the buffer is empty. Default is :code:`0`.
"""
send_timeout: float = field(
validator=validators.instance_of(float), converter=float, default=0
)
"""The maximum time in seconds to wait for an answer from the broker on polling.
Default is :code:`0`."""
kafka_config: Optional[MappingProxyType] = field(
validator=[
validators.instance_of(MappingProxyType),
Expand Down Expand Up @@ -286,7 +298,7 @@ def store_custom(self, document: dict, target: str) -> None:
self._producer.poll(self._config.send_timeout)
self.metrics.number_of_processed_events += 1
except BufferError:
# block program until buffer is empty
# block program until buffer is empty or timeout is reached
self._producer.flush(timeout=self._config.flush_timeout)
except BaseException as error:
raise CriticalOutputError(
Expand Down Expand Up @@ -327,9 +339,22 @@ def store_failed(
self._producer.flush(timeout=self._config.flush_timeout)

def shut_down(self) -> None:
"""ensures that all messages are flushed"""
if self._producer is not None:
self._producer.flush(self._config.flush_timeout)
"""ensures that all messages are flushed. According to
https://confluent-kafka-python.readthedocs.io/en/latest/#confluent_kafka.Producer.flush
flush without the timeout parameter will block until all messages are delivered.
This ensures no messages will get lost on shutdown.
"""
if self._producer is None:
return
remaining_messages = self._producer.flush()
if remaining_messages:
self.metrics.number_of_errors += 1
logger.error(
"Flushing producer timed out. %s messages are still in the buffer.",
remaining_messages,
)
else:
logger.info("Producer flushed successfully. %s messages remaining.", remaining_messages)

def health(self) -> bool:
"""Check the health of kafka producer."""
Expand Down
9 changes: 9 additions & 0 deletions tests/unit/connector/test_confluent_kafka_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,3 +195,12 @@ def test_health_counts_metrics_on_kafka_exception(self):
self.object._producer.list_topics.side_effect = KafkaException("test error")
assert not self.object.health()
assert self.object.metrics.number_of_errors == 1

def test_shutdown_logs_and_counts_error_if_queue_not_fully_flushed(self):
self.object.metrics.number_of_errors = 0
self.object._producer = mock.MagicMock()
self.object._producer.flush.return_value = 1
with mock.patch("logging.Logger.error") as mock_error:
self.object.shut_down()
mock_error.assert_called()
self.object.metrics.number_of_errors = 1

0 comments on commit 78d4b24

Please sign in to comment.