diff --git a/cmd/health_cmd.go b/cmd/health_cmd.go index 9a694517..53f7a5a2 100644 --- a/cmd/health_cmd.go +++ b/cmd/health_cmd.go @@ -9,6 +9,7 @@ import ( "github.com/xataio/pgstream/cmd/config" "github.com/xataio/pgstream/internal/health" + "github.com/xataio/pgstream/internal/phase" pglib "github.com/xataio/pgstream/internal/postgres" loglib "github.com/xataio/pgstream/pkg/log" ) @@ -28,8 +29,9 @@ const ( // fails the command immediately rather than silently degrading. // // sourcePostgresURL, when non-empty, wires the /ready endpoint to ping the -// source database. -func startHealthServer(ctx context.Context, logger loglib.Logger, sourcePostgresURL string) (func(), error) { +// source database. phaseTracker, when non-nil, wires the /status endpoint to +// report the current pipeline phase. +func startHealthServer(ctx context.Context, logger loglib.Logger, sourcePostgresURL string, phaseTracker *phase.Tracker) (func(), error) { cfg, err := config.ParseHealthConfig() if err != nil { return nil, fmt.Errorf("parsing health config: %w", err) @@ -43,6 +45,12 @@ func startHealthServer(ctx context.Context, logger loglib.Logger, sourcePostgres health.WithVersion(version()), } + if phaseTracker != nil { + opts = append(opts, health.WithPhaseProvider(func() string { + return string(phaseTracker.Get()) + })) + } + var pool *pglib.Pool if sourcePostgresURL != "" { pool, err = pglib.NewConnPool(ctx, sourcePostgresURL, pglib.WithMaxConnections(healthReadinessPoolMaxConns)) diff --git a/cmd/run_cmd.go b/cmd/run_cmd.go index 4e36906a..aa4fcc90 100644 --- a/cmd/run_cmd.go +++ b/cmd/run_cmd.go @@ -11,6 +11,7 @@ import ( "github.com/spf13/viper" "github.com/xataio/pgstream/cmd/config" "github.com/xataio/pgstream/internal/log/zerolog" + "github.com/xataio/pgstream/internal/phase" "github.com/xataio/pgstream/pkg/stream" ) @@ -71,7 +72,9 @@ func run(ctx context.Context) error { } stdLogger := zerolog.NewStdLogger(logger) - stopHealth, err := startHealthServer(ctx, stdLogger, streamConfig.SourcePostgresURL()) + phaseTracker := phase.NewTracker() + opts = append(opts, stream.WithPhaseTracker(phaseTracker)) + stopHealth, err := startHealthServer(ctx, stdLogger, streamConfig.SourcePostgresURL(), phaseTracker) if err != nil { return err } diff --git a/cmd/snapshot_cmd.go b/cmd/snapshot_cmd.go index 306e4a29..95408eb8 100644 --- a/cmd/snapshot_cmd.go +++ b/cmd/snapshot_cmd.go @@ -10,6 +10,7 @@ import ( "github.com/spf13/viper" "github.com/xataio/pgstream/cmd/config" "github.com/xataio/pgstream/internal/log/zerolog" + "github.com/xataio/pgstream/internal/phase" "github.com/xataio/pgstream/pkg/stream" ) @@ -44,13 +45,14 @@ func snapshot(ctx context.Context) error { defer provider.Close() stdLogger := zerolog.NewStdLogger(logger) - stopHealth, err := startHealthServer(ctx, stdLogger, streamConfig.SourcePostgresURL()) + phaseTracker := phase.NewTracker() + stopHealth, err := startHealthServer(ctx, stdLogger, streamConfig.SourcePostgresURL(), phaseTracker) if err != nil { return err } defer stopHealth() - return stream.Snapshot(ctx, stdLogger, streamConfig, provider.NewInstrumentation("snapshot")) + return stream.Snapshot(ctx, stdLogger, streamConfig, provider.NewInstrumentation("snapshot"), stream.WithPhaseTracker(phaseTracker)) } func snapshotFlagBinding(cmd *cobra.Command, args []string) error { diff --git a/config_template.yaml b/config_template.yaml index 3baf922c..abfd7e59 100644 --- a/config_template.yaml +++ b/config_template.yaml @@ -12,7 +12,7 @@ instrumentation: endpoint: "0.0.0.0:4317" sample_ratio: 0.5 # ratio of traces that will be sampled. Must be between 0.0-1.0, where 0 is no traces sampled, and 1 is all traces sampled. health: - enabled: false # whether to expose the /health (liveness) and /ready (readiness) HTTP endpoints. Only honoured by the run and snapshot commands. Defaults to false. + enabled: false # whether to expose the /health (liveness), /ready (readiness), and /status (pipeline phase) HTTP endpoints. Only honoured by the run and snapshot commands. Defaults to false. address: "localhost:9910" # address the health server binds to. Defaults to localhost:9910. source: diff --git a/docs/Observability.md b/docs/Observability.md index 0ceb370e..f845145f 100644 --- a/docs/Observability.md +++ b/docs/Observability.md @@ -64,6 +64,18 @@ All metrics follow the `pgstream.*` naming convention and include relevant attri **⚠️ Important:** This metric only tracks pgstream's consumer lag. It's strongly recommended to also monitor your source PostgreSQL metrics, particularly the built-in replication lag metrics (`pg_stat_replication.flush_lag`, `pg_stat_replication.replay_lag`) to get a complete picture of replication health. +### Pipeline Phase + +| Metric | Type | Unit | Description | +| -------------------------- | --------------- | ---- | --------------------------------------------------------------------------- | +| `pgstream.pipeline.phase` | ObservableGauge | 1 | Reports `1` for the active pipeline phase and `0` for the others (`phase=snapshot\|replication`) | + +**Attributes:** + +- `phase`: Phase this data point refers to (`snapshot` or `replication`) + +**Usage:** Distinguish whether a running process is performing the initial snapshot or streaming logical replication, without scraping logs. Both `phase="snapshot"` and `phase="replication"` series are reported on every collection, with `1` marking the active phase and `0` the inactive one, so queries like `sum by (phase)` stay unambiguous across transitions. Nothing is reported until the first phase transition. The gauge flips to `replication` when logical replication starts. The same value is also exposed via the health server `GET /status` endpoint when health checks are enabled. + ### WAL Event Processing | Metric | Type | Unit | Description | diff --git a/docs/configuration.md b/docs/configuration.md index a7c91ee2..0320c200 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -440,7 +440,7 @@ One of exponential/constant/disable retries retry policies can be provided for t
Health endpoint -Exposes `/health` (liveness, always 200) and `/ready` (readiness, pings the source postgres database when configured). Only the `run` and `snapshot` commands start the server. Responses are JSON. +Exposes `/health` (liveness, always 200), `/ready` (readiness, pings the source postgres database when configured), and `/status` (current pipeline phase: `snapshot` or `replication`). Only the `run` and `snapshot` commands start the server. Responses are JSON. | Environment Variable | Default | Required | Description | | ------------------------------ | ---------------- | -------- | ------------------------------------------------------------------------------------------------------------ | diff --git a/internal/health/health.go b/internal/health/health.go index c8927d1d..cebd8ec5 100644 --- a/internal/health/health.go +++ b/internal/health/health.go @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 -// Package health exposes a small HTTP server with /health (liveness) and -// /ready (readiness) endpoints so that operators and orchestrators can -// monitor a running pgstream process. +// Package health exposes a small HTTP server with /health (liveness), +// /ready (readiness), and /status (pipeline phase) endpoints so that +// operators and orchestrators can monitor a running pgstream process. package health import ( @@ -29,12 +29,13 @@ type Config struct { Address string } -// Server serves /health and /ready over HTTP. +// Server serves /health, /ready, and /status over HTTP. type Server struct { logger loglib.Logger address string version string readinessCheck func(ctx context.Context) error + phaseProvider func() string listener net.Listener httpServer *http.Server } @@ -84,6 +85,16 @@ func WithReadinessCheck(fn func(ctx context.Context) error) Option { } } +// WithPhaseProvider registers a function the /status handler will invoke to +// report the current pipeline phase (e.g. "snapshot" or "replication"). +// If no provider is registered, /status still returns status/version with an +// empty phase. +func WithPhaseProvider(fn func() string) Option { + return func(s *Server) { + s.phaseProvider = fn + } +} + // Listen binds the configured address synchronously and prepares the HTTP // server. After it returns successfully, Serve must be called (typically // from a goroutine) to start accepting requests. Splitting bind from serve @@ -98,6 +109,7 @@ func (s *Server) Listen() error { mux := http.NewServeMux() mux.HandleFunc("/health", s.handleHealth) mux.HandleFunc("/ready", s.handleReady) + mux.HandleFunc("/status", s.handleStatus) s.listener = ln s.httpServer = &http.Server{ @@ -133,6 +145,10 @@ func (s *Server) handleHealth(w http.ResponseWriter, _ *http.Request) { s.writeJSON(w, http.StatusOK, s.okPayload()) } +func (s *Server) handleStatus(w http.ResponseWriter, _ *http.Request) { + s.writeJSON(w, http.StatusOK, s.statusPayload()) +} + func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) { if s.readinessCheck != nil { ctx, cancel := context.WithTimeout(r.Context(), readinessCheckTimeout) @@ -159,6 +175,16 @@ func (s *Server) okPayload() map[string]string { return payload } +func (s *Server) statusPayload() map[string]string { + payload := s.okPayload() + phase := "" + if s.phaseProvider != nil { + phase = s.phaseProvider() + } + payload["phase"] = phase + return payload +} + func (s *Server) writeJSON(w http.ResponseWriter, code int, body any) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(code) diff --git a/internal/health/health_test.go b/internal/health/health_test.go index c07e111c..51fbdcbb 100644 --- a/internal/health/health_test.go +++ b/internal/health/health_test.go @@ -101,6 +101,49 @@ func TestHandleReady_CheckFails(t *testing.T) { require.Equal(t, "source unreachable", failure["error"]) } +func TestHandleStatus_NoProvider(t *testing.T) { + t.Parallel() + s := NewServer(Config{}, WithVersion("v1.2.3")) + + req := httptest.NewRequest(http.MethodGet, "/status", nil) + rec := httptest.NewRecorder() + s.handleStatus(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "application/json", rec.Header().Get("Content-Type")) + body, _ := io.ReadAll(rec.Body) + var got map[string]string + require.NoError(t, json.Unmarshal(body, &got)) + require.Equal(t, "ok", got["status"]) + require.Equal(t, "v1.2.3", got["version"]) + require.Equal(t, "", got["phase"]) +} + +func TestHandleStatus_WithProvider(t *testing.T) { + t.Parallel() + phase := "snapshot" + s := NewServer(Config{}, WithVersion("v1.2.3"), WithPhaseProvider(func() string { + return phase + })) + + req := httptest.NewRequest(http.MethodGet, "/status", nil) + rec := httptest.NewRecorder() + s.handleStatus(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + body, _ := io.ReadAll(rec.Body) + var got map[string]string + require.NoError(t, json.Unmarshal(body, &got)) + require.Equal(t, "snapshot", got["phase"]) + + phase = "replication" + rec = httptest.NewRecorder() + s.handleStatus(rec, req) + body, _ = io.ReadAll(rec.Body) + require.NoError(t, json.Unmarshal(body, &got)) + require.Equal(t, "replication", got["phase"]) +} + func TestListenServeShutdown(t *testing.T) { t.Parallel() s := NewServer(Config{Address: "127.0.0.1:0"}) diff --git a/internal/phase/phase.go b/internal/phase/phase.go new file mode 100644 index 00000000..df4fa982 --- /dev/null +++ b/internal/phase/phase.go @@ -0,0 +1,50 @@ +// SPDX-License-Identifier: Apache-2.0 + +// Package phase tracks whether a running pgstream process is in the initial +// snapshot or logical replication (streaming) phase. +package phase + +import "sync/atomic" + +// Phase identifies the current pipeline stage. +type Phase string + +const ( + // Snapshot is the initial data/schema snapshot stage. + Snapshot Phase = "snapshot" + // Replication is the logical replication / streaming stage. + Replication Phase = "replication" +) + +// Tracker holds the current pipeline phase. It is safe for concurrent use. +type Tracker struct { + current atomic.Value // stores Phase +} + +// NewTracker returns a tracker with an empty phase until the first Set. +func NewTracker() *Tracker { + t := &Tracker{} + t.current.Store(Phase("")) + return t +} + +// Set updates the current phase. +func (t *Tracker) Set(p Phase) { + if t == nil { + return + } + t.current.Store(p) +} + +// Get returns the current phase, or empty if unset / tracker is nil. +func (t *Tracker) Get() Phase { + if t == nil { + return "" + } + v := t.current.Load() + p, ok := v.(Phase) + if !ok { + return "" + } + return p +} diff --git a/internal/phase/phase_test.go b/internal/phase/phase_test.go new file mode 100644 index 00000000..9a6932fd --- /dev/null +++ b/internal/phase/phase_test.go @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 + +package phase + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestTracker_GetSet(t *testing.T) { + t.Parallel() + + tr := NewTracker() + require.Equal(t, Phase(""), tr.Get()) + + tr.Set(Snapshot) + require.Equal(t, Snapshot, tr.Get()) + + tr.Set(Replication) + require.Equal(t, Replication, tr.Get()) +} + +func TestTracker_NilSafe(t *testing.T) { + t.Parallel() + + var tr *Tracker + require.Equal(t, Phase(""), tr.Get()) + tr.Set(Snapshot) // must not panic + require.Equal(t, Phase(""), tr.Get()) +} diff --git a/pkg/stream/integration/helper_test.go b/pkg/stream/integration/helper_test.go index 15877eb9..1761209a 100644 --- a/pkg/stream/integration/helper_test.go +++ b/pkg/stream/integration/helper_test.go @@ -8,6 +8,7 @@ import ( "encoding/json" "errors" "fmt" + "net" "net/http" "net/http/httptest" "sync/atomic" @@ -15,7 +16,9 @@ import ( "time" "github.com/stretchr/testify/require" + "github.com/xataio/pgstream/internal/health" "github.com/xataio/pgstream/internal/log/zerolog" + "github.com/xataio/pgstream/internal/phase" pglib "github.com/xataio/pgstream/internal/postgres" "github.com/xataio/pgstream/pkg/backoff" kafkalib "github.com/xataio/pgstream/pkg/kafka" @@ -132,6 +135,60 @@ func runSnapshot(t *testing.T, ctx context.Context, cfg *stream.Config) { }) } +func startPhaseHealthServer(t *testing.T, tracker *phase.Tracker) (statusURL string, stop func()) { + t.Helper() + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := ln.Addr().String() + require.NoError(t, ln.Close()) + + srv := health.NewServer(health.Config{Address: addr}, + health.WithVersion("test"), + health.WithPhaseProvider(func() string { + return string(tracker.Get()) + }), + ) + require.NoError(t, srv.Listen()) + + go func() { + if err := srv.Serve(); err != nil { + t.Logf("health server stopped: %v", err) + } + }() + + return "http://" + addr + "/status", func() { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + require.NoError(t, srv.Shutdown(shutdownCtx)) + } +} + +func fetchStatusPhase(t *testing.T, statusURL string) string { + t.Helper() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, statusURL, nil) + require.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return "" + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return "" + } + + var body map[string]string + if err := json.NewDecoder(resp.Body).Decode(&body); err != nil { + return "" + } + return body["phase"] +} + func initStream(t *testing.T, ctx context.Context, url string) { err := stream.Init(ctx, &stream.InitConfig{ PostgresURL: url, diff --git a/pkg/stream/integration/phase_status_integration_test.go b/pkg/stream/integration/phase_status_integration_test.go new file mode 100644 index 00000000..01863550 --- /dev/null +++ b/pkg/stream/integration/phase_status_integration_test.go @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 + +package integration + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "github.com/stretchr/testify/require" + "github.com/xataio/pgstream/internal/phase" + pglib "github.com/xataio/pgstream/internal/postgres" + "github.com/xataio/pgstream/internal/testcontainers" + "github.com/xataio/pgstream/pkg/stream" +) + +func Test_PhaseStatus_SnapshotAndReplication(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + var snapshotPGURL string + pgcleanup, err := testcontainers.SetupPostgresContainer(context.Background(), &snapshotPGURL, testcontainers.Postgres14, "config/postgresql.conf") + require.NoError(t, err) + defer pgcleanup() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + testTable := "phase_status_integration_test" + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf("create table %s(id serial primary key, name text)", testTable)) + // bulk rows lengthen the snapshot window so /status can observe phase=snapshot + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + "insert into %s(name) select 'row_' || g from generate_series(1, 5000) g", testTable)) + + cfg := &stream.Config{ + Listener: testPostgresListenerCfgWithSnapshot(snapshotPGURL, targetPGURL, []string{testTable}), + Processor: testPostgresProcessorCfg(), + } + initStream(t, ctx, snapshotPGURL) + + tracker := phase.NewTracker() + statusURL, stopHealth := startPhaseHealthServer(t, tracker) + defer stopHealth() + + runCtx, runCancel := context.WithCancel(ctx) + defer runCancel() + + done := make(chan error, 1) + go func() { + done <- stream.Run(runCtx, testLogger(), cfg, false, nil, stream.WithPhaseTracker(tracker)) + }() + + var sawSnapshot, sawReplication bool + require.Eventually(t, func() bool { + switch fetchStatusPhase(t, statusURL) { + case string(phase.Snapshot): + sawSnapshot = true + case string(phase.Replication): + sawReplication = true + } + return sawReplication + }, 90*time.Second, 25*time.Millisecond, "timed out waiting for /status to report replication phase") + + require.True(t, sawSnapshot, "expected /status to report snapshot phase during initial snapshot") + + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf("insert into %s(name) values('live')", testTable)) + + targetConn, err := pglib.NewConn(ctx, targetPGURL) + require.NoError(t, err) + + require.Eventually(t, func() bool { + rows, err := targetConn.Query(ctx, fmt.Sprintf("select name from %s where name = 'live'", testTable)) + if err != nil { + return false + } + defer rows.Close() + return rows.Next() + }, 20*time.Second, time.Second, "replication did not deliver post-snapshot insert to target") + + runCancel() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(15 * time.Second): + t.Fatal("timeout waiting for stream to stop") + } +} diff --git a/pkg/stream/phase_metrics.go b/pkg/stream/phase_metrics.go new file mode 100644 index 00000000..086e211e --- /dev/null +++ b/pkg/stream/phase_metrics.go @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 + +package stream + +import ( + "context" + "fmt" + + "github.com/xataio/pgstream/internal/phase" + "github.com/xataio/pgstream/pkg/otel" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +// registerPhaseMetric registers an observable gauge that reports the current +// pipeline phase. No-op when instrumentation or the tracker is nil. +func registerPhaseMetric(instrumentation *otel.Instrumentation, tracker *phase.Tracker) error { + if instrumentation == nil || !instrumentation.IsEnabled() || tracker == nil { + return nil + } + if instrumentation.Meter == nil { + return nil + } + + gauge, err := instrumentation.Meter.Int64ObservableGauge("pgstream.pipeline.phase", + metric.WithUnit("1"), + metric.WithDescription("Reports 1 for the active pipeline phase and 0 for the others (phase=snapshot|replication)")) + if err != nil { + return fmt.Errorf("creating pipeline phase gauge: %w", err) + } + + _, err = instrumentation.Meter.RegisterCallback(func(_ context.Context, o metric.Observer) error { + current := tracker.Get() + if current == "" { + return nil // no phase entered yet — emit nothing + } + for _, p := range []phase.Phase{phase.Snapshot, phase.Replication} { + var v int64 + if p == current { + v = 1 + } + o.ObserveInt64(gauge, v, metric.WithAttributes(attribute.String("phase", string(p)))) + } + return nil + }, gauge) + if err != nil { + return fmt.Errorf("registering pipeline phase callback: %w", err) + } + return nil +} diff --git a/pkg/stream/phase_metrics_test.go b/pkg/stream/phase_metrics_test.go new file mode 100644 index 00000000..9abba99a --- /dev/null +++ b/pkg/stream/phase_metrics_test.go @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 + +package stream + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "github.com/xataio/pgstream/internal/phase" + "github.com/xataio/pgstream/pkg/otel" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" +) + +func TestRegisterPhaseMetric_NoopCases(t *testing.T) { + t.Parallel() + + tracker := phase.NewTracker() + tracker.Set(phase.Snapshot) + + require.NoError(t, registerPhaseMetric(nil, tracker)) + require.NoError(t, registerPhaseMetric(&otel.Instrumentation{}, tracker)) + require.NoError(t, registerPhaseMetric(&otel.Instrumentation{Meter: nil}, tracker)) + _, instrumentation := testInstrumentation(t) + require.NoError(t, registerPhaseMetric(instrumentation, nil)) +} + +func TestRegisterPhaseMetric_ObservesCurrentPhase(t *testing.T) { + t.Parallel() + + reader, instrumentation := testInstrumentation(t) + tracker := phase.NewTracker() + + require.NoError(t, registerPhaseMetric(instrumentation, tracker)) + + tracker.Set(phase.Snapshot) + points := collectPipelinePhaseMetric(t, reader) + require.Len(t, points, 2) + require.EqualValues(t, 1, valueForPhase(t, points, "snapshot")) + require.EqualValues(t, 0, valueForPhase(t, points, "replication")) + + tracker.Set(phase.Replication) + points = collectPipelinePhaseMetric(t, reader) + require.Len(t, points, 2) + require.EqualValues(t, 0, valueForPhase(t, points, "snapshot")) + require.EqualValues(t, 1, valueForPhase(t, points, "replication")) +} + +func valueForPhase(t *testing.T, points []metricdata.DataPoint[int64], phase string) int64 { + t.Helper() + + for _, p := range points { + if p.Attributes == attribute.NewSet(attribute.String("phase", phase)) { + return p.Value + } + } + t.Fatalf("no data point found for phase %q", phase) + return 0 +} + +func TestRegisterPhaseMetric_EmptyPhaseNotObserved(t *testing.T) { + t.Parallel() + + reader, instrumentation := testInstrumentation(t) + tracker := phase.NewTracker() + + require.NoError(t, registerPhaseMetric(instrumentation, tracker)) + + points := collectPipelinePhaseMetric(t, reader) + require.Empty(t, points) +} + +func testInstrumentation(t *testing.T) (*metric.ManualReader, *otel.Instrumentation) { + t.Helper() + + reader := metric.NewManualReader() + provider := metric.NewMeterProvider(metric.WithReader(reader)) + return reader, &otel.Instrumentation{ + Meter: provider.Meter("test"), + } +} + +func collectPipelinePhaseMetric(t *testing.T, reader *metric.ManualReader) []metricdata.DataPoint[int64] { + t.Helper() + + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + if m.Name != "pgstream.pipeline.phase" { + continue + } + gauge, ok := m.Data.(metricdata.Gauge[int64]) + require.True(t, ok) + return gauge.DataPoints + } + } + return nil +} diff --git a/pkg/stream/stream_init.go b/pkg/stream/stream_init.go index 0e7daf44..091f2464 100644 --- a/pkg/stream/stream_init.go +++ b/pkg/stream/stream_init.go @@ -8,6 +8,7 @@ import ( "fmt" migratorlib "github.com/xataio/pgstream/internal/migrator" + "github.com/xataio/pgstream/internal/phase" pglib "github.com/xataio/pgstream/internal/postgres" _ "github.com/golang-migrate/migrate/v4/database/postgres" @@ -22,6 +23,8 @@ type InitConfig struct { InjectorMigrationsEnabled bool MigrationsOnly bool Upgrade bool + // PhaseTracker is optional runtime state for stream.Run; Init/Destroy ignore it. + PhaseTracker *phase.Tracker } type InitOption func(*InitConfig) @@ -38,6 +41,14 @@ func WithUpgrade() InitOption { } } +// WithPhaseTracker registers a tracker updated as the pipeline moves between +// snapshot and replication phases. Used by stream.Run; ignored by Init/Destroy. +func WithPhaseTracker(t *phase.Tracker) InitOption { + return func(cfg *InitConfig) { + cfg.PhaseTracker = t + } +} + const ( pgstreamSchema = "pgstream" ) diff --git a/pkg/stream/stream_run.go b/pkg/stream/stream_run.go index 32e25cd8..e1185995 100644 --- a/pkg/stream/stream_run.go +++ b/pkg/stream/stream_run.go @@ -28,13 +28,21 @@ import ( ) // Run will run the configured pgstream processes. This call is blocking. +// Pass WithPhaseTracker to expose snapshot/replication phase via /status and metrics. func Run(ctx context.Context, logger loglib.Logger, config *Config, init bool, instrumentation *otel.Instrumentation, opts ...InitOption) error { if err := config.IsValid(); err != nil { return fmt.Errorf("incompatible configuration: %w", err) } + initConfig := config.GetInitConfig(opts...) + phaseTracker := initConfig.PhaseTracker + + if err := registerPhaseMetric(instrumentation, phaseTracker); err != nil { + return fmt.Errorf("registering pipeline phase metric: %w", err) + } + if init { - if err := Init(ctx, config.GetInitConfig(opts...)); err != nil { + if err := Init(ctx, initConfig); err != nil { return err } } @@ -125,6 +133,7 @@ func Run(ctx context.Context, logger loglib.Logger, config *Config, init bool, i logger.Info("postgres listener configured") opts := []pglistener.Option{ pglistener.WithLogger(logger), + pglistener.WithPhaseTracker(phaseTracker), } if config.Listener.Postgres.Snapshot != nil { logger.Info("initial snapshot enabled") @@ -161,7 +170,8 @@ func Run(ctx context.Context, logger loglib.Logger, config *Config, init bool, i listener, err = kafkalistener.NewWALReader( kafkaReader, processor.ProcessWALEvent, - kafkalistener.WithLogger(logger)) + kafkalistener.WithLogger(logger), + kafkalistener.WithPhaseTracker(phaseTracker)) if err != nil { return err } diff --git a/pkg/stream/stream_snapshot.go b/pkg/stream/stream_snapshot.go index ee2e9025..33cf6b89 100644 --- a/pkg/stream/stream_snapshot.go +++ b/pkg/stream/stream_snapshot.go @@ -14,7 +14,9 @@ import ( "golang.org/x/sync/errgroup" ) -func Snapshot(ctx context.Context, logger loglib.Logger, config *Config, instrumentation *otel.Instrumentation) error { +// Snapshot performs a one-time data snapshot. This call is blocking. +// Pass WithPhaseTracker to expose the snapshot phase via /status and metrics. +func Snapshot(ctx context.Context, logger loglib.Logger, config *Config, instrumentation *otel.Instrumentation, opts ...InitOption) error { if config.Listener.Postgres == nil { return errors.New("source postgres snapshot not configured: ensure source.postgres is set") } @@ -23,6 +25,12 @@ func Snapshot(ctx context.Context, logger loglib.Logger, config *Config, instrum return fmt.Errorf("incompatible configuration: %w", err) } + tracker := config.GetInitConfig(opts...).PhaseTracker + + if err := registerPhaseMetric(instrumentation, tracker); err != nil { + return fmt.Errorf("registering pipeline phase metric: %w", err) + } + eg, ctx := errgroup.WithContext(ctx) // Processor @@ -53,7 +61,7 @@ func Snapshot(ctx context.Context, logger loglib.Logger, config *Config, instrum if err != nil { return err } - listener := snapshotlistener.New(snapshotGenerator) + listener := snapshotlistener.New(snapshotGenerator, snapshotlistener.WithPhaseTracker(tracker)) defer listener.Close() eg.Go(func() error { diff --git a/pkg/wal/listener/kafka/wal_kafka_reader.go b/pkg/wal/listener/kafka/wal_kafka_reader.go index 72f02220..2854b2d8 100644 --- a/pkg/wal/listener/kafka/wal_kafka_reader.go +++ b/pkg/wal/listener/kafka/wal_kafka_reader.go @@ -8,6 +8,7 @@ import ( "fmt" "github.com/xataio/pgstream/internal/json" + "github.com/xataio/pgstream/internal/phase" "github.com/xataio/pgstream/pkg/kafka" loglib "github.com/xataio/pgstream/pkg/log" "github.com/xataio/pgstream/pkg/wal" @@ -19,6 +20,7 @@ type Reader struct { unmarshaler func([]byte, any) error logger loglib.Logger offsetParser kafka.OffsetParser + phaseTracker *phase.Tracker // processRecord is called for a new record. processRecord payloadProcessor @@ -58,7 +60,15 @@ func WithLogger(logger loglib.Logger) Option { } } +// WithPhaseTracker registers a tracker set to replication when Listen starts. +func WithPhaseTracker(t *phase.Tracker) Option { + return func(r *Reader) { + r.phaseTracker = t + } +} + func (r *Reader) Listen(ctx context.Context) error { + r.phaseTracker.Set(phase.Replication) for { select { case <-ctx.Done(): diff --git a/pkg/wal/listener/postgres/wal_pg_listener.go b/pkg/wal/listener/postgres/wal_pg_listener.go index 19a088e4..950a6c03 100644 --- a/pkg/wal/listener/postgres/wal_pg_listener.go +++ b/pkg/wal/listener/postgres/wal_pg_listener.go @@ -8,6 +8,7 @@ import ( "fmt" "github.com/xataio/pgstream/internal/json" + "github.com/xataio/pgstream/internal/phase" loglib "github.com/xataio/pgstream/pkg/log" "github.com/xataio/pgstream/pkg/wal" "github.com/xataio/pgstream/pkg/wal/replication" @@ -20,6 +21,7 @@ type Listener struct { logger loglib.Logger lsnParser replication.LSNParser snapshotGenerator snapshotGenerator + phaseTracker *phase.Tracker // Function called for processing WAL events. processEvent listenerProcessWalEvent @@ -77,6 +79,14 @@ func WithInitialSnapshot(sg snapshotGenerator) Option { } } +// WithPhaseTracker registers a tracker updated as the listener moves between +// snapshot and replication phases. +func WithPhaseTracker(t *phase.Tracker) Option { + return func(l *Listener) { + l.phaseTracker = t + } +} + // Listen starts the subscription process to listen for updates from PG. func (l *Listener) Listen(ctx context.Context) error { if l.snapshotGenerator != nil { @@ -84,11 +94,15 @@ func (l *Listener) Listen(ctx context.Context) error { l.logger.Error(err, "pg snapshot and listen") return err } + // snapshotAndListen already entered the listen loop; only reached on + // clean return (context cancel) or if snapshot path returned early. + return nil } if err := l.replicationHandler.StartReplication(ctx); err != nil { return fmt.Errorf("start replication: %w", err) } + l.phaseTracker.Set(phase.Replication) return l.listen(ctx) } @@ -99,6 +113,8 @@ func (l *Listener) Close() error { } func (l *Listener) snapshotAndListen(ctx context.Context) error { + l.phaseTracker.Set(phase.Snapshot) + lsn, err := l.replicationHandler.GetCurrentLSN(ctx) if err != nil { return err @@ -111,6 +127,7 @@ func (l *Listener) snapshotAndListen(ctx context.Context) error { if err := l.replicationHandler.StartReplicationFromLSN(ctx, lsn); err != nil { return fmt.Errorf("start replication from LSN %s: %w", l.lsnParser.ToString(lsn), err) } + l.phaseTracker.Set(phase.Replication) return l.listen(ctx) } diff --git a/pkg/wal/listener/postgres/wal_pg_listener_test.go b/pkg/wal/listener/postgres/wal_pg_listener_test.go index b328da20..f7e26bfd 100644 --- a/pkg/wal/listener/postgres/wal_pg_listener_test.go +++ b/pkg/wal/listener/postgres/wal_pg_listener_test.go @@ -11,6 +11,7 @@ import ( "time" "github.com/stretchr/testify/require" + "github.com/xataio/pgstream/internal/phase" "github.com/xataio/pgstream/pkg/log" "github.com/xataio/pgstream/pkg/wal" "github.com/xataio/pgstream/pkg/wal/replication" @@ -472,3 +473,98 @@ func TestListener_Listen(t *testing.T) { }) } } + +func TestListener_Listen_PhaseTransitions(t *testing.T) { + t.Parallel() + + t.Run("replication only sets replication after start", func(t *testing.T) { + t.Parallel() + + tracker := phase.NewTracker() + doneChan := make(chan struct{}, 1) + + h := newMockReplicationHandler() + h.StartReplicationFn = func(ctx context.Context) error { + require.Equal(t, phase.Phase(""), tracker.Get()) + return nil + } + h.ReceiveMessageFn = func(ctx context.Context, i uint64) (*replication.Message, error) { + require.Equal(t, phase.Replication, tracker.Get()) + doneChan <- struct{}{} + <-ctx.Done() + return nil, ctx.Err() + } + + l := New(h, func(context.Context, *wal.Event) error { return nil }, + WithLogger(log.NewNoopLogger()), + WithPhaseTracker(tracker)) + l.lsnParser = newMockLSNParser() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + errCh := make(chan error, 1) + go func() { errCh <- l.Listen(ctx) }() + + select { + case <-doneChan: + cancel() + case <-ctx.Done(): + t.Fatal("timeout waiting for receive") + } + require.ErrorIs(t, <-errCh, context.Canceled) + require.Equal(t, phase.Replication, tracker.Get()) + }) + + t.Run("snapshot and replication flips at StartReplicationFromLSN", func(t *testing.T) { + t.Parallel() + + tracker := phase.NewTracker() + doneChan := make(chan struct{}, 1) + + h := newMockReplicationHandler() + h.GetCurrentLSNFn = func(context.Context) (replication.LSN, error) { + require.Equal(t, phase.Snapshot, tracker.Get()) + return testLSN, nil + } + h.StartReplicationFromLSNFn = func(ctx context.Context, lsn replication.LSN) error { + require.Equal(t, phase.Snapshot, tracker.Get()) + require.Equal(t, testLSN, lsn) + return nil + } + h.ReceiveMessageFn = func(ctx context.Context, i uint64) (*replication.Message, error) { + require.Equal(t, phase.Replication, tracker.Get()) + doneChan <- struct{}{} + <-ctx.Done() + return nil, ctx.Err() + } + + gen := &mockGenerator{ + createSnapshotFn: func(ctx context.Context) error { + require.Equal(t, phase.Snapshot, tracker.Get()) + return nil + }, + } + + l := New(h, func(context.Context, *wal.Event) error { return nil }, + WithLogger(log.NewNoopLogger()), + WithInitialSnapshot(gen), + WithPhaseTracker(tracker)) + l.lsnParser = newMockLSNParser() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + errCh := make(chan error, 1) + go func() { errCh <- l.Listen(ctx) }() + + select { + case <-doneChan: + cancel() + case <-ctx.Done(): + t.Fatal("timeout waiting for receive") + } + require.ErrorIs(t, <-errCh, context.Canceled) + require.Equal(t, phase.Replication, tracker.Get()) + }) +} diff --git a/pkg/wal/listener/snapshot/wal_snapshot_listener.go b/pkg/wal/listener/snapshot/wal_snapshot_listener.go index e18ad55d..344ee1f7 100644 --- a/pkg/wal/listener/snapshot/wal_snapshot_listener.go +++ b/pkg/wal/listener/snapshot/wal_snapshot_listener.go @@ -4,6 +4,8 @@ package snapshot import ( "context" + + "github.com/xataio/pgstream/internal/phase" ) type Generator interface { @@ -12,17 +14,32 @@ type Generator interface { } type Listener struct { - generator Generator + generator Generator + phaseTracker *phase.Tracker } -func New(generator Generator) *Listener { - return &Listener{ +type Option func(*Listener) + +func New(generator Generator, opts ...Option) *Listener { + l := &Listener{ generator: generator, } + for _, opt := range opts { + opt(l) + } + return l +} + +// WithPhaseTracker registers a tracker set to snapshot when Listen starts. +func WithPhaseTracker(t *phase.Tracker) Option { + return func(l *Listener) { + l.phaseTracker = t + } } // Listen starts the snapshot generation process. func (l *Listener) Listen(ctx context.Context) error { + l.phaseTracker.Set(phase.Snapshot) return l.generator.CreateSnapshot(ctx) }