Files
clickstream-ch-kafka-supers…/generator/generator.py
T
ddadminandDmitry Dementiev 731d991f94 refactor(generator): история batch в Kafka вместо ClickHouse
- Зачем:
  - ревью rev5: ClickHouse-интеграция была проблемной (порт 9000 native vs HTTP,
    неработающий fallback, отсутствие DDL для базы meta)
  - архитектурно чище: генератор остаётся pure Kafka producer,
    история доступна для аналитики через стандартный ingestion
- Что:
  - удален ClickHouseBatchHistory, clickhouse-connect зависимость
  - добавлен KafkaBatchHistory с записью в топик generator_batch_history
  - добавлен BatchRecord.to_dict() для JSON-сериализации
  - добавлен рабочий fallback: Kafka → InMemory при недоступности
  - удален pytest-asyncio (не использовался)
  - добавлены тесты test_kafka_history.py (15 тестов) и test_service.py (6 тестов)
  - обновлена документация: топик вместо таблицы ClickHouse
- Проверка:
  - make generator-test: 44/44 тестов пройдено
  - docker-compose валиден, генератор не зависит от clickhouse
2026-06-09 17:27:16 +03:00

693 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""
Автономный генератор событий для Kafka (MVP rev5).
Режим 'steady-stream': публикуем постепенно, короткими тиками (1-10 сек),
держим целевую интенсивность events/min без крупных минутных batch.
"""
import json
import logging
import math
import os
import random
import sys
import time
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
# Prometheus метрики
from prometheus_client import Counter, Gauge, Histogram, start_http_server
# Kafka импортируем lazy для возможности тестирования без Kafka
_kafka_imported = False
KafkaProducer = None
KafkaError = None
def _import_kafka():
global _kafka_imported, KafkaProducer, KafkaError
if not _kafka_imported:
from kafka import KafkaProducer
from kafka.errors import KafkaError
_kafka_imported = True
return KafkaProducer, KafkaError
# ---------------------------------------------------------------------------
# Настройка логирования
# ---------------------------------------------------------------------------
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("generator")
# ---------------------------------------------------------------------------
# Prometheus метрики
# ---------------------------------------------------------------------------
METRICS_EVENTS_TOTAL = Counter(
"generator_events_total",
"Total number of events sent to Kafka",
["topic"]
)
METRICS_ERRORS_TOTAL = Counter(
"generator_publish_errors_total",
"Total number of publish errors",
["topic"]
)
METRICS_TICK_DURATION = Histogram(
"generator_tick_duration_seconds",
"Duration of generator tick in seconds"
)
METRICS_LAST_SUCCESS = Gauge(
"generator_last_success_timestamp",
"Unix timestamp of last successful tick"
)
# ---------------------------------------------------------------------------
# Конфигурация через env
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class Config:
"""Конфигурация генератора из переменных окружения."""
# Подключение к Kafka
kafka_bootstrap_servers: str = field(
default_factory=lambda: os.getenv("KAFKA_BOOTSTRAP_SERVERS", "kafka:29092")
)
# Параметры генерации
tick_seconds: int = field(
default_factory=lambda: int(os.getenv("GEN_TICK_SECONDS", "5"))
)
lambda_base_per_min: int = field(
default_factory=lambda: int(os.getenv("GEN_LAMBDA_BASE_PER_MIN", "200"))
)
jitter_pct: int = field(
default_factory=lambda: int(os.getenv("GEN_JITTER_PCT", "20"))
)
min_events_per_tick: int = field(
default_factory=lambda: int(os.getenv("GEN_MIN_EVENTS_PER_TICK", "5"))
)
max_events_per_tick: int = field(
default_factory=lambda: int(os.getenv("GEN_MAX_EVENTS_PER_TICK", "50"))
)
# Пути к данным
data_dir: Path = field(
default_factory=lambda: Path(os.getenv("GEN_DATA_DIR", "/data"))
)
# Сид для воспроизводимости
seed: int | None = field(
default_factory=lambda: int(os.getenv("GEN_SEED"))
if os.getenv("GEN_SEED")
else None
)
# Включение/выключение генерации
enabled: bool = field(
default_factory=lambda: os.getenv("GEN_ENABLED", "true").lower() == "true"
)
# Порт для Prometheus метрик
metrics_port: int = field(
default_factory=lambda: int(os.getenv("GEN_METRICS_PORT", "9109"))
)
# Топик для истории batch
history_topic: str = field(
default_factory=lambda: os.getenv("GEN_HISTORY_TOPIC", "generator_batch_history")
)
def __post_init__(self):
# Валидация параметров
if self.tick_seconds < 1:
raise ValueError("GEN_TICK_SECONDS must be >= 1")
if self.lambda_base_per_min < 1:
raise ValueError("GEN_LAMBDA_BASE_PER_MIN must be >= 1")
if not self.data_dir.exists():
raise ValueError(f"Data directory does not exist: {self.data_dir}")
# ---------------------------------------------------------------------------
# Загрузка базового словаря событий
# ---------------------------------------------------------------------------
@dataclass
class EventDictionary:
"""Базовый словарь событий из JSONL файлов."""
browser_events: list[dict[str, Any]]
location_events: list[dict[str, Any]]
device_events: list[dict[str, Any]]
geo_events: list[dict[str, Any]]
# Индексы для быстрого поиска
location_by_event_id: dict[str, dict] = field(default_factory=dict)
device_by_click_id: dict[str, dict] = field(default_factory=dict)
geo_by_click_id: dict[str, dict] = field(default_factory=dict)
def __post_init__(self):
# Строим индексы для связности
for loc in self.location_events:
self.location_by_event_id[loc["event_id"]] = loc
for dev in self.device_events:
self.device_by_click_id[dev["click_id"]] = dev
for geo in self.geo_events:
self.geo_by_click_id[geo["click_id"]] = geo
@classmethod
def load(cls, data_dir: Path) -> "EventDictionary":
"""Загружает события из JSONL файлов."""
logger.info(f"Loading event dictionary from {data_dir}")
def load_jsonl(filename: str) -> list[dict]:
path = data_dir / filename
events = []
with open(path, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
events.append(json.loads(line))
logger.info(f" Loaded {len(events)} events from {filename}")
return events
return cls(
browser_events=load_jsonl("browser_events.jsonl"),
location_events=load_jsonl("location_events.jsonl"),
device_events=load_jsonl("device_events.jsonl"),
geo_events=load_jsonl("geo_events.jsonl"),
)
# ---------------------------------------------------------------------------
# Генерация событий
# ---------------------------------------------------------------------------
class EventGenerator:
"""Генератор событий с сохранением связности."""
def __init__(self, dictionary: EventDictionary, config: Config):
self.dictionary = dictionary
self.config = config
self.rng = random.Random(config.seed)
def _new_uuid(self) -> str:
"""Генерирует новый UUID."""
return str(uuid.uuid4())
def _current_timestamp(self) -> str:
"""Возвращает текущую метку времени в формате JSONL."""
return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S.%f")
def _hour_factor(self) -> float:
"""Возвращает коэффициент интенсивности в зависимости от часа дня."""
hour = datetime.now(timezone.utc).hour
# Дневное окно (9-18): 1.2
# Ночное окно (0-5): 0.7
# Остальное время: 1.0
if 9 <= hour <= 18:
return 1.2
elif 0 <= hour <= 5:
return 0.7
return 1.0
def _calculate_events_count(self) -> int:
"""Вычисляет количество событий для текущего тика (Poisson + jitter)."""
# Базовая интенсивность с учётом часа
lambda_minute = self.config.lambda_base_per_min * self._hour_factor()
# Масштабируем на длительность тика
lambda_tick = lambda_minute * (self.config.tick_seconds / 60.0)
# Генерируем Poisson
count = 0
L = math.exp(-lambda_tick)
p = 1.0
while p > L:
p *= self.rng.random()
count += 1
count -= 1
# Применяем jitter (вариативность)
if self.config.jitter_pct > 0:
jitter_factor = 1.0 + self.rng.uniform(
-self.config.jitter_pct / 100.0,
self.config.jitter_pct / 100.0
)
count = int(count * jitter_factor)
# Применяем границы
count = max(self.config.min_events_per_tick, min(count, self.config.max_events_per_tick))
return count
def generate_batch(self, batch_size: int) -> dict[str, list[dict]]:
"""
Генерирует батч событий с сохранением связей.
Возвращает словарь {topic: [events]}
"""
batch = {
"browser_events": [],
"location_events": [],
"device_events": [],
"geo_events": [],
}
for _ in range(batch_size):
# Выбираем случайное браузерное событие как базу
base_browser = self.rng.choice(self.dictionary.browser_events)
base_location = self.dictionary.location_by_event_id.get(base_browser["event_id"])
base_device = self.dictionary.device_by_click_id.get(base_browser["click_id"])
base_geo = self.dictionary.geo_by_click_id.get(base_browser["click_id"])
# Генерируем новые ID
new_event_id = self._new_uuid()
new_click_id = self._new_uuid()
new_timestamp = self._current_timestamp()
# Создаём новое браузерное событие
browser_event = {
**base_browser,
"event_id": new_event_id,
"click_id": new_click_id,
"event_timestamp": new_timestamp,
}
batch["browser_events"].append(browser_event)
# Связанное location событие
if base_location:
location_event = {
**base_location,
"event_id": new_event_id,
}
batch["location_events"].append(location_event)
# Связанное device событие
if base_device:
device_event = {
**base_device,
"click_id": new_click_id,
}
batch["device_events"].append(device_event)
# Связанное geo событие
if base_geo:
geo_event = {
**base_geo,
"click_id": new_click_id,
}
batch["geo_events"].append(geo_event)
return batch
# ---------------------------------------------------------------------------
# Batch record для истории
# ---------------------------------------------------------------------------
@dataclass
class BatchRecord:
"""Запись об отправленном батче."""
batch_id: str
started_at: datetime
finished_at: datetime
sent_total: int
sent_browser: int
sent_location: int
sent_device: int
sent_geo: int
status: str # 'success', 'partial', 'error'
error_message: str | None = None
def to_dict(self) -> dict:
"""Конвертирует в словарь для сериализации."""
return {
"batch_id": self.batch_id,
"started_at": self.started_at.isoformat(),
"finished_at": self.finished_at.isoformat(),
"sent_total": self.sent_total,
"sent_browser": self.sent_browser,
"sent_location": self.sent_location,
"sent_device": self.sent_device,
"sent_geo": self.sent_geo,
"status": self.status,
"error_message": self.error_message,
}
# ---------------------------------------------------------------------------
# In-memory fallback для истории
# ---------------------------------------------------------------------------
class InMemoryBatchHistory:
"""Fallback хранение истории batch в памяти."""
def __init__(self):
self.batches: list[BatchRecord] = []
def add(self, record: BatchRecord):
self.batches.append(record)
if len(self.batches) > 1000:
self.batches = self.batches[-1000:]
def get_stats(self) -> dict:
if not self.batches:
return {}
total = len(self.batches)
success = sum(1 for b in self.batches if b.status == "success")
return {
"total_batches": total,
"success_rate": success / total if total > 0 else 0,
"last_batch_status": self.batches[-1].status if self.batches else None,
}
# ---------------------------------------------------------------------------
# Kafka history - пишет историю в отдельный топик
# ---------------------------------------------------------------------------
class KafkaBatchHistory:
"""Хранение истории batch в Kafka (отдельный топик)."""
HISTORY_TOPIC = "generator_batch_history"
def __init__(self, bootstrap_servers: str):
self.bootstrap_servers = bootstrap_servers
self.producer = None
self._initialized = False
self._connect()
def _connect(self):
"""Устанавливает соединение с Kafka."""
KafkaProducerCls, KafkaErrorCls = _import_kafka()
logger.info(f"Connecting to Kafka for history at {self.bootstrap_servers}")
try:
self.producer = KafkaProducerCls(
bootstrap_servers=self.bootstrap_servers,
value_serializer=lambda v: json.dumps(v).encode("utf-8"),
key_serializer=lambda k: k.encode("utf-8") if k else None,
retries=3,
retry_backoff_ms=1000,
)
self._initialized = True
logger.info("Connected to Kafka for history successfully")
except Exception as e:
logger.warning(f"Failed to connect to Kafka for history: {e}")
self._initialized = False
def add(self, record: BatchRecord):
"""Добавляет запись в историю (топик Kafka)."""
if not self._initialized or self.producer is None:
logger.debug("Kafka history not available, skipping")
return
try:
key = record.batch_id
value = record.to_dict()
self.producer.send(self.HISTORY_TOPIC, key=key, value=value)
except Exception as e:
logger.warning(f"Failed to write batch history to Kafka: {e}")
def flush(self):
"""Сбрасывает буфер."""
if self.producer:
self.producer.flush()
def close(self):
"""Закрывает соединение."""
if self.producer:
self.producer.close()
def get_stats(self) -> dict:
"""Возвращает статус подключения."""
return {
"initialized": self._initialized,
"topic": self.HISTORY_TOPIC,
}
# ---------------------------------------------------------------------------
# Kafka publisher
# ---------------------------------------------------------------------------
class KafkaPublisher:
"""Публикация событий в Kafka."""
def __init__(self, bootstrap_servers: str):
self.bootstrap_servers = bootstrap_servers
self.producer = None
self._connect()
def _connect(self):
"""Устанавливает соединение с Kafka."""
KafkaProducerCls, KafkaErrorCls = _import_kafka()
logger.info(f"Connecting to Kafka at {self.bootstrap_servers}")
try:
self.producer = KafkaProducerCls(
bootstrap_servers=self.bootstrap_servers,
value_serializer=lambda v: json.dumps(v).encode("utf-8"),
key_serializer=lambda k: k.encode("utf-8") if k else None,
batch_size=16384,
linger_ms=100,
retries=3,
retry_backoff_ms=1000,
)
logger.info("Connected to Kafka successfully")
except KafkaErrorCls as e:
logger.error(f"Failed to connect to Kafka: {e}")
raise
def publish(self, topic: str, events: list[dict]) -> tuple[int, int]:
"""
Публикует события в топик.
Returns:
(sent_count, error_count)
"""
if not self.producer:
raise RuntimeError("Producer not connected")
_, KafkaErrorCls = _import_kafka()
sent = 0
errors = 0
futures = []
for event in events:
key = event.get("event_id") or event.get("click_id")
try:
future = self.producer.send(topic, key=key, value=event)
futures.append(future)
except Exception as e:
logger.error(f"Failed to send message to {topic}: {e}")
errors += 1
METRICS_ERRORS_TOTAL.labels(topic=topic).inc()
# Ждём подтверждений
for future in futures:
try:
future.get(timeout=10)
sent += 1
METRICS_EVENTS_TOTAL.labels(topic=topic).inc()
except Exception as e:
logger.error(f"Failed to confirm message delivery: {e}")
errors += 1
METRICS_ERRORS_TOTAL.labels(topic=topic).inc()
return sent, errors
def flush(self):
"""Сбрасывает буфер."""
if self.producer:
self.producer.flush()
def close(self):
"""Закрывает соединение."""
if self.producer:
self.producer.close()
# ---------------------------------------------------------------------------
# Основной цикл генератора
# ---------------------------------------------------------------------------
class GeneratorService:
"""Основной сервис генератора."""
def __init__(self, config: Config):
self.config = config
self.dictionary = EventDictionary.load(config.data_dir)
self.generator = EventGenerator(self.dictionary, config)
self.publisher: KafkaPublisher | None = None
self.history: KafkaBatchHistory | InMemoryBatchHistory | None = None
self._running = False
def start(self):
"""Запускает основной цикл."""
if not self.config.enabled:
logger.warning("Generator is disabled (GEN_ENABLED=false)")
return
# Запускаем HTTP-сервер для Prometheus метрик
logger.info(f"Starting metrics server on port {self.config.metrics_port}")
start_http_server(self.config.metrics_port)
logger.info("Starting generator service...")
logger.info(f"Configuration: tick={self.config.tick_seconds}s, "
f"lambda_base={self.config.lambda_base_per_min}/min, "
f"jitter={self.config.jitter_pct}%")
# Подключаемся к Kafka для публикации событий
self.publisher = KafkaPublisher(self.config.kafka_bootstrap_servers)
# Инициализируем историю (Kafka с fallback на in-memory)
self.history = KafkaBatchHistory(self.config.kafka_bootstrap_servers)
if not self.history._initialized:
logger.warning("Falling back to InMemoryBatchHistory")
self.history = InMemoryBatchHistory()
self._running = True
try:
self._main_loop()
except KeyboardInterrupt:
logger.info("Received shutdown signal")
finally:
self.stop()
def stop(self):
"""Останавливает сервис."""
logger.info("Stopping generator service...")
self._running = False
if self.publisher:
self.publisher.close()
if isinstance(self.history, KafkaBatchHistory):
self.history.close()
def _main_loop(self):
"""Основной цикл тиков."""
tick = 0
while self._running:
tick += 1
tick_start = time.time()
batch_id = str(uuid.uuid4())[:8]
with METRICS_TICK_DURATION.time():
logger.info(f"=== Tick {tick} (batch_id={batch_id}) ===")
try:
# Вычисляем количество событий
events_count = self.generator._calculate_events_count()
logger.info(f"Generating ~{events_count} base events")
# Генерируем батч
gen_start = time.time()
batch = self.generator.generate_batch(events_count)
gen_duration = time.time() - gen_start
# Публикуем в Kafka
pub_start = time.time()
total_sent = 0
total_errors = 0
sent_counts = {}
for topic, events in batch.items():
if events:
sent, errors = self.publisher.publish(topic, events)
sent_counts[topic] = {"sent": sent, "errors": errors}
total_sent += sent
total_errors += errors
self.publisher.flush()
if isinstance(self.history, KafkaBatchHistory):
self.history.flush()
pub_duration = time.time() - pub_start
# Определяем статус
if total_errors == 0:
status = "success"
elif total_sent > 0:
status = "partial"
else:
status = "error"
# Обновляем метрику последнего успешного тика
if status in ("success", "partial"):
METRICS_LAST_SUCCESS.set_to_current_time()
# Сохраняем в историю
batch_record = BatchRecord(
batch_id=batch_id,
started_at=datetime.fromtimestamp(tick_start, tz=timezone.utc),
finished_at=datetime.now(timezone.utc),
sent_total=total_sent,
sent_browser=sent_counts.get("browser_events", {}).get("sent", 0),
sent_location=sent_counts.get("location_events", {}).get("sent", 0),
sent_device=sent_counts.get("device_events", {}).get("sent", 0),
sent_geo=sent_counts.get("geo_events", {}).get("sent", 0),
status=status,
error_message=None if status == "success" else f"Errors: {total_errors}",
)
self.history.add(batch_record)
# Логируем результат
tick_duration = time.time() - tick_start
logger.info(
f"Batch {batch_id} completed: "
f"sent={total_sent}, errors={total_errors}, "
f"gen_time={gen_duration:.3f}s, pub_time={pub_duration:.3f}s, "
f"total_time={tick_duration:.3f}s"
)
for topic, counts in sent_counts.items():
if counts["sent"] > 0:
logger.info(f" {topic}: {counts['sent']} sent")
except Exception as e:
logger.exception(f"Error in tick {tick}: {e}")
self.history.add(
BatchRecord(
batch_id=batch_id,
started_at=datetime.fromtimestamp(tick_start, tz=timezone.utc),
finished_at=datetime.now(timezone.utc),
sent_total=0,
sent_browser=0,
sent_location=0,
sent_device=0,
sent_geo=0,
status="error",
error_message=str(e),
)
)
# Ждём до следующего тика
elapsed = time.time() - tick_start
sleep_time = max(0, self.config.tick_seconds - elapsed)
if sleep_time > 0:
logger.debug(f"Sleeping for {sleep_time:.1f}s until next tick")
time.sleep(sleep_time)
# ---------------------------------------------------------------------------
# Точка входа
# ---------------------------------------------------------------------------
def main():
try:
config = Config()
service = GeneratorService(config)
service.start()
except Exception as e:
logger.exception(f"Fatal error: {e}")
sys.exit(1)
if __name__ == "__main__":
main()