mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-09-29 05:56:38 +08:00
feat(db): add ClickHouse batchwriter framework and skill
Introduce internal/db/batchwriter as a reusable generic buffered writer for per-domain ClickHouse flush pipelines, with unit tests and default batch tuning aligned with audit log ingestion. Add clickhouse-batchwriter agent skill and cross-references in AGENTS.md and database-migration. Business layers are not wired yet.
This commit is contained in:
@@ -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/<domain>/` | 采集、入队、背压;`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/<domain>/` 或 `internal/repository/analytics/<domain>_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`
|
||||
@@ -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/<domain>/` 编排采集与入队;高频写入通过 `internal/db/batchwriter` 各域独立实例异步 flush(详见 `clickhouse-batchwriter` 技能),`FlushFunc` 只调 repository `BatchInsert*`;管理端统计 API 只读 repository,不触达 DDL。
|
||||
|
||||
### ClickHouse 验证
|
||||
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -24,6 +24,7 @@ sidebar: false
|
||||
|
||||
### 新增
|
||||
|
||||
- 新增 `internal/db/batchwriter` 通用批量写入框架,支持各业务域独立队列实例、按条数/时间 flush、非阻塞入队与优雅停机;业务层尚未接入。
|
||||
|
||||
### 变更
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user