diff --git a/config.yaml b/config.yaml index c47dfbd7..8bd8ae20 100644 --- a/config.yaml +++ b/config.yaml @@ -25,6 +25,14 @@ pipeline: backend: auto batch_size: auto + resilience: + retry: + initial_delay_seconds: 1.0 + max_delay_seconds: 30.0 + backoff_multiplier: 2.0 + jitter_seconds: 0.25 + log_every_attempts: 5 + scaling: defaults: executor: thread diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 6d7b8d59..3d409010 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -177,7 +177,7 @@ services: condition: service_healthy kafka3: condition: service_healthy - profiles: ["dev"] + profiles: ["dev", "dev-gpu"] inspector: extends: @@ -190,71 +190,7 @@ services: condition: service_healthy kafka3: condition: service_healthy - profiles: ["prod"] - - inspector-dev-gpu: - extends: - file: "docker-compose/dev/docker-compose.pipeline.yml" - service: inspector - depends_on: - kafka1: - condition: service_healthy - kafka2: - condition: service_healthy - kafka3: - condition: service_healthy - deploy: - mode: "replicated" - replicas: 1 - resources: - limits: - cpus: '2' - memory: 512m - reservations: - cpus: '1' - memory: 256m - devices: - - driver: nvidia - count: all - capabilities: [gpu] - environment: - - GROUP_ID=data_inspection - - NUMBER_OF_INSTANCES=1 - - NVIDIA_VISIBLE_DEVICES=all - - NVIDIA_DRIVER_CAPABILITIES=compute,utility - profiles: ["dev-gpu"] - - inspector-gpu: - extends: - file: "docker-compose/prod/docker-compose.pipeline.yml" - service: inspector - depends_on: - kafka1: - condition: service_healthy - kafka2: - condition: service_healthy - kafka3: - condition: service_healthy - deploy: - mode: "replicated" - replicas: 1 - resources: - limits: - cpus: '2' - memory: 512m - reservations: - cpus: '1' - memory: 256m - devices: - - driver: nvidia - count: all - capabilities: [gpu] - environment: - - GROUP_ID=data_inspection - - NUMBER_OF_INSTANCES=1 - - NVIDIA_VISIBLE_DEVICES=all - - NVIDIA_DRIVER_CAPABILITIES=compute,utility - profiles: ["prod-gpu"] + profiles: ["prod", "prod-gpu"] detector-dev: @@ -270,10 +206,10 @@ services: condition: service_healthy profiles: ["dev"] - detector: + detector-dev-gpu: extends: - file: "docker-compose/prod/docker-compose.pipeline.yml" - service: detector + file: "docker-compose/dev/docker-compose.pipeline.yml" + service: detector-gpu depends_on: kafka1: condition: service_healthy @@ -281,11 +217,11 @@ services: condition: service_healthy kafka3: condition: service_healthy - profiles: ["prod"] + profiles: ["dev-gpu"] - detector-dev-gpu: + detector: extends: - file: "docker-compose/dev/docker-compose.pipeline.yml" + file: "docker-compose/prod/docker-compose.pipeline.yml" service: detector depends_on: kafka1: @@ -294,25 +230,12 @@ services: condition: service_healthy kafka3: condition: service_healthy - deploy: - mode: "replicated" - replicas: 1 - resources: - reservations: - devices: - - driver: nvidia - count: all - capabilities: [gpu] - environment: - - GROUP_ID=data_analysis - - NVIDIA_VISIBLE_DEVICES=all - - NVIDIA_DRIVER_CAPABILITIES=compute,utility - profiles: ["dev-gpu"] + profiles: ["prod"] detector-gpu: extends: file: "docker-compose/prod/docker-compose.pipeline.yml" - service: detector + service: detector-gpu depends_on: kafka1: condition: service_healthy @@ -320,19 +243,6 @@ services: condition: service_healthy kafka3: condition: service_healthy - deploy: - mode: "replicated" - replicas: 1 - resources: - reservations: - devices: - - driver: nvidia - count: all - capabilities: [gpu] - environment: - - GROUP_ID=data_analysis - - NVIDIA_VISIBLE_DEVICES=all - - NVIDIA_DRIVER_CAPABILITIES=compute,utility profiles: ["prod-gpu"] alerter-dev: diff --git a/docker/docker-compose/base/docker-compose.kafka.yml b/docker/docker-compose/base/docker-compose.kafka.yml index f917072c..5b0bd59a 100644 --- a/docker/docker-compose/base/docker-compose.kafka.yml +++ b/docker/docker-compose/base/docker-compose.kafka.yml @@ -47,6 +47,10 @@ services: KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 3 KAFKA_AUTHORIZER_CLASS_NAME: kafka.security.authorizer.AclAuthorizer KAFKA_ALLOW_EVERYONE_IF_NO_ACL_FOUND: "true" + KAFKA_LOG_RETENTION_HOURS: 4 + KAFKA_LOG_RETENTION_BYTES: 10737418240 + KAFKA_LOG_SEGMENT_BYTES: 1073741824 + KAFKA_LOG_CLEANUP_POLICY: delete volumes: - kafka-data1:/var/lib/kafka/data @@ -80,6 +84,10 @@ services: KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 3 KAFKA_AUTHORIZER_CLASS_NAME: kafka.security.authorizer.AclAuthorizer KAFKA_ALLOW_EVERYONE_IF_NO_ACL_FOUND: "true" + KAFKA_LOG_RETENTION_HOURS: 4 + KAFKA_LOG_RETENTION_BYTES: 10737418240 + KAFKA_LOG_SEGMENT_BYTES: 1073741824 + KAFKA_LOG_CLEANUP_POLICY: delete volumes: - kafka-data2:/var/lib/kafka/data @@ -113,6 +121,10 @@ services: KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 3 KAFKA_AUTHORIZER_CLASS_NAME: kafka.security.authorizer.AclAuthorizer KAFKA_ALLOW_EVERYONE_IF_NO_ACL_FOUND: "true" + KAFKA_LOG_RETENTION_HOURS: 4 + KAFKA_LOG_RETENTION_BYTES: 10737418240 + KAFKA_LOG_SEGMENT_BYTES: 1073741824 + KAFKA_LOG_CLEANUP_POLICY: delete volumes: - kafka-data3:/var/lib/kafka/data diff --git a/docker/docker-compose/dev/docker-compose.pipeline.yml b/docker/docker-compose/dev/docker-compose.pipeline.yml index eb67a7d2..946789ce 100644 --- a/docker/docker-compose/dev/docker-compose.pipeline.yml +++ b/docker/docker-compose/dev/docker-compose.pipeline.yml @@ -80,6 +80,29 @@ services: environment: - GROUP_ID=data_analysis + detector-gpu: + build: + context: ../../.. + dockerfile: docker/dockerfiles/Dockerfile.detector + restart: "unless-stopped" + volumes: + - ../../../config.yaml:/app/config.yaml + networks: + hamstring: + deploy: + mode: "replicated" + replicas: 1 + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + environment: + - GROUP_ID=data_analysis + - NVIDIA_VISIBLE_DEVICES=all + - NVIDIA_DRIVER_CAPABILITIES=compute,utility + alerter: build: context: ../../.. diff --git a/docker/docker-compose/prod/docker-compose.monitoring.yml b/docker/docker-compose/prod/docker-compose.monitoring.yml index 9d3a61d3..7dbfa0d9 100644 --- a/docker/docker-compose/prod/docker-compose.monitoring.yml +++ b/docker/docker-compose/prod/docker-compose.monitoring.yml @@ -1,6 +1,6 @@ services: monitoring_agent: - image: ghcr.io/astraos-de/hamstring-monitoring:v2.0.0 + image: ghcr.io/astraos-de/hamstring-monitoring:v2.2.0-dev restart: "unless-stopped" volumes: - ../../../config.yaml:/app/config.yaml diff --git a/docker/docker-compose/prod/docker-compose.pipeline.yml b/docker/docker-compose/prod/docker-compose.pipeline.yml index 8019b1c4..2ea6803c 100644 --- a/docker/docker-compose/prod/docker-compose.pipeline.yml +++ b/docker/docker-compose/prod/docker-compose.pipeline.yml @@ -1,6 +1,6 @@ services: logserver: - image: ghcr.io/astraos-de/hamstring-logserver:v2.0.0 + image: ghcr.io/astraos-de/hamstring-logserver:v2.2.0-dev restart: "unless-stopped" networks: hamstring: @@ -10,7 +10,7 @@ services: environment: - GROUP_ID=log_storage logcollector: - image: ghcr.io/astraos-de/hamstring-logcollector:v2.0.0 + image: ghcr.io/astraos-de/hamstring-logcollector:v2.2.0-dev restart: "unless-stopped" volumes: - ../../../config.yaml:/app/config.yaml @@ -20,7 +20,7 @@ services: - GROUP_ID=log_collection prefilter: - image: ghcr.io/astraos-de/hamstring-prefilter:v2.0.0 + image: ghcr.io/astraos-de/hamstring-prefilter:v2.2.0-dev restart: "unless-stopped" volumes: - ../../../config.yaml:/app/config.yaml @@ -33,7 +33,7 @@ services: - GROUP_ID=log_filtering inspector: - image: ghcr.io/astraos-de/hamstring-inspector:v2.0.0 + image: ghcr.io/astraos-de/hamstring-inspector:v2.2.0-dev restart: "unless-stopped" volumes: - ../../../config.yaml:/app/config.yaml @@ -54,7 +54,7 @@ services: - NUMBER_OF_INSTANCES=1 detector: - image: ghcr.io/astraos-de/hamstring-detector:v2.0.0 + image: ghcr.io/astraos-de/hamstring-detector:v2.2.0-dev restart: "unless-stopped" volumes: - ../../../config.yaml:/app/config.yaml @@ -66,8 +66,29 @@ services: environment: - GROUP_ID=data_analysis + detector-gpu: + image: ghcr.io/astraos-de/hamstring-detector:v2.2.0-dev + restart: "unless-stopped" + volumes: + - ../../../config.yaml:/app/config.yaml + networks: + hamstring: + deploy: + mode: "replicated" + replicas: 1 + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + environment: + - GROUP_ID=data_analysis + - NVIDIA_VISIBLE_DEVICES=all + - NVIDIA_DRIVER_CAPABILITIES=compute,utility + alerter: - image: ghcr.io/astraos-de/hamstring-alerter:v2.0.0 + image: ghcr.io/astraos-de/hamstring-alerter:v2.2.0-dev restart: "unless-stopped" volumes: - ../../../config.yaml:/app/config.yaml diff --git a/docker/docker_swarm/docker-compose.swarm.yml b/docker/docker_swarm/docker-compose.swarm.yml new file mode 100644 index 00000000..52edefe4 --- /dev/null +++ b/docker/docker_swarm/docker-compose.swarm.yml @@ -0,0 +1,394 @@ +# Compose file to deploy directly with docker stack deploy. +--- +version: "3.8" + +services: + zookeeper: + image: confluentinc/cp-zookeeper:7.9.3 + networks: + - hamstring + environment: + ZOOKEEPER_CLIENT_PORT: "2181" + ZOOKEEPER_TICK_TIME: "2000" + ZOOKEEPER_SERVER_ID: "1" + volumes: + - zk-data:/var/lib/zookeeper/data + - zk-txn-logs:/var/lib/zookeeper/log + healthcheck: + test: ["CMD-SHELL", "curl -s localhost:8080/commands | grep ruok"] + interval: 10s + timeout: 5s + retries: 3 + deploy: + placement: + constraints: + - "node.hostname == workstation" + restart_policy: + condition: any + + kafka1: + image: confluentinc/cp-kafka:7.9.3 + networks: + - hamstring + ports: + - target: 8097 + published: 8097 + protocol: tcp + mode: ingress + environment: + KAFKA_BROKER_ID: "1" + KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181 + KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: INTERNAL:PLAINTEXT,EXTERNAL:PLAINTEXT + KAFKA_ADVERTISED_LISTENERS: INTERNAL://kafka1:19092,EXTERNAL://workstation:8097 + KAFKA_LISTENERS: INTERNAL://0.0.0.0:19092,EXTERNAL://0.0.0.0:8097 + KAFKA_INTER_BROKER_LISTENER_NAME: INTERNAL + KAFKA_AUTO_CREATE_TOPICS_ENABLE: "false" + KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR: "3" + KAFKA_TRANSACTION_STATE_LOG_MIN_ISR: "2" + KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: "3" + KAFKA_AUTHORIZER_CLASS_NAME: kafka.security.authorizer.AclAuthorizer + KAFKA_ALLOW_EVERYONE_IF_NO_ACL_FOUND: "true" + KAFKA_LOG_RETENTION_HOURS: 4 + KAFKA_LOG_RETENTION_BYTES: 10737418240 + KAFKA_LOG_SEGMENT_BYTES: 1073741824 + KAFKA_LOG_CLEANUP_POLICY: delete + volumes: + - kafka-data1:/var/lib/kafka/data + healthcheck: + test: ["CMD-SHELL", "nc -z localhost 19092"] + interval: 30s + timeout: 10s + retries: 10 + deploy: + placement: + constraints: + - "node.hostname == workstation" + restart_policy: + condition: any + + kafka2: + image: confluentinc/cp-kafka:7.9.3 + networks: + - hamstring + ports: + - target: 8098 + published: 8098 + protocol: tcp + mode: ingress + environment: + KAFKA_BROKER_ID: "2" + KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181 + KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: INTERNAL:PLAINTEXT,EXTERNAL:PLAINTEXT + KAFKA_ADVERTISED_LISTENERS: INTERNAL://kafka2:19093,EXTERNAL://workstation:8098 + KAFKA_LISTENERS: INTERNAL://0.0.0.0:19093,EXTERNAL://0.0.0.0:8098 + KAFKA_INTER_BROKER_LISTENER_NAME: INTERNAL + KAFKA_AUTO_CREATE_TOPICS_ENABLE: "false" + KAFKA_LOG4J_LOGGERS: kafka.controller=INFO,kafka.producer.async.DefaultEventHandler=INFO,state.change.logger=INFO + KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR: "3" + KAFKA_TRANSACTION_STATE_LOG_MIN_ISR: "2" + KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: "3" + KAFKA_AUTHORIZER_CLASS_NAME: kafka.security.authorizer.AclAuthorizer + KAFKA_ALLOW_EVERYONE_IF_NO_ACL_FOUND: "true" + KAFKA_LOG_RETENTION_HOURS: 4 + KAFKA_LOG_RETENTION_BYTES: 10737418240 + KAFKA_LOG_SEGMENT_BYTES: 1073741824 + KAFKA_LOG_CLEANUP_POLICY: delete + volumes: + - kafka-data2:/var/lib/kafka/data + healthcheck: + test: ["CMD-SHELL", "nc -z localhost 19093"] + interval: 30s + timeout: 10s + retries: 10 + deploy: + placement: + constraints: + - "node.hostname == workstation" + restart_policy: + condition: any + + kafka3: + image: confluentinc/cp-kafka:7.9.3 + networks: + - hamstring + ports: + - target: 8099 + published: 8099 + protocol: tcp + mode: ingress + environment: + KAFKA_BROKER_ID: "3" + KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181 + KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: INTERNAL:PLAINTEXT,EXTERNAL:PLAINTEXT + KAFKA_ADVERTISED_LISTENERS: INTERNAL://kafka3:19094,EXTERNAL://workstation:8099 + KAFKA_LISTENERS: INTERNAL://0.0.0.0:19094,EXTERNAL://0.0.0.0:8099 + KAFKA_INTER_BROKER_LISTENER_NAME: INTERNAL + KAFKA_AUTO_CREATE_TOPICS_ENABLE: "false" + KAFKA_LOG4J_LOGGERS: kafka.controller=INFO,kafka.producer.async.DefaultEventHandler=INFO,state.change.logger=INFO + KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR: "3" + KAFKA_TRANSACTION_STATE_LOG_MIN_ISR: "2" + KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: "3" + KAFKA_AUTHORIZER_CLASS_NAME: kafka.security.authorizer.AclAuthorizer + KAFKA_ALLOW_EVERYONE_IF_NO_ACL_FOUND: "true" + KAFKA_LOG_RETENTION_HOURS: 4 + KAFKA_LOG_RETENTION_BYTES: 10737418240 + KAFKA_LOG_SEGMENT_BYTES: 1073741824 + KAFKA_LOG_CLEANUP_POLICY: delete + volumes: + - kafka-data3:/var/lib/kafka/data + healthcheck: + test: ["CMD-SHELL", "nc -z localhost 19094"] + interval: 30s + timeout: 10s + retries: 10 + deploy: + placement: + constraints: + - "node.hostname == workstation" + restart_policy: + condition: any + + clickhouse-server: + image: clickhouse/clickhouse-server:24.3.12.75-alpine + networks: + - hamstring + ports: + - target: 8123 + published: 8123 + protocol: tcp + mode: ingress + - target: 9000 + published: 9000 + protocol: tcp + mode: ingress + volumes: + - "/home/max/HAMSTRING/hamstring/docker/create_tables:/docker-entrypoint-initdb.d" + - ch_data:/var/lib/clickhouse/ + - ch_logs:/var/log/clickhouse-server/ + healthcheck: + test: ["CMD-SHELL", "nc -z 127.0.0.1 8123"] + interval: 10s + timeout: 5s + retries: 3 + deploy: + placement: + constraints: + - "node.hostname == workstation" + restart_policy: + condition: any + + grafana: + image: grafana/grafana:11.2.2-security-01 + networks: + - hamstring + ports: + - target: 3000 + published: 3000 + protocol: tcp + mode: ingress + volumes: + - "/home/max/HAMSTRING/hamstring/docker/grafana-provisioning/dashboards:/etc/grafana/provisioning/dashboards" + - "/home/max/HAMSTRING/hamstring/docker/grafana-provisioning/dashboards/dashboards.yaml:/etc/grafana/provisioning/dashboards/dashboards.yaml" + - "/home/max/HAMSTRING/hamstring/docker/grafana-provisioning/datasources.yaml:/etc/grafana/provisioning/datasources/datasources.yaml" + environment: + GF_SECURITY_ADMIN_USER: admin + GF_SECURITY_ADMIN_PASSWORD: admin + GF_INSTALL_PLUGINS: grafana-clickhouse-datasource + healthcheck: + test: ["CMD-SHELL", "nc -z localhost 3000"] + interval: 10s + timeout: 5s + retries: 3 + deploy: + placement: + constraints: + - "node.hostname == workstation" + restart_policy: + condition: any + + prometheus: + image: prom/prometheus + networks: + - hamstring + ports: + - target: 9090 + published: 9088 + protocol: tcp + mode: ingress + volumes: + - "/home/max/HAMSTRING/hamstring/docker/prometheus/prometheus.yml:/etc/prometheus/prometheus.yml" + deploy: + placement: + constraints: + - "node.hostname == workstation" + restart_policy: + condition: any + + kafka-exporter: + image: danielqsj/kafka-exporter + command: + - --kafka.server=kafka1:19092 + - --kafka.server=kafka2:19093 + - --kafka.server=kafka3:19094 + networks: + - hamstring + ports: + - target: 9308 + published: 9308 + protocol: tcp + mode: ingress + deploy: + restart_policy: + condition: any + + monitoring_agent: + image: ghcr.io/astraos-de/hamstring-monitoring:v2.2.0-dev + networks: + - hamstring + volumes: + - "/home/max/HAMSTRING/hamstring/config.yaml:/app/config.yaml" + environment: + GROUP_ID: monitoring_agent + HAMSTRING_CONFIG_CHECKSUM: "d9bc812b76c2b58c633b412a6c07f981f3c4d787" + deploy: + restart_policy: + condition: any + + logserver: + image: ghcr.io/astraos-de/hamstring-logserver:v2.2.0-dev + networks: + - hamstring + volumes: + - "/home/max/HAMSTRING/hamstring/docker/default.txt:/opt/file.txt" + - "/home/max/HAMSTRING/hamstring/config.yaml:/app/config.yaml" + environment: + GROUP_ID: log_storage + HAMSTRING_CONFIG_CHECKSUM: "d9bc812b76c2b58c633b412a6c07f981f3c4d787" + deploy: + restart_policy: + condition: any + + logcollector: + image: ghcr.io/astraos-de/hamstring-logcollector:v2.2.0-dev + networks: + - hamstring + volumes: + - "/home/max/HAMSTRING/hamstring/config.yaml:/app/config.yaml" + environment: + GROUP_ID: log_collection + HAMSTRING_CONFIG_CHECKSUM: "d9bc812b76c2b58c633b412a6c07f981f3c4d787" + deploy: + restart_policy: + condition: any + + prefilter: + image: ghcr.io/astraos-de/hamstring-prefilter:v2.2.0-dev + networks: + - hamstring + volumes: + - "/home/max/HAMSTRING/hamstring/config.yaml:/app/config.yaml" + environment: + GROUP_ID: log_filtering + HAMSTRING_CONFIG_CHECKSUM: "d9bc812b76c2b58c633b412a6c07f981f3c4d787" + deploy: + mode: replicated + replicas: 1 + restart_policy: + condition: any + + inspector: + image: ghcr.io/astraos-de/hamstring-inspector:v2.2.0-dev + networks: + - hamstring + volumes: + - "/home/max/HAMSTRING/hamstring/config.yaml:/app/config.yaml" + environment: + GROUP_ID: data_inspection + NUMBER_OF_INSTANCES: "1" + HAMSTRING_CONFIG_CHECKSUM: "d9bc812b76c2b58c633b412a6c07f981f3c4d787" + deploy: + mode: replicated + replicas: 1 + resources: + limits: + cpus: "2" + memory: 512M + reservations: + cpus: "1" + memory: 256M + restart_policy: + condition: any + + detector: + image: ghcr.io/astraos-de/hamstring-detector:v2.2.0-dev + networks: + - hamstring + volumes: + - "/home/max/HAMSTRING/hamstring/config.yaml:/app/config.yaml" + environment: + GROUP_ID: data_analysis + NVIDIA_DRIVER_CAPABILITIES: compute,utility + NVIDIA_VISIBLE_DEVICES: all + HAMSTRING_CONFIG_CHECKSUM: "d9bc812b76c2b58c633b412a6c07f981f3c4d787" + deploy: + mode: replicated + replicas: 1 + # Use when nvidia gpu is available to you + # resources: + # reservations: + # generic_resources: + # - discrete_resource_spec: + # kind: NVIDIA-GPU + # value: 1 + restart_policy: + condition: any + + alerter: + image: ghcr.io/astraos-de/hamstring-alerter:v2.2.0-dev + networks: + - hamstring + volumes: + - "/home/max/HAMSTRING/hamstring/config.yaml:/app/config.yaml" + - "/home/max/HAMSTRING/hamstring/docker/logs:/opt/logs" + environment: + GROUP_ID: data_alerting + HAMSTRING_CONFIG_CHECKSUM: "d9bc812b76c2b58c633b412a6c07f981f3c4d787" + deploy: + mode: replicated + replicas: 1 + restart_policy: + condition: any + + zeek: + image: ghcr.io/astraos-de/hamstring-zeek:v2.2.0-dev + networks: + - hamstring + volumes: + - "/home/max/HAMSTRING/hamstring/config.yaml:/opt/config.yaml" + - "/home/max/HAMSTRING/hamstring/docker/data/test_pcaps:/opt/static_files" + environment: + CONTAINER_NAME: zeek + HAMSTRING_CONFIG_CHECKSUM: "d9bc812b76c2b58c633b412a6c07f981f3c4d787" + deploy: + restart_policy: + condition: any + +networks: + hamstring: + driver: overlay + attachable: true + ipam: + driver: default + config: + - subnet: 172.27.0.0/16 + +volumes: + ch_data: + ch_logs: + kafka-data1: + kafka-data2: + kafka-data3: + zk-data: + zk-txn-logs: + +# code: language=yaml diff --git a/docs/index.rst b/docs/index.rst index 91d73fd2..06875993 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,8 +1,9 @@ Welcome to HAMSTRING's documentation! ====================================== -**HAMSTRING** is a machine learning-based DNS classifier for detecting Domain Generation Algorithms (DGAs), tunneling, and -data exfiltration by malicious actors. +**HAMSTRING** is a CIDS framework to run signature and machine learning-based IDS classifiers. It employs several string and anomaly-based filtering techniques +to maximize detection efficiency. Currently the tool focueses on DNS attacks, as it incorporates heiDGAF (DGA detection) & Domainator (DNS Tunneling Detector) +:cite:p:`petrov_domainator_2025` :cite:p:`machmeier_heidgaf`. Check out the :doc:`usage` section for further information on how to use the software, including how to :ref:`install ` and :ref:`configure ` the project. For more details on the implementation diff --git a/docs/refs.bib b/docs/refs.bib index c1f26a69..3ea9bdc4 100644 --- a/docs/refs.bib +++ b/docs/refs.bib @@ -13,6 +13,20 @@ @inproceedings{schuppen_fanci_2018 pages = {1165--1181}, } + +Machmeier, S., Ludwig, M., Fuchs, M. and Heuveline, V., "", in 2025 15th SPRING graduate workshop, pp. 42, Nuremberg, Germany, 2025 + +@inproceedings{machmeier_heidgaf, + title = {{heiDGAF}: {A} {Novel} {Open}-{Source} {CIDS} {Solution} for {DGA} {Detection}}, + shorttitle = {{heiDGAF}}, + urldate = {2025-06-05}, + booktitle = {15th {SPRING} graduate workshop}, + author = {Machmeier, Stefan and Ludwig, Maximilian and Fuchs, Manuel and Heuveline, Vincent}, + year = {2025}, + pages = {42}, + file = {Available Version (via Google Scholar):/home/max/Zotero/storage/4Q8ZB2YB/Machmeier et al. - 2025 - heiDGAF A Novel Open-Source CIDS Solution for DGA Detection.pdf:application/pdf}, +} + @inproceedings{petrov_domainator_2025, address = {Cham}, title = {Domainator: Detecting and Identifying DNS-Tunneling Malware Using Metadata Sequences}, diff --git a/src/base/kafka_handler.py b/src/base/kafka_handler.py index 5ddf3c2c..007f9553 100644 --- a/src/base/kafka_handler.py +++ b/src/base/kafka_handler.py @@ -11,7 +11,7 @@ import time import uuid from abc import abstractmethod -from typing import Optional +from typing import Callable, Optional import marshmallow_dataclass from confluent_kafka import ( @@ -25,8 +25,8 @@ sys.path.append(os.getcwd()) from src.base.data_classes.batch import Batch from src.base.log_config import get_logger +from src.base.retry import retry_forever from src.base.utils import kafka_delivery_report, setup_config -import uuid logger = get_logger() @@ -188,6 +188,32 @@ def _wait_for_admin_futures(futures: dict, operation: str) -> None: raise +def _is_retriable_kafka_exception(exception: Exception) -> bool: + if isinstance(exception, (KafkaException, BufferError, RuntimeError, OSError)): + return True + return False + + +def _is_retriable_kafka_error(error) -> bool: + retriable = getattr(error, "retriable", None) + if callable(retriable) and retriable(): + return True + + retriable_codes = { + getattr(KafkaError, name) + for name in ( + "_ALL_BROKERS_DOWN", + "_TRANSPORT", + "_TIMED_OUT", + "_MSG_TIMED_OUT", + "_RESOLVE", + "_WAIT_COORD", + ) + if hasattr(KafkaError, name) + } + return hasattr(error, "code") and error.code() in retriable_codes + + def ensure_topics( admin_client: AdminClient, topics: str | list[str], @@ -210,7 +236,10 @@ def ensure_topics( else _as_bool(auto_expand_partitions) ) - cluster_metadata = admin_client.list_topics(timeout=10) + cluster_metadata = retry_forever( + lambda: admin_client.list_topics(timeout=10), + "Kafka metadata lookup", + ) topics_metadata = getattr(cluster_metadata, "topics", {}) existing_topics = ( set(topics_metadata.keys()) @@ -226,22 +255,30 @@ def ensure_topics( "Creating Kafka topics %s.", missing_topics, ) - futures = admin_client.create_topics( - [ - NewTopic( - topic, - target_partitions_by_topic[topic], - replication_factor_by_topic[topic], - ) - for topic in missing_topics - ] + retry_forever( + lambda: _wait_for_admin_futures( + admin_client.create_topics( + [ + NewTopic( + topic, + target_partitions_by_topic[topic], + replication_factor_by_topic[topic], + ) + for topic in missing_topics + ] + ), + "create topic", + ), + f"Kafka topic creation for {missing_topics}", ) - _wait_for_admin_futures(futures, "create topic") if not auto_expand_partitions: return target_partitions_by_topic - cluster_metadata = admin_client.list_topics(timeout=10) + cluster_metadata = retry_forever( + lambda: admin_client.list_topics(timeout=10), + "Kafka metadata lookup after topic creation", + ) topics_to_expand = [] for topic in normalized_topics: current_partition_count = _topic_partition_count(cluster_metadata, topic) @@ -258,8 +295,13 @@ def ensure_topics( topics_to_expand.append(NewPartitions(topic, target_partitions)) if topics_to_expand: - futures = admin_client.create_partitions(topics_to_expand) - _wait_for_admin_futures(futures, "expand partitions") + retry_forever( + lambda: _wait_for_admin_futures( + admin_client.create_partitions(topics_to_expand), + "expand partitions", + ), + f"Kafka partition expansion for {[str(topic) for topic in topics_to_expand]}", + ) return target_partitions_by_topic @@ -334,7 +376,39 @@ def __init__(self, conf): Should contain broker settings and producer-specific options. """ super().__init__() - self.producer = Producer(conf) + self.conf = conf + self.producer = self._new_producer() + + def _new_producer(self): + return retry_forever( + lambda: Producer(self.conf), + "Kafka producer creation", + ) + + def _reset_producer(self) -> None: + try: + if self.producer: + self.producer.flush(5) + except Exception as exception: + logger.warning("Ignoring Kafka producer flush failure during reconnect: %s", exception) + self.producer = self._new_producer() + + def _with_producer_retry(self, description: str, operation: Callable[[], None]) -> None: + def attempt(): + try: + operation() + except Exception as exception: + if not _is_retriable_kafka_exception(exception): + raise + logger.warning("%s failed, recreating Kafka producer: %s", description, exception) + self._reset_producer() + raise + + retry_forever( + attempt, + description, + retryable=(KafkaException, BufferError, RuntimeError, OSError), + ) @abstractmethod def produce(self, *args, **kwargs): @@ -358,7 +432,8 @@ def __del__(self) -> None: Ensures that all pending messages are flushed before the producer is destroyed, preventing message loss. """ - self.producer.flush() + if self.producer: + self.producer.flush() class SimpleKafkaProduceHandler(KafkaProduceHandler): @@ -410,13 +485,30 @@ def produce(self, topic: str, data: str, key: None | str = None) -> None: """ if not data: return - self.producer.flush() - self.producer.produce( - topic=topic, - key=key, - value=data, - callback=kafka_delivery_report, - ) + + def operation(): + delivery_errors = [] + + def delivery_callback(err, msg): + kafka_delivery_report(err, msg) + if err: + delivery_errors.append(err) + + self.producer.flush() + self.producer.produce( + topic=topic, + key=key, + value=data, + callback=delivery_callback, + ) + self.producer.flush() + if delivery_errors: + delivery_error = delivery_errors[0] + if _is_retriable_kafka_error(delivery_error): + raise KafkaException(delivery_error) + raise ValueError(f"Kafka delivery failed: {delivery_error}") + + self._with_producer_retry(f"Kafka produce to {topic}", operation) class ExactlyOnceKafkaProduceHandler(KafkaProduceHandler): @@ -458,7 +550,17 @@ def __init__(self): } super().__init__(conf) - self.producer.init_transactions() + self._init_transactions_with_retry() + + def _reset_producer(self) -> None: + super()._reset_producer() + self._init_transactions_with_retry() + + def _init_transactions_with_retry(self) -> None: + retry_forever( + lambda: self.producer.init_transactions(), + "Kafka transactional producer initialization", + ) def produce(self, topic: str, data: str, key: None | str = None) -> None: """Produce a message to the specified Kafka topic with exactly-once semantics. @@ -480,24 +582,29 @@ def produce(self, topic: str, data: str, key: None | str = None) -> None: if not data: return - self.producer.flush() - self.producer.begin_transaction() + def operation(): + self.producer.flush() + self.producer.begin_transaction() - try: - self.producer.produce( - topic=topic, - key=key, - value=data, - callback=kafka_delivery_report, - ) + try: + self.producer.produce( + topic=topic, + key=key, + value=data, + callback=kafka_delivery_report, + ) + self.commit_transaction_with_retry() + except Exception as e: + logger.info(f"aborted for topic {topic}") + try: + self.producer.abort_transaction() + except Exception as abort_exception: + logger.warning("Kafka transaction abort failed: %s", abort_exception) + logger.error("Transaction aborted.") + logger.error(e) + raise - self.commit_transaction_with_retry() - except Exception as e: - logger.info(f"aborted for topic {topic}") - self.producer.abort_transaction() - logger.error("Transaction aborted.") - logger.error(e) - raise + self._with_producer_retry(f"Kafka transactional produce to {topic}", operation) def commit_transaction_with_retry( self, max_retries: int = 3, retry_interval_ms: int = 1000 @@ -572,6 +679,7 @@ def __init__(self, topics: str | list[str]) -> None: if isinstance(topics, str): topics = [topics] + self.topics = topics # get brokers self.brokers = ",".join( @@ -580,37 +688,58 @@ def __init__(self, topics: str | list[str]) -> None: for broker in KAFKA_BROKERS ] ) + self.conf = self._build_consumer_conf() + self._connect_consumer() - # create consumer - conf = { + def _build_consumer_conf(self) -> dict: + return { "bootstrap.servers": self.brokers, - "group.id": build_consumer_group_id(topics), + "group.id": build_consumer_group_id(self.topics), "enable.auto.commit": False, "auto.offset.reset": "earliest", "enable.partition.eof": True, "max.poll.interval.ms": KAFKA_CONSUMER_MAX_POLL_INTERVAL_MS, } - self.consumer = Consumer(conf) - # create topics - admin_client = AdminClient( - { - "bootstrap.servers": self.brokers, - } - ) - target_partitions_by_topic = ensure_topics(admin_client, topics) + def _connect_consumer(self) -> None: + def connect(): + consumer = Consumer(self.conf) + admin_client = AdminClient( + { + "bootstrap.servers": self.brokers, + } + ) + target_partitions_by_topic = ensure_topics(admin_client, self.topics) + + if not self._all_topics_created(self.topics, target_partitions_by_topic, consumer): + try: + consumer.close() + except Exception: + pass + raise TooManyFailedAttemptsError("Not all topics were created.") - # check if topics are created - if not self._all_topics_created(topics, target_partitions_by_topic): - raise TooManyFailedAttemptsError("Not all topics were created.") + consumer.subscribe(self.topics) + return consumer - # subscribe to the topics - self.consumer.subscribe(topics) + self.consumer = retry_forever(connect, f"Kafka consumer setup for {self.topics}") + + def _reset_consumer(self) -> None: + try: + if self.consumer: + self.consumer.close() + except Exception as exception: + logger.warning("Ignoring Kafka consumer close failure during reconnect: %s", exception) + self._last_consumed_message = None + self._connect_consumer() def commit(self) -> None: """Commit the last message returned by ``consume``.""" if self.consumer and self._last_consumed_message is not None: - self.consumer.commit(self._last_consumed_message) + retry_forever( + lambda: self.consumer.commit(self._last_consumed_message), + "Kafka consumer offset commit", + retryable=(KafkaException, RuntimeError, OSError), + ) self._last_consumed_message = None @abstractmethod @@ -660,7 +789,10 @@ def consume_as_json(self) -> tuple[Optional[str], dict]: raise ValueError("Unknown data format") def _all_topics_created( - self, topics: list[str], min_partitions: int | dict[str, int] = 1 + self, + topics: list[str], + min_partitions: int | dict[str, int] = 1, + consumer=None, ) -> bool: """Verify that all specified topics have been created successfully. @@ -676,8 +808,13 @@ def _all_topics_created( """ number_of_retries_left = 30 all_topics_created = False + consumer = consumer or self.consumer while not all_topics_created: # try for 15 seconds - assigned_topics = self.consumer.list_topics(timeout=10) + assigned_topics = retry_forever( + lambda: consumer.list_topics(timeout=10), + "Kafka topic visibility check", + retryable=(KafkaException, RuntimeError, OSError), + ) all_topics_created = True for topic in topics: @@ -700,6 +837,28 @@ def _all_topics_created( return True + def _poll_message(self): + while True: + try: + msg = self.consumer.poll(timeout=1.0) + except (KafkaException, RuntimeError, OSError) as exception: + logger.warning("Kafka consumer poll failed, reconnecting: %s", exception) + self._reset_consumer() + continue + + if msg is None: + return None + + if msg.error(): + if msg.error().code() == KafkaError._PARTITION_EOF: + return None + if _is_retriable_kafka_error(msg.error()): + logger.warning("Kafka consumer error is retriable, reconnecting: %s", msg.error()) + self._reset_consumer() + return None + + return msg + def __del__(self) -> None: """Cleanup method called when the object is destroyed @@ -795,7 +954,7 @@ def consume(self) -> tuple[Optional[str], Optional[str], Optional[str]]: try: while True: - msg = self.consumer.poll(timeout=1.0) + msg = self._poll_message() if msg is None: if not empty_data_retrieved: @@ -804,11 +963,8 @@ def consume(self) -> tuple[Optional[str], Optional[str], Optional[str]]: empty_data_retrieved = True continue if msg.error(): - if msg.error().code() == KafkaError._PARTITION_EOF: - continue - else: - logger.error(f"Consumer error: {msg.error()}") - raise ValueError("Message is invalid") + logger.error(f"Consumer error: {msg.error()}") + raise ValueError("Message is invalid") # unpack message key = msg.key().decode("utf-8") if msg.key() else None @@ -860,7 +1016,7 @@ def consume(self) -> tuple[Optional[str], Optional[str], Optional[str]]: try: while True: - msg = self.consumer.poll(timeout=1.0) + msg = self._poll_message() if msg is None: if not empty_data_retrieved: @@ -870,11 +1026,8 @@ def consume(self) -> tuple[Optional[str], Optional[str], Optional[str]]: continue if msg.error(): - if msg.error().code() == KafkaError._PARTITION_EOF: - continue - else: - logger.error(f"Consumer error: {msg.error()}") - raise ValueError("Message is invalid") + logger.error(f"Consumer error: {msg.error()}") + raise ValueError("Message is invalid") # unpack message key = msg.key().decode("utf-8") if msg.key() else None diff --git a/src/base/retry.py b/src/base/retry.py new file mode 100644 index 00000000..779798e5 --- /dev/null +++ b/src/base/retry.py @@ -0,0 +1,83 @@ +import os +import random +import time +from typing import Any, Callable, TypeVar + +from src.base.log_config import get_logger +from src.base.utils import setup_config + +logger = get_logger("base.retry") + +T = TypeVar("T") + +_DEFAULT_CONFIG = { + "initial_delay_seconds": 1.0, + "max_delay_seconds": 30.0, + "backoff_multiplier": 2.0, + "jitter_seconds": 0.25, + "log_every_attempts": 5, +} + + +def resilience_config() -> dict[str, Any]: + config = setup_config() + retry_config = ( + config.get("pipeline", {}) + .get("resilience", {}) + .get("retry", {}) + ) + merged = dict(_DEFAULT_CONFIG) + merged.update(retry_config) + return merged + + +def retry_forever( + operation: Callable[[], T], + description: str, + retry_config: dict[str, Any] | None = None, + retryable: tuple[type[BaseException], ...] = (Exception,), +) -> T: + config = retry_config or resilience_config() + initial_delay = _float_setting(config, "initial_delay_seconds") + max_delay = _float_setting(config, "max_delay_seconds") + multiplier = max(1.0, _float_setting(config, "backoff_multiplier")) + jitter = max(0.0, _float_setting(config, "jitter_seconds")) + log_every = max(1, _int_setting(config, "log_every_attempts")) + + delay = initial_delay + attempt = 0 + + while True: + try: + return operation() + except retryable as exception: + attempt += 1 + if attempt == 1 or attempt % log_every == 0: + logger.warning( + "%s failed on attempt %d: %s. Retrying in %.1fs.", + description, + attempt, + exception, + delay, + ) + sleep_for = delay + (random.uniform(0, jitter) if jitter else 0) + time.sleep(sleep_for) + delay = min(max_delay, delay * multiplier) + + +def _float_setting(config: dict[str, Any], key: str) -> float: + env_key = f"HAMSTRING_RETRY_{key.upper()}" + value = os.getenv(env_key, config.get(key, _DEFAULT_CONFIG[key])) + try: + return float(value) + except (TypeError, ValueError): + return float(_DEFAULT_CONFIG[key]) + + +def _int_setting(config: dict[str, Any], key: str) -> int: + env_key = f"HAMSTRING_RETRY_{key.upper()}" + value = os.getenv(env_key, config.get(key, _DEFAULT_CONFIG[key])) + try: + return int(value) + except (TypeError, ValueError): + return int(_DEFAULT_CONFIG[key]) diff --git a/src/monitoring/clickhouse_batch_sender.py b/src/monitoring/clickhouse_batch_sender.py index 4df8c85f..e4f73c4d 100644 --- a/src/monitoring/clickhouse_batch_sender.py +++ b/src/monitoring/clickhouse_batch_sender.py @@ -10,6 +10,7 @@ sys.path.append(os.getcwd()) from src.base.log_config import get_logger +from src.base.retry import retry_forever from src.base.utils import setup_config logger = get_logger() @@ -221,11 +222,23 @@ def __init__(self): self.timer = None self.batch = {key: [] for key in self.tables} - self._client = clickhouse_connect.get_client( - host=CLICKHOUSE_HOSTNAME, - ) + self._client = self._connect_client() self.lock = Lock() + def _connect_client(self): + return retry_forever( + lambda: clickhouse_connect.get_client(host=CLICKHOUSE_HOSTNAME), + "ClickHouse client connection", + ) + + def _reset_client(self) -> None: + try: + if self._client: + self._client.close() + except Exception as exception: + logger.warning("Ignoring ClickHouse client close failure during reconnect: %s", exception) + self._client = self._connect_client() + def __del__(self): self.insert_all() @@ -266,13 +279,28 @@ def insert(self, table_name: str): """ if self.batch[table_name]: with self.lock: - self._client.insert( - table_name, - self.batch.get(table_name), - column_names=list(self.tables.get(table_name).columns), - ) + pending_rows = self.batch.get(table_name) + column_names = list(self.tables.get(table_name).columns) + + def insert_batch(): + try: + self._client.insert( + table_name, + pending_rows, + column_names=column_names, + ) + except Exception as exception: + logger.warning( + "ClickHouse insert for table '%s' failed, reconnecting: %s", + table_name, + exception, + ) + self._reset_client() + raise + + retry_forever(insert_batch, f"ClickHouse insert for table '{table_name}'") logger.debug( - f"Inserted {table_name=},{self.batch.get(table_name)=},{list(self.tables.get(table_name).columns)=}" + f"Inserted {table_name=},{pending_rows=},{column_names=}" ) self.batch[table_name] = [] diff --git a/src/monitoring/monitoring_agent.py b/src/monitoring/monitoring_agent.py index 948a9e31..5000d351 100644 --- a/src/monitoring/monitoring_agent.py +++ b/src/monitoring/monitoring_agent.py @@ -12,6 +12,7 @@ from src.base.log_config import get_logger from src.base.utils import setup_config from src.base.execution import create_pipeline_executor +from src.base.retry import retry_forever logger = get_logger() module_name = "monitoring.agent" @@ -48,10 +49,17 @@ def _iter_statements(sql_content: str): file_path = os.path.join(CREATE_TABLES_DIRECTORY, filename) sql_content = _load_contents(file_path) - with clickhouse_connect.get_client(host=CLICKHOUSE_HOSTNAME) as client: + client = retry_forever( + lambda: clickhouse_connect.get_client(host=CLICKHOUSE_HOSTNAME), + "ClickHouse table preparation connection", + ) + with client: for statement in _iter_statements(sql_content): try: - client.command(statement) + retry_forever( + lambda statement=statement: client.command(statement), + "ClickHouse CREATE TABLE statement", + ) except Exception as e: logger.critical("Error in CREATE TABLE statement") raise e diff --git a/tests/clickhouse/test_clickhouse_batch_sender.py b/tests/clickhouse/test_clickhouse_batch_sender.py index 1afb5243..c24f1964 100644 --- a/tests/clickhouse/test_clickhouse_batch_sender.py +++ b/tests/clickhouse/test_clickhouse_batch_sender.py @@ -75,6 +75,23 @@ def test_successful(self): self.assertIsNotNone(sut.lock) self.assertEqual({key: [] for key in sut.tables}, sut.batch) + @patch("src.base.retry.time.sleep", return_value=None) + def test_retries_until_clickhouse_is_available(self, mock_sleep): + client = Mock() + with patch( + "src.monitoring.clickhouse_batch_sender.clickhouse_connect" + ) as mock_clickhouse_connect: + mock_clickhouse_connect.get_client.side_effect = [ + RuntimeError("clickhouse unavailable"), + client, + ] + + sut = ClickHouseBatchSender() + + self.assertEqual(client, sut._client) + self.assertEqual(2, mock_clickhouse_connect.get_client.call_count) + mock_sleep.assert_called() + class TestDel(unittest.TestCase): def setUp(self): @@ -213,7 +230,33 @@ def test_filled_batch(self): ["one", "two", "three"], column_names=["col_1", "col_2"], ) - self.assertEquals([], self.sut.batch[test_table_name]) + self.assertEqual([], self.sut.batch[test_table_name]) + + @patch("src.base.retry.time.sleep", return_value=None) + def test_filled_batch_retries_without_dropping_rows(self, mock_sleep): + # Arrange + test_table_name = "test_table" + first_client = Mock() + second_client = Mock() + first_client.insert.side_effect = RuntimeError("clickhouse unavailable") + + self.sut.tables = { + test_table_name: Table(test_table_name, {"col_1": str, "col_2": str}) + } + self.sut.batch = {test_table_name: ["one", "two", "three"]} + self.sut._client = first_client + + with patch.object(self.sut, "_connect_client", return_value=second_client): + self.sut.insert(test_table_name) + + first_client.insert.assert_called_once() + second_client.insert.assert_called_once_with( + test_table_name, + ["one", "two", "three"], + column_names=["col_1", "col_2"], + ) + self.assertEqual([], self.sut.batch[test_table_name]) + mock_sleep.assert_called() def test_empty_batch(self): # Arrange @@ -230,7 +273,7 @@ def test_empty_batch(self): # Assert self.sut._client.insert.assert_not_called() - self.assertEquals([], self.sut.batch[test_table_name]) + self.assertEqual([], self.sut.batch[test_table_name]) class TestInsertAll(unittest.TestCase): diff --git a/tests/kafka/test_exactly_once_kafka_consume_handler.py b/tests/kafka/test_exactly_once_kafka_consume_handler.py index 55a5c431..d5958513 100644 --- a/tests/kafka/test_exactly_once_kafka_consume_handler.py +++ b/tests/kafka/test_exactly_once_kafka_consume_handler.py @@ -78,9 +78,12 @@ def test_init(self, mock_consumer, mock_admin_client, mock_all_topics_created): "src.base.kafka_handler.KafkaConsumeHandler._all_topics_created", return_value=True, ) + @patch("src.base.retry.time.sleep", return_value=None) @patch("src.base.kafka_handler.AdminClient") @patch("src.base.kafka_handler.Consumer") - def test_init_fail(self, mock_consumer, mock_admin_client, mock_all_topics_created): + def test_init_retries_until_subscribe_succeeds( + self, mock_consumer, mock_admin_client, mock_sleep, mock_all_topics_created + ): mock_consumer_instance = Mock() mock_consumer.return_value = mock_consumer_instance @@ -93,16 +96,15 @@ def test_init_fail(self, mock_consumer, mock_admin_client, mock_all_topics_creat "max.poll.interval.ms": 1800000, } - with patch.object( - mock_consumer_instance, "subscribe", side_effect=KafkaException - ): - with self.assertRaises(KafkaException): - sut = ExactlyOnceKafkaConsumeHandler(topics="test_topic") + mock_consumer_instance.subscribe.side_effect = [KafkaException(), None] - self.assertEqual(mock_consumer_instance, sut.consumer) + sut = ExactlyOnceKafkaConsumeHandler(topics="test_topic") - mock_consumer.assert_called_once_with(expected_conf) - mock_consumer_instance.assign.assert_called_once() + self.assertEqual(mock_consumer_instance, sut.consumer) + self.assertEqual(2, mock_consumer.call_count) + mock_consumer.assert_any_call(expected_conf) + self.assertEqual(2, mock_consumer_instance.subscribe.call_count) + mock_sleep.assert_called() class TestConsume(unittest.TestCase): @@ -164,7 +166,8 @@ def test_consumer_error_partition_eof(self): def test_consumer_raises_other_error(self): other_error = Mock() - other_error.code.return_value = KafkaError._ALL_BROKERS_DOWN + other_error.retriable.return_value = False + other_error.code.return_value = 123456 msg = Mock() msg.error.return_value = other_error diff --git a/tests/kafka/test_exactly_once_kafka_produce_handler.py b/tests/kafka/test_exactly_once_kafka_produce_handler.py index b60401fc..02ddf47b 100644 --- a/tests/kafka/test_exactly_once_kafka_produce_handler.py +++ b/tests/kafka/test_exactly_once_kafka_produce_handler.py @@ -46,6 +46,7 @@ def test_init(self, mock_producer, mock_uuid): mock_producer.assert_called_once_with(expected_conf) mock_producer_instance.init_transactions.assert_called_once() + @patch("src.base.retry.time.sleep", return_value=None) @patch("src.base.kafka_handler.uuid") @patch("src.base.kafka_handler.logger") @patch( @@ -66,7 +67,9 @@ def test_init(self, mock_producer, mock_uuid): ], ) @patch("src.base.kafka_handler.Producer") - def test_init_fail(self, mock_producer, mock_logger, mock_uuid): + def test_init_retries_until_transactions_initialize( + self, mock_producer, mock_logger, mock_uuid, mock_sleep + ): mock_producer_instance = MagicMock() mock_producer.return_value = mock_producer_instance mock_uuid.uuid4.return_value = "fixed‑uuid‑1234‑abcd‑5678‑90ef" @@ -78,14 +81,17 @@ def test_init_fail(self, mock_producer, mock_logger, mock_uuid): "message.max.bytes": 1000000000, } - with patch.object( - mock_producer_instance, "init_transactions", side_effect=KafkaException - ): - with self.assertRaises(KafkaException): - sut = ExactlyOnceKafkaProduceHandler() + mock_producer_instance.init_transactions.side_effect = [ + KafkaException(), + None, + ] - mock_producer.assert_called_once_with(expected_conf) - mock_producer_instance.init_transactions.assert_called_once() + sut = ExactlyOnceKafkaProduceHandler() + + self.assertEqual(mock_producer_instance, sut.producer) + mock_producer.assert_called_once_with(expected_conf) + self.assertEqual(2, mock_producer_instance.init_transactions.call_count) + mock_sleep.assert_called() class TestSend(unittest.TestCase): diff --git a/tests/kafka/test_kafka_consume_handler.py b/tests/kafka/test_kafka_consume_handler.py index 0f258ebd..ac61cd23 100644 --- a/tests/kafka/test_kafka_consume_handler.py +++ b/tests/kafka/test_kafka_consume_handler.py @@ -7,7 +7,6 @@ ensure_topics, KafkaConsumeHandler, KafkaMessageFetchException, - TooManyFailedAttemptsError, _desired_topic_partitions, _topic_replication_factor, _topic_config, @@ -111,6 +110,26 @@ def test_existing_topic_with_enough_partitions_is_left_unchanged(self): admin_client.create_topics.assert_not_called() admin_client.create_partitions.assert_not_called() + @patch("src.base.retry.time.sleep", return_value=None) + def test_metadata_lookup_retries_until_kafka_is_available(self, mock_sleep): + admin_client = MagicMock() + admin_client.list_topics.side_effect = [ + RuntimeError("broker unavailable"), + _metadata({"test_topic": 4}), + _metadata({"test_topic": 4}), + ] + + target_partitions_by_topic = ensure_topics( + admin_client, + ["test_topic"], + target_partitions=4, + replication_factor=2, + ) + + self.assertEqual({"test_topic": 4}, target_partitions_by_topic) + self.assertEqual(3, admin_client.list_topics.call_count) + mock_sleep.assert_called() + def test_auto_expand_can_be_disabled(self): admin_client = MagicMock() admin_client.list_topics.return_value = _metadata({"test_topic": 2}) @@ -238,12 +257,13 @@ def test_init_successful( ) @patch( "src.base.kafka_handler.KafkaConsumeHandler._all_topics_created", - return_value=False, + side_effect=[False, True], ) + @patch("src.base.retry.time.sleep", return_value=None) @patch("src.base.kafka_handler.AdminClient") @patch("src.base.kafka_handler.Consumer") - def test_init_unsuccessful( - self, mock_consumer, mock_admin_client, mock_all_topics_created + def test_init_retries_until_topics_are_visible( + self, mock_consumer, mock_admin_client, mock_sleep, mock_all_topics_created ): # Arrange mock_consumer_instance = MagicMock() @@ -259,12 +279,15 @@ def test_init_unsuccessful( } # Act - with self.assertRaises(TooManyFailedAttemptsError): - KafkaConsumeHandler(topics="test_topic") + sut = KafkaConsumeHandler(topics="test_topic") # Assert - mock_consumer.assert_called_once_with(expected_conf) - mock_consumer_instance.subscribe.assert_not_called() + self.assertEqual(mock_consumer_instance, sut.consumer) + self.assertEqual(2, mock_consumer.call_count) + mock_consumer.assert_any_call(expected_conf) + mock_consumer_instance.close.assert_called_once() + mock_consumer_instance.subscribe.assert_called_once() + mock_sleep.assert_called() @patch("src.base.kafka_handler.CONSUMER_GROUP_ID", "test_group_id") @patch( diff --git a/tests/kafka/test_simple_kafka_consume_handler.py b/tests/kafka/test_simple_kafka_consume_handler.py index eb454e9b..0b060085 100644 --- a/tests/kafka/test_simple_kafka_consume_handler.py +++ b/tests/kafka/test_simple_kafka_consume_handler.py @@ -114,7 +114,8 @@ def test_consumer_error_partition_eof(self): def test_consumer_raises_other_error(self): other_error = Mock() - other_error.code.return_value = KafkaError._ALL_BROKERS_DOWN + other_error.retriable.return_value = False + other_error.code.return_value = 123456 msg = Mock() msg.error.return_value = other_error diff --git a/tests/kafka/test_simple_kafka_produce_handler.py b/tests/kafka/test_simple_kafka_produce_handler.py index 281997f2..7919d722 100644 --- a/tests/kafka/test_simple_kafka_produce_handler.py +++ b/tests/kafka/test_simple_kafka_produce_handler.py @@ -1,8 +1,9 @@ import unittest -from unittest.mock import patch, Mock +from unittest.mock import ANY, patch, Mock + +from confluent_kafka import KafkaError from src.base.kafka_handler import SimpleKafkaProduceHandler -from src.base.utils import kafka_delivery_report class TestInit(unittest.TestCase): @@ -58,13 +59,60 @@ def test_with_data(self): sut.produce("test_topic", "test_data") # Assert - mock_producer_instance.flush.assert_called_once() + self.assertEqual(2, mock_producer_instance.flush.call_count) mock_producer_instance.produce.assert_called_once_with( topic="test_topic", key=None, value="test_data", - callback=kafka_delivery_report, + callback=ANY, + ) + + @patch("src.base.retry.time.sleep", return_value=None) + def test_with_data_recreates_producer_after_transient_failure(self, mock_sleep): + with patch("src.base.kafka_handler.Producer") as mock_producer: + first_producer = Mock() + second_producer = Mock() + first_producer.flush.side_effect = BufferError("queue full") + mock_producer.side_effect = [first_producer, second_producer] + + sut = SimpleKafkaProduceHandler() + sut.produce("test_topic", "test_data") + + self.assertEqual(2, mock_producer.call_count) + first_producer.flush.assert_called() + second_producer.produce.assert_called_once_with( + topic="test_topic", + key=None, + value="test_data", + callback=ANY, + ) + mock_sleep.assert_called() + + @patch("src.base.retry.time.sleep", return_value=None) + def test_with_data_retries_delivery_callback_error(self, mock_sleep): + with patch("src.base.kafka_handler.Producer") as mock_producer: + first_producer = Mock() + second_producer = Mock() + delivery_error = KafkaError(KafkaError._ALL_BROKERS_DOWN) + + def fail_delivery(**kwargs): + kwargs["callback"](delivery_error, None) + + first_producer.produce.side_effect = fail_delivery + mock_producer.side_effect = [first_producer, second_producer] + + sut = SimpleKafkaProduceHandler() + sut.produce("test_topic", "test_data") + + self.assertEqual(2, mock_producer.call_count) + first_producer.produce.assert_called_once() + second_producer.produce.assert_called_once_with( + topic="test_topic", + key=None, + value="test_data", + callback=ANY, ) + mock_sleep.assert_called() def test_without_data(self): with patch("src.base.kafka_handler.Producer") as mock_producer: