From 9491b2a744862cfd6c08d6e72182eb440444858d Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 19 Jun 2026 17:13:46 +0800 Subject: [PATCH] feat(clickhouse): wire batchwriter into business ingestion paths Migrate risk_control audit logs to internal/db/batchwriter and add openflare/chwriter with per-table async flush for observability timeseries and node access logs. Replace single-row ClickHouse inserts and pre-insert SELECT count() dedup with repository BatchInsert* APIs, in-process TTL dedup for metric/report snapshots, and bootstrap initialization on API startup. --- .agent/skills/clickhouse-batchwriter/SKILL.md | 13 +- docs/changelog/index.md | 3 +- internal/apps/openflare/chwriter/dedup.go | 42 ++ .../apps/openflare/chwriter/dedup_test.go | 24 ++ internal/apps/openflare/chwriter/writer.go | 213 ++++++++++ internal/apps/risk_control/logics.go | 136 ++++--- internal/apps/risk_control/middleware_test.go | 65 ++- internal/bootstrap/bootstrap.go | 2 + internal/model/openflare_access_log_store.go | 6 +- .../model/openflare_observability_store.go | 26 +- .../analytics/node_observability_writer.go | 379 ++++++++++-------- 11 files changed, 633 insertions(+), 276 deletions(-) create mode 100644 internal/apps/openflare/chwriter/dedup.go create mode 100644 internal/apps/openflare/chwriter/dedup_test.go create mode 100644 internal/apps/openflare/chwriter/writer.go diff --git a/.agent/skills/clickhouse-batchwriter/SKILL.md b/.agent/skills/clickhouse-batchwriter/SKILL.md index 038f10b7..ce67d278 100644 --- a/.agent/skills/clickhouse-batchwriter/SKILL.md +++ b/.agent/skills/clickhouse-batchwriter/SKILL.md @@ -58,9 +58,9 @@ writer.Stop(stopCtx) // close 队列 + drain + 最终 flush | 域 | 表 | 现状 | 目标形态 | | :--- | :--- | :--- | :--- | -| 管理端审计 | `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 查询 | +| 管理端审计 | `w_user_access_logs` | `risk_control` → `batchwriter` + `analyticsrepo.BatchInsert` | 已接入 | +| 边缘访问日志 | `of_node_access_logs` | `openflare/chwriter` 异步 flush | 已接入 | +| 可观测时序 | `of_node_metric_snapshots` 等 5 表 | `openflare/chwriter` 五表独立 writer + 进程内短 TTL 去重 | 已接入 | **不要**把 audit、access log、observability 并入同一 channel。 @@ -154,7 +154,8 @@ make code-check - 框架：`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` +- 审计写入：`internal/apps/risk_control/logics.go` +- OpenFlare 写入胶水：`internal/apps/openflare/chwriter/writer.go` +- 节点访问日志 repository：`internal/repository/analytics/node_access_log_writer.go` +- 可观测 repository：`internal/repository/analytics/node_observability_writer.go` - Bootstrap：`internal/bootstrap/bootstrap.go` \ No newline at end of file diff --git a/docs/changelog/index.md b/docs/changelog/index.md index 819c0a34..e115b4bc 100644 --- a/docs/changelog/index.md +++ b/docs/changelog/index.md @@ -24,7 +24,8 @@ sidebar: false ### 新增 -- 新增 `internal/db/batchwriter` 通用批量写入框架，支持各业务域独立队列实例、按条数/时间 flush、非阻塞入队与优雅停机；业务层尚未接入。 +- 新增 `internal/db/batchwriter` 通用批量写入框架，支持各业务域独立队列实例、按条数/时间 flush、非阻塞入队与优雅停机。 +- 业务层接入批量写入：`risk_control` 审计日志迁移至 `batchwriter`；OpenFlare 可观测时序与节点访问日志通过 `internal/apps/openflare/chwriter` 异步 flush，移除写前 `SELECT count()` 去重。 ### 变更 diff --git a/internal/apps/openflare/chwriter/dedup.go b/internal/apps/openflare/chwriter/dedup.go new file mode 100644 index 00000000..d9b8127c --- /dev/null +++ b/internal/apps/openflare/chwriter/dedup.go @@ -0,0 +1,42 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package chwriter + +import ( + "sync" + "time" +) + +const dedupTTL = 2 * time.Minute + +type dedupSet struct { + mu sync.Mutex + keys map[string]time.Time +} + +func newDedupSet() *dedupSet { + return &dedupSet{keys: make(map[string]time.Time)} +} + +// markIfNew records key when it has not been seen within dedupTTL. +func (s *dedupSet) markIfNew(key string) bool { + if key == "" { + return false + } + + now := time.Now() + s.mu.Lock() + defer s.mu.Unlock() + + for existing, expiresAt := range s.keys { + if now.After(expiresAt) { + delete(s.keys, existing) + } + } + if expiresAt, exists := s.keys[key]; exists && now.Before(expiresAt) { + return false + } + s.keys[key] = now.Add(dedupTTL) + return true +} \ No newline at end of file diff --git a/internal/apps/openflare/chwriter/dedup_test.go b/internal/apps/openflare/chwriter/dedup_test.go new file mode 100644 index 00000000..e36542e0 --- /dev/null +++ b/internal/apps/openflare/chwriter/dedup_test.go @@ -0,0 +1,24 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package chwriter + +import "testing" + +func TestDedupSetMarkIfNew(t *testing.T) { + t.Parallel() + + set := newDedupSet() + if !set.markIfNew("node-a|1") { + t.Fatal("markIfNew() = false, want true on first key") + } + if set.markIfNew("node-a|1") { + t.Fatal("markIfNew() = true, want false on duplicate key") + } + if !set.markIfNew("node-b|1") { + t.Fatal("markIfNew() = false, want true on different key") + } + if set.markIfNew("") { + t.Fatal("markIfNew() = true, want false on empty key") + } +} \ No newline at end of file diff --git a/internal/apps/openflare/chwriter/writer.go b/internal/apps/openflare/chwriter/writer.go new file mode 100644 index 00000000..eee83561 --- /dev/null +++ b/internal/apps/openflare/chwriter/writer.go @@ -0,0 +1,213 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +// Package chwriter queues OpenFlare ClickHouse writes and flushes them through +// internal/db/batchwriter with per-table writer instances. +package chwriter + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/Rain-kl/Wavelet/internal/config" + "github.com/Rain-kl/Wavelet/internal/db/batchwriter" + analyticsmodel "github.com/Rain-kl/Wavelet/internal/model/analytics" + analyticsrepo "github.com/Rain-kl/Wavelet/internal/repository/analytics" + "github.com/Rain-kl/Wavelet/pkg/logger" +) + +const ( + observabilityQueueSize = 5_000 + observabilityMaxBatchSize = 200 + observabilityFlushEvery = 2 * time.Second + + nodeAccessLogQueueSize = 10_000 + nodeAccessLogMaxBatchSize = 1_000 + nodeAccessLogFlushEvery = time.Second +) + +var ( + initOnce sync.Once + + metricSnapshotWriter *batchwriter.Writer[analyticsmodel.NodeMetricSnapshot] + requestReportWriter *batchwriter.Writer[analyticsmodel.NodeRequestReport] + openrestyWriter *batchwriter.Writer[analyticsmodel.NodeObsOpenresty] + frpsWriter *batchwriter.Writer[analyticsmodel.NodeObsFrps] + frpcWriter *batchwriter.Writer[analyticsmodel.NodeObsFrpc] + nodeAccessLogWriter *batchwriter.Writer[analyticsmodel.NodeAccessLog] + + metricSnapshotDedup *dedupSet + requestReportDedup *dedupSet +) + +// Init starts OpenFlare ClickHouse batch writers. Safe to call multiple times. +func Init(ctx context.Context) { + if !config.Config.ClickHouse.Enabled { + return + } + + initOnce.Do(func() { + metricSnapshotDedup = newDedupSet() + requestReportDedup = newDedupSet() + + metricSnapshotWriter = mustNewObservabilityWriter("metric_snapshots", analyticsrepo.BatchInsertNodeMetricSnapshots) + requestReportWriter = mustNewObservabilityWriter("request_reports", analyticsrepo.BatchInsertNodeRequestReports) + openrestyWriter = mustNewObservabilityWriter("openresty_obs", analyticsrepo.BatchInsertNodeObsOpenresty) + frpsWriter = mustNewObservabilityWriter("frps_obs", analyticsrepo.BatchInsertNodeObsFrps) + frpcWriter = mustNewObservabilityWriter("frpc_obs", analyticsrepo.BatchInsertNodeObsFrpc) + nodeAccessLogWriter = mustNewNodeAccessLogWriter() + + metricSnapshotWriter.Start(ctx) + requestReportWriter.Start(ctx) + openrestyWriter.Start(ctx) + frpsWriter.Start(ctx) + frpcWriter.Start(ctx) + nodeAccessLogWriter.Start(ctx) + }) +} + +// Stop drains all OpenFlare ClickHouse writers. +func Stop(ctx context.Context) error { + if !running() { + return nil + } + + var firstErr error + for _, writer := range []batchStopper{ + metricSnapshotWriter, + requestReportWriter, + openrestyWriter, + frpsWriter, + frpcWriter, + nodeAccessLogWriter, + } { + if writer == nil { + continue + } + if err := writer.Stop(ctx); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} + +// QueueMetricSnapshot enqueues a metric snapshot for asynchronous flush. +func QueueMetricSnapshot(snapshot analyticsmodel.NodeMetricSnapshot) { + if metricSnapshotWriter == nil { + return + } + key := fmt.Sprintf("%s|%d", snapshot.NodeID, snapshot.CapturedAt.UTC().UnixNano()) + if !metricSnapshotDedup.markIfNew(key) { + return + } + metricSnapshotWriter.TryEnqueue(snapshot) +} + +// QueueRequestReport enqueues a request report for asynchronous flush. +func QueueRequestReport(report analyticsmodel.NodeRequestReport) { + if requestReportWriter == nil { + return + } + key := fmt.Sprintf( + "%s|%d|%d", + report.NodeID, + report.WindowStartedAt.UTC().UnixNano(), + report.WindowEndedAt.UTC().UnixNano(), + ) + if !requestReportDedup.markIfNew(key) { + return + } + requestReportWriter.TryEnqueue(report) +} + +// QueueOpenrestyObservation enqueues an OpenResty observation for asynchronous flush. +func QueueOpenrestyObservation(observation analyticsmodel.NodeObsOpenresty) { + if openrestyWriter == nil { + return + } + openrestyWriter.TryEnqueue(observation) +} + +// QueueFrpsObservation enqueues an FRPS observation for asynchronous flush. +func QueueFrpsObservation(observation analyticsmodel.NodeObsFrps) { + if frpsWriter == nil { + return + } + frpsWriter.TryEnqueue(observation) +} + +// QueueFrpcObservation enqueues an FRPC observation for asynchronous flush. +func QueueFrpcObservation(observation analyticsmodel.NodeObsFrpc) { + if frpcWriter == nil { + return + } + frpcWriter.TryEnqueue(observation) +} + +// QueueNodeAccessLogs enqueues node access logs for asynchronous flush. +func QueueNodeAccessLogs(logs []analyticsmodel.NodeAccessLog) { + if nodeAccessLogWriter == nil || len(logs) == 0 { + return + } + for _, logItem := range logs { + nodeAccessLogWriter.TryEnqueue(logItem) + } +} + +func mustNewObservabilityWriter[T any](name string, flush batchwriter.FlushFunc[T]) *batchwriter.Writer[T] { + cfg := batchwriter.Config{ + Name: name, + QueueSize: observabilityQueueSize, + MaxBatchSize: observabilityMaxBatchSize, + FlushInterval: observabilityFlushEvery, + } + writer, err := batchwriter.New( + cfg, + flush, + withObservabilityDropHandler[T](name), + batchwriter.WithFlushErrorHandler[T](func(ctx context.Context, batchSize int, err error) { + logger.ErrorF(ctx, "[OpenFlare] flush %s failed (batch=%d): %v", name, batchSize, err) + }), + ) + if err != nil { + panic(fmt.Sprintf("openflare chwriter %s: %v", name, err)) + } + return writer +} + +func mustNewNodeAccessLogWriter() *batchwriter.Writer[analyticsmodel.NodeAccessLog] { + cfg := batchwriter.Config{ + Name: "node_access_logs", + QueueSize: nodeAccessLogQueueSize, + MaxBatchSize: nodeAccessLogMaxBatchSize, + FlushInterval: nodeAccessLogFlushEvery, + } + writer, err := batchwriter.New[analyticsmodel.NodeAccessLog](cfg, analyticsrepo.BatchInsertNodeAccessLogs, + batchwriter.WithDropHandler[analyticsmodel.NodeAccessLog](func(item analyticsmodel.NodeAccessLog) { + logger.WarnF(context.Background(), "[OpenFlare] node access log queue full, dropping log for node %s path %s", item.NodeID, item.Path) + }), + batchwriter.WithFlushErrorHandler[analyticsmodel.NodeAccessLog](func(ctx context.Context, batchSize int, err error) { + logger.ErrorF(ctx, "[OpenFlare] flush node access logs failed (batch=%d): %v", batchSize, err) + }), + ) + if err != nil { + panic(fmt.Sprintf("openflare chwriter node_access_logs: %v", err)) + } + return writer +} + +func withObservabilityDropHandler[T any](name string) batchwriter.Option[T] { + return batchwriter.WithDropHandler(func(_ T) { + logger.WarnF(context.Background(), "[OpenFlare] %s queue full, dropping observability item", name) + }) +} + +type batchStopper interface { + Stop(ctx context.Context) error +} + +func running() bool { + return metricSnapshotWriter != nil && metricSnapshotWriter.Running() +} \ No newline at end of file diff --git a/internal/apps/risk_control/logics.go b/internal/apps/risk_control/logics.go index 067ad23c..b2f969f9 100644 --- a/internal/apps/risk_control/logics.go +++ b/internal/apps/risk_control/logics.go @@ -5,88 +5,96 @@ package risk_control import ( "context" - "time" + "sync" "github.com/Rain-kl/Wavelet/internal/config" + "github.com/Rain-kl/Wavelet/internal/db/batchwriter" "github.com/Rain-kl/Wavelet/internal/model/analytics" analyticsrepo "github.com/Rain-kl/Wavelet/internal/repository/analytics" "github.com/Rain-kl/Wavelet/pkg/logger" ) -var logChan chan *analytics.UserAccessLog - -const ( - defaultQueueSize = 10000 - maxBatchSize = 1000 - flushInterval = 1 * time.Second +var ( + logWriterMu sync.RWMutex + logWriter *batchwriter.Writer[*analytics.UserAccessLog] ) -// InitLogWriter 初始化日志写入通道和后台写入协程 +// InitLogWriter initializes the ClickHouse access-log batch writer. func InitLogWriter(ctx context.Context) { if !config.Config.ClickHouse.Enabled { return } - logChan = make(chan *analytics.UserAccessLog, defaultQueueSize) - go startBatchWorker(context.WithoutCancel(ctx)) -} - -// IsBufferFull 检查当前本地缓冲队列是否已满 -// 如果没有启用 ClickHouse，默认返回 false，不触发限流 -func IsBufferFull() bool { - if !config.Config.ClickHouse.Enabled || logChan == nil { - return false - } - return len(logChan) >= cap(logChan) -} - -// QueueAccessLog 异步非阻塞地将日志推入缓冲队列 -func QueueAccessLog(logItem *analytics.UserAccessLog) { - if !config.Config.ClickHouse.Enabled || logChan == nil { + logWriterMu.Lock() + defer logWriterMu.Unlock() + if logWriter != nil { return } - select { - case logChan <- logItem: - default: - logger.WarnF(context.Background(), "[RiskControl] Log queue full, dropping log item for path: %s", logItem.Path) + cfg := batchwriter.DefaultConfig() + writer, err := batchwriter.New[*analytics.UserAccessLog](cfg, func(ctx context.Context, items []*analytics.UserAccessLog) error { + rows := make([]analytics.UserAccessLog, 0, len(items)) + for _, item := range items { + if item == nil { + continue + } + rows = append(rows, *item) + } + return analyticsrepo.BatchInsert(ctx, rows) + }, + batchwriter.WithDropHandler[*analytics.UserAccessLog](func(item *analytics.UserAccessLog) { + path := "" + if item != nil { + path = item.Path + } + logger.WarnF(context.Background(), "[RiskControl] Log queue full, dropping log item for path: %s", path) + }), + batchwriter.WithFlushErrorHandler[*analytics.UserAccessLog](func(ctx context.Context, batchSize int, err error) { + logger.ErrorF(ctx, "[RiskControl] Send ClickHouse batch failed (batch=%d): %v", batchSize, err) + }), + ) + if err != nil { + logger.ErrorF(ctx, "[RiskControl] init log writer failed: %v", err) + return + } + + writer.Start(ctx) + logWriter = writer +} + +// IsBufferFull reports whether the access-log queue has no remaining capacity. +func IsBufferFull() bool { + writer := currentLogWriter() + if writer == nil { + return false + } + return writer.IsFull() +} + +// QueueAccessLog enqueues an access log without blocking. +func QueueAccessLog(logItem *analytics.UserAccessLog) { + writer := currentLogWriter() + if writer == nil || logItem == nil { + return + } + writer.TryEnqueue(logItem) +} + +// SetLogWriterForTest swaps the access-log writer for unit tests. +func SetLogWriterForTest(writer *batchwriter.Writer[*analytics.UserAccessLog]) func() { + logWriterMu.Lock() + previous := logWriter + logWriter = writer + logWriterMu.Unlock() + return func() { + logWriterMu.Lock() + logWriter = previous + logWriterMu.Unlock() } } -func startBatchWorker(ctx context.Context) { - ticker := time.NewTicker(flushInterval) - defer ticker.Stop() - - var batch []*analytics.UserAccessLog - - flush := func() { - if len(batch) == 0 { - return - } - - items := make([]analytics.UserAccessLog, len(batch)) - for i, item := range batch { - items[i] = *item - } - if err := analyticsrepo.BatchInsert(ctx, items); err != nil { - logger.ErrorF(ctx, "[RiskControl] Send ClickHouse batch failed: %v", err) - } - batch = nil - } - - for { - select { - case item, ok := <-logChan: - if !ok { - flush() - return - } - batch = append(batch, item) - if len(batch) >= maxBatchSize { - flush() - } - case <-ticker.C: - flush() - } - } +func currentLogWriter() *batchwriter.Writer[*analytics.UserAccessLog] { + logWriterMu.RLock() + defer logWriterMu.RUnlock() + return logWriter } \ No newline at end of file diff --git a/internal/apps/risk_control/middleware_test.go b/internal/apps/risk_control/middleware_test.go index c20965d1..f1381faf 100644 --- a/internal/apps/risk_control/middleware_test.go +++ b/internal/apps/risk_control/middleware_test.go @@ -4,6 +4,7 @@ package risk_control import ( + "context" "encoding/json" "net/http" "net/http/httptest" @@ -12,6 +13,7 @@ import ( "github.com/Rain-kl/Wavelet/internal/apps/oauth" "github.com/Rain-kl/Wavelet/internal/config" + "github.com/Rain-kl/Wavelet/internal/db/batchwriter" "github.com/Rain-kl/Wavelet/internal/model" "github.com/Rain-kl/Wavelet/internal/model/analytics" "github.com/Rain-kl/Wavelet/internal/testhelper" @@ -19,6 +21,35 @@ import ( "github.com/stretchr/testify/assert" ) +const testLogQueueSize = 10_000 + +func newTestLogWriter(t *testing.T, queueSize int) (*batchwriter.Writer[*analytics.UserAccessLog], chan *analytics.UserAccessLog) { + t.Helper() + + received := make(chan *analytics.UserAccessLog, queueSize) + cfg := batchwriter.Config{ + QueueSize: queueSize, + MaxBatchSize: 1, + FlushInterval: 5 * time.Millisecond, + } + writer, err := batchwriter.New[*analytics.UserAccessLog](cfg, func(_ context.Context, items []*analytics.UserAccessLog) error { + for _, item := range items { + if item != nil { + received <- item + } + } + return nil + }) + assert.NoError(t, err) + writer.Start(context.Background()) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = writer.Stop(stopCtx) + }) + return writer, received +} + func TestRiskControlMiddleware(t *testing.T) { gin.SetMode(gin.TestMode) @@ -41,15 +72,15 @@ func TestRiskControlMiddleware(t *testing.T) { t.Run("ClickHouse enabled - Normal Authenticated Request", func(t *testing.T) { config.Config.ClickHouse.Enabled = true - logChan = make(chan *analytics.UserAccessLog, defaultQueueSize) + writer, received := newTestLogWriter(t, testLogQueueSize) + resetWriter := SetLogWriterForTest(writer) defer func() { config.Config.ClickHouse.Enabled = false - logChan = nil + resetWriter() }() r := gin.New() r.Use(func(c *gin.Context) { - // Mock authentication middleware placing user in context user := &model.User{ID: 12345} oauth.SetToContext(c, oauth.UserObjKey, user) c.Next() @@ -68,9 +99,8 @@ func TestRiskControlMiddleware(t *testing.T) { assert.Equal(t, http.StatusOK, w.Code) assert.Equal(t, "ok", w.Body.String()) - // Verify log is enqueued select { - case logItem := <-logChan: + case logItem := <-received: assert.Equal(t, uint64(12345), logItem.UserID) assert.Equal(t, "/test", logItem.Path) assert.Equal(t, http.MethodGet, logItem.Method) @@ -78,17 +108,18 @@ func TestRiskControlMiddleware(t *testing.T) { assert.NotEmpty(t, logItem.Headers) assert.Contains(t, logItem.Headers, "X-Test-Header") assert.NotContains(t, logItem.Headers, "Cookie") - case <-time.After(100 * time.Millisecond): - t.Fatal("expected log item in logChan, but got none") + case <-time.After(200 * time.Millisecond): + t.Fatal("expected flushed log item, but got none") } }) t.Run("ClickHouse enabled - Unauthenticated Request", func(t *testing.T) { config.Config.ClickHouse.Enabled = true - logChan = make(chan *analytics.UserAccessLog, defaultQueueSize) + writer, received := newTestLogWriter(t, testLogQueueSize) + resetWriter := SetLogWriterForTest(writer) defer func() { config.Config.ClickHouse.Enabled = false - logChan = nil + resetWriter() }() r := testhelper.NewTestGinEngine(RiskControlMiddleware()) @@ -103,26 +134,24 @@ func TestRiskControlMiddleware(t *testing.T) { assert.Equal(t, http.StatusOK, w.Code) assert.Equal(t, "ok", w.Body.String()) - // Verify no log is enqueued select { - case <-logChan: + case <-received: t.Fatal("expected no log item for unauthenticated request") case <-time.After(50 * time.Millisecond): - // Success } }) t.Run("ClickHouse enabled - Buffer Full Rate Limiting", func(t *testing.T) { config.Config.ClickHouse.Enabled = true - logChan = make(chan *analytics.UserAccessLog, 2) // small capacity for quick fill + writer, _ := newTestLogWriter(t, 2) + resetWriter := SetLogWriterForTest(writer) defer func() { config.Config.ClickHouse.Enabled = false - logChan = nil + resetWriter() }() - // fill logChan up to cap to simulate buffer full - for len(logChan) < cap(logChan) { - logChan <- &analytics.UserAccessLog{} + for range 2 { + assert.True(t, writer.TryEnqueue(&analytics.UserAccessLog{})) } r := testhelper.NewTestGinEngine(RiskControlMiddleware()) @@ -141,4 +170,4 @@ func TestRiskControlMiddleware(t *testing.T) { assert.NoError(t, err) assert.Contains(t, resp["error_msg"], "系统繁忙") }) -} +} \ No newline at end of file diff --git a/internal/bootstrap/bootstrap.go b/internal/bootstrap/bootstrap.go index b21af1fd..b7cf8796 100644 --- a/internal/bootstrap/bootstrap.go +++ b/internal/bootstrap/bootstrap.go @@ -11,6 +11,7 @@ import ( admin_push "github.com/Rain-kl/Wavelet/internal/apps/admin/push" "github.com/Rain-kl/Wavelet/internal/apps/admin/push/custom_events" + "github.com/Rain-kl/Wavelet/internal/apps/openflare/chwriter" "github.com/Rain-kl/Wavelet/internal/apps/risk_control" taskhandlers "github.com/Rain-kl/Wavelet/internal/task/handlers" "github.com/Rain-kl/Wavelet/pkg/logger" @@ -83,6 +84,7 @@ func Init(ctx context.Context, opts Options) { } if opts.API { risk_control.InitLogWriter(ctx) + chwriter.Init(ctx) } }) } diff --git a/internal/model/openflare_access_log_store.go b/internal/model/openflare_access_log_store.go index 6d135ec9..c6ee4e57 100644 --- a/internal/model/openflare_access_log_store.go +++ b/internal/model/openflare_access_log_store.go @@ -8,6 +8,7 @@ import ( "sync" "time" + "github.com/Rain-kl/Wavelet/internal/apps/openflare/chwriter" analyticsmodel "github.com/Rain-kl/Wavelet/internal/model/analytics" analyticsrepo "github.com/Rain-kl/Wavelet/internal/repository/analytics" ) @@ -63,7 +64,7 @@ func NewMemoryAccessLogStore() accessLogStore { type clickhouseAccessLogStore struct{} -func (clickhouseAccessLogStore) InsertBatch(ctx context.Context, records []*OpenFlareAccessLog) error { +func (clickhouseAccessLogStore) InsertBatch(_ context.Context, records []*OpenFlareAccessLog) error { logs := make([]analyticsmodel.NodeAccessLog, 0, len(records)) for _, record := range records { if record == nil { @@ -71,7 +72,8 @@ func (clickhouseAccessLogStore) InsertBatch(ctx context.Context, records []*Open } logs = append(logs, toAnalyticsNodeAccessLog(record)) } - return analyticsrepo.BatchInsertNodeAccessLogs(ctx, logs) + chwriter.QueueNodeAccessLogs(logs) + return nil } func (clickhouseAccessLogStore) List(ctx context.Context, query OpenFlareAccessLogQuery) ([]*OpenFlareAccessLog, error) { diff --git a/internal/model/openflare_observability_store.go b/internal/model/openflare_observability_store.go index 9c9ad95d..25136adf 100644 --- a/internal/model/openflare_observability_store.go +++ b/internal/model/openflare_observability_store.go @@ -9,6 +9,7 @@ import ( "sync" "time" + "github.com/Rain-kl/Wavelet/internal/apps/openflare/chwriter" analyticsmodel "github.com/Rain-kl/Wavelet/internal/model/analytics" analyticsrepo "github.com/Rain-kl/Wavelet/internal/repository/analytics" ) @@ -74,11 +75,12 @@ func NewMemoryObservabilityStore() observabilityStore { type clickhouseObservabilityStore struct{} -func (clickhouseObservabilityStore) InsertMetricSnapshot(ctx context.Context, record *OpenFlareMetricSnapshot) error { +func (clickhouseObservabilityStore) InsertMetricSnapshot(_ context.Context, record *OpenFlareMetricSnapshot) error { if record == nil { return nil } - return analyticsrepo.InsertNodeMetricSnapshot(ctx, toAnalyticsNodeMetricSnapshot(record)) + chwriter.QueueMetricSnapshot(toAnalyticsNodeMetricSnapshot(record)) + return nil } func (clickhouseObservabilityStore) ListMetricSnapshots(ctx context.Context, nodeID string, since time.Time, limit int) ([]*OpenFlareMetricSnapshot, error) { @@ -97,11 +99,12 @@ func (clickhouseObservabilityStore) DeleteMetricSnapshotsBefore(ctx context.Cont return analyticsrepo.DeleteNodeMetricSnapshotsBefore(ctx, cutoff) } -func (clickhouseObservabilityStore) InsertRequestReport(ctx context.Context, record *OpenFlareRequestReport) error { +func (clickhouseObservabilityStore) InsertRequestReport(_ context.Context, record *OpenFlareRequestReport) error { if record == nil { return nil } - return analyticsrepo.InsertNodeRequestReport(ctx, toAnalyticsNodeRequestReport(record)) + chwriter.QueueRequestReport(toAnalyticsNodeRequestReport(record)) + return nil } func (clickhouseObservabilityStore) ListRequestReports(ctx context.Context, nodeID string, since time.Time, limit int) ([]*OpenFlareRequestReport, error) { @@ -120,11 +123,12 @@ func (clickhouseObservabilityStore) DeleteRequestReportsBefore(ctx context.Conte return analyticsrepo.DeleteNodeRequestReportsBefore(ctx, cutoff) } -func (clickhouseObservabilityStore) InsertNodeObservationOpenresty(ctx context.Context, record *OpenFlareNodeObservationOpenresty) error { +func (clickhouseObservabilityStore) InsertNodeObservationOpenresty(_ context.Context, record *OpenFlareNodeObservationOpenresty) error { if record == nil { return nil } - return analyticsrepo.InsertNodeObsOpenresty(ctx, toAnalyticsNodeObsOpenresty(record)) + chwriter.QueueOpenrestyObservation(toAnalyticsNodeObsOpenresty(record)) + return nil } func (clickhouseObservabilityStore) ListNodeObservationOpenresty(ctx context.Context, nodeID string, since time.Time, limit int) ([]*OpenFlareNodeObservationOpenresty, error) { @@ -143,11 +147,12 @@ func (clickhouseObservabilityStore) DeleteNodeObservationOpenrestyBefore(ctx con return analyticsrepo.DeleteNodeObsOpenrestyBefore(ctx, cutoff) } -func (clickhouseObservabilityStore) InsertNodeObservationFrps(ctx context.Context, record *OpenFlareNodeObservationFrps) error { +func (clickhouseObservabilityStore) InsertNodeObservationFrps(_ context.Context, record *OpenFlareNodeObservationFrps) error { if record == nil { return nil } - return analyticsrepo.InsertNodeObsFrps(ctx, toAnalyticsNodeObsFrps(record)) + chwriter.QueueFrpsObservation(toAnalyticsNodeObsFrps(record)) + return nil } func (clickhouseObservabilityStore) ListNodeObservationFrps(ctx context.Context, nodeID string, since time.Time, limit int) ([]*OpenFlareNodeObservationFrps, error) { @@ -166,11 +171,12 @@ func (clickhouseObservabilityStore) DeleteNodeObservationFrpsBefore(ctx context. return analyticsrepo.DeleteNodeObsFrpsBefore(ctx, cutoff) } -func (clickhouseObservabilityStore) InsertNodeObservationFrpc(ctx context.Context, record *OpenFlareNodeObservationFrpc) error { +func (clickhouseObservabilityStore) InsertNodeObservationFrpc(_ context.Context, record *OpenFlareNodeObservationFrpc) error { if record == nil { return nil } - return analyticsrepo.InsertNodeObsFrpc(ctx, toAnalyticsNodeObsFrpc(record)) + chwriter.QueueFrpcObservation(toAnalyticsNodeObsFrpc(record)) + return nil } func (clickhouseObservabilityStore) ListNodeObservationFrpc(ctx context.Context, nodeID string, since time.Time, limit int) ([]*OpenFlareNodeObservationFrpc, error) { diff --git a/internal/repository/analytics/node_observability_writer.go b/internal/repository/analytics/node_observability_writer.go index fd0d41e6..6e0688b8 100644 --- a/internal/repository/analytics/node_observability_writer.go +++ b/internal/repository/analytics/node_observability_writer.go @@ -14,26 +14,19 @@ import ( analyticsmodel "github.com/Rain-kl/Wavelet/internal/model/analytics" ) -// InsertNodeMetricSnapshot writes a metric snapshot when no row exists for node_id+captured_at. +// InsertNodeMetricSnapshot writes a single metric snapshot via the batch API. func InsertNodeMetricSnapshot(ctx context.Context, snapshot analyticsmodel.NodeMetricSnapshot) error { - nodeID := strings.TrimSpace(snapshot.NodeID) - if nodeID == "" { + if strings.TrimSpace(snapshot.NodeID) == "" { return nil } - capturedAt := snapshot.CapturedAt.UTC() + return BatchInsertNodeMetricSnapshots(ctx, []analyticsmodel.NodeMetricSnapshot{snapshot}) +} - exists, err := nodeObservabilityRowExists( - ctx, - fmt.Sprintf("SELECT count() FROM %s WHERE node_id = ? AND captured_at = ?", nodeMetricSnapshotTableName()), - nodeID, capturedAt, - ) - if err != nil { - return err - } - if exists { +// BatchInsertNodeMetricSnapshots writes metric snapshots to ClickHouse. +func BatchInsertNodeMetricSnapshots(ctx context.Context, snapshots []analyticsmodel.NodeMetricSnapshot) error { + if len(snapshots) == 0 { return nil } - if db.ChConn == nil { return fmt.Errorf("clickhouse connection is not initialized") } @@ -44,30 +37,40 @@ func InsertNodeMetricSnapshot(ctx context.Context, snapshot analyticsmodel.NodeM } now := time.Now().UTC() - id := snapshot.ID - if id == 0 { - id = idgen.NextUint64ID() + for _, snapshot := range snapshots { + nodeID := strings.TrimSpace(snapshot.NodeID) + if nodeID == "" { + continue + } + id := snapshot.ID + if id == 0 { + id = idgen.NextUint64ID() + } + createdAt := snapshot.CreatedAt + if createdAt.IsZero() { + createdAt = now + } + if err := batch.Append( + id, + nodeID, + snapshot.CapturedAt.UTC(), + snapshot.CPUUsagePercent, + snapshot.MemoryUsedBytes, + snapshot.MemoryTotalBytes, + snapshot.StorageUsedBytes, + snapshot.StorageTotalBytes, + snapshot.DiskReadBytes, + snapshot.DiskWriteBytes, + snapshot.NetworkRxBytes, + snapshot.NetworkTxBytes, + createdAt.UTC(), + ); err != nil { + return fmt.Errorf("append node metric snapshot to batch: %w", err) + } } - createdAt := snapshot.CreatedAt - if createdAt.IsZero() { - createdAt = now - } - if err := batch.Append( - id, - nodeID, - capturedAt, - snapshot.CPUUsagePercent, - snapshot.MemoryUsedBytes, - snapshot.MemoryTotalBytes, - snapshot.StorageUsedBytes, - snapshot.StorageTotalBytes, - snapshot.DiskReadBytes, - snapshot.DiskWriteBytes, - snapshot.NetworkRxBytes, - snapshot.NetworkTxBytes, - createdAt.UTC(), - ); err != nil { - return fmt.Errorf("append node metric snapshot to batch: %w", err) + + if batch.Rows() == 0 { + return nil } if err := batch.Send(); err != nil { return fmt.Errorf("send clickhouse batch: %w", err) @@ -75,30 +78,19 @@ func InsertNodeMetricSnapshot(ctx context.Context, snapshot analyticsmodel.NodeM return nil } -// InsertNodeRequestReport writes a request report when no row exists for node_id+window bounds. +// InsertNodeRequestReport writes a single request report via the batch API. func InsertNodeRequestReport(ctx context.Context, report analyticsmodel.NodeRequestReport) error { - nodeID := strings.TrimSpace(report.NodeID) - if nodeID == "" { + if strings.TrimSpace(report.NodeID) == "" { return nil } - windowStartedAt := report.WindowStartedAt.UTC() - windowEndedAt := report.WindowEndedAt.UTC() + return BatchInsertNodeRequestReports(ctx, []analyticsmodel.NodeRequestReport{report}) +} - exists, err := nodeObservabilityRowExists( - ctx, - fmt.Sprintf( - "SELECT count() FROM %s WHERE node_id = ? AND window_started_at = ? AND window_ended_at = ?", - nodeRequestReportTableName(), - ), - nodeID, windowStartedAt, windowEndedAt, - ) - if err != nil { - return err - } - if exists { +// BatchInsertNodeRequestReports writes request reports to ClickHouse. +func BatchInsertNodeRequestReports(ctx context.Context, reports []analyticsmodel.NodeRequestReport) error { + if len(reports) == 0 { return nil } - if db.ChConn == nil { return fmt.Errorf("clickhouse connection is not initialized") } @@ -109,28 +101,38 @@ func InsertNodeRequestReport(ctx context.Context, report analyticsmodel.NodeRequ } now := time.Now().UTC() - id := report.ID - if id == 0 { - id = idgen.NextUint64ID() + for _, report := range reports { + nodeID := strings.TrimSpace(report.NodeID) + if nodeID == "" { + continue + } + id := report.ID + if id == 0 { + id = idgen.NextUint64ID() + } + createdAt := report.CreatedAt + if createdAt.IsZero() { + createdAt = now + } + if err := batch.Append( + id, + nodeID, + report.WindowStartedAt.UTC(), + report.WindowEndedAt.UTC(), + report.RequestCount, + report.ErrorCount, + report.UniqueVisitorCount, + report.StatusCodesJSON, + report.TopDomainsJSON, + report.SourceCountriesJSON, + createdAt.UTC(), + ); err != nil { + return fmt.Errorf("append node request report to batch: %w", err) + } } - createdAt := report.CreatedAt - if createdAt.IsZero() { - createdAt = now - } - if err := batch.Append( - id, - nodeID, - windowStartedAt, - windowEndedAt, - report.RequestCount, - report.ErrorCount, - report.UniqueVisitorCount, - report.StatusCodesJSON, - report.TopDomainsJSON, - report.SourceCountriesJSON, - createdAt.UTC(), - ); err != nil { - return fmt.Errorf("append node request report to batch: %w", err) + + if batch.Rows() == 0 { + return nil } if err := batch.Send(); err != nil { return fmt.Errorf("send clickhouse batch: %w", err) @@ -138,34 +140,19 @@ func InsertNodeRequestReport(ctx context.Context, report analyticsmodel.NodeRequ return nil } -// InsertNodeObsOpenresty writes an OpenResty observability snapshot. +// InsertNodeObsOpenresty writes a single OpenResty observation via the batch API. func InsertNodeObsOpenresty(ctx context.Context, obs analyticsmodel.NodeObsOpenresty) error { - nodeID := strings.TrimSpace(obs.NodeID) - if nodeID == "" { + if strings.TrimSpace(obs.NodeID) == "" { return nil } - return insertNodeObsOpenrestyBatch(ctx, obs, nodeID) + return BatchInsertNodeObsOpenresty(ctx, []analyticsmodel.NodeObsOpenresty{obs}) } -// InsertNodeObsFrps writes an FRPS observability snapshot. -func InsertNodeObsFrps(ctx context.Context, obs analyticsmodel.NodeObsFrps) error { - nodeID := strings.TrimSpace(obs.NodeID) - if nodeID == "" { +// BatchInsertNodeObsOpenresty writes OpenResty observations to ClickHouse. +func BatchInsertNodeObsOpenresty(ctx context.Context, observations []analyticsmodel.NodeObsOpenresty) error { + if len(observations) == 0 { return nil } - return insertNodeObsFrpsBatch(ctx, obs, nodeID) -} - -// InsertNodeObsFrpc writes an FRPC observability snapshot. -func InsertNodeObsFrpc(ctx context.Context, obs analyticsmodel.NodeObsFrpc) error { - nodeID := strings.TrimSpace(obs.NodeID) - if nodeID == "" { - return nil - } - return insertNodeObsFrpcBatch(ctx, obs, nodeID) -} - -func insertNodeObsOpenrestyBatch(ctx context.Context, obs analyticsmodel.NodeObsOpenresty, nodeID string) error { if db.ChConn == nil { return fmt.Errorf("clickhouse connection is not initialized") } @@ -176,28 +163,38 @@ func insertNodeObsOpenrestyBatch(ctx context.Context, obs analyticsmodel.NodeObs } now := time.Now().UTC() - id := obs.ID - if id == 0 { - id = idgen.NextUint64ID() + for _, obs := range observations { + nodeID := strings.TrimSpace(obs.NodeID) + if nodeID == "" { + continue + } + id := obs.ID + if id == 0 { + id = idgen.NextUint64ID() + } + createdAt := obs.CreatedAt + if createdAt.IsZero() { + createdAt = now + } + capturedAt := obs.CapturedAt.UTC() + if capturedAt.IsZero() { + capturedAt = now + } + if err := batch.Append( + id, + nodeID, + capturedAt, + obs.OpenrestyRxBytes, + obs.OpenrestyTxBytes, + obs.OpenrestyConnections, + createdAt.UTC(), + ); err != nil { + return fmt.Errorf("append node openresty observation to batch: %w", err) + } } - createdAt := obs.CreatedAt - if createdAt.IsZero() { - createdAt = now - } - capturedAt := obs.CapturedAt.UTC() - if capturedAt.IsZero() { - capturedAt = now - } - if err := batch.Append( - id, - nodeID, - capturedAt, - obs.OpenrestyRxBytes, - obs.OpenrestyTxBytes, - obs.OpenrestyConnections, - createdAt.UTC(), - ); err != nil { - return fmt.Errorf("append node openresty observation to batch: %w", err) + + if batch.Rows() == 0 { + return nil } if err := batch.Send(); err != nil { return fmt.Errorf("send clickhouse batch: %w", err) @@ -205,7 +202,19 @@ func insertNodeObsOpenrestyBatch(ctx context.Context, obs analyticsmodel.NodeObs return nil } -func insertNodeObsFrpsBatch(ctx context.Context, obs analyticsmodel.NodeObsFrps, nodeID string) error { +// InsertNodeObsFrps writes a single FRPS observation via the batch API. +func InsertNodeObsFrps(ctx context.Context, obs analyticsmodel.NodeObsFrps) error { + if strings.TrimSpace(obs.NodeID) == "" { + return nil + } + return BatchInsertNodeObsFrps(ctx, []analyticsmodel.NodeObsFrps{obs}) +} + +// BatchInsertNodeObsFrps writes FRPS observations to ClickHouse. +func BatchInsertNodeObsFrps(ctx context.Context, observations []analyticsmodel.NodeObsFrps) error { + if len(observations) == 0 { + return nil + } if db.ChConn == nil { return fmt.Errorf("clickhouse connection is not initialized") } @@ -216,29 +225,39 @@ func insertNodeObsFrpsBatch(ctx context.Context, obs analyticsmodel.NodeObsFrps, } now := time.Now().UTC() - id := obs.ID - if id == 0 { - id = idgen.NextUint64ID() + for _, obs := range observations { + nodeID := strings.TrimSpace(obs.NodeID) + if nodeID == "" { + continue + } + id := obs.ID + if id == 0 { + id = idgen.NextUint64ID() + } + createdAt := obs.CreatedAt + if createdAt.IsZero() { + createdAt = now + } + capturedAt := obs.CapturedAt.UTC() + if capturedAt.IsZero() { + capturedAt = now + } + if err := batch.Append( + id, + nodeID, + capturedAt, + obs.FrpsConnections, + obs.FrpsProxyCount, + obs.FrpsClientCount, + obs.FrpsProxies, + createdAt.UTC(), + ); err != nil { + return fmt.Errorf("append node frps observation to batch: %w", err) + } } - createdAt := obs.CreatedAt - if createdAt.IsZero() { - createdAt = now - } - capturedAt := obs.CapturedAt.UTC() - if capturedAt.IsZero() { - capturedAt = now - } - if err := batch.Append( - id, - nodeID, - capturedAt, - obs.FrpsConnections, - obs.FrpsProxyCount, - obs.FrpsClientCount, - obs.FrpsProxies, - createdAt.UTC(), - ); err != nil { - return fmt.Errorf("append node frps observation to batch: %w", err) + + if batch.Rows() == 0 { + return nil } if err := batch.Send(); err != nil { return fmt.Errorf("send clickhouse batch: %w", err) @@ -246,7 +265,19 @@ func insertNodeObsFrpsBatch(ctx context.Context, obs analyticsmodel.NodeObsFrps, return nil } -func insertNodeObsFrpcBatch(ctx context.Context, obs analyticsmodel.NodeObsFrpc, nodeID string) error { +// InsertNodeObsFrpc writes a single FRPC observation via the batch API. +func InsertNodeObsFrpc(ctx context.Context, obs analyticsmodel.NodeObsFrpc) error { + if strings.TrimSpace(obs.NodeID) == "" { + return nil + } + return BatchInsertNodeObsFrpc(ctx, []analyticsmodel.NodeObsFrpc{obs}) +} + +// BatchInsertNodeObsFrpc writes FRPC observations to ClickHouse. +func BatchInsertNodeObsFrpc(ctx context.Context, observations []analyticsmodel.NodeObsFrpc) error { + if len(observations) == 0 { + return nil + } if db.ChConn == nil { return fmt.Errorf("clickhouse connection is not initialized") } @@ -257,42 +288,40 @@ func insertNodeObsFrpcBatch(ctx context.Context, obs analyticsmodel.NodeObsFrpc, } now := time.Now().UTC() - id := obs.ID - if id == 0 { - id = idgen.NextUint64ID() + for _, obs := range observations { + nodeID := strings.TrimSpace(obs.NodeID) + if nodeID == "" { + continue + } + id := obs.ID + if id == 0 { + id = idgen.NextUint64ID() + } + createdAt := obs.CreatedAt + if createdAt.IsZero() { + createdAt = now + } + capturedAt := obs.CapturedAt.UTC() + if capturedAt.IsZero() { + capturedAt = now + } + if err := batch.Append( + id, + nodeID, + capturedAt, + obs.TunnelStatus, + obs.ConnectedRelaysCount, + createdAt.UTC(), + ); err != nil { + return fmt.Errorf("append node frpc observation to batch: %w", err) + } } - createdAt := obs.CreatedAt - if createdAt.IsZero() { - createdAt = now - } - capturedAt := obs.CapturedAt.UTC() - if capturedAt.IsZero() { - capturedAt = now - } - if err := batch.Append( - id, - nodeID, - capturedAt, - obs.TunnelStatus, - obs.ConnectedRelaysCount, - createdAt.UTC(), - ); err != nil { - return fmt.Errorf("append node frpc observation to batch: %w", err) + + if batch.Rows() == 0 { + return nil } if err := batch.Send(); err != nil { return fmt.Errorf("send clickhouse batch: %w", err) } return nil -} - -func nodeObservabilityRowExists(ctx context.Context, countSQL string, args ...any) (bool, error) { - conn, err := observabilityConn() - if err != nil { - return false, err - } - var count int64 - if err := conn.QueryRow(ctx, countSQL, args...).Scan(&count); err != nil { - return false, fmt.Errorf("check observability row exists: %w", err) - } - return count > 0, nil -} +} \ No newline at end of file