diff --git a/controlplane/telemetry/cmd/geoprobe-target/main.go b/controlplane/telemetry/cmd/geoprobe-target/main.go index dc8db060a4..3bdf25ceb6 100644 --- a/controlplane/telemetry/cmd/geoprobe-target/main.go +++ b/controlplane/telemetry/cmd/geoprobe-target/main.go @@ -91,22 +91,8 @@ func main() { var chWriter *geoprobe.ClickhouseWriter if chCfg := geoprobe.ClickhouseConfigFromEnv(); chCfg != nil { log.Info("clickhouse enabled", "addr", chCfg.Addr, "db", chCfg.Database) - - if err := geoprobe.RunMigrations(*chCfg, log); err != nil { - fmt.Fprintf(os.Stderr, "clickhouse migration failed: %v\n", err) - os.Exit(1) - } - - chConn, err := geoprobe.NewClickhouseConn(*chCfg) - if err != nil { - fmt.Fprintf(os.Stderr, "clickhouse connect failed: %v\n", err) - os.Exit(1) - } - defer chConn.Close() - - chWriter = geoprobe.NewClickhouseWriter(chConn, chCfg.Database, log) + chWriter = geoprobe.NewClickhouseWriter(*chCfg, log) go chWriter.Run(ctx) - log.Info("clickhouse writer started") } errCh := make(chan error, 2) diff --git a/controlplane/telemetry/internal/geoprobe/clickhouse.go b/controlplane/telemetry/internal/geoprobe/clickhouse.go index ccb9e3b962..0bcfd6ffc3 100644 --- a/controlplane/telemetry/internal/geoprobe/clickhouse.go +++ b/controlplane/telemetry/internal/geoprobe/clickhouse.go @@ -7,6 +7,7 @@ import ( "fmt" "log/slog" "os" + "strings" "sync" "time" @@ -28,6 +29,8 @@ func ClickhouseConfigFromEnv() *ClickhouseConfig { if addr == "" { return nil } + addr = strings.TrimPrefix(addr, "https://") + addr = strings.TrimPrefix(addr, "http://") db := os.Getenv("CLICKHOUSE_DB") if db == "" { db = "default" @@ -47,7 +50,8 @@ func ClickhouseConfigFromEnv() *ClickhouseConfig { func NewClickhouseConn(cfg ClickhouseConfig) (driver.Conn, error) { opts := &clickhouse.Options{ - Addr: []string{cfg.Addr}, + Protocol: clickhouse.HTTP, + Addr: []string{cfg.Addr}, Auth: clickhouse.Auth{ Database: cfg.Database, Username: cfg.Username, @@ -119,41 +123,76 @@ func OffsetRowFromLocationOffset(offset *LocationOffset, sourceAddr string, sigV return row } +const maxBufferedRows = 10_000 + type ClickhouseWriter struct { + cfg ClickhouseConfig conn driver.Conn - db string buf []OffsetRow mu sync.Mutex log *slog.Logger } -func NewClickhouseWriter(conn driver.Conn, db string, log *slog.Logger) *ClickhouseWriter { +func NewClickhouseWriter(cfg ClickhouseConfig, log *slog.Logger) *ClickhouseWriter { return &ClickhouseWriter{ - conn: conn, - db: db, - buf: make([]OffsetRow, 0, 64), - log: log, + cfg: cfg, + buf: make([]OffsetRow, 0, 64), + log: log, } } func (w *ClickhouseWriter) Record(row OffsetRow) { w.mu.Lock() + if len(w.buf) >= maxBufferedRows { + w.mu.Unlock() + w.log.Warn("clickhouse buffer full, dropping row", "max", maxBufferedRows) + return + } w.buf = append(w.buf, row) w.mu.Unlock() } +func (w *ClickhouseWriter) connect(ctx context.Context) error { + if err := RunMigrations(w.cfg, w.log); err != nil { + return fmt.Errorf("migrations: %w", err) + } + conn, err := NewClickhouseConn(w.cfg) + if err != nil { + return fmt.Errorf("connect: %w", err) + } + w.conn = conn + w.log.Info("clickhouse connected", "addr", w.cfg.Addr, "db", w.cfg.Database) + return nil +} + +func (w *ClickhouseWriter) Close() { + if w.conn != nil { + _ = w.conn.Close() + w.conn = nil + } +} + func (w *ClickhouseWriter) Run(ctx context.Context) { ticker := time.NewTicker(10 * time.Second) defer ticker.Stop() + defer w.Close() for { select { case <-ctx.Done(): - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - w.flush(shutdownCtx) - cancel() + if w.conn != nil { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + w.flush(shutdownCtx) + cancel() + } return case <-ticker.C: + if w.conn == nil { + if err := w.connect(ctx); err != nil { + w.log.Error("clickhouse connection failed, will retry", "error", err) + continue + } + } w.flush(ctx) } } @@ -170,10 +209,11 @@ func (w *ClickhouseWriter) flush(ctx context.Context) { w.mu.Unlock() batch, err := w.conn.PrepareBatch(ctx, fmt.Sprintf( - `INSERT INTO "%s".location_offsets`, w.db, + `INSERT INTO "%s".location_offsets`, w.cfg.Database, )) if err != nil { w.log.Error("failed to prepare batch", "error", err, "dropped_rows", len(rows)) + w.Close() return } @@ -200,6 +240,7 @@ func (w *ClickhouseWriter) flush(ctx context.Context) { ); err != nil { w.log.Error("failed to append row", "error", err, "dropped_rows", len(rows)) _ = batch.Abort() + w.Close() return } } @@ -207,6 +248,7 @@ func (w *ClickhouseWriter) flush(ctx context.Context) { if err := batch.Send(); err != nil { w.log.Error("failed to send batch", "error", err, "dropped_rows", len(rows)) _ = batch.Close() + w.Close() return } _ = batch.Close() diff --git a/controlplane/telemetry/internal/geoprobe/clickhouse_test.go b/controlplane/telemetry/internal/geoprobe/clickhouse_test.go index 943ade42df..c05cc2392b 100644 --- a/controlplane/telemetry/internal/geoprobe/clickhouse_test.go +++ b/controlplane/telemetry/internal/geoprobe/clickhouse_test.go @@ -1,6 +1,8 @@ package geoprobe import ( + "fmt" + "log/slog" "testing" "time" @@ -61,10 +63,42 @@ func TestOffsetRowFromLocationOffset(t *testing.T) { require.WithinDuration(t, time.Now(), row.ReceivedAt, 2*time.Second) } -func TestClickhouseWriterRecordBuffers(t *testing.T) { - w := &ClickhouseWriter{ - buf: make([]OffsetRow, 0), +func TestClickhouseConfigFromEnv(t *testing.T) { + tests := []struct { + name string + addr string + wantAddr string + }{ + { + name: "plain host:port", + addr: "clickhouse.example.com:8443", + wantAddr: "clickhouse.example.com:8443", + }, + { + name: "strips https:// scheme prefix", + addr: "https://clickhouse.example.com:8443", + wantAddr: "clickhouse.example.com:8443", + }, + { + name: "strips http:// scheme prefix", + addr: "http://localhost:8123", + wantAddr: "localhost:8123", + }, } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("CLICKHOUSE_ADDR", tt.addr) + t.Setenv("CLICKHOUSE_DB", "testdb") + cfg := ClickhouseConfigFromEnv() + require.NotNil(t, cfg) + require.Equal(t, tt.wantAddr, cfg.Addr) + }) + } +} + +func TestClickhouseWriterRecordBuffers(t *testing.T) { + w := NewClickhouseWriter(ClickhouseConfig{Addr: "unused"}, slog.Default()) w.Record(OffsetRow{SourceAddr: "a"}) w.Record(OffsetRow{SourceAddr: "b"}) @@ -72,3 +106,14 @@ func TestClickhouseWriterRecordBuffers(t *testing.T) { require.Len(t, w.buf, 2) w.mu.Unlock() } + +func TestClickhouseWriterRecordBufferCap(t *testing.T) { + w := NewClickhouseWriter(ClickhouseConfig{Addr: "unused"}, slog.Default()) + for i := range maxBufferedRows + 100 { + w.Record(OffsetRow{SourceAddr: fmt.Sprintf("addr-%d", i)}) + } + + w.mu.Lock() + require.Len(t, w.buf, maxBufferedRows) + w.mu.Unlock() +} diff --git a/controlplane/telemetry/internal/geoprobe/migrations.go b/controlplane/telemetry/internal/geoprobe/migrations.go index 0b92fc7d46..19467e9aa6 100644 --- a/controlplane/telemetry/internal/geoprobe/migrations.go +++ b/controlplane/telemetry/internal/geoprobe/migrations.go @@ -13,7 +13,8 @@ import ( func RunMigrations(cfg ClickhouseConfig, log *slog.Logger) error { opts := &clickhouse.Options{ - Addr: []string{cfg.Addr}, + Protocol: clickhouse.HTTP, + Addr: []string{cfg.Addr}, Auth: clickhouse.Auth{ Database: cfg.Database, Username: cfg.Username, diff --git a/e2e/geoprobe_test.go b/e2e/geoprobe_test.go index ed505ece2a..03b8bbf0a7 100644 --- a/e2e/geoprobe_test.go +++ b/e2e/geoprobe_test.go @@ -236,7 +236,7 @@ func TestE2E_GeoprobeDiscovery(t *testing.T) { // since we need to generate the sender keypair inside it. log.Debug("==> Starting geoprobe target container") targetContainerID := startGeoprobeTarget(t, log, dn, targetIPStr, &geoprobeTargetOpts{ - clickhouseAddr: "clickhouse:9000", + clickhouseAddr: "clickhouse:8123", clickhousePass: "test", })