telemetry: geoprobe-target clickhouse fixes

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)

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,13 +240,15 @@ 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
 		}
 	}
 
 	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()

controlplane/telemetry/internal/geoprobe/clickhouse_test.go

@@ -1,6 +1,8 @@
 package geoprobe
 
 import (
+	"fmt"
+	"log/slog"
 	"testing"
 	"time"
 
@@ -61,14 +63,57 @@ 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"})
 
 	w.mu.Lock()
 	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()
+}

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,

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",
 	})