233 lines
8.5 KiB
Python
233 lines
8.5 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""批量写入缓冲 + 幂等去重 —— 「数据不丢不重」的实现保障（PRD 5.2 / Issue #4）。
|
|||
|
|
|
|||
|
|
设计：
|
|||
|
|
- push() 单条/批量入队：按 (device_id, point_id, ts) 幂等去重（**不重**）；
|
|||
|
|
- 触发 flush：缓冲达到 batch_size 或距上次 flush 超过 flush_interval
|
|||
|
|
（**批量写入**，默认 5000 条 / 0.1s，对齐 PRD 5.2「5k 条/100ms」基线）；
|
|||
|
|
- at-least-once：sink 写成功才清空缓冲；失败整批保留、下次 flush 重发（**不丢**）；
|
|||
|
|
- sink 侧按去重键幂等（MemorySink / TdengineSink 均已实现）：失败重发时
|
|||
|
|
即使发生部分写入也不会重复落库，以 sink 返回的成功条数为准。
|
|||
|
|
|
|||
|
|
StoreSink 抽象：
|
|||
|
|
- MemorySink —— 内存落库（本地联调 / 端到端测试）；
|
|||
|
|
- TdengineSink —— 把样本批转为按子表分组的批量 INSERT，交客户端执行
|
|||
|
|
（不依赖 taospy；executor 由调用方注入连接适配器或联调函数）。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
from abc import ABC, abstractmethod
|
|||
|
|
from typing import Callable, Dict, List, Optional, Tuple
|
|||
|
|
|
|||
|
|
from .tdengine_schema import build_batch_insert
|
|||
|
|
|
|||
|
|
# 幂等去重键：同一设备-测点-时间戳视为同一条样本（重复推送只落一次）
|
|||
|
|
DEDUP_KEYS: Tuple[str, ...] = ("device_id", "point_id", "ts")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _dedup_key(row: dict, keys: Tuple[str, ...]) -> tuple:
|
|||
|
|
"""从样本行提取去重键（缺失字段统一为字符串 "None"，保证键可哈希稳定）。"""
|
|||
|
|
return tuple(str(row.get(k)) for k in keys)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class StoreSink(ABC):
|
|||
|
|
"""存储端抽象：批量写入的落库实现。"""
|
|||
|
|
|
|||
|
|
@abstractmethod
|
|||
|
|
def write(self, rows: List[dict]) -> int:
|
|||
|
|
"""写入一批样本。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
rows: 样本列表（含 device_id / point_id / value / ts / quality）。
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
成功写入的条数。
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
写入失败抛异常 —— 调用方（BatchWriter）保留缓冲待重发（不丢）。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def close(self) -> None:
|
|||
|
|
"""释放资源（可选覆写；幂等，可多次调用）。"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MemorySink(StoreSink):
|
|||
|
|
"""内存落库：按去重键幂等、行保序（本地联调 / 端到端测试用）。"""
|
|||
|
|
|
|||
|
|
def __init__(self, dedup_keys: Tuple[str, ...] = DEDUP_KEYS):
|
|||
|
|
self.dedup_keys = dedup_keys
|
|||
|
|
self._rows: List[dict] = []
|
|||
|
|
self._seen: set = set()
|
|||
|
|
self.write_count: int = 0
|
|||
|
|
self.closed: bool = False
|
|||
|
|
|
|||
|
|
def write(self, rows: List[dict]) -> int:
|
|||
|
|
n = 0
|
|||
|
|
for row in rows:
|
|||
|
|
key = _dedup_key(row, self.dedup_keys)
|
|||
|
|
if key in self._seen:
|
|||
|
|
continue
|
|||
|
|
self._seen.add(key)
|
|||
|
|
self._rows.append(dict(row))
|
|||
|
|
n += 1
|
|||
|
|
self.write_count += n
|
|||
|
|
return n
|
|||
|
|
|
|||
|
|
def close(self) -> None:
|
|||
|
|
self.closed = True
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def rows(self) -> List[dict]:
|
|||
|
|
"""已落库样本（保序副本，供验收断言）。"""
|
|||
|
|
return list(self._rows)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TdengineSink(StoreSink):
|
|||
|
|
"""TDengine 批量落库：样本批 → 按子表分组的批量 INSERT，交 executor 执行。
|
|||
|
|
|
|||
|
|
不依赖 taospy：executor(statements) 由调用方注入（连接适配器执行 SQL，
|
|||
|
|
或打印 SQL 的联调函数）。子表命名由模板命名器推导（点位字典驱动）。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
naming: "TemplateNaming",
|
|||
|
|
executor: Callable[[List[str]], None],
|
|||
|
|
dedup_keys: Tuple[str, ...] = DEDUP_KEYS,
|
|||
|
|
):
|
|||
|
|
self.naming = naming
|
|||
|
|
self.executor = executor
|
|||
|
|
self.dedup_keys = dedup_keys
|
|||
|
|
self._seen: set = set()
|
|||
|
|
self.write_count: int = 0
|
|||
|
|
|
|||
|
|
def write(self, rows: List[dict]) -> int:
|
|||
|
|
fresh = []
|
|||
|
|
for row in rows:
|
|||
|
|
key = _dedup_key(row, self.dedup_keys)
|
|||
|
|
if key in self._seen:
|
|||
|
|
continue
|
|||
|
|
self._seen.add(key)
|
|||
|
|
fresh.append(row)
|
|||
|
|
if fresh:
|
|||
|
|
# 批量 INSERT：每子表一条 VALUES 多值语句（对齐 PRD 5.2 批量写入）
|
|||
|
|
self.executor(build_batch_insert(self.naming, fresh))
|
|||
|
|
self.write_count += len(fresh)
|
|||
|
|
return len(fresh)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BatchWriter:
|
|||
|
|
"""批量写入缓冲：批量聚合 + 幂等去重 + 失败重试（**不丢不重**）。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
sink: StoreSink 落库实现（MemorySink / TdengineSink / 自定义）。
|
|||
|
|
batch_size: 缓冲达到该条数触发一次 flush（默认 5000，5k/100ms 基线）。
|
|||
|
|
flush_interval: 距上次 flush 超过该秒数，push 时自动 flush（默认 0.1s）。
|
|||
|
|
dedup_keys: 幂等去重键（默认 设备-测点-时间戳）。
|
|||
|
|
clock: 时钟注入（测试可传假时钟），默认 time.monotonic。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
sink: StoreSink,
|
|||
|
|
batch_size: int = 5000,
|
|||
|
|
flush_interval: float = 0.1,
|
|||
|
|
dedup_keys: Tuple[str, ...] = DEDUP_KEYS,
|
|||
|
|
clock: Callable[[], float] = time.monotonic,
|
|||
|
|
):
|
|||
|
|
self.sink = sink
|
|||
|
|
self.batch_size = max(1, int(batch_size))
|
|||
|
|
self.flush_interval = max(0.0, float(flush_interval))
|
|||
|
|
self.dedup_keys = dedup_keys
|
|||
|
|
self._clock = clock
|
|||
|
|
self._buffer: List[dict] = []
|
|||
|
|
self._seen: set = set() # 已成功提交的去重键（跨批幂等窗口）
|
|||
|
|
self._buffered_keys: set = set() # 当前缓冲内去重键（缓冲未提交前即拦截重复）
|
|||
|
|
self._last_flush = self._clock()
|
|||
|
|
self._stats: Dict[str, int] = {
|
|||
|
|
"received": 0,
|
|||
|
|
"written": 0,
|
|||
|
|
"duplicates": 0,
|
|||
|
|
"flush_count": 0,
|
|||
|
|
"failed_flushes": 0,
|
|||
|
|
}
|
|||
|
|
self._lock = threading.Lock()
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
def push(self, row: dict) -> bool:
|
|||
|
|
"""入队一条样本。
|
|||
|
|
|
|||
|
|
自动刷盘（达到 batch_size / 超 flush_interval）失败时**不中断入队**：
|
|||
|
|
失败批保留缓冲待重发（不丢），统计计入 failed_flushes，可经 stats() 观察。
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
True=接受入队；False=与已入队/已提交样本重复（幂等丢弃，计入 duplicates）。
|
|||
|
|
"""
|
|||
|
|
with self._lock:
|
|||
|
|
key = _dedup_key(row, self.dedup_keys)
|
|||
|
|
if key in self._seen or key in self._buffered_keys:
|
|||
|
|
self._stats["duplicates"] += 1
|
|||
|
|
return False
|
|||
|
|
self._buffer.append(dict(row))
|
|||
|
|
self._buffered_keys.add(key)
|
|||
|
|
self._stats["received"] += 1
|
|||
|
|
if len(self._buffer) >= self.batch_size or (
|
|||
|
|
self._clock() - self._last_flush >= self.flush_interval
|
|||
|
|
):
|
|||
|
|
try:
|
|||
|
|
self._flush_locked()
|
|||
|
|
except Exception:
|
|||
|
|
pass # 失败保留缓冲，由下次 push / flush 重试
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
def push_many(self, rows: List[dict]) -> int:
|
|||
|
|
"""批量入队，返回接受条数（重复自动过滤）。"""
|
|||
|
|
accepted = 0
|
|||
|
|
for row in rows:
|
|||
|
|
if self.push(row):
|
|||
|
|
accepted += 1
|
|||
|
|
return accepted
|
|||
|
|
|
|||
|
|
def flush(self) -> int:
|
|||
|
|
"""主动刷新：把缓冲整批写入 sink（失败抛异常并保留缓冲）。"""
|
|||
|
|
with self._lock:
|
|||
|
|
return self._flush_locked()
|
|||
|
|
|
|||
|
|
def _flush_locked(self) -> int:
|
|||
|
|
if not self._buffer:
|
|||
|
|
return 0
|
|||
|
|
batch = list(self._buffer)
|
|||
|
|
try:
|
|||
|
|
n = self.sink.write(batch)
|
|||
|
|
except Exception:
|
|||
|
|
# 写入失败：缓冲整体保留，下次 flush 重发（不丢）
|
|||
|
|
self._stats["failed_flushes"] += 1
|
|||
|
|
raise
|
|||
|
|
# 写入成功：清缓冲；sink 侧已按去重键幂等，重发不会重复落库
|
|||
|
|
self._buffer.clear()
|
|||
|
|
self._buffered_keys.clear()
|
|||
|
|
for row in batch:
|
|||
|
|
self._seen.add(_dedup_key(row, self.dedup_keys))
|
|||
|
|
self._stats["written"] += n
|
|||
|
|
self._stats["flush_count"] += 1
|
|||
|
|
self._last_flush = self._clock()
|
|||
|
|
return n
|
|||
|
|
|
|||
|
|
def pending(self) -> int:
|
|||
|
|
"""当前缓冲中待写入条数。"""
|
|||
|
|
with self._lock:
|
|||
|
|
return len(self._buffer)
|
|||
|
|
|
|||
|
|
def stats(self) -> dict:
|
|||
|
|
"""累计统计（received/written/duplicates/flush_count/failed_flushes）。"""
|
|||
|
|
with self._lock:
|
|||
|
|
return dict(self._stats)
|
|||
|
|
|
|||
|
|
def close(self) -> None:
|
|||
|
|
"""收尾：刷新剩余缓冲并关闭 sink（写失败异常上抛，由调用方处理）。"""
|
|||
|
|
with self._lock:
|
|||
|
|
self._flush_locked()
|
|||
|
|
self.sink.close()
|