Give chwriter.Registry periodic refresh, matching Tantivy's tracker

Closing search's active-tenant gap last commit surfaced a real asymmetry by comparison: chwriter.Registry's per-tenant writer map was still a snapshot built once at enterprise-ingest startup with no refresh at all, while search's new ActiveTenantTracker refreshes every minute. A tenant deprovisioned after enterprise-ingest started would keep writing successfully to ClickHouse until the next restart -- a real, disclosed staleness gap, not matched by anything on the Tantivy side anymore. Registry.StartRefreshing spawns a goroutine that re-lists active tenants every minute (dataSourceRefreshInterval, same interval as search's tracker) via a new SourceLister callback and reconciles the writer map: opens a connection for a newly-active tenant, closes and removes one no longer active. New connections are dialed before taking the write lock, so a slow/unreachable ClickHouse for one newly-active tenant never blocks WriteBatch's read lock. A refresh failure (lister error, or one tenant's connection failing to open) logs and leaves the existing map untouched for that tick -- the same last-known-good posture ActiveTenantTracker already uses, so a transient rbacstore/Postgres blip doesn't evict every other tenant's already-working writer. WriteBatch now takes a read lock and Close takes a write lock -- the writer map was safe unsynchronized before only because it was immutable after New() returned; StartRefreshing makes it mutable at runtime. enterprise-ingest/main.go extracts the existing rbacstore-row-to- DataSource adaptation into tenantDataSourceLister, reused for both the initial synchronous load and StartRefreshing's periodic calls, so the two can't drift into checking different things. Verified: the lister-error-keeps-last-known-good path is Docker-free (same "construct a Registry directly, bypass New" trick the existing fail-closed tests use). The actual add/remove reconciliation against real ClickHouse connections (TestRefreshAddsNewlyActiveTenant, TestRefreshRemovesNoLongerActiveTenant) are skip-gated live-ClickHouse tests, same CHWRITER_TEST_CLICKHOUSE_ADDR convention as this package's existing integration tests -- not run against a live database in this environment. This closes the last disclosed gap from Phase 4's write-routing work: both storage engines now share the same one-minute active-tenant staleness bound instead of one being materially staler than the other.
2026-08-14 23:55:40 -07:00
parent 088677643f
commit 2e8ab1ed6a
7 changed files with 379 additions and 90 deletions
@@ -24,6 +24,9 @@ package chwriter
 import (
 	"context"
 	"fmt"
+	"log/slog"
+	"sync"
+	"time"

 	"github.com/sentry/sentry/ingest/clickhousewriter"
 	"github.com/sentry/sentry/ingest/consumer"
@@ -44,11 +47,17 @@ type DataSource struct {

 // Registry implements ingest/consumer's chWriter interface
 // (WriteBatch(ctx, []consumer.Record) error) by routing each record to
-// its tenant's dedicated connection. Immutable after New returns -- see
-// this file's doc comment.
+// its tenant's dedicated connection. The writer map used to be
+// immutable after New returned; StartRefreshing (below) makes it
+// mutable at runtime, guarded by mu -- WriteBatch takes a read lock (the
+// common case, and concurrent reads don't block each other), a refresh
+// takes a write lock only for the brief final swap, never while
+// actually dialing ClickHouse (see refresh's comment).
 type Registry struct {
+	addr string
+
+	mu      sync.RWMutex
 	writers map[string]*clickhousewriter.Writer
-	closers []func()
 }

 // New opens one real ClickHouse connection per DataSource (same native
@@ -58,7 +67,7 @@ type Registry struct {
 // Registry fails to construct rather than silently running with a
 // partial tenant set.
 func New(ctx context.Context, addr string, sources []DataSource) (*Registry, error) {
-	reg := &Registry{writers: make(map[string]*clickhousewriter.Writer, len(sources))}
+	reg := &Registry{addr: addr, writers: make(map[string]*clickhousewriter.Writer, len(sources))}
 	for _, src := range sources {
 		w, err := clickhousewriter.New(ctx, clickhousewriter.Config{
 			Addr: addr, Database: src.Database, Username: src.Username, Password: src.Password,
@@ -68,16 +77,21 @@ func New(ctx context.Context, addr string, sources []DataSource) (*Registry, err
 			return nil, fmt.Errorf("chwriter: opening connection for tenant %q: %w", src.TenantID, err)
 		}
 		reg.writers[src.TenantID] = w
-		reg.closers = append(reg.closers, func() { _ = w.Close() })
 	}
 	return reg, nil
 }

 // Close releases every underlying connection -- call once at process
-// shutdown, same lifecycle as chrunner.Registry.Close.
+// shutdown, same lifecycle as chrunner.Registry.Close. Safe to call
+// even with StartRefreshing's goroutine still running (it only ever
+// adds/removes individual writers under mu, never assumes the whole map
+// survives), though callers should still cancel that goroutine's
+// context first to stop it from reopening what Close just shut down.
 func (r *Registry) Close() {
-	for _, c := range r.closers {
-		c()
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	for _, w := range r.writers {
+		_ = w.Close()
 	}
 }

@@ -105,6 +119,8 @@ func (r *Registry) WriteBatch(ctx context.Context, records []consumer.Record) er
 		byTenant[rec.TenantID] = append(byTenant[rec.TenantID], rec)
 	}

+	r.mu.RLock()
+	defer r.mu.RUnlock()
 	for tenantID, group := range byTenant {
 		if tenantID == "" {
 			return fmt.Errorf("chwriter: %d record(s) in this batch have no tenant_id, refusing to write any of it", len(group))
@@ -123,3 +139,91 @@ func (r *Registry) WriteBatch(ctx context.Context, records []consumer.Record) er
 	}
 	return nil
 }
+
+// SourceLister re-lists the data sources a Registry should have a
+// writer for -- a narrow function type, not an rbacstore dependency,
+// same reasoning DataSource's doc comment gives for not importing
+// rbacstore directly here. enterprise-ingest's main.go supplies one
+// backed by rbacstore.ListProvisionedDataSources (the same query New's
+// caller already runs once at startup).
+type SourceLister func(ctx context.Context) ([]DataSource, error)
+
+// StartRefreshing closes the staleness gap disclosed in
+// /docs/security/threat-model.md as an asymmetry with search's
+// tenants.ActiveTenantTracker (Tantivy's write-side active-tenant gate,
+// which already refreshes every 60s): spawns a goroutine that
+// periodically re-lists data sources via lister and reconciles the
+// writer map -- opens a connection for any newly-active tenant, closes
+// and removes any tenant no longer present (deprovisioned or suspended
+// since the last refresh). Stops when ctx is cancelled; call at most
+// once per Registry. A refresh failure (lister error, or one tenant's
+// new connection failing to open) logs via logger and leaves the
+// existing map alone for that tick -- a transient rbacstore/Postgres
+// blip, or one bad tenant's connection, must not evict every other
+// tenant's already-working writer, the same "last-known-good" posture
+// ActiveTenantTracker's periodic refresh uses.
+func (r *Registry) StartRefreshing(ctx context.Context, lister SourceLister, interval time.Duration, logger *slog.Logger) {
+	go func() {
+		ticker := time.NewTicker(interval)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				r.refresh(ctx, lister, logger)
+			}
+		}
+	}()
+}
+
+// refresh dials any newly-needed connections *before* taking the write
+// lock, so a slow/unreachable ClickHouse for one newly-active tenant
+// never blocks WriteBatch's read lock for longer than the map swap
+// itself takes.
+func (r *Registry) refresh(ctx context.Context, lister SourceLister, logger *slog.Logger) {
+	sources, err := lister(ctx)
+	if err != nil {
+		logger.Error("chwriter: refreshing data sources failed, keeping last-known-good writer set", "error", err)
+		return
+	}
+	fresh := make(map[string]DataSource, len(sources))
+	for _, src := range sources {
+		fresh[src.TenantID] = src
+	}
+
+	r.mu.RLock()
+	var toOpen []DataSource
+	for tenantID, src := range fresh {
+		if _, ok := r.writers[tenantID]; !ok {
+			toOpen = append(toOpen, src)
+		}
+	}
+	r.mu.RUnlock()
+
+	newWriters := make(map[string]*clickhousewriter.Writer, len(toOpen))
+	for _, src := range toOpen {
+		w, err := clickhousewriter.New(ctx, clickhousewriter.Config{
+			Addr: r.addr, Database: src.Database, Username: src.Username, Password: src.Password,
+		})
+		if err != nil {
+			logger.Error("chwriter: opening connection for newly-active tenant failed, will retry next refresh", "tenant_id", src.TenantID, "error", err)
+			continue
+		}
+		newWriters[src.TenantID] = w
+	}
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	for tenantID, w := range newWriters {
+		r.writers[tenantID] = w
+		logger.Info("chwriter: added writer for newly-active tenant", "tenant_id", tenantID)
+	}
+	for tenantID, w := range r.writers {
+		if _, ok := fresh[tenantID]; !ok {
+			_ = w.Close()
+			delete(r.writers, tenantID)
+			logger.Info("chwriter: removed writer for tenant no longer active/provisioned", "tenant_id", tenantID)
+		}
+	}
+}
@@ -14,7 +14,10 @@ package chwriter

 import (
 	"context"
+	"errors"
 	"fmt"
+	"io"
+	"log/slog"
 	"os"
 	"testing"

@@ -27,6 +30,10 @@ import (
 	logsv1 "github.com/sentry/sentry/proto/sentry/logs/v1"
 )

+func discardLogger() *slog.Logger {
+	return slog.New(slog.NewTextHandler(io.Discard, nil))
+}
+
 // TestWriteBatchRefusesEmptyTenantID and
 // TestWriteBatchRefusesUnknownTenantWithEmptyRegistry construct a
 // Registry directly (bypassing New, which would dial ClickHouse) so
@@ -52,6 +59,31 @@ func TestWriteBatchRefusesUnknownTenantWithEmptyRegistry(t *testing.T) {
 	}
 }

+// TestRefreshListerErrorLeavesRegistryUnchanged is Docker-free the same
+// way the two tests above are: refresh's early-return on a lister error
+// happens before anything touches ClickHouse, so this genuinely
+// exercises the "keep last-known-good" path -- see refresh's doc
+// comment on StartRefreshing.
+func TestRefreshListerErrorLeavesRegistryUnchanged(t *testing.T) {
+	reg := &Registry{writers: map[string]*clickhousewriter.Writer{}}
+	lister := func(context.Context) ([]DataSource, error) {
+		return nil, errors.New("rbacstore unreachable")
+	}
+
+	reg.refresh(context.Background(), lister, discardLogger())
+
+	// Still refuses -- refresh must not have added a writer for "acme"
+	// (there's nothing a failed lister call could have legitimately
+	// learned), and must not have panicked reaching into a nil/partial
+	// state either.
+	err := reg.WriteBatch(context.Background(), []consumer.Record{
+		{TenantID: "acme", Record: &logsv1.LogRecord{Message: "m"}},
+	})
+	if err == nil {
+		t.Fatal("expected WriteBatch to still refuse tenant acme after a failed refresh")
+	}
+}
+
 func testAddr(t *testing.T) string {
 	t.Helper()
 	addr := os.Getenv("CHWRITER_TEST_CLICKHOUSE_ADDR")
@@ -155,3 +187,74 @@ func TestRegistryRefusesUnprovisionedTenant(t *testing.T) {
 		t.Fatal("expected WriteBatch to refuse a tenant with no provisioned connection, not silently drop or misroute it")
 	}
 }
+
+// TestRefreshAddsNewlyActiveTenant is the live counterpart to
+// TestRefreshListerErrorLeavesRegistryUnchanged: proves refresh actually
+// opens a real, usable connection for a tenant that appears in a later
+// lister call but wasn't present at New() time -- the scenario
+// StartRefreshing exists to handle (a tenant provisioned after
+// enterprise-ingest already started).
+func TestRefreshAddsNewlyActiveTenant(t *testing.T) {
+	addr := testAddr(t)
+	ctx := context.Background()
+	tenantA, credsA := provisionTestTenant(t, addr)
+
+	reg, err := New(ctx, addr, nil) // starts with zero tenants, same as a cold start before any tenant exists
+	if err != nil {
+		t.Fatalf("New: %v", err)
+	}
+	defer reg.Close()
+
+	if err := reg.WriteBatch(ctx, []consumer.Record{
+		{TenantID: tenantA, Record: &logsv1.LogRecord{Message: "m", RecordId: uuid.NewString()}},
+	}); err == nil {
+		t.Fatal("expected WriteBatch to refuse tenantA before the first refresh has run")
+	}
+
+	lister := func(context.Context) ([]DataSource, error) {
+		return []DataSource{{TenantID: tenantA, Database: tenantA, Username: credsA.Username, Password: credsA.Password}}, nil
+	}
+	reg.refresh(ctx, lister, discardLogger())
+
+	if err := reg.WriteBatch(ctx, []consumer.Record{
+		{TenantID: tenantA, Record: &logsv1.LogRecord{Host: "h1", Message: "after-refresh", RecordId: uuid.NewString()}},
+	}); err != nil {
+		t.Fatalf("expected WriteBatch to succeed for tenantA after refresh added it, got: %v", err)
+	}
+}
+
+// TestRefreshRemovesNoLongerActiveTenant is TestRefreshAddsNewlyActiveTenant's
+// mirror image: a tenant present at New() time that a later lister call
+// no longer returns (deprovisioned or suspended) must lose its writer,
+// not keep writing indefinitely until process restart -- the exact
+// staleness gap this whole mechanism exists to close.
+func TestRefreshRemovesNoLongerActiveTenant(t *testing.T) {
+	addr := testAddr(t)
+	ctx := context.Background()
+	tenantA, credsA := provisionTestTenant(t, addr)
+
+	reg, err := New(ctx, addr, []DataSource{
+		{TenantID: tenantA, Database: tenantA, Username: credsA.Username, Password: credsA.Password},
+	})
+	if err != nil {
+		t.Fatalf("New: %v", err)
+	}
+	defer reg.Close()
+
+	if err := reg.WriteBatch(ctx, []consumer.Record{
+		{TenantID: tenantA, Record: &logsv1.LogRecord{Host: "h1", Message: "before-removal", RecordId: uuid.NewString()}},
+	}); err != nil {
+		t.Fatalf("expected WriteBatch to succeed for tenantA before refresh removes it, got: %v", err)
+	}
+
+	lister := func(context.Context) ([]DataSource, error) {
+		return nil, nil // tenantA no longer active/provisioned as of this refresh
+	}
+	reg.refresh(ctx, lister, discardLogger())
+
+	if err := reg.WriteBatch(ctx, []consumer.Record{
+		{TenantID: tenantA, Record: &logsv1.LogRecord{Message: "after-removal", RecordId: uuid.NewString()}},
+	}); err == nil {
+		t.Fatal("expected WriteBatch to refuse tenantA after refresh removed it, not keep writing with a stale connection")
+	}
+}