fix: use context for readiness probes

ory · Nov 21, 2024 · d5c96ba · d5c96ba
1 parent 05c5e48
commit d5c96ba
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 28 deletions.
diff --git a/driver/factory_test.go b/driver/factory_test.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/ory/x/servicelocatorx"
 
@@ -35,7 +36,9 @@ func TestDriverNew(t *testing.T) {
 	require.NoError(t, err)
 
 	assert.EqualValues(t, config.DefaultSQLiteMemoryDSN, r.Config().DSN(ctx))
-	require.NoError(t, r.Persister().Ping())
+	pingCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	t.Cleanup(cancel)
+	require.NoError(t, r.Persister().Ping(pingCtx))
 
 	assert.NotEqual(t, uuid.Nil.String(), r.Persister().NetworkID(context.Background()).String())
 

diff --git a/driver/registry_default.go b/driver/registry_default.go
@@ -253,8 +253,8 @@ func (m *RegistryDefault) HealthHandler(_ context.Context) *healthx.Handler {
 	if m.healthxHandler == nil {
 		m.healthxHandler = healthx.NewHandler(m.Writer(), config.Version,
 			healthx.ReadyCheckers{
-				"database": func(_ *http.Request) error {
-					return m.Ping()
+				"database": func(r *http.Request) error {
+					return m.PingContext(r.Context())
 				},
 				"migrations": func(r *http.Request) error {
 					if m.migrationStatus != nil && !m.migrationStatus.HasPending() {
@@ -683,7 +683,9 @@ func (m *RegistryDefault) Init(ctx context.Context, ctxer contextx.Contextualize
 			return err
 		}
 
-		if err := p.Ping(); err != nil {
+		ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+		defer cancel()
+		if err := c.Store.SQLDB().PingContext(ctx); err != nil {
 			m.Logger().WithError(err).Warnf("Unable to ping database, retrying.")
 			return err
 		}
@@ -810,8 +812,12 @@ func (m *RegistryDefault) Persister() persistence.Persister {
 	return m.persister
 }
 
+func (m *RegistryDefault) PingContext(ctx context.Context) error {
+	return m.persister.Ping(ctx)
+}
+
 func (m *RegistryDefault) Ping() error {
-	return m.persister.Ping()
+	return m.persister.Ping(context.Background())
 }
 
 func (m *RegistryDefault) WithCSRFTokenGenerator(cg x.CSRFToken) {

diff --git a/go.mod b/go.mod
@@ -1,10 +1,10 @@
 module github.com/ory/kratos
 
-go 1.22
+go 1.23
 
 replace (
 	// https://github.com/gobuffalo/pop/pull/833
-	github.com/gobuffalo/pop/v6 => github.com/ory/pop/v6 v6.2.0
+	github.com/gobuffalo/pop/v6 => github.com/ory/pop/v6 v6.2.1-0.20241121111754-e5dfc0f3344b
 
 	github.com/gorilla/sessions => github.com/ory/sessions v1.2.2-0.20220110165800-b09c17334dc2
 	github.com/mattn/go-sqlite3 => github.com/mattn/go-sqlite3 v1.14.22
@@ -70,7 +70,7 @@ require (
 	github.com/ory/jsonschema/v3 v3.0.8
 	github.com/ory/mail/v3 v3.0.0
 	github.com/ory/nosurf v1.2.7
-	github.com/ory/x v0.0.665-0.20241031130226-ae5097122246
+	github.com/ory/x v0.0.669
 	github.com/peterhellberg/link v1.2.0
 	github.com/phayes/freeport v0.0.0-20220201140144-74d24b5ae9f5
 	github.com/pkg/errors v0.9.1

diff --git a/go.sum b/go.sum
@@ -252,8 +252,6 @@ github.com/gobuffalo/github_flavored_markdown v1.1.4 h1:WacrEGPXUDX+BpU1GM/Y0ADg
 github.com/gobuffalo/github_flavored_markdown v1.1.4/go.mod h1:Vl9686qrVVQou4GrHRK/KOG3jCZOKLUqV8MMOAYtlso=
 github.com/gobuffalo/helpers v0.6.7 h1:C9CedoRSfgWg2ZoIkVXgjI5kgmSpL34Z3qdnzpfNVd8=
 github.com/gobuffalo/helpers v0.6.7/go.mod h1:j0u1iC1VqlCaJEEVkZN8Ia3TEzfj/zoXANqyJExTMTA=
-github.com/gobuffalo/here v0.6.7 h1:hpfhh+kt2y9JLDfhYUxxCRxQol540jsVfKUZzjlbp8o=
-github.com/gobuffalo/here v0.6.7/go.mod h1:vuCfanjqckTuRlqAitJz6QC4ABNnS27wLb816UhsPcc=
 github.com/gobuffalo/httptest v1.5.2 h1:GpGy520SfY1QEmyPvaqmznTpG4gEQqQ82HtHqyNEreM=
 github.com/gobuffalo/httptest v1.5.2/go.mod h1:FA23yjsWLGj92mVV74Qtc8eqluc11VqcWr8/C1vxt4g=
 github.com/gobuffalo/nulls v0.4.2 h1:GAqBR29R3oPY+WCC7JL9KKk9erchaNuV6unsOSZGQkw=
@@ -547,8 +545,6 @@ github.com/mailhog/storage v1.0.1 h1:uut2nlG5hIxbsl6f8DGznPAHwQLf3/7Na2t4gmrIais
 github.com/mailhog/storage v1.0.1/go.mod h1:4EAUf5xaEVd7c/OhvSxOOwQ66jT6q2er+BDBQ0EVrew=
 github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
 github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
-github.com/markbates/pkger v0.17.1 h1:/MKEtWqtc0mZvu9OinB9UzVN9iYCwLWuyUv4Bw+PCno=
-github.com/markbates/pkger v0.17.1/go.mod h1:0JoVlrol20BSywW79rN3kdFFsE5xYM+rSCQDXbLhiuI=
 github.com/matryer/is v1.4.0/go.mod h1:8I/i5uYgLzgsgEloJE1U6xx5HkBQpAZvepWuujKwMRU=
 github.com/mattn/go-colorable v0.1.8/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
 github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
@@ -638,12 +634,12 @@ github.com/ory/mail/v3 v3.0.0 h1:8LFMRj473vGahFD/ntiotWEd4S80FKYFtiZTDfOQ+sM=
 github.com/ory/mail/v3 v3.0.0/go.mod h1:JGAVeZF8YAlxbaFDUHqRZAKBCSeW2w1vuxf28hFbZAw=
 github.com/ory/nosurf v1.2.7 h1:YrHrbSensQyU6r6HT/V5+HPdVEgrOTMJiLoJABSBOp4=
 github.com/ory/nosurf v1.2.7/go.mod h1:d4L3ZBa7Amv55bqxCBtCs63wSlyaiCkWVl4vKf3OUxA=
-github.com/ory/pop/v6 v6.2.0 h1:hRFOGAOEHw91kUHQ32k5NHqCkcHrRou/romvrJP1w0E=
-github.com/ory/pop/v6 v6.2.0/go.mod h1:okVAYKGtgunD/wbW3NGhZTndJCS+6FqO+cA89rQ4doc=
+github.com/ory/pop/v6 v6.2.1-0.20241121111754-e5dfc0f3344b h1:BIzoOe2/wynZBQak1po0tzgvARseIKsR2bF6b+SZoKE=
+github.com/ory/pop/v6 v6.2.1-0.20241121111754-e5dfc0f3344b/go.mod h1:okVAYKGtgunD/wbW3NGhZTndJCS+6FqO+cA89rQ4doc=
 github.com/ory/sessions v1.2.2-0.20220110165800-b09c17334dc2 h1:zm6sDvHy/U9XrGpixwHiuAwpp0Ock6khSVHkrv6lQQU=
 github.com/ory/sessions v1.2.2-0.20220110165800-b09c17334dc2/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM=
-github.com/ory/x v0.0.665-0.20241031130226-ae5097122246 h1:h6Jt8glJkehgQxr7MP3q5gmR4Ub0RWqlVgPXTnJU5rs=
-github.com/ory/x v0.0.665-0.20241031130226-ae5097122246/go.mod h1:7SCTki3N0De3ZpqlxhxU/94ZrOCfNEnXwVtd0xVt+L8=
+github.com/ory/x v0.0.669 h1:pBrju8B5Oie9RjebOwWf1Sj+6dPNIPI3nkVeC8rjUno=
+github.com/ory/x v0.0.669/go.mod h1:0Av1u/Gh7WXCrEDJJnySAJrDzluaWllOfl5zqf9Dky8=
 github.com/pelletier/go-toml v1.9.5 h1:4yBQzkHv+7BHq2PQUZF3Mx0IYxG7LsP222s7Agd3ve8=
 github.com/pelletier/go-toml v1.9.5/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
 github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM=

diff --git a/persistence/reference.go b/persistence/reference.go
@@ -58,13 +58,13 @@ type Persister interface {
 
 	CleanupDatabase(context.Context, time.Duration, time.Duration, int) error
 	Close(context.Context) error
-	Ping() error
-	MigrationStatus(c context.Context) (popx.MigrationStatuses, error)
-	MigrateDown(c context.Context, steps int) error
-	MigrateUp(c context.Context) error
+	Ping(context.Context) error
+	MigrationStatus(context.Context) (popx.MigrationStatuses, error)
+	MigrateDown(ctx context.Context, steps int) error
+	MigrateUp(context.Context) error
 	Migrator() *popx.Migrator
 	MigrationBox() *popx.MigrationBox
-	GetConnection(ctx context.Context) *pop.Connection
+	GetConnection(context.Context) *pop.Connection
 	x.TransactionalPersister
 	Networker
 }

diff --git a/persistence/sql/persister.go b/persistence/sql/persister.go
@@ -178,13 +178,8 @@ func (p *Persister) Close(ctx context.Context) error {
 	return errors.WithStack(p.GetConnection(ctx).Close())
 }
 
-func (p *Persister) Ping() error {
-	type pinger interface {
-		Ping() error
-	}
-
-	// This can not be contextualized because of some gobuffalo/pop limitations.
-	return errors.WithStack(p.c.Store.(pinger).Ping())
+func (p *Persister) Ping(ctx context.Context) error {
+	return errors.WithStack(p.c.Store.SQLDB().PingContext(ctx))
 }
 
 func (p *Persister) CleanupDatabase(ctx context.Context, wait time.Duration, older time.Duration, batchSize int) error {