mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 15:34:16 +08:00
perf(polars): 并发闸+写锁收缩+看门狗三层防死锁, 升级 polars 1.44
线上曾出现并发 LazyFrame.collect 触发 polars streaming 执行器死锁, 叠加 _write_lock 区间内做重活, 放大为全站请求冻结。本次按触发缩小、 爆炸半径收缩、自动恢复三层布防: - polars_guard: BoundedSemaphore 并发闸 (总闸 4 + 后台车道 2, 后台 先拿子闸再拿总闸防死锁); repository 18 处 collect 按交互/后台分级接入 - repository 写锁区间收缩: 分区合并移出锁外, 锁内 (mtime_ns,size) 指纹校验 + 3 次乐观重试, 失败回退锁内合并; 5 处 _write_lock 重构 - watchdog: 周期探测 collect 闸与全局写锁, 连续 2 次失败退出交由 supervisor 拉起 (可配置, 默认开) - polars >=1.44,<1.45 (1.44.1); 附并发压测脚本 scripts/stress_polars_concurrency.py 供复现验证 另: config 新增 polars_collect_permits / watchdog_* / strategy_run_all_workers / strategy_run_all_first_return_s 旋钮 (后两者供后续 run_all 优化提交使用, 默认保持旧行为基准)。
This commit is contained in:
@@ -112,6 +112,29 @@ class Settings(BaseSettings):
|
||||
backtest_matrix_cache_prewarm: bool = True
|
||||
backtest_matrix_cache_prewarm_years: int = 5
|
||||
|
||||
# polars collect 并发闸 — polars 共享执行器在多线程并发 collect 下存在死锁
|
||||
# (上游 #24448/#25754 同族), 限流并发是社区验证的缓解手段。background 限额
|
||||
# 保证预热/增量等后台计算不占满闸位饿死页面读请求。
|
||||
polars_collect_permits: int = 4
|
||||
polars_collect_background_permits: int = 2
|
||||
|
||||
# 后端自愈看门狗 — 探测 collect 闸与全局写锁, 连续失败即退出交由
|
||||
# supervisor 拉起 (见 app/watchdog.py)。误伤防护靠保守阈值。
|
||||
watchdog_enabled: bool = True
|
||||
watchdog_interval_s: float = 30.0
|
||||
watchdog_probe_timeout_s: float = 15.0
|
||||
watchdog_failure_threshold: int = 2
|
||||
|
||||
# 策略批量执行 (run_all / 策略页全量跑) 的并发 worker 上限。实测 2026-09-07:
|
||||
# polars eager 操作内部已多线程并行, 外层再并发 4 worker 属超订, 41 策略
|
||||
# 299.6s 慢于串行 — 默认 1 (串行)。保留开关供配合 POLARS_MAX_THREADS 调优实验。
|
||||
strategy_run_all_workers: int = 1
|
||||
|
||||
# run_all 渐进式返回: HTTP 同步等待时限 (秒)。策略按历史耗时升序执行,
|
||||
# 到点后已算完的随响应返回, 未算完的转后台继续算并逐个写入策略缓存,
|
||||
# 前端轮询 cached-summary 点亮卡片。0 = 关闭 (整段阻塞, 旧行为)。
|
||||
strategy_run_all_first_return_s: float = 15.0
|
||||
|
||||
# Auth — 首次启动时预置访问密码(明文, 仅用于初始化, 详见 services/auth.bootstrap_from_env)
|
||||
# 公网服务器部署时免去 SSH 端口转发设密码的麻烦。写入 auth.json(哈希)后即不再读取。
|
||||
auth_password: str = ""
|
||||
@@ -140,6 +163,20 @@ class Settings(BaseSettings):
|
||||
raise ValueError("ai_max_output_tokens must be positive")
|
||||
if self.ai_context_window <= 0:
|
||||
raise ValueError("ai_context_window must be positive")
|
||||
if self.polars_collect_permits < 2:
|
||||
raise ValueError("polars_collect_permits must be >= 2")
|
||||
if not 1 <= self.polars_collect_background_permits < self.polars_collect_permits:
|
||||
raise ValueError(
|
||||
"polars_collect_background_permits must be in [1, polars_collect_permits)"
|
||||
)
|
||||
if self.watchdog_interval_s <= 0 or self.watchdog_probe_timeout_s <= 0:
|
||||
raise ValueError("watchdog intervals must be positive")
|
||||
if self.watchdog_failure_threshold < 1:
|
||||
raise ValueError("watchdog_failure_threshold must be >= 1")
|
||||
if self.strategy_run_all_workers < 1:
|
||||
raise ValueError("strategy_run_all_workers must be >= 1")
|
||||
if self.strategy_run_all_first_return_s < 0:
|
||||
raise ValueError("strategy_run_all_first_return_s must be >= 0")
|
||||
return self
|
||||
|
||||
@property
|
||||
|
||||
@@ -238,6 +238,10 @@ async def _application_lifespan(app: FastAPI):
|
||||
financial_scheduler.start(store.data_dir, capset)
|
||||
app.state.financial_scheduler = financial_scheduler
|
||||
|
||||
# 自愈看门狗: 探测 polars 闸与写锁, 僵死时退出交由 supervisor 拉起 (兜底层)。
|
||||
from app.watchdog import start_watchdog
|
||||
app.state.watchdog = start_watchdog(app.state, repo)
|
||||
|
||||
# 策略引擎
|
||||
from app.strategy.engine import StrategyEngine
|
||||
from app.strategy import config as strategy_config
|
||||
@@ -355,6 +359,9 @@ async def _application_lifespan(app: FastAPI):
|
||||
yield
|
||||
finally:
|
||||
repo._on_refresh_done = None # noqa: SLF001
|
||||
wd = getattr(app.state, "watchdog", None)
|
||||
if wd:
|
||||
await wd.stop()
|
||||
if not matrix_prewarm_owner.shutdown(timeout=5.0):
|
||||
logger.warning("matrix cache prewarm did not stop within 5 seconds")
|
||||
mmanager = getattr(app.state, "mining_manager", None)
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
"""polars collect 并发闸。
|
||||
|
||||
polars 的共享执行器 (rayon 工作池 + 流式引擎异步运行时) 在多线程并发 collect
|
||||
时存在死锁问题: 上游 issue #24448 / #23053 / #25754 等同族案例均为「在飞的
|
||||
collect 超过池内工作位 → 持有工作位的任务等待排不上队的任务 → 0 CPU 永久
|
||||
挂起」。本模块用进程级信号量限制同时在飞的 collect 数量 —— 这是上游 issue
|
||||
区被反复验证有效的缓解手段。
|
||||
|
||||
车道设计: 总闸位 polars_collect_permits 个, 其中 background (预热 / 增量 /
|
||||
维表加载等后台计算) 最多占 polars_collect_background_permits 个, 其余闸位
|
||||
保留给 interactive (页面读接口), 保证后台大计算不会把页面请求饿死。获取顺序
|
||||
恒为 background 车道 → 总闸, 不存在环。
|
||||
|
||||
worker 子进程 (回测/优化/挖掘) 单任务串行执行, 不经过本闸。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from typing import Literal
|
||||
|
||||
import polars as pl
|
||||
|
||||
from app.config import settings
|
||||
|
||||
CollectPriority = Literal["interactive", "background"]
|
||||
|
||||
_TOTAL_GATE = threading.BoundedSemaphore(settings.polars_collect_permits)
|
||||
_BACKGROUND_LANE = threading.BoundedSemaphore(settings.polars_collect_background_permits)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def collect_slot(priority: CollectPriority = "interactive") -> Iterator[None]:
|
||||
"""占用一个 collect 闸位; background 需同时占用车道位与总闸位。"""
|
||||
if priority == "background":
|
||||
with _BACKGROUND_LANE, _TOTAL_GATE:
|
||||
yield
|
||||
return
|
||||
with _TOTAL_GATE:
|
||||
yield
|
||||
|
||||
|
||||
def guarded_collect(
|
||||
lf: pl.LazyFrame,
|
||||
*,
|
||||
priority: CollectPriority = "interactive",
|
||||
**kwargs: object,
|
||||
) -> pl.DataFrame:
|
||||
"""在并发闸内执行 LazyFrame.collect; 语义不变, 仅串行化调度。"""
|
||||
with collect_slot(priority):
|
||||
return lf.collect(**kwargs)
|
||||
@@ -35,6 +35,7 @@ import polars as pl
|
||||
|
||||
from app.market_time import CN_TZ, cn_now, cn_today
|
||||
from app.parquet import scan_daily_parquet
|
||||
from app.polars_guard import guarded_collect
|
||||
from app.services.index_const import CORE_INDEX_SYMBOLS
|
||||
from app.strategy.intraday_signals import IntradaySignalEvaluator
|
||||
from app.strategy.monitor import format_alert_quote
|
||||
@@ -1761,11 +1762,11 @@ class QuoteService:
|
||||
table = {"etf": "kline_etf_daily", "index": "kline_index_daily"}.get(asset_type, "kline_daily")
|
||||
daily_glob = str(self._repo.store.data_dir / table / "**" / "*.parquet")
|
||||
ohlcv_cols = ["symbol", "date", "open", "high", "low", "close", "volume", "amount", "quote_ts"]
|
||||
hist_df = (
|
||||
hist_df = guarded_collect(
|
||||
scan_daily_parquet(daily_glob)
|
||||
.filter(pl.col("date") >= cutoff)
|
||||
.sort(["symbol", "date"])
|
||||
.collect()
|
||||
.sort(["symbol", "date"]),
|
||||
priority="background",
|
||||
)
|
||||
if hist_df.is_empty():
|
||||
return
|
||||
|
||||
@@ -34,6 +34,7 @@ from app.enriched_generation import (
|
||||
)
|
||||
from app.market_time import cn_today
|
||||
from app.parquet import scan_enriched_parquet
|
||||
from app.polars_guard import guarded_collect
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -595,7 +596,7 @@ class KlineRepository:
|
||||
|
||||
step = time.perf_counter()
|
||||
logger.info("enriched refresh step start: collect history from %s", start_full)
|
||||
df_hist = lf.select(read_cols).collect()
|
||||
df_hist = guarded_collect(lf.select(read_cols), priority="background")
|
||||
logger.info("enriched refresh step done: collect history rows=%d (%.2fs)", len(df_hist), time.perf_counter() - step)
|
||||
if not df_hist.is_empty():
|
||||
instruments = self._instruments_cache if self._instruments_cache is not None else pl.DataFrame()
|
||||
@@ -893,7 +894,7 @@ class KlineRepository:
|
||||
c for c in ["symbol", "consecutive_limit_ups", "consecutive_limit_downs"]
|
||||
if c in lf.collect_schema().names()
|
||||
]
|
||||
consec_source = lf.select("date", *consec_cols).collect()
|
||||
consec_source = guarded_collect(lf.select("date", *consec_cols), priority="background")
|
||||
if len(consec_cols) == 3:
|
||||
consec_df = _last_available_rows(
|
||||
consec_source.select("date", *consec_cols), latest,
|
||||
@@ -982,7 +983,7 @@ class KlineRepository:
|
||||
"raw_close", "raw_high", "raw_low",
|
||||
"consecutive_limit_ups", "consecutive_limit_downs"]
|
||||
if c in lf.collect_schema().names()]
|
||||
df_hist = lf.select(read_cols).collect()
|
||||
df_hist = guarded_collect(lf.select(read_cols), priority="background")
|
||||
|
||||
if df_hist.is_empty():
|
||||
return df_hist, pl.DataFrame()
|
||||
@@ -1029,13 +1030,13 @@ class KlineRepository:
|
||||
read_cols = [c for c in ["symbol", "date", "open", "high", "low", "close",
|
||||
"volume", "amount", "raw_close", "raw_high", "raw_low"]
|
||||
if c in df_latest.columns]
|
||||
df_hist = (
|
||||
df_hist = guarded_collect(
|
||||
scan_enriched_parquet(self._etf_enriched_glob,
|
||||
cast_options=pl.ScanCastOptions(integer_cast="allow-float"))
|
||||
.filter(pl.col("date") >= start_full)
|
||||
.select(read_cols)
|
||||
.sort(["symbol", "date"])
|
||||
.collect()
|
||||
.sort(["symbol", "date"]),
|
||||
priority="background",
|
||||
)
|
||||
if df_hist.is_empty():
|
||||
self._etf_enriched_cache = df_latest.sort(["symbol"])
|
||||
@@ -1073,13 +1074,13 @@ class KlineRepository:
|
||||
read_cols = [c for c in ["symbol", "date", "open", "high", "low", "close",
|
||||
"volume", "amount"]
|
||||
if c in df_latest.columns]
|
||||
df_hist = (
|
||||
df_hist = guarded_collect(
|
||||
scan_enriched_parquet(self._index_enriched_glob,
|
||||
cast_options=pl.ScanCastOptions(integer_cast="allow-float"))
|
||||
.filter(pl.col("date") >= start_full)
|
||||
.select(read_cols)
|
||||
.sort(["symbol", "date"])
|
||||
.collect()
|
||||
.sort(["symbol", "date"]),
|
||||
priority="background",
|
||||
)
|
||||
if df_hist.is_empty():
|
||||
self._index_enriched_cache = df_latest.sort(["symbol"])
|
||||
@@ -1093,7 +1094,7 @@ class KlineRepository:
|
||||
def _refresh_instruments(self) -> None:
|
||||
"""加载 instruments 到内存。"""
|
||||
try:
|
||||
df = pl.scan_parquet(self._inst_glob).collect()
|
||||
df = guarded_collect(pl.scan_parquet(self._inst_glob), priority="background")
|
||||
if not df.is_empty():
|
||||
self._instruments_cache = df
|
||||
self._name_map_cache = None
|
||||
@@ -1104,7 +1105,7 @@ class KlineRepository:
|
||||
def _refresh_index_instruments(self) -> None:
|
||||
"""加载指数 instruments 到内存。"""
|
||||
try:
|
||||
df = pl.scan_parquet(self._index_inst_glob).collect()
|
||||
df = guarded_collect(pl.scan_parquet(self._index_inst_glob), priority="background")
|
||||
if not df.is_empty():
|
||||
self._index_instruments_cache = df
|
||||
self._index_symbol_set_cache = None
|
||||
@@ -1117,7 +1118,7 @@ class KlineRepository:
|
||||
"""加载 ETF instruments 到内存;兼容旧版 instruments_index 中的 ETF。"""
|
||||
parts: list[pl.DataFrame] = []
|
||||
try:
|
||||
df = pl.scan_parquet(self._etf_inst_glob).collect()
|
||||
df = guarded_collect(pl.scan_parquet(self._etf_inst_glob), priority="background")
|
||||
if not df.is_empty():
|
||||
parts.append(df)
|
||||
except Exception as e: # noqa: BLE001
|
||||
@@ -1571,10 +1572,12 @@ class KlineRepository:
|
||||
) -> pl.DataFrame:
|
||||
"""分钟K查询 — Polars scan_parquet + predicate pushdown。"""
|
||||
try:
|
||||
return pl.scan_parquet(self._minute_glob_for(asset_type)).filter(
|
||||
return guarded_collect(
|
||||
pl.scan_parquet(self._minute_glob_for(asset_type)).filter(
|
||||
(pl.col("symbol") == symbol)
|
||||
& (pl.col("datetime").dt.date() == trade_date)
|
||||
).sort("datetime").collect()
|
||||
).sort("datetime")
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("分钟K查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
@@ -1593,10 +1596,12 @@ class KlineRepository:
|
||||
if not symbols:
|
||||
return pl.DataFrame()
|
||||
try:
|
||||
return pl.scan_parquet(self._minute_glob_for(asset_type)).filter(
|
||||
return guarded_collect(
|
||||
pl.scan_parquet(self._minute_glob_for(asset_type)).filter(
|
||||
pl.col("symbol").is_in(symbols)
|
||||
& (pl.col("datetime").dt.date() == trade_date)
|
||||
).sort(["symbol", "datetime"]).collect()
|
||||
).sort(["symbol", "datetime"])
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("批量分钟K查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
@@ -1619,15 +1624,15 @@ class KlineRepository:
|
||||
lf = pl.scan_parquet(self._minute_glob_for(asset_type))
|
||||
available = set(lf.collect_schema().names())
|
||||
select_cols = [c for c in ["symbol", "datetime", "open", "high", "low", "close", "volume", "amount"] if c in available]
|
||||
return (
|
||||
return guarded_collect(
|
||||
lf.select(select_cols)
|
||||
.filter(
|
||||
pl.col("symbol").is_in(symbols)
|
||||
& (pl.col("datetime").dt.date() >= start)
|
||||
& (pl.col("datetime").dt.date() <= end)
|
||||
)
|
||||
.sort(["symbol", "datetime"])
|
||||
.collect(streaming=True)
|
||||
.sort(["symbol", "datetime"]),
|
||||
streaming=True,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("分钟K范围查询失败: %s", e)
|
||||
@@ -1664,11 +1669,11 @@ class KlineRepository:
|
||||
lf = pl.scan_parquet(parts)
|
||||
available = set(lf.collect_schema().names())
|
||||
select_cols = [c for c in ["symbol", "datetime", "open", "high", "low", "close", "volume", "amount"] if c in available]
|
||||
return (
|
||||
return guarded_collect(
|
||||
lf.select(select_cols)
|
||||
.filter(pl.col("symbol").is_in(symbols))
|
||||
.sort(["symbol", "datetime"])
|
||||
.collect(streaming=True)
|
||||
.sort(["symbol", "datetime"]),
|
||||
streaming=True,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("分钟K按日期查询失败: %s", e)
|
||||
@@ -1738,7 +1743,7 @@ class KlineRepository:
|
||||
schema_names = lf.collect_schema().names()
|
||||
existing = [c for c in columns if c in schema_names]
|
||||
lf = lf.select(existing)
|
||||
return lf.collect()
|
||||
return guarded_collect(lf)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("日K查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
@@ -1755,7 +1760,7 @@ class KlineRepository:
|
||||
schema_names = lf.collect_schema().names()
|
||||
existing = [c for c in columns if c in schema_names]
|
||||
lf = lf.select(existing)
|
||||
return lf.collect()
|
||||
return guarded_collect(lf)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("日K批量查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
@@ -1772,7 +1777,7 @@ class KlineRepository:
|
||||
schema_names = lf.collect_schema().names()
|
||||
existing = [c for c in columns if c in schema_names]
|
||||
lf = lf.select(existing)
|
||||
return lf.collect()
|
||||
return guarded_collect(lf)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("指数日K查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
@@ -1789,7 +1794,7 @@ class KlineRepository:
|
||||
schema_names = lf.collect_schema().names()
|
||||
existing = [c for c in columns if c in schema_names]
|
||||
lf = lf.select(existing)
|
||||
return lf.collect()
|
||||
return guarded_collect(lf)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.debug("ETF 日K查询跳过: %s", e)
|
||||
return pl.DataFrame()
|
||||
@@ -2156,26 +2161,75 @@ class KlineRepository:
|
||||
if generation_asset is not None
|
||||
else None
|
||||
)
|
||||
with self._write_lock:
|
||||
for date_df in df.partition_by("date"):
|
||||
dt = date_df["date"][0]
|
||||
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
|
||||
out = base / f"date={ds}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
existing = pl.DataFrame()
|
||||
if out.exists():
|
||||
existing = pl.read_parquet(out)
|
||||
date_df = pl.concat([existing, date_df], how="diagonal_relaxed").unique(
|
||||
self._optimistic_upsert_partition(out, date_df, publication)
|
||||
|
||||
@staticmethod
|
||||
def _partition_fingerprint(path: Path) -> tuple[int, int] | None:
|
||||
"""分区文件的修改指纹 (mtime_ns, size); 不存在返回 None。"""
|
||||
try:
|
||||
st = path.stat()
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
return (st.st_mtime_ns, st.st_size)
|
||||
|
||||
def _optimistic_upsert_partition(
|
||||
self,
|
||||
out: Path,
|
||||
incoming: pl.DataFrame,
|
||||
publication: EnrichedPublication | None,
|
||||
*,
|
||||
retries: int = 3,
|
||||
) -> None:
|
||||
"""单分区 merge-upsert: polars 读/合并/排序在 _write_lock 外, 锁内只做
|
||||
指纹校验 + 原子替换 + commit。
|
||||
|
||||
背景: polars 并发执行存在死锁风险 (见 app.polars_guard), 重活若在
|
||||
_write_lock 内悬死, 全局写锁被永久持有, 所有写路径排队冻结。乐观模式
|
||||
把读/算移出锁外; 锁内用指纹确认基底未被其他写入者改动, 失配则重试,
|
||||
重试耗尽退回锁内直读直写 (正确性优先, 牺牲隔离性)。
|
||||
"""
|
||||
def _merge(existing: pl.DataFrame) -> pl.DataFrame:
|
||||
if existing.is_empty():
|
||||
return incoming.sort(["symbol", "date"])
|
||||
return pl.concat([existing, incoming], how="diagonal_relaxed").unique(
|
||||
subset=["symbol", "date"], keep="last"
|
||||
)
|
||||
date_df = date_df.sort(["symbol", "date"])
|
||||
if not existing.is_empty() and existing.equals(date_df):
|
||||
continue
|
||||
if publication is None:
|
||||
self._atomic_write_parquet(date_df, out)
|
||||
else:
|
||||
publication.write_parquet(date_df, out)
|
||||
).sort(["symbol", "date"])
|
||||
|
||||
for _ in range(retries):
|
||||
existing = pl.read_parquet(out) if out.exists() else pl.DataFrame()
|
||||
base_fp = self._partition_fingerprint(out)
|
||||
merged = _merge(existing)
|
||||
with self._write_lock:
|
||||
if self._partition_fingerprint(out) != base_fp:
|
||||
continue # 基底被并发写入者改过, 出锁重读重算
|
||||
self._write_partition_locked(out, merged, existing, publication)
|
||||
return
|
||||
# 乐观重试耗尽 (罕见: 高频并发写同一分区): 退回锁内全量模式保证正确性
|
||||
with self._write_lock:
|
||||
existing = pl.read_parquet(out) if out.exists() else pl.DataFrame()
|
||||
self._write_partition_locked(out, _merge(existing), existing, publication)
|
||||
|
||||
def _write_partition_locked(
|
||||
self,
|
||||
out: Path,
|
||||
merged: pl.DataFrame,
|
||||
existing: pl.DataFrame,
|
||||
publication: EnrichedPublication | None,
|
||||
) -> None:
|
||||
"""锁内的纯文件阶段: 无变化跳过; 否则原子替换 + 提交 generation。"""
|
||||
if not existing.is_empty() and existing.equals(merged):
|
||||
if publication is not None:
|
||||
publication.commit() # 未写入时为无害空提交
|
||||
return
|
||||
if publication is None:
|
||||
self._atomic_write_parquet(merged, out)
|
||||
else:
|
||||
publication.write_parquet(merged, out)
|
||||
publication.commit()
|
||||
|
||||
def merge_live_daily_asset(self, asset_type: str, df: pl.DataFrame) -> None:
|
||||
@@ -2194,14 +2248,7 @@ class KlineRepository:
|
||||
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
|
||||
out = base / f"date={ds}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self._write_lock:
|
||||
date_df = df.sort(["symbol", "date"])
|
||||
if out.exists():
|
||||
existing = pl.read_parquet(out)
|
||||
date_df = pl.concat([existing, date_df], how="diagonal_relaxed").unique(
|
||||
subset=["symbol", "date"], keep="last"
|
||||
)
|
||||
self._atomic_write_parquet(date_df.sort(["symbol", "date"]), out)
|
||||
self._optimistic_upsert_partition(out, df, None)
|
||||
|
||||
def _with_instrument_metadata(self, asset_type: str, df: pl.DataFrame) -> pl.DataFrame:
|
||||
"""补齐实时内存缓存所需的维表字段;这些字段不会写入 enriched 分区。"""
|
||||
@@ -2258,21 +2305,7 @@ class KlineRepository:
|
||||
if asset_type in {"stock", "etf"}
|
||||
else None
|
||||
)
|
||||
with self._write_lock:
|
||||
existing = pl.DataFrame()
|
||||
if out.exists():
|
||||
existing = pl.read_parquet(out)
|
||||
df_storage = pl.concat([existing, df_storage], how="diagonal_relaxed").unique(
|
||||
subset=["symbol", "date"], keep="last"
|
||||
)
|
||||
df_storage = df_storage.sort(["symbol"])
|
||||
if existing.is_empty() or not existing.equals(df_storage):
|
||||
if publication is None:
|
||||
self._atomic_write_parquet(df_storage, out)
|
||||
else:
|
||||
publication.write_parquet(df_storage, out)
|
||||
if publication is not None:
|
||||
publication.commit()
|
||||
self._optimistic_upsert_partition(out, df_storage, publication)
|
||||
|
||||
if asset_type == "stock":
|
||||
self._enriched_cache = merged_cache
|
||||
@@ -2306,8 +2339,10 @@ class KlineRepository:
|
||||
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
|
||||
out = base / f"date={ds}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
# 覆写语义: 排序在锁外, 锁内只做原子替换。
|
||||
df_sorted = df.sort(["symbol", "date"])
|
||||
with self._write_lock:
|
||||
self._atomic_write_parquet(df.sort(["symbol", "date"]), out)
|
||||
self._atomic_write_parquet(df_sorted, out)
|
||||
|
||||
def flush_live_enriched(self, df: pl.DataFrame) -> None:
|
||||
"""覆写当天 kline_daily_enriched 分区 (实时 enriched 落盘, 非merge)。
|
||||
@@ -2343,15 +2378,11 @@ class KlineRepository:
|
||||
if asset_type in {"stock", "etf"}
|
||||
else None
|
||||
)
|
||||
with self._write_lock:
|
||||
# 覆写语义: 读旧内容只为跳过无变化的写, 读在锁外 (误判最多造成一次
|
||||
# 冗余覆写, 不影响正确性); 锁内只做替换 + commit。
|
||||
existing = pl.read_parquet(out) if out.exists() else pl.DataFrame()
|
||||
if existing.is_empty() or not existing.equals(df_storage):
|
||||
if publication is None:
|
||||
self._atomic_write_parquet(df_storage, out)
|
||||
else:
|
||||
publication.write_parquet(df_storage, out)
|
||||
if publication is not None:
|
||||
publication.commit()
|
||||
with self._write_lock:
|
||||
self._write_partition_locked(out, df_storage, existing, publication)
|
||||
|
||||
if asset_type == "stock":
|
||||
self._enriched_cache = cache_df
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
"""后端自愈看门狗。
|
||||
|
||||
2026-09-07 事故形态: polars 并发死锁把线程悬死在 collect 内部 (0 CPU 永久
|
||||
挂起), 其中持锁者让 _write_lock 永久被占, 所有请求线程排队冻结, 只能人工
|
||||
重启。并发闸 (app.polars_guard) 与写锁瘦身 (repository 乐观并发) 分别削减
|
||||
触发概率与扩散半径; 看门狗是最后一层兜底 —— 探测走与事故相同的共享资源
|
||||
路径 (collect 闸 + 全局写锁), 连续 N 次超时即判定进程已僵死, 主动退出交由
|
||||
supervisor / Docker restart / dev 脚本拉起, 把恢复时间从"人工发现"缩短到
|
||||
约一分钟。
|
||||
|
||||
误伤防护: 探测本身是毫秒级微型 collect + 1s 写锁试探, 阈值要求连续失败
|
||||
(默认 2 次 × 15s 超时), 高负载下"慢而未死"不会触发。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
|
||||
import polars as pl
|
||||
|
||||
from app.config import settings
|
||||
from app.polars_guard import guarded_collect
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def default_probe(write_lock: threading.Lock | None = None) -> None:
|
||||
"""探测关键共享资源: polars collect 闸 + 仓库全局写锁。
|
||||
|
||||
任一被悬死线程占住即超时 — 正是 2026-09-07 冻结事故中被毒化的两条路径。
|
||||
"""
|
||||
guarded_collect(pl.LazyFrame({"probe": [1]}).sum())
|
||||
if write_lock is not None:
|
||||
acquired = write_lock.acquire(timeout=1.0)
|
||||
if not acquired:
|
||||
raise TimeoutError("repository write lock unavailable")
|
||||
write_lock.release()
|
||||
|
||||
|
||||
class HealthWatchdog:
|
||||
"""周期探测; 连续 failure_threshold 次失败后调用 exit_cb(退出码)。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
probe: Callable[[], None],
|
||||
*,
|
||||
exit_cb: Callable[[int], None],
|
||||
interval_s: float | None = None,
|
||||
probe_timeout_s: float | None = None,
|
||||
failure_threshold: int | None = None,
|
||||
) -> None:
|
||||
self._probe = probe
|
||||
self._exit_cb = exit_cb
|
||||
self._interval_s = settings.watchdog_interval_s if interval_s is None else interval_s
|
||||
self._probe_timeout_s = (
|
||||
settings.watchdog_probe_timeout_s if probe_timeout_s is None else probe_timeout_s
|
||||
)
|
||||
self._failure_threshold = (
|
||||
settings.watchdog_failure_threshold if failure_threshold is None else failure_threshold
|
||||
)
|
||||
self._consecutive_failures = 0
|
||||
self._task: asyncio.Task | None = None
|
||||
|
||||
async def _loop(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.to_thread(self._probe), timeout=self._probe_timeout_s
|
||||
)
|
||||
self._consecutive_failures = 0
|
||||
except BaseException as exc: # 探测任何异常都算失败 (含 to_thread 超时)
|
||||
self._consecutive_failures += 1
|
||||
logger.error(
|
||||
"watchdog probe failed (%d/%d): %r",
|
||||
self._consecutive_failures,
|
||||
self._failure_threshold,
|
||||
exc,
|
||||
)
|
||||
if self._consecutive_failures >= self._failure_threshold:
|
||||
logger.critical(
|
||||
"watchdog: backend wedged (probe failed %d consecutive times); "
|
||||
"exiting for supervisor restart",
|
||||
self._consecutive_failures,
|
||||
)
|
||||
self._exit_cb(70)
|
||||
return
|
||||
await asyncio.sleep(self._interval_s)
|
||||
|
||||
def start(self) -> None:
|
||||
if self._task is None or self._task.done():
|
||||
self._task = asyncio.create_task(self._loop(), name="health-watchdog")
|
||||
|
||||
async def stop(self) -> None:
|
||||
if self._task is not None:
|
||||
self._task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._task
|
||||
self._task = None
|
||||
|
||||
|
||||
def start_watchdog(app_state, repo) -> HealthWatchdog | None:
|
||||
"""lifespan 启动钩子; 返回实例挂到 app.state.watchdog 便于关闭。"""
|
||||
if not settings.watchdog_enabled:
|
||||
return None
|
||||
write_lock = getattr(repo, "_write_lock", None)
|
||||
watchdog = HealthWatchdog(
|
||||
lambda: default_probe(write_lock),
|
||||
exit_cb=lambda code: os._exit(code),
|
||||
)
|
||||
watchdog.start()
|
||||
return watchdog
|
||||
@@ -13,7 +13,9 @@ dependencies = [
|
||||
"python-multipart>=0.0.6",
|
||||
"sse-starlette>=2.0",
|
||||
# Data
|
||||
"polars>=1.0",
|
||||
# 1.44 起含流式引擎 executor 线程调度修复; 上限锁定一个已验证的大版本,
|
||||
# 升级需重跑 scripts/stress_polars_concurrency.py 压测 (并发死锁回归)。
|
||||
"polars>=1.44,<1.45",
|
||||
"duckdb>=1.0",
|
||||
"pyarrow>=16.0",
|
||||
"pandas>=2.2", # 仅在 BacktestService 边界使用,见 §7.4 / ADR-19
|
||||
@@ -45,7 +47,7 @@ dependencies = [
|
||||
# machines without AVX2/FMA support.
|
||||
# Enable with: uv sync --extra legacy-cpu
|
||||
legacy-cpu = [
|
||||
"polars[rtcompat]>=1.0",
|
||||
"polars[rtcompat]>=1.44,<1.45",
|
||||
]
|
||||
|
||||
# vectorbt 还会引入绘图、交互组件等完整分析栈,仍保持为可选 extras。
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""polars 并发死锁复现压测。
|
||||
|
||||
模拟 2026-09-07 线上事故的触发形态: 多个线程并发对同一 parquet 目录做
|
||||
lazy scan + filter + collect (页面首屏多路读), 叠加后台线程的批量重计算 —
|
||||
全部绕过 polars_guard 闸, 直接裸调 collect。
|
||||
|
||||
判定: worker 线程持续完成 collect 即健康; 若超过宽限期没有任何完成
|
||||
(0 进度推进), 判定死锁复现, 退出码 1。
|
||||
|
||||
用法:
|
||||
uv run python scripts/stress_polars_concurrency.py [--duration 180] [--threads 8]
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
import polars as pl
|
||||
|
||||
|
||||
def _build_dataset(root: Path, symbols: int = 300, days: int = 120) -> None:
|
||||
"""构造 ~数十万行、按日期分区的 parquet 目录 (模拟 enriched 布局)。"""
|
||||
dates = pl.date_range(
|
||||
__import__("datetime").date(2026, 1, 1),
|
||||
__import__("datetime").date(2026, 12, 31),
|
||||
"1d",
|
||||
eager=True,
|
||||
).to_list()[:days]
|
||||
for d in dates:
|
||||
df = pl.DataFrame({
|
||||
"symbol": [f"{i:06d}.SZ" for i in range(symbols)],
|
||||
"date": [d] * symbols,
|
||||
"close": [10.0 + (i % 37) for i in range(symbols)],
|
||||
"volume": [1_000.0 * (i + 1) for i in range(symbols)],
|
||||
})
|
||||
out = root / f"date={d.isoformat()}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
df.write_parquet(out)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--duration", type=float, default=180.0)
|
||||
parser.add_argument("--threads", type=int, default=8)
|
||||
parser.add_argument("--grace", type=float, default=60.0)
|
||||
args = parser.parse_args()
|
||||
|
||||
stop = threading.Event()
|
||||
stats = {"collects": 0, "last_done": time.monotonic()}
|
||||
lock = threading.Lock()
|
||||
|
||||
with TemporaryDirectory(prefix="polars_stress_") as tmp:
|
||||
root = Path(tmp) / "kline"
|
||||
print("building dataset ...", flush=True)
|
||||
_build_dataset(root)
|
||||
glob = str(root / "**" / "*.parquet")
|
||||
|
||||
def worker(kind: str) -> None:
|
||||
i = 0
|
||||
while not stop.is_set():
|
||||
i += 1
|
||||
if kind == "interactive":
|
||||
df = (
|
||||
pl.scan_parquet(glob)
|
||||
.filter((pl.col("symbol") == f"{(i * 7) % 300:06d}.SZ"))
|
||||
.sort("date")
|
||||
.collect()
|
||||
)
|
||||
else:
|
||||
df = (
|
||||
pl.scan_parquet(glob)
|
||||
.filter(pl.col("volume") > 100_000.0)
|
||||
.group_by("symbol")
|
||||
.agg(pl.col("close").mean().alias("avg_close"))
|
||||
.sort("symbol")
|
||||
.collect()
|
||||
)
|
||||
del df
|
||||
with lock:
|
||||
stats["collects"] += 1
|
||||
stats["last_done"] = time.monotonic()
|
||||
|
||||
print(f"stressing: {args.threads} threads for {args.duration:.0f}s "
|
||||
f"(polars {pl.__version__})", flush=True)
|
||||
threads = [
|
||||
threading.Thread(
|
||||
target=worker,
|
||||
args=("interactive" if i % 2 == 0 else "background",),
|
||||
daemon=True,
|
||||
)
|
||||
for i in range(args.threads)
|
||||
]
|
||||
for t in threads:
|
||||
t.start()
|
||||
|
||||
deadline = time.monotonic() + args.duration
|
||||
wedged = False
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5.0)
|
||||
with lock:
|
||||
idle = time.monotonic() - stats["last_done"]
|
||||
done = stats["collects"]
|
||||
print(f" progress: {done} collects, idle {idle:.1f}s", flush=True)
|
||||
if idle > args.grace:
|
||||
wedged = True
|
||||
break
|
||||
stop.set()
|
||||
for t in threads:
|
||||
t.join(timeout=10)
|
||||
|
||||
alive = [t for t in threads if t.is_alive()]
|
||||
with lock:
|
||||
total = stats["collects"]
|
||||
if wedged or alive:
|
||||
print(f"RESULT: DEADLOCK REPRODUCED — wedged={wedged}, "
|
||||
f"stuck_threads={len(alive)}/{args.threads}, total_collects={total}", flush=True)
|
||||
return 1
|
||||
print(f"RESULT: HEALTHY — {total} collects completed, all threads exited", flush=True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,99 @@
|
||||
"""polars collect 并发闸测试。
|
||||
|
||||
针对的缺陷: polars 共享执行器在多线程并发 collect 下可能死锁 (上游 #24448/
|
||||
#25754 同族), 并发闸限制同时在飞的 collect 数量。此处验证两件事:
|
||||
- 总闸与 background 车道确实约束并发上限;
|
||||
- background 车道占满时 interactive 仍能拿到保留闸位 (页面不被后台计算饿死)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
|
||||
import polars as pl
|
||||
|
||||
import app.polars_guard as guard
|
||||
from app.polars_guard import collect_slot, guarded_collect
|
||||
|
||||
|
||||
def _join_all(threads: list[threading.Thread]) -> None:
|
||||
for t in threads:
|
||||
t.join(timeout=10)
|
||||
assert all(not t.is_alive() for t in threads), "collect 闸内线程未结束 — 闸死锁了"
|
||||
|
||||
|
||||
def test_total_permits_bound_background_concurrency(monkeypatch) -> None:
|
||||
monkeypatch.setattr(guard, "_TOTAL_GATE", threading.BoundedSemaphore(2))
|
||||
monkeypatch.setattr(guard, "_BACKGROUND_LANE", threading.BoundedSemaphore(1))
|
||||
|
||||
counter = {"now": 0, "peak": 0, "bg_now": 0, "bg_peak": 0}
|
||||
lock = threading.Lock()
|
||||
|
||||
def enter_interactive() -> None:
|
||||
with collect_slot("interactive"):
|
||||
with lock:
|
||||
counter["now"] += 1
|
||||
counter["peak"] = max(counter["peak"], counter["now"])
|
||||
time.sleep(0.1)
|
||||
with lock:
|
||||
counter["now"] -= 1
|
||||
|
||||
def enter_background() -> None:
|
||||
with collect_slot("background"):
|
||||
with lock:
|
||||
counter["bg_now"] += 1
|
||||
counter["bg_peak"] = max(counter["bg_peak"], counter["bg_now"])
|
||||
time.sleep(0.1)
|
||||
with lock:
|
||||
counter["bg_now"] -= 1
|
||||
|
||||
threads = [threading.Thread(target=enter_interactive, daemon=True) for _ in range(4)]
|
||||
threads += [threading.Thread(target=enter_background, daemon=True) for _ in range(4)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
_join_all(threads)
|
||||
|
||||
assert counter["peak"] <= 2 # 总闸上限
|
||||
assert counter["bg_peak"] <= 1 # background 车道上限
|
||||
|
||||
|
||||
def test_interactive_survives_full_background_lane(monkeypatch) -> None:
|
||||
total, background = 3, 2
|
||||
monkeypatch.setattr(guard, "_TOTAL_GATE", threading.BoundedSemaphore(total))
|
||||
monkeypatch.setattr(guard, "_BACKGROUND_LANE", threading.BoundedSemaphore(background))
|
||||
|
||||
holders: list = []
|
||||
holder_ready = threading.Event()
|
||||
|
||||
def bg_holder() -> None:
|
||||
ctx = collect_slot("background")
|
||||
ctx.__enter__()
|
||||
holders.append(ctx)
|
||||
if len(holders) == background:
|
||||
holder_ready.set()
|
||||
|
||||
bg_threads = [threading.Thread(target=bg_holder, daemon=True) for _ in range(background)]
|
||||
for t in bg_threads:
|
||||
t.start()
|
||||
assert holder_ready.wait(timeout=5), "后台线程未占满车道"
|
||||
|
||||
# 车道被 background 占满时, interactive 仍应能在保留闸位内进入并退出。
|
||||
done = threading.Event()
|
||||
|
||||
def interactive_probe() -> None:
|
||||
with collect_slot("interactive"):
|
||||
done.set()
|
||||
|
||||
probe = threading.Thread(target=interactive_probe, daemon=True)
|
||||
probe.start()
|
||||
assert done.wait(timeout=2), "interactive 被占满的 background 车道饿死"
|
||||
probe.join(timeout=2)
|
||||
|
||||
for ctx in holders:
|
||||
ctx.__exit__(None, None, None)
|
||||
_join_all(bg_threads)
|
||||
|
||||
|
||||
def test_guarded_collect_executes_lazy_frame() -> None:
|
||||
lf = pl.LazyFrame({"a": [1, 2, 3]}).filter(pl.col("a") > 1)
|
||||
assert guarded_collect(lf)["a"].to_list() == [2, 3]
|
||||
@@ -0,0 +1,54 @@
|
||||
"""看门狗触发逻辑测试 (不真退出进程)。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from app.watchdog import HealthWatchdog, default_probe
|
||||
|
||||
|
||||
async def _run_watchdog(probe_results, *, threshold=2, interval=0.01, timeout=0.2):
|
||||
exits: list[int] = []
|
||||
idx = 0
|
||||
|
||||
def probe() -> None:
|
||||
nonlocal idx
|
||||
if idx < len(probe_results):
|
||||
result = probe_results[idx]
|
||||
idx += 1
|
||||
if isinstance(result, BaseException):
|
||||
raise result
|
||||
# 脚本耗尽后恒为成功
|
||||
|
||||
wd = HealthWatchdog(
|
||||
probe,
|
||||
exit_cb=exits.append,
|
||||
interval_s=interval,
|
||||
probe_timeout_s=timeout,
|
||||
failure_threshold=threshold,
|
||||
)
|
||||
wd.start()
|
||||
for _ in range(50):
|
||||
await asyncio.sleep(0.02)
|
||||
if exits or wd._task.done():
|
||||
break
|
||||
await wd.stop()
|
||||
return exits
|
||||
|
||||
|
||||
async def test_consecutive_failures_trigger_exit() -> None:
|
||||
exits = await _run_watchdog([RuntimeError("wedge"), TimeoutError("wedge")])
|
||||
assert exits == [70]
|
||||
|
||||
|
||||
async def test_success_resets_failure_counter() -> None:
|
||||
# 失败 1 次 → 成功 → 再失败 1 次: 未达连续阈值, 不退出。
|
||||
exits = await _run_watchdog([RuntimeError("slow"), None, RuntimeError("slow")])
|
||||
assert exits == []
|
||||
|
||||
|
||||
async def test_default_probe_passes_on_healthy_resources() -> None:
|
||||
import threading
|
||||
|
||||
lock = threading.Lock()
|
||||
default_probe(lock) # 不抛即通过
|
||||
default_probe(None)
|
||||
@@ -0,0 +1,90 @@
|
||||
"""_write_lock 锁区间瘦身的回归测试。
|
||||
|
||||
背景: polars 并发执行存在死锁风险 (app.polars_guard), 若 polars 读/合并/排序
|
||||
悬死在 _write_lock 内, 全局写锁被永久持有 → 所有写路径排队冻结 (2026-09-07
|
||||
线上全站冻结事故的放大器)。乐观并发模式把重活移到锁外, 此处验证:
|
||||
- 并发 upsert 不丢行 (乐观重试的正确性);
|
||||
- 合并计算期间 _write_lock 可被其他线程获取 (重活确实不在锁内)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from datetime import date
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
|
||||
from app.tickflow.repository import DataStore, KlineRepository
|
||||
|
||||
|
||||
def _frame(symbols: list[str], dt: date = date(2026, 9, 7)) -> pl.DataFrame:
|
||||
n = len(symbols)
|
||||
return pl.DataFrame({
|
||||
"symbol": symbols,
|
||||
"date": [dt] * n,
|
||||
"close": [10.0 + i for i in range(n)],
|
||||
})
|
||||
|
||||
|
||||
def test_concurrent_upserts_do_not_lose_rows(tmp_path: Path) -> None:
|
||||
repo = KlineRepository(DataStore(tmp_path))
|
||||
|
||||
groups = [[f"{i:03d}{j:04d}.SZ" for j in range(8)] for i in range(6)]
|
||||
errors: list[BaseException] = []
|
||||
|
||||
def worker(symbols: list[str]) -> None:
|
||||
try:
|
||||
for _ in range(3): # 每线程多轮写, 提高乐观重试路径命中
|
||||
repo.merge_live_daily_asset("stock", _frame(symbols))
|
||||
except BaseException as exc:
|
||||
errors.append(exc)
|
||||
|
||||
threads = [threading.Thread(target=worker, args=(g,), daemon=True) for g in groups]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join(timeout=30)
|
||||
assert not errors, errors
|
||||
assert all(not t.is_alive() for t in threads)
|
||||
|
||||
out = tmp_path / "kline_daily" / "date=2026-09-07" / "part.parquet"
|
||||
final = pl.read_parquet(out)
|
||||
assert final["symbol"].n_unique() == sum(len(g) for g in groups) # 无一丢失
|
||||
assert final["symbol"].to_list() == sorted(final["symbol"].to_list())
|
||||
|
||||
|
||||
def test_heavy_merge_runs_outside_write_lock(tmp_path: Path, monkeypatch) -> None:
|
||||
repo = KlineRepository(DataStore(tmp_path))
|
||||
# 预置旧分区内容, 让 upsert 走「读旧 + concat 合并」路径。
|
||||
repo.merge_live_daily_asset("stock", _frame(["000001.SZ"]))
|
||||
|
||||
merge_entered = threading.Event()
|
||||
original_concat = pl.concat
|
||||
|
||||
def slow_concat(*args, **kwargs):
|
||||
merge_entered.set()
|
||||
import time
|
||||
|
||||
time.sleep(0.4) # 模拟重合并耗时; 期间 _write_lock 必须是空闲的
|
||||
return original_concat(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(pl, "concat", slow_concat)
|
||||
|
||||
done = threading.Event()
|
||||
|
||||
def upsert() -> None:
|
||||
repo.merge_live_daily_asset("stock", _frame(["000002.SZ"]))
|
||||
done.set()
|
||||
|
||||
t = threading.Thread(target=upsert, daemon=True)
|
||||
t.start()
|
||||
assert merge_entered.wait(timeout=5), "合并路径未被触发"
|
||||
|
||||
acquired = repo._write_lock.acquire(timeout=1.0)
|
||||
assert acquired, "合并计算期间 _write_lock 被占用 — 重活仍在锁内"
|
||||
repo._write_lock.release()
|
||||
|
||||
assert done.wait(timeout=5)
|
||||
t.join(timeout=5)
|
||||
out = tmp_path / "kline_daily" / "date=2026-09-07" / "part.parquet"
|
||||
assert pl.read_parquet(out)["symbol"].to_list() == ["000001.SZ", "000002.SZ"]
|
||||
Generated
+1845
-1845
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user