fix(kafka): обновлено хранение state генератора без pickle

- Зачем:
  - упростить и сделать безопаснее сериализацию состояния генератора
  - повысить устойчивость старта при временной недоступности Kafka
- Что:
  - заменена сериализация rng state на JSON-safe формат без pickle/base64
  - добавлено восстановление tuple-структуры rng state после json
  - добавлен общий retry с exponential backoff для ensure_topics
  - обновлена документация по continuity и ограничениям после рестарта
  - расширены тесты state-сценариев и retry-логики
- Проверка:
  - make generator-test
  - 50 passed
This commit is contained in:
2026-06-09 17:27:16 +03:00
committed by Dmitry Dementiev
parent 31e8901eb0
commit e9a46e47ae
3 changed files with 192 additions and 56 deletions
+4 -2
View File
@@ -116,9 +116,11 @@ docker compose exec kafka /opt/kafka/bin/kafka-console-consumer.sh \
Генератор сохраняет своё состояние между перезапусками в Kafka-топик `generator_state` (compact topic). Это позволяет: Генератор сохраняет своё состояние между перезапусками в Kafka-топик `generator_state` (compact topic). Это позволяет:
- Продолжить нумерацию тиков с места остановки - Продолжить нумерацию тиков с места остановки (continuity)
- Сохранить последовательность случайных чисел (RNG state) - Сохранить последовательность случайных чисел (RNG state)
- Избежать дублирования при рестарте - Восстановить интенсивность генерации после рестарта
**Важно:** восстанавливается continuity по номеру тика и интенсивности, но не гарантируется отсутствие дублирования событий — `event_id` и `click_id` всегда генерируются заново (`uuid4()`).
### Как работает ### Как работает
+62 -18
View File
@@ -6,12 +6,10 @@
держим целевую интенсивность events/min без крупных минутных batch. держим целевую интенсивность events/min без крупных минутных batch.
""" """
import base64
import json import json
import logging import logging
import math import math
import os import os
import pickle
import random import random
import sys import sys
import time import time
@@ -368,25 +366,33 @@ class BatchRecord:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# State record для сохранения состояния генератора # State record для сохранения состояния генератора
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def _nested_list_to_tuple(obj):
"""Рекурсивно преобразует list в tuple (для восстановления RNG state после JSON)."""
if isinstance(obj, list):
return tuple(_nested_list_to_tuple(x) for x in obj)
return obj
@dataclass @dataclass
class GeneratorState: class GeneratorState:
"""Состояние генератора для восстановления после рестарта.""" """Состояние генератора для восстановления после рестарта.
Используем JSON-safe сериализацию:
- rng_state от random.getstate() - кортеж из простых типов (int, tuple),
безопасно сериализуется в JSON напрямую без pickle
"""
tick: int tick: int
rng_state: tuple # результат random.getstate() rng_state: tuple # результат random.getstate() - JSON-serializable
last_batch_id: str last_batch_id: str
last_timestamp: datetime last_timestamp: datetime
version: str = "1.0" version: str = "1.0"
def to_dict(self) -> dict: def to_dict(self) -> dict:
"""Конвертирует в словарь для сериализации.""" """Конвертирует в словарь для JSON-сериализации."""
# Сериализуем rng_state через pickle + base64
rng_state_bytes = pickle.dumps(self.rng_state)
rng_state_b64 = base64.b64encode(rng_state_bytes).decode("utf-8")
return { return {
"tick": self.tick, "tick": self.tick,
"rng_state": rng_state_b64, "rng_state": self.rng_state, # tuple из int - JSON-serializable
"last_batch_id": self.last_batch_id, "last_batch_id": self.last_batch_id,
"last_timestamp": self.last_timestamp.isoformat(), "last_timestamp": self.last_timestamp.isoformat(),
"version": self.version, "version": self.version,
@@ -394,14 +400,10 @@ class GeneratorState:
@classmethod @classmethod
def from_dict(cls, data: dict) -> "GeneratorState": def from_dict(cls, data: dict) -> "GeneratorState":
"""Создаёт состояние из словаря.""" """Создаёт состояние из словаря (JSON-only, без pickle)."""
# Десериализуем rng_state
rng_state_bytes = base64.b64decode(data["rng_state"])
rng_state = pickle.loads(rng_state_bytes)
return cls( return cls(
tick=data["tick"], tick=data["tick"],
rng_state=rng_state, rng_state=_nested_list_to_tuple(data["rng_state"]), # рекурсивно list->tuple
last_batch_id=data["last_batch_id"], last_batch_id=data["last_batch_id"],
last_timestamp=datetime.fromisoformat(data["last_timestamp"]), last_timestamp=datetime.fromisoformat(data["last_timestamp"]),
version=data.get("version", "1.0"), version=data.get("version", "1.0"),
@@ -445,7 +447,11 @@ class KafkaStateManager:
self.producer.close() self.producer.close()
def load(self) -> GeneratorState | None: def load(self) -> GeneratorState | None:
"""Загружает последнее состояние из топика.""" """Загружает последнее состояние из топика.
Для compact topic хранится только последнее значение для ключа,
поэтому читаем все сообщения и берём последнее с нужным ключом.
"""
from kafka import KafkaConsumer from kafka import KafkaConsumer
logger.info(f"Loading state from topic {self.STATE_TOPIC}") logger.info(f"Loading state from topic {self.STATE_TOPIC}")
@@ -479,15 +485,51 @@ class KafkaStateManager:
return None return None
# ---------------------------------------------------------------------------
# Утилиты для работы с Kafka с retry/backoff
# ---------------------------------------------------------------------------
def _with_retry(operation, max_retries: int = 5, base_delay: float = 1.0, max_delay: float = 30.0):
"""Выполняет операцию с экспоненциальным backoff и ограниченным числом попыток.
Args:
operation: функция для выполнения
max_retries: максимальное число попыток
base_delay: начальная задержка между попытками (сек)
max_delay: максимальная задержка между попытками (сек)
Returns:
результат операции
Raises:
последнее исключение после исчерпания попыток
"""
import time
last_exception = None
for attempt in range(max_retries):
try:
return operation()
except Exception as e:
last_exception = e
if attempt < max_retries - 1:
delay = min(base_delay * (2 ** attempt), max_delay)
logger.warning(f"Operation failed (attempt {attempt + 1}/{max_retries}): {e}. Retrying in {delay:.1f}s...")
time.sleep(delay)
else:
logger.error(f"Operation failed after {max_retries} attempts: {e}")
raise last_exception
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Создание топиков для истории и состояния # Создание топиков для истории и состояния
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def ensure_topics(bootstrap_servers: str) -> None: def ensure_topics(bootstrap_servers: str) -> None:
"""Создаёт необходимые топики если они не существуют.""" """Создаёт необходимые топики если они не существуют (с retry на подключение)."""
from kafka import KafkaAdminClient from kafka import KafkaAdminClient
from kafka.admin import NewTopic from kafka.admin import NewTopic
from kafka.errors import TopicAlreadyExistsError from kafka.errors import TopicAlreadyExistsError
def _create_topics():
admin_client = KafkaAdminClient(bootstrap_servers=bootstrap_servers) admin_client = KafkaAdminClient(bootstrap_servers=bootstrap_servers)
try: try:
# Топик для истории батчей (обычный, с retention) # Топик для истории батчей (обычный, с retention)
@@ -518,6 +560,8 @@ def ensure_topics(bootstrap_servers: str) -> None:
finally: finally:
admin_client.close() admin_client.close()
_with_retry(_create_topics, max_retries=5, base_delay=1.0)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Kafka history - пишет историю в отдельный топик # Kafka history - пишет историю в отдельный топик
+99 -9
View File
@@ -1,8 +1,7 @@
""" """
Тесты сохранения и восстановления состояния генератора. Тесты сохранения и восстановления состояния генератора.
""" """
import base64 import json
import pickle
from datetime import datetime, timezone from datetime import datetime, timezone
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
@@ -48,7 +47,7 @@ class TestGeneratorState:
assert state.version == "1.0" assert state.version == "1.0"
def test_to_dict_serialization(self): def test_to_dict_serialization(self):
"""Сериализация в словарь.""" """Сериализация в словарь (JSON-safe, без pickle)."""
now = datetime.now(timezone.utc) now = datetime.now(timezone.utc)
rng_state = (3, (1, 2, 3), None) rng_state = (3, (1, 2, 3), None)
@@ -66,13 +65,14 @@ class TestGeneratorState:
assert data["last_timestamp"] == now.isoformat() assert data["last_timestamp"] == now.isoformat()
assert data["version"] == "1.0" assert data["version"] == "1.0"
# Проверяем что rng_state сериализован через pickle+base64 # Проверяем что rng_state сериализован как tuple (JSON-safe, без pickle)
assert "rng_state" in data assert "rng_state" in data
assert isinstance(data["rng_state"], str) assert data["rng_state"] == rng_state
# Проверяем что можно десериализовать # Проверяем что можно сериализовать в JSON и восстановить
decoded = base64.b64decode(data["rng_state"]) json_str = json.dumps(data)
restored_rng = pickle.loads(decoded) restored_data = json.loads(json_str)
assert restored_rng == rng_state restored_state = GeneratorState.from_dict(restored_data)
assert restored_state.rng_state == rng_state
def test_from_dict_deserialization(self): def test_from_dict_deserialization(self):
"""Десериализация из словаря.""" """Десериализация из словаря."""
@@ -289,3 +289,93 @@ class TestKafkaStateManager:
manager.close() manager.close()
mock_producer.close.assert_called_once() mock_producer.close.assert_called_once()
class TestJsonSafeState:
"""Тесты JSON-safe сериализации state (без pickle)."""
def test_json_roundtrip_with_nested_tuples(self):
"""Проверка что nested tuple корректно восстанавливается после JSON."""
from generator import _nested_list_to_tuple
# Симулируем что получаем после json.loads() - все tuple становятся list
json_loaded = [3, [1, 2, 3], None]
result = _nested_list_to_tuple(json_loaded)
assert result == (3, (1, 2, 3), None)
assert isinstance(result, tuple)
assert isinstance(result[1], tuple)
def test_json_roundtrip_rng_state(self):
"""Полный цикл: rng.getstate() -> JSON -> from_dict -> setstate."""
import random
rng = random.Random(42)
# Делаем несколько вызовов
values_before = [rng.random() for _ in range(10)]
# Сохраняем state
state = GeneratorState(
tick=100,
rng_state=rng.getstate(),
last_batch_id="test123",
last_timestamp=datetime.now(timezone.utc),
)
# Сериализуем через JSON (как в Kafka)
data = state.to_dict()
json_str = json.dumps(data)
restored_data = json.loads(json_str)
# Восстанавливаем
restored_state = GeneratorState.from_dict(restored_data)
# Проверяем что RNG state восстановлен корректно
rng2 = random.Random()
rng2.setstate(restored_state.rng_state)
# Проверяем что следующие значения совпадают
values_after = [rng2.random() for _ in range(5)]
# Оригинальный RNG должен дать те же значения
values_expected = [rng.random() for _ in range(5)]
assert values_after == values_expected
class TestWithRetry:
"""Тесты функции _with_retry."""
def test_success_on_first_attempt(self):
"""Успех с первой попытки."""
from generator import _with_retry
operation = MagicMock(return_value="success")
result = _with_retry(operation, max_retries=3, base_delay=0.01)
assert result == "success"
assert operation.call_count == 1
def test_success_after_retries(self):
"""Успех после нескольких попыток."""
from generator import _with_retry
operation = MagicMock(side_effect=[Exception("fail1"), Exception("fail2"), "success"])
result = _with_retry(operation, max_retries=3, base_delay=0.01)
assert result == "success"
assert operation.call_count == 3
def test_failure_after_all_retries(self):
"""Исчерпание всех попыток."""
from generator import _with_retry
operation = MagicMock(side_effect=Exception("always fails"))
with pytest.raises(Exception, match="always fails"):
_with_retry(operation, max_retries=3, base_delay=0.01)
assert operation.call_count == 3