diff --git a/.agent/skills/clickhouse-batchwriter/SKILL.md b/.agent/skills/clickhouse-batchwriter/SKILL.md new file mode 100644 index 00000000..038f10b7 --- /dev/null +++ b/.agent/skills/clickhouse-batchwriter/SKILL.md @@ -0,0 +1,160 @@ +--- +name: "clickhouse-batchwriter" +description: "Wavelet 项目专用:当新增或修改 ClickHouse 批量写入、接入 internal/db/batchwriter、将业务域异步 flush 到分析表、迁移 risk_control/节点访问日志/可观测时序写入、或评估 async_insert 与背压策略时必须使用。本技能指导分层职责、各域独立 Writer 实例、repository 批量 API 与禁止写法。" +--- + +# ClickHouse 批量写入开发 + +开始前阅读根目录 `AGENTS.md`。ClickHouse 是辅助 OLAP 存储,**厌恶高频单条写入**(过多小 part);写入路径必须优先批量或异步聚合。 + +DDL 与表结构变更见 `database-migration` 技能;本技能只覆盖**运行时写入架构**。 + +## 分层职责 + +| 层级 | 路径 | 职责 | +| :--- | :--- | :--- | +| 连接 | `internal/db/clickhouse.go` | `ChConn`(原生批量写)、`ChDB`(GORM 查询);禁止在业务包直接 `clickhouse.Open` | +| 批量框架 | `internal/db/batchwriter/` | 泛型队列 + 按条数/时间 flush + 非阻塞入队 + 优雅停机;**各业务域独立实例** | +| Model | `internal/model/analytics/` | 列定义、`TableName()`、`BatchInsertSQL()`(及可选 `InsertColumns()`) | +| Repository | `internal/repository/analytics/` | `BatchInsert*` / `BatchInsertNodeAccessLogs` 等;`PrepareBatch` + 多行 `Append` + 一次 `Send` | +| Apps | `internal/apps//` | 采集、入队、背压;`FlushFunc` 只调 repository,不写 SQL、不 `PrepareBatch` | +| 装配 | `internal/bootstrap/bootstrap.go` | 进程启动时 `Writer.Start`、停机时 `Writer.Stop`(与 `risk_control.InitLogWriter` 同级) | + +**禁止**在 Handler / middleware 内直接 `db.ChConn.PrepareBatch`;**禁止**在 repository 内启动 goroutine 或维护全局 channel(队列生命周期由 apps + bootstrap 或专用 writer 包负责)。 + +## batchwriter 框架契约 + +```go +writer, err := batchwriter.New[YourType](cfg, flushFunc, opts...) +writer.Start(ctx) +writer.TryEnqueue(item) // 非阻塞;满则 false +writer.IsFull() // 背压探测 +writer.Stop(stopCtx) // close 队列 + drain + 最终 flush +``` + +### Config 默认值(`batchwriter.DefaultConfig()`) + +- `QueueSize`: 10_000 +- `MaxBatchSize`: 1_000 +- `FlushInterval`: 1s + +各域可独立覆盖;可观测低频指标可用更小 `MaxBatchSize`(如 100)与更长 `FlushInterval`(如 2–5s),但**不要**退化为逐条 `Send`。 + +### 可选回调 + +- `WithFlushErrorHandler[T]`:flush 失败时记录日志;批次丢弃后 worker 继续 +- `WithDropHandler[T]`:队列满或未 `Start` 时丢弃项 + +### FlushFunc 规范 + +- 签名:`func(ctx context.Context, items []T) error` +- 内部调用 `internal/repository/analytics` 的 `BatchInsert*`(传入 `[]analyticsmodel.X`) +- 在 flush 边界记录一次错误日志,不要把 DB 驱动错误直接暴露给 HTTP 客户端 +- `Start` 使用 `context.WithoutCancel(parent)`,避免请求 ctx 取消中断后台 flush + +## 各域独立实例(不共享队列) + +每个业务域拥有自己的 `Writer`、配置与 `FlushFunc`: + +| 域 | 表 | 现状 | 目标形态 | +| :--- | :--- | :--- | :--- | +| 管理端审计 | `w_user_access_logs` | `risk_control` 手写 channel worker | 迁移至 `batchwriter` + `analyticsrepo.BatchInsert` | +| 边缘访问日志 | `of_node_access_logs` | 心跳内 `BatchInsertNodeAccessLogs` | 高 QPS 时增加独立 `batchwriter`;低规模可维持心跳内批量 | +| 可观测时序 | `of_node_metric_snapshots` 等 5 表 | repository 逐条 `PrepareBatch` + 部分写前 `SELECT count()` | **优先改造**:`batchwriter` 或心跳内按表聚合 + 去掉 exists 查询 | + +**不要**把 audit、access log、observability 并入同一 channel。 + +## 新增 ClickHouse 写入工作流 + +1. **Model**:在 `internal/model/analytics/` 定义 struct 与 `BatchInsertSQL()`(列顺序与 goose DDL 一致)。 +2. **Goose DDL**:在 `internal/db/migrator/goose/clickhouse/` 新增迁移(见 `database-migration`)。 +3. **Repository**:实现 `BatchInsertX(ctx, []analyticsmodel.X) error`: + - `len(items)==0` 直接返回 + - `db.ChConn == nil` 返回明确错误 + - 一次 `PrepareBatch` → 循环 `Append` → 一次 `Send` +4. **Writer 胶水**(`internal/apps//` 或 `internal/repository/analytics/_writer.go`): + - `New` + `Start` 在 bootstrap 注册 + - 业务路径 `TryEnqueue`;HTTP 背压用 `IsFull()` +5. **测试**: + - repository:mock `ChConn` 验证 `BatchInsertSQL` 与 append 列数 + - batchwriter:`go test ./internal/db/batchwriter` +6. 运行 `make code-check`;有 API 变更时 `make swagger`。 + +## 背压与丢弃策略 + +| 场景 | 推荐策略 | +| :--- | :--- | +| 管理端 API 审计 | 队列满 → `IsFull()` 触发 429(见 `risk_control` middleware) | +| Agent 心跳指标 | 队列满 → `WithDropHandler` 记 warn;不阻塞心跳响应 | +| 边缘 access log | 优先扩大队列与 batch;必要时丢弃最旧或采样 | + +## 禁止写法 + +```go +// ❌ 单条伪批量:每条都 PrepareBatch + Send +batch.Append(oneRow) +batch.Send() + +// ❌ 写前 OLTP 式去重(高 RTT + 仍产生小 part) +SELECT count() FROM ... WHERE node_id = ? AND captured_at = ? + +// ❌ Handler 内直接写 ClickHouse +db.ChConn.PrepareBatch(...) + +// ❌ 全局单队列承载所有分析表 +var globalChan chan any +``` + +去重应使用:`ReplacingMergeTree`、查询侧 `argMax`、或进程内短 TTL 去重缓存——**不要**在每次 insert 前 `SELECT count()`。 + +## async_insert(补充,非主方案) + +可在 `internal/db/clickhouse.go` 的 `Settings` 增加服务端异步写入作为第二层防护: + +```go +"async_insert": 1, +"wait_for_async_insert": 1, +``` + +**不能替代**应用层批量;接入前需评估丢失可观测性与服务端负载。优先完成 `batchwriter` 接入后再考虑。 + +## Bootstrap 装配示例 + +```go +// internal/bootstrap/bootstrap.go(示意) +var userAccessLogWriter *batchwriter.Writer[*analytics.UserAccessLog] + +func RegisterAPI(ctx context.Context) { + // ... + if config.Config.ClickHouse.Enabled { + initUserAccessLogWriter(ctx) // Start writer + risk_control.BindWriter(userAccessLogWriter) // 或逐步替换 InitLogWriter + } +} +``` + +- `RegisterAPI` / `RegisterAll`:`Start` +- 进程优雅停机:带超时的 `Stop(ctx)` +- 使用 `sync.Once` 保证幂等 + +## 验证清单 + +```bash +go test ./internal/db/batchwriter +go test ./internal/repository/analytics +make code-check +``` + +- flush 按 `MaxBatchSize` 与 `FlushInterval` 触发 +- `Stop` 能 drain 队列内剩余项 +- repository 层无 goroutine、无 channel +- `clickhouse.enabled: false` 时不 `Start` writer、不入队 + +## 相关文件速查 + +- 框架:`internal/db/batchwriter/{config,writer,errs}.go` +- 连接:`internal/db/clickhouse.go` +- 审计写入(待迁移):`internal/apps/risk_control/logics.go` +- 节点访问日志:`internal/repository/analytics/node_access_log_writer.go` +- 可观测写入(待改造):`internal/repository/analytics/node_observability_writer.go` +- Bootstrap:`internal/bootstrap/bootstrap.go` \ No newline at end of file diff --git a/.agent/skills/database-migration/SKILL.md b/.agent/skills/database-migration/SKILL.md index 84feb061..43852c7b 100644 --- a/.agent/skills/database-migration/SKILL.md +++ b/.agent/skills/database-migration/SKILL.md @@ -117,8 +117,8 @@ ClickHouse 是**辅助 OLAP 存储**,与 PostgreSQL/SQLite 主库**完全独 1. **Model**:在 `internal/model/analytics/` 定义 struct,`gorm:"column:..."` 与 DDL 列名一一对应;实现 `TableName()`,批量写入表可提供 `InsertColumns()` / `BatchInsertSQL()`。 2. **Goose SQL**:在 `internal/db/migrator/goose/clickhouse/` 新增递增版本文件(格式同主库,如 `YYYYMMDDNNNN_create_xxx.sql`),编写 `-- +goose Up` / `-- +goose Down`。 -3. **Repository**:在 `internal/repository/analytics/` 实现写入(优先 `db.ChConn` 批量)与查询(`db.ChDB`);连接未初始化时返回明确错误,**不要**在 handler 写 SQL。 -4. **Apps**:在 `internal/apps/` 编排业务(如中间件采集、管理端统计 API),只调用 repository,不触达 DDL。 +3. **Repository**:在 `internal/repository/analytics/` 实现 `BatchInsert*`(`db.ChConn` 一次 `PrepareBatch` + 多行 `Append` + 一次 `Send`)与查询(`db.ChDB`);连接未初始化时返回明确错误,**不要**在 handler 写 SQL,**不要**在 repository 内维护 channel/goroutine。 +4. **Apps**:在 `internal/apps//` 编排采集与入队;高频写入通过 `internal/db/batchwriter` 各域独立实例异步 flush(详见 `clickhouse-batchwriter` 技能),`FlushFunc` 只调 repository `BatchInsert*`;管理端统计 API 只读 repository,不触达 DDL。 ### ClickHouse 验证 diff --git a/AGENTS.md b/AGENTS.md index 4f1da2fc..a2342f99 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -43,6 +43,7 @@ | `new-async-task` | 添加或修改 Asynq 任务、定时任务、TaskHandler、任务元数据 | | `new-setting` | 添加或修改系统/业务/公开设置、`/admin/system` 参数或 `/admin/settings` 图形化设置 | | `database-migration` | 数据库表结构变更、goose SQL 迁移(PG/SQLite/ClickHouse)、seed 数据 | +| `clickhouse-batchwriter` | ClickHouse 批量写入、`internal/db/batchwriter` 接入、分析表异步 flush、背压与写入路径改造 | | `file-upload` | 业务上传文件、Worker 程序化摄取、`upload.Ingest` 策略选型、文件访问与 `w_uploads` / 统计排查 | | `push-notification` | 系统通知推送事件、统一触发器投递、带消息推送的业务功能 | | `release-guide` | 根据自上一正式版本 Tag 以来的提交整理 Version Bump 提交信息以触发双语 Release | diff --git a/docs/changelog/index.md b/docs/changelog/index.md index 7e100119..819c0a34 100644 --- a/docs/changelog/index.md +++ b/docs/changelog/index.md @@ -24,6 +24,7 @@ sidebar: false ### 新增 +- 新增 `internal/db/batchwriter` 通用批量写入框架,支持各业务域独立队列实例、按条数/时间 flush、非阻塞入队与优雅停机;业务层尚未接入。 ### 变更 diff --git a/go.mod b/go.mod index 14bccfad..b910302f 100644 --- a/go.mod +++ b/go.mod @@ -20,6 +20,7 @@ require ( github.com/glebarez/sqlite v1.11.0 github.com/go-acme/lego/v4 v4.35.2 github.com/go-jose/go-jose/v4 v4.1.4 + github.com/google/go-cmp v0.7.0 github.com/google/uuid v1.6.0 github.com/gorilla/websocket v1.5.3 github.com/hibiken/asynq v0.25.1 diff --git a/internal/db/batchwriter/config.go b/internal/db/batchwriter/config.go new file mode 100644 index 00000000..32972ecb --- /dev/null +++ b/internal/db/batchwriter/config.go @@ -0,0 +1,52 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package batchwriter + +import ( + "fmt" + "time" +) + +const ( + defaultQueueSize = 10_000 + defaultMaxBatchSize = 1_000 + defaultFlushEvery = time.Second +) + +// Config controls queue capacity and flush thresholds for a Writer instance. +type Config struct { + // Name identifies the writer in logs and diagnostics. Optional. + Name string + + // QueueSize is the buffered channel capacity. + QueueSize int + + // MaxBatchSize triggers a flush when the in-memory batch reaches this count. + MaxBatchSize int + + // FlushInterval triggers a time-based flush even when the batch is smaller. + FlushInterval time.Duration +} + +// DefaultConfig returns production-friendly defaults aligned with audit log batching. +func DefaultConfig() Config { + return Config{ + QueueSize: defaultQueueSize, + MaxBatchSize: defaultMaxBatchSize, + FlushInterval: defaultFlushEvery, + } +} + +func (c Config) validate() error { + if c.QueueSize <= 0 { + return fmt.Errorf("batchwriter: queue size must be positive") + } + if c.MaxBatchSize <= 0 { + return fmt.Errorf("batchwriter: max batch size must be positive") + } + if c.FlushInterval <= 0 { + return fmt.Errorf("batchwriter: flush interval must be positive") + } + return nil +} \ No newline at end of file diff --git a/internal/db/batchwriter/errs.go b/internal/db/batchwriter/errs.go new file mode 100644 index 00000000..f6242925 --- /dev/null +++ b/internal/db/batchwriter/errs.go @@ -0,0 +1,8 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package batchwriter + +import "errors" + +var errNilFlushFunc = errors.New("batchwriter: flush func is required") \ No newline at end of file diff --git a/internal/db/batchwriter/writer.go b/internal/db/batchwriter/writer.go new file mode 100644 index 00000000..ac8c24e4 --- /dev/null +++ b/internal/db/batchwriter/writer.go @@ -0,0 +1,215 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +// Package batchwriter provides a reusable buffered batch writer for high-throughput +// append-only sinks such as ClickHouse. Each business domain should own an independent +// Writer instance with its own queue, flush callback, and tuning parameters. +package batchwriter + +import ( + "context" + "sync" + "time" +) + +// FlushFunc persists a batch of queued items. It is invoked from the worker goroutine. +type FlushFunc[T any] func(ctx context.Context, items []T) error + +// FlushErrorHandler is called when FlushFunc returns an error. The batch is discarded +// after the handler returns; the worker continues processing. +type FlushErrorHandler func(ctx context.Context, batchSize int, err error) + +// Writer buffers items and flushes them by size or interval. +type Writer[T any] struct { + cfg Config + flush FlushFunc[T] + + onFlushError FlushErrorHandler + onDrop func(T) + + startOnce sync.Once + stopOnce sync.Once + + mu sync.RWMutex + ch chan T + workerCtx context.Context + done chan struct{} +} + +// Option configures optional Writer callbacks. +type Option[T any] func(*Writer[T]) + +// WithFlushErrorHandler registers a callback for flush failures. +func WithFlushErrorHandler[T any](handler FlushErrorHandler) Option[T] { + return func(w *Writer[T]) { + w.onFlushError = handler + } +} + +// WithDropHandler registers a callback when TryEnqueue cannot accept an item. +func WithDropHandler[T any](handler func(T)) Option[T] { + return func(w *Writer[T]) { + w.onDrop = handler + } +} + +// New creates a Writer. Call Start before enqueueing items. +func New[T any](cfg Config, flush FlushFunc[T], opts ...Option[T]) (*Writer[T], error) { + if flush == nil { + return nil, errNilFlushFunc + } + if err := cfg.validate(); err != nil { + return nil, err + } + + w := &Writer[T]{ + cfg: cfg, + flush: flush, + done: make(chan struct{}), + } + for _, opt := range opts { + opt(w) + } + return w, nil +} + +// Start launches the background worker. It is safe to call at most once. +func (w *Writer[T]) Start(parent context.Context) { + w.startOnce.Do(func() { + w.mu.Lock() + defer w.mu.Unlock() + + w.ch = make(chan T, w.cfg.QueueSize) + w.workerCtx = context.WithoutCancel(parent) + go w.run() + }) +} + +// Stop closes the queue and waits until the worker drains pending items and exits. +func (w *Writer[T]) Stop(ctx context.Context) error { + w.mu.RLock() + ch := w.ch + done := w.done + w.mu.RUnlock() + + if ch == nil { + return nil + } + + w.stopOnce.Do(func() { + close(ch) + }) + + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// Running reports whether Start has been called and Stop has not completed. +func (w *Writer[T]) Running() bool { + w.mu.RLock() + defer w.mu.RUnlock() + if w.ch == nil { + return false + } + select { + case <-w.done: + return false + default: + return true + } +} + +// TryEnqueue adds one item without blocking. It returns false when the writer is not +// running or the queue is full. +func (w *Writer[T]) TryEnqueue(item T) bool { + w.mu.RLock() + ch := w.ch + w.mu.RUnlock() + if ch == nil { + w.notifyDrop(item) + return false + } + + select { + case ch <- item: + return true + default: + w.notifyDrop(item) + return false + } +} + +// IsFull reports whether the queue has no remaining capacity. +func (w *Writer[T]) IsFull() bool { + w.mu.RLock() + defer w.mu.RUnlock() + if w.ch == nil { + return false + } + return len(w.ch) >= cap(w.ch) +} + +// Len returns the current queue depth. +func (w *Writer[T]) Len() int { + w.mu.RLock() + defer w.mu.RUnlock() + if w.ch == nil { + return 0 + } + return len(w.ch) +} + +// Cap returns the queue capacity. +func (w *Writer[T]) Cap() int { + return w.cfg.QueueSize +} + +func (w *Writer[T]) run() { + ticker := time.NewTicker(w.cfg.FlushInterval) + defer ticker.Stop() + + batch := make([]T, 0, w.cfg.MaxBatchSize) + flush := func() { + if len(batch) == 0 { + return + } + items := append([]T(nil), batch...) + if err := w.flush(w.workerCtx, items); err != nil { + if w.onFlushError != nil { + w.onFlushError(w.workerCtx, len(items), err) + } + } + batch = batch[:0] + } + + defer func() { + flush() + close(w.done) + }() + + for { + select { + case item, ok := <-w.ch: + if !ok { + return + } + batch = append(batch, item) + if len(batch) >= w.cfg.MaxBatchSize { + flush() + } + case <-ticker.C: + flush() + } + } +} + +func (w *Writer[T]) notifyDrop(item T) { + if w.onDrop == nil { + return + } + w.onDrop(item) +} \ No newline at end of file diff --git a/internal/db/batchwriter/writer_test.go b/internal/db/batchwriter/writer_test.go new file mode 100644 index 00000000..0fd410df --- /dev/null +++ b/internal/db/batchwriter/writer_test.go @@ -0,0 +1,284 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package batchwriter + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/google/go-cmp/cmp" +) + +func TestNewRejectsInvalidConfig(t *testing.T) { + t.Parallel() + + _, err := New[int](Config{}, func(context.Context, []int) error { return nil }) + if err == nil { + t.Fatal("New() = nil, want validation error") + } +} + +func TestNewRejectsNilFlushFunc(t *testing.T) { + t.Parallel() + + cfg := DefaultConfig() + _, err := New[int](cfg, nil) + if !errors.Is(err, errNilFlushFunc) { + t.Fatalf("New() error = %v, want %v", err, errNilFlushFunc) + } +} + +func TestWriterFlushesOnMaxBatchSize(t *testing.T) { + t.Parallel() + + var ( + mu sync.Mutex + batches [][]int + ) + cfg := DefaultConfig() + cfg.MaxBatchSize = 3 + cfg.FlushInterval = time.Hour + + writer, err := New[int](cfg, func(_ context.Context, items []int) error { + mu.Lock() + defer mu.Unlock() + batches = append(batches, append([]int(nil), items...)) + return nil + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + writer.Start(context.Background()) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := writer.Stop(stopCtx); err != nil { + t.Fatalf("Stop() error = %v", err) + } + }) + + for i := range 3 { + if !writer.TryEnqueue(i + 1) { + t.Fatalf("TryEnqueue(%d) = false, want true", i+1) + } + } + + deadline := time.Now().Add(time.Second) + for { + mu.Lock() + ready := len(batches) == 1 + mu.Unlock() + if ready || time.Now().After(deadline) { + break + } + time.Sleep(10 * time.Millisecond) + } + + mu.Lock() + got := batches + mu.Unlock() + + want := [][]int{{1, 2, 3}} + if diff := cmp.Diff(want, got); diff != "" { + t.Fatalf("flush batches mismatch (-want +got):\n%s", diff) + } +} + +func TestWriterFlushesOnInterval(t *testing.T) { + t.Parallel() + + var ( + mu sync.Mutex + batch []int + ) + cfg := DefaultConfig() + cfg.MaxBatchSize = 100 + cfg.FlushInterval = 20 * time.Millisecond + + writer, err := New[int](cfg, func(_ context.Context, items []int) error { + mu.Lock() + defer mu.Unlock() + batch = append([]int(nil), items...) + return nil + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + writer.Start(context.Background()) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := writer.Stop(stopCtx); err != nil { + t.Fatalf("Stop() error = %v", err) + } + }) + + if !writer.TryEnqueue(42) { + t.Fatal("TryEnqueue() = false, want true") + } + + deadline := time.Now().Add(time.Second) + for { + mu.Lock() + ready := len(batch) == 1 + mu.Unlock() + if ready || time.Now().After(deadline) { + break + } + time.Sleep(5 * time.Millisecond) + } + + mu.Lock() + got := batch + mu.Unlock() + + want := []int{42} + if diff := cmp.Diff(want, got); diff != "" { + t.Fatalf("interval flush mismatch (-want +got):\n%s", diff) + } +} + +func TestWriterTryEnqueueDropsWhenFull(t *testing.T) { + t.Parallel() + + cfg := DefaultConfig() + cfg.QueueSize = 1 + cfg.MaxBatchSize = 10 + cfg.FlushInterval = time.Hour + + var dropped int + writer, err := New[int](cfg, func(context.Context, []int) error { return nil }, WithDropHandler[int](func(int) { + dropped++ + })) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + writer.Start(context.Background()) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = writer.Stop(stopCtx) + }) + + if !writer.TryEnqueue(1) { + t.Fatal("TryEnqueue(1) = false, want true") + } + if writer.TryEnqueue(2) { + t.Fatal("TryEnqueue(2) = true, want false") + } + if !writer.IsFull() { + t.Fatal("IsFull() = false, want true") + } + if dropped != 1 { + t.Fatalf("dropped = %d, want 1", dropped) + } +} + +func TestWriterStopDrainsQueuedItems(t *testing.T) { + t.Parallel() + + cfg := DefaultConfig() + cfg.MaxBatchSize = 10 + cfg.FlushInterval = time.Hour + + var flushed []int + writer, err := New[int](cfg, func(_ context.Context, items []int) error { + flushed = append(flushed, items...) + return nil + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + writer.Start(context.Background()) + for i := range 2 { + if !writer.TryEnqueue(i + 1) { + t.Fatalf("TryEnqueue(%d) = false, want true", i+1) + } + } + + stopCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := writer.Stop(stopCtx); err != nil { + t.Fatalf("Stop() error = %v", err) + } + + want := []int{1, 2} + if diff := cmp.Diff(want, flushed); diff != "" { + t.Fatalf("Stop() drain mismatch (-want +got):\n%s", diff) + } + if writer.Running() { + t.Fatal("Running() = true after Stop(), want false") + } +} + +func TestWriterInvokesFlushErrorHandler(t *testing.T) { + t.Parallel() + + cfg := DefaultConfig() + cfg.MaxBatchSize = 1 + cfg.FlushInterval = time.Hour + + flushErr := errors.New("flush failed") + var ( + mu sync.Mutex + errCount int + batchSize int + ) + + writer, err := New[int](cfg, func(context.Context, []int) error { + return flushErr + }, WithFlushErrorHandler[int](func(_ context.Context, size int, err error) { + mu.Lock() + defer mu.Unlock() + errCount++ + batchSize = size + if !errors.Is(err, flushErr) { + t.Errorf("flush error = %v, want %v", err, flushErr) + } + })) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + writer.Start(context.Background()) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = writer.Stop(stopCtx) + }) + + if !writer.TryEnqueue(7) { + t.Fatal("TryEnqueue() = false, want true") + } + + deadline := time.Now().Add(time.Second) + for { + mu.Lock() + ready := errCount == 1 + mu.Unlock() + if ready || time.Now().After(deadline) { + break + } + time.Sleep(5 * time.Millisecond) + } + + mu.Lock() + gotCount := errCount + gotSize := batchSize + mu.Unlock() + + if gotCount != 1 { + t.Fatalf("flush error handler count = %d, want 1", gotCount) + } + if gotSize != 1 { + t.Fatalf("flush error handler batch size = %d, want 1", gotSize) + } +} \ No newline at end of file