fix(datastores): ReadyState respects context cancellation/deadline (#3262)

Author: miparnisari

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Postgres: read replicas no longer intermittently return `object definition not found` under load. The strict read-replica guard now verifies that the replica's snapshot has caught up to the revision being read (snapshot domination) instead of checking a single transaction id, and raises from within the read itself rather than from a trailing assertion, so a replica that catches up mid-query can no longer let an incomplete read through. In both cases the read correctly falls back to the primary. (https://github.com/authzed/spicedb/pull/3243)
 - Prevent ReadRelationships from doing work that's immediately discarded when the `optional_limit` parameter is used (https://github.com/authzed/spicedb/pull/3253)
 - MemDB: overlapping write transactions could violate every snapshot-consistency invariant of the datastore — a committed write could be invisible at its own returned revision (breaking read-your-writes, e.g. an at-exact-snapshot `Check` right after `WriteRelationships`), a later commit could leak into reads at an earlier revision, a write visible at one revision could be missing at a later one (including head), and two concurrent transactions could be assigned the same revision. Revisions are now assigned at write-transaction acquisition, where writer serialization makes the two orders identical. (https://github.com/authzed/spicedb/pull/3239)
+- All datastores now answer Ready (or not) within a bounded time (https://github.com/authzed/spicedb/pull/3262)
 
 ## [1.56.0] - 2026-07-24
 ### Added

internal/datastore/crdb/crdb.go

@@ -90,17 +90,17 @@ func newCRDBDatastore(ctx context.Context, url string, options ...Option) (datas
 		return nil, common.RedactAndLogSensitiveConnString(ctx, errUnableToInstantiate, err, url)
 	}
 
-	initCtx, initCancel := context.WithTimeout(context.Background(), 5*time.Minute)
-	defer initCancel()
-
 	healthChecker, err := pool.NewNodeHealthChecker(url)
 	if err != nil {
 		return nil, common.RedactAndLogSensitiveConnString(ctx, errUnableToInstantiate, err, url)
 	}
 
 	// The initPool is a 1-connection pool that is only used for setup tasks.
-	// The actual pools are not given the initCtx, since cancellation can
-	// interfere with pool setup.
+	// If the database is completely unreachable (e.g. wrong credentials),
+	// this will block for 15 seconds and then error out.
+	// TODO(miparnisari): remove this initCtx once spicedb has a k8s startup probe. It will become unnecessary.
+	initCtx, initCancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer initCancel()
 	initPoolConfig := readPoolConfig.Copy()
 	initPoolConfig.MinConns = 1
 	initPool, err := pool.NewRetryPool(initCtx, "init", initPoolConfig, healthChecker, config.maxRetries, config.connectRate)
@@ -208,7 +208,8 @@ func newCRDBDatastore(ctx context.Context, url string, options ...Option) (datas
 	ds.SetNowFunc(ds.headRevisionInternal)
 	ds.SetNowOnlyFunc(ds.headRevisionInternalNoHash)
 
-	// this ctx and cancel is tied to the lifetime of the datastore
+	// The actual pools are not given the initCtx.
+	// This ctx and cancel is tied to the lifetime of the datastore
 	ds.ctx, ds.cancel = context.WithCancel(context.Background())
 	ds.writePool, err = pool.NewRetryPool(ds.ctx, "write", writePoolConfig, healthChecker, config.maxRetries, config.connectRate)
 	if err != nil {
@@ -426,7 +427,7 @@ func wrapError(err error) error {
 // to be ready to receive traffic, and total connections counts connections in the constructing
 // state, which cannot receive traffic.
 func (cds *crdbDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, error) {
-	currentRevision, err := migrations.NewCRDBDriver(cds.dburl)
+	currentRevision, err := migrations.NewCRDBDriver(ctx, cds.dburl)
 	if err != nil {
 		return datastore.ReadyState{}, err
 	}

internal/datastore/crdb/crdb_test.go

@@ -67,7 +67,7 @@ func crdbTestVersion() string {
 func TestCRDBDatastoreWithoutIntegrity(t *testing.T) {
 	t.Parallel()
 	b := testdatastore.RunCRDBForTesting(t, crdbTestVersion())
-	test.All(t, crdbFactory.NewTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
+	test.All(t, crdbFactory.NewTester(test.PausableTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
 		ctx := t.Context()
 		ds := b.NewDatastore(t, func(engine, uri string) datastore.Datastore {
 			ds, err := NewCRDBDatastore(
@@ -88,7 +88,7 @@ func TestCRDBDatastoreWithoutIntegrity(t *testing.T) {
 		})
 
 		return ds, nil
-	})))
+	}), b)))
 
 	t.Run("TestWatchStreaming", createDatastoreTest(
 		b,
@@ -209,7 +209,7 @@ func TestCRDBDatastoreWithIntegrity(t *testing.T) { //nolint:tparallel
 	t.Parallel()
 	b := testdatastore.RunCRDBForTesting(t, crdbTestVersion())
 
-	test.AllWithExceptions(t, crdbFactory.NewTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
+	test.AllWithExceptions(t, crdbFactory.NewTester(test.PausableTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
 		ctx := t.Context()
 		ds := b.NewDatastore(t, func(engine, uri string) datastore.Datastore {
 			ds, err := NewCRDBDatastore(
@@ -234,7 +234,7 @@ func TestCRDBDatastoreWithIntegrity(t *testing.T) { //nolint:tparallel
 		})
 
 		return ds, nil
-	})), test.WithCategories(test.MigrationCategory))
+	}), b)), test.WithCategories(test.MigrationCategory))
 
 	unwrappedTester := test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
 		ctx := t.Context()
@@ -311,7 +311,7 @@ func TestWatchFeatureDetection(t *testing.T) {
 			ctx := t.Context()
 			adminConn, connStrings := newCRDBWithUser(t)
 
-			migrationDriver, err := crdbmigrations.NewCRDBDriver(connStrings[testuser])
+			migrationDriver, err := crdbmigrations.NewCRDBDriver(ctx, connStrings[testuser])
 			require.NoError(t, err)
 			require.NoError(t, crdbmigrations.CRDBMigrations.Run(ctx, migrationDriver, migrate.Head, migrate.LiveRun))
 

internal/datastore/crdb/enginebuilder.go

@@ -18,8 +18,8 @@ func init() {
 	migration.RegisterMigratableEngine(Engine, migrations.CRDBMigrations, newMigrationDriverFromConfig, "add-schema-tables")
 }
 
-func newMigrationDriverFromConfig(_ context.Context, cfg *migration.Config) (*migrations.CRDBDriver, error) {
-	return migrations.NewCRDBDriver(cfg.DatastoreURI)
+func newMigrationDriverFromConfig(ctx context.Context, cfg *migration.Config) (*migrations.CRDBDriver, error) {
+	return migrations.NewCRDBDriver(ctx, cfg.DatastoreURI)
 }
 
 func newDatastoreFromConfig(ctx context.Context, opts datastorecfg.Config) (datastore.Datastore, error) {

internal/datastore/crdb/migrations/driver.go

@@ -29,15 +29,15 @@ type CRDBDriver struct {
 
 // NewCRDBDriver creates a new driver with active connections to the database
 // specified.
-func NewCRDBDriver(url string) (*CRDBDriver, error) {
+func NewCRDBDriver(ctx context.Context, url string) (*CRDBDriver, error) {
 	connConfig, err := pgx.ParseConfig(url)
 	if err != nil {
 		return nil, fmt.Errorf(errUnableToInstantiate, err)
 	}
 	pgxcommon.ConfigurePGXLogger(connConfig)
 	pgxcommon.ConfigureOTELTracer(connConfig, false)
 
-	db, err := pgx.ConnectConfig(context.Background(), connConfig)
+	db, err := pgx.ConnectConfig(ctx, connConfig)
 	if err != nil {
 		return nil, fmt.Errorf(errUnableToInstantiate, err)
 	}

internal/datastore/mysql/datastore_test.go

@@ -120,7 +120,7 @@ func TestMySQLDatastoreDSNWithoutParseTime(t *testing.T) {
 func TestMySQL8Datastore(t *testing.T) {
 	b := testdatastore.RunMySQLForTestingWithOptions(t, testdatastore.MySQLTesterOptions{MigrateForNewDatastore: true})
 	dst := datastoreTester{b: b}
-	test.AllWithExceptions(t, mysqlFactory.NewTester(test.DatastoreTesterFunc(dst.createDatastore)), test.WithCategories(test.WatchSchemaCategory))
+	test.AllWithExceptions(t, mysqlFactory.NewTester(test.PausableTester(test.DatastoreTesterFunc(dst.createDatastore), b)), test.WithCategories(test.WatchSchemaCategory))
 	additionalMySQLTests(t, b)
 }
 

internal/datastore/postgres/postgres_shared_test.go

@@ -110,7 +110,7 @@ func testPostgresDatastore(t *testing.T, config postgresTestConfig) {
 		b := testdatastore.RunPostgresForTesting(t, config.pgVersion, config.pgbouncer)
 		ctx := t.Context()
 
-		test.AllWithExceptions(t, pgFactory.NewTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
+		test.AllWithExceptions(t, pgFactory.NewTester(test.PausableTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
 			ds := b.NewDatastore(t, func(engine, uri string) datastore.Datastore {
 				ds, err := newPostgresDatastore(ctx, uri, primaryInstanceID,
 					RevisionQuantization(revisionParameters.Quantization),
@@ -125,7 +125,7 @@ func testPostgresDatastore(t *testing.T, config postgresTestConfig) {
 				return indexcheck.WrapWithIndexCheckingDatastoreProxyIfApplicable(ds)
 			})
 			return ds, nil
-		})), test.WithCategories(test.GCCategory))
+		}), b)), test.WithCategories(test.GCCategory))
 
 		t.Run("TransactionTimestamps", createDatastoreTest(
 			b,
@@ -320,7 +320,7 @@ func testPostgresDatastoreWithoutCommitTimestamps(t *testing.T, config postgresT
 
 		// NOTE: watch API requires the commit timestamps, so we skip those tests here.
 		// NOTE: gc tests take exclusive locks, so they are run under non-parallel.
-		test.AllWithExceptions(t, pgFactory.NewTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
+		test.AllWithExceptions(t, pgFactory.NewTester(test.PausableTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
 			ds := b.NewDatastore(t, func(engine, uri string) datastore.Datastore {
 				ds, err := newPostgresDatastore(ctx, uri, primaryInstanceID,
 					RevisionQuantization(revisionParameters.Quantization),
@@ -334,7 +334,7 @@ func testPostgresDatastoreWithoutCommitTimestamps(t *testing.T, config postgresT
 				return ds
 			})
 			return ds, nil
-		})), test.WithCategories(test.WatchCategory, test.GCCategory, test.MigrationCategory))
+		}), b)), test.WithCategories(test.WatchCategory, test.GCCategory, test.MigrationCategory))
 	})
 
 	t.Run(fmt.Sprintf("postgres-%s-gc", pgVersion), func(t *testing.T) {

internal/datastore/spanner/spanner_test.go

@@ -29,7 +29,7 @@ func TestSpannerDatastore(t *testing.T) {
 	b := testdatastore.RunSpannerForTesting(t)
 
 	// Transaction tests are excluded because, for reasons unknown, one cannot read its own write in one transaction in the Spanner emulator.
-	test.AllWithExceptions(t, spannerFactory.NewTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
+	test.AllWithExceptions(t, spannerFactory.NewTester(test.PausableTester(test.DatastoreTesterFunc(func(t testing.TB, revisionParameters test.RevisionParameters, watchBufferLength uint16) (datastore.Datastore, error) {
 		ds := b.NewDatastore(t, func(engine, uri string) datastore.Datastore {
 			ds, err := NewSpannerDatastore(ctx, uri,
 				RevisionQuantization(revisionParameters.Quantization),
@@ -43,7 +43,7 @@ func TestSpannerDatastore(t *testing.T) {
 			return ds
 		})
 		return ds, nil
-	})), test.WithCategories(test.GCCategory, test.StatsCategory, test.TransactionCategory))
+	}), b)), test.WithCategories(test.GCCategory, test.StatsCategory, test.TransactionCategory))
 
 	t.Run("TestFakeStats", createDatastoreTest(
 		b,

internal/services/health/health.go

@@ -28,6 +28,7 @@ func NewHealthManager(dispatcher dispatch.Dispatcher, dsc DatastoreChecker) Mana
 // traffic.
 type DatastoreChecker interface {
 	// ReadyState returns whether the datastore is ready to be used.
+	// This call must respect context cancellation and deadlines.
 	ReadyState(ctx context.Context) (datastore.ReadyState, error)
 }
 
@@ -40,8 +41,15 @@ type Manager interface {
 	// HealthSvc is the health service this manager is managing.
 	HealthSvc() *grpcutil.AuthlessHealthServer
 
-	// Checker blocks until the status is SERVING.
+	// Checker blocks until the status is SERVING or until the context is done.
 	Checker(ctx context.Context) error
+
+	// right now, Checker is serving both as a startup probe and a readiness probe.
+	// However, its implementation is really only functioning as a startup probe because the function exits quickly.
+	// TODO(miparnisari): Split Checker into two functions: Startup and Readiness (to match kubernetes probes)
+	// Startup can continue with this implementation. Readiness can ... be something else (TBD).
+	// ALso: we don't need a per-service probe, just one a general one for the entire service.
+	// Also: we need a function Close() that sets status to NOT_SERVING.
 }
 
 type healthManager struct {
@@ -60,7 +68,11 @@ func (hm *healthManager) RegisterReportedService(serviceName string) {
 	hm.healthSvc.SetServingStatus(serviceName, healthpb.HealthCheckResponse_NOT_SERVING)
 }
 
+// Checker blocks until the underlying dependencies are ready.
+// When they are, it marks the service as "serving" and returns.
+// If they never become ready, it just returns.
 func (hm *healthManager) Checker(ctx context.Context) error {
+	log.Ctx(ctx).Info().Msg("HEALTHCHECK")
 	// Run immediately for the initial check
 	backoffInterval := backoff.NewExponentialBackOff()
 
@@ -96,6 +108,7 @@ func (hm *healthManager) Checker(ctx context.Context) error {
 	}
 }
 
+// checkIsReady returns true if both the datastore and the dispatcher are ready
 func (hm *healthManager) checkIsReady(ctx context.Context) bool {
 	log.Ctx(ctx).Debug().Msg("checking if datastore and dispatcher are ready")
 

internal/testserver/datastore/crdb.go

@@ -18,11 +18,13 @@ import (
 
 // crdbTester is safe for concurrent use by tests.
 type crdbTester struct {
+	pausableContainer
+
 	// endpoint is the host:port of the cockroach container.
 	endpoint string
 }
 
-var _ RunningEngineForTest = (*crdbTester)(nil)
+var _ PausableEngineForTest = (*crdbTester)(nil)
 
 // RunCRDBForTesting returns a RunningEngineForTest for CRDB
 func RunCRDBForTesting(t testing.TB, crdbVersion string, opts ...testcontainers.ContainerCustomizer) *crdbTester {
@@ -56,7 +58,8 @@ func RunCRDBForTesting(t testing.TB, crdbVersion string, opts ...testcontainers.
 	require.NoError(t, err)
 
 	return &crdbTester{
-		endpoint: net.JoinHostPort(host, mappedPort.Port()),
+		pausableContainer: pausableContainer{container: container},
+		endpoint:          net.JoinHostPort(host, mappedPort.Port()),
 	}
 }
 
@@ -88,7 +91,7 @@ func (r *crdbTester) NewDatabase(t testing.TB) string {
 func (r *crdbTester) NewDatastore(t testing.TB, initFunc InitFunc) datastore.Datastore {
 	connectStr := r.NewDatabase(t)
 
-	migrationDriver, err := crdbmigrations.NewCRDBDriver(connectStr)
+	migrationDriver, err := crdbmigrations.NewCRDBDriver(t.Context(), connectStr)
 	require.NoError(t, err)
 	require.NoError(t, crdbmigrations.CRDBMigrations.Run(t.Context(), migrationDriver, migrate.Head, migrate.LiveRun))
 	t.Cleanup(func() {