From 0cfb570019b789623834394a60f98f63e9146215 Mon Sep 17 00:00:00 2001 From: Nikita Kharashun Date: Tue, 5 May 2026 01:06:33 +0300 Subject: [PATCH 1/2] fix connection churn in mongo (cherry picked from commit ca40192262175986795ab23af1a44d7ca7bf541c) --- checks/mongo/check.go | 32 ++++++++++++++++++++++++++++++++ checks/mongo/check_test.go | 31 +++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/checks/mongo/check.go b/checks/mongo/check.go index 4df6f63..8d3d0d7 100644 --- a/checks/mongo/check.go +++ b/checks/mongo/check.go @@ -2,6 +2,7 @@ package mongo import ( "context" + "errors" "fmt" "time" @@ -16,6 +17,8 @@ const ( defaultTimeoutPing = 5 * time.Second ) +var errNilClient = errors.New("mongoDB health check failed on ping: nil mongo client") + // Config is the MongoDB checker configuration settings container. type Config struct { // DSN is the MongoDB instance connection DSN. Required. @@ -29,9 +32,38 @@ type Config struct { TimeoutPing time.Duration } +// NewPingCheck creates a MongoDB health check using an existing client. +// The caller is responsible for managing the client's lifecycle. +func NewPingCheck(client *mongo.Client, timeout time.Duration) func(ctx context.Context) error { + if client == nil { + return func(context.Context) error { + return errNilClient + } + } + + if timeout == 0 { + timeout = defaultTimeoutPing + } + + return func(ctx context.Context) error { + ctxPing, cancelPing := context.WithTimeout(ctx, timeout) + defer cancelPing() + + err := client.Ping(ctxPing, readpref.Primary()) + if err != nil { + return fmt.Errorf("mongoDB health check failed on ping: %w", err) + } + + return nil + } +} + // New creates new MongoDB health check that verifies the following: // - connection establishing // - doing the ping command +// +// Deprecated: use NewPingCheck instead. It uses an existing *mongo.Client +// to avoid connection churn. func New(config Config) func(ctx context.Context) error { if config.TimeoutConnect == 0 { config.TimeoutConnect = defaultTimeoutConnect diff --git a/checks/mongo/check_test.go b/checks/mongo/check_test.go index c391870..376c1a6 100644 --- a/checks/mongo/check_test.go +++ b/checks/mongo/check_test.go @@ -2,15 +2,46 @@ package mongo import ( "context" + "go.mongodb.org/mongo-driver/mongo" + "go.mongodb.org/mongo-driver/mongo/options" "os" "strings" "testing" + "time" "github.com/stretchr/testify/require" ) const mgDSNEnv = "HEALTH_GO_MG_DSN" +func TestNewPingCheck(t *testing.T) { + ctx := context.Background() + + t.Run("success check with valid client", func(t *testing.T) { + dsn := getDSN(t) + + client, err := mongo.Connect(ctx, options.Client().ApplyURI(dsn)) + require.NoError(t, err) + + defer func() { + errDisc := client.Disconnect(ctx) + require.NoError(t, errDisc) + }() + + check := NewPingCheck(client, 0) + + err = check(ctx) + require.NoError(t, err) + }) + + t.Run("fails on nil client", func(t *testing.T) { + check := NewPingCheck(nil, 5*time.Second) + + err := check(ctx) + require.ErrorIs(t, err, errNilClient) + }) +} + func TestNew(t *testing.T) { check := New(Config{ DSN: getDSN(t), From a87fd89c855774ac88fa000f842f1392cf6b16ff Mon Sep 17 00:00:00 2001 From: Nikita Kharashun Date: Fri, 3 Jul 2026 19:28:40 +0300 Subject: [PATCH 2/2] support reusing an existing mongo client via Config.Client Replace NewPingCheck with an optional Client field in Config, mirroring the cassandra check's Session pattern, so the check can reuse the application's connection pool instead of establishing and tearing down a new connection on every run. --- checks/mongo/check.go | 97 +++++++++++++++++++------------------- checks/mongo/check_test.go | 52 ++++++++++---------- 2 files changed, 73 insertions(+), 76 deletions(-) diff --git a/checks/mongo/check.go b/checks/mongo/check.go index 8d3d0d7..c588f8a 100644 --- a/checks/mongo/check.go +++ b/checks/mongo/check.go @@ -17,13 +17,17 @@ const ( defaultTimeoutPing = 5 * time.Second ) -var errNilClient = errors.New("mongoDB health check failed on ping: nil mongo client") - // Config is the MongoDB checker configuration settings container. type Config struct { - // DSN is the MongoDB instance connection DSN. Required. + // DSN is the MongoDB instance connection DSN. Optional if Client is supplied. DSN string + // Client is an existing mongo client and can be used in place of DSN. Recommended, + // since it reuses the application's connection pool instead of establishing and + // tearing down a new connection on every check. The caller is responsible for + // managing the client's lifecycle. Optional if DSN is supplied. + Client *mongo.Client + // TimeoutConnect defines timeout for establishing mongo connection, if not set - default value is used TimeoutConnect time.Duration // TimeoutDisconnect defines timeout for closing connection, if not set - default value is used @@ -32,38 +36,9 @@ type Config struct { TimeoutPing time.Duration } -// NewPingCheck creates a MongoDB health check using an existing client. -// The caller is responsible for managing the client's lifecycle. -func NewPingCheck(client *mongo.Client, timeout time.Duration) func(ctx context.Context) error { - if client == nil { - return func(context.Context) error { - return errNilClient - } - } - - if timeout == 0 { - timeout = defaultTimeoutPing - } - - return func(ctx context.Context) error { - ctxPing, cancelPing := context.WithTimeout(ctx, timeout) - defer cancelPing() - - err := client.Ping(ctxPing, readpref.Primary()) - if err != nil { - return fmt.Errorf("mongoDB health check failed on ping: %w", err) - } - - return nil - } -} - // New creates new MongoDB health check that verifies the following: -// - connection establishing +// - connection establishing (skipped when an existing Client is supplied) // - doing the ping command -// -// Deprecated: use NewPingCheck instead. It uses an existing *mongo.Client -// to avoid connection churn. func New(config Config) func(ctx context.Context) error { if config.TimeoutConnect == 0 { config.TimeoutConnect = defaultTimeoutConnect @@ -78,28 +53,16 @@ func New(config Config) func(ctx context.Context) error { } return func(ctx context.Context) (checkErr error) { - client, err := mongo.NewClient(options.Client().ApplyURI(config.DSN)) + shutdown, client, err := initClient(ctx, config) if err != nil { - checkErr = fmt.Errorf("mongoDB health check failed on client creation: %w", err) - return - } - - ctxConn, cancelConn := context.WithTimeout(ctx, config.TimeoutConnect) - defer cancelConn() - - err = client.Connect(ctxConn) - if err != nil { - checkErr = fmt.Errorf("mongoDB health check failed on connect: %w", err) + checkErr = err return } defer func() { - ctxDisc, cancelDisc := context.WithTimeout(ctx, config.TimeoutDisconnect) - defer cancelDisc() - // override checkErr only if there were no other errors - if err := client.Disconnect(ctxDisc); err != nil && checkErr == nil { - checkErr = fmt.Errorf("mongoDB health check failed on closing connection: %w", err) + if err := shutdown(ctx); err != nil && checkErr == nil { + checkErr = err } }() @@ -115,3 +78,39 @@ func New(config Config) func(ctx context.Context) error { return } } + +func initClient(ctx context.Context, c Config) (func(ctx context.Context) error, *mongo.Client, error) { + if c.Client != nil { + return func(context.Context) error { return nil }, c.Client, nil + } + + if c.DSN == "" { + return nil, nil, errors.New("mongoDB DSN or an existing client is required to initialize mongoDB health check") + } + + client, err := mongo.NewClient(options.Client().ApplyURI(c.DSN)) + if err != nil { + return nil, nil, fmt.Errorf("mongoDB health check failed on client creation: %w", err) + } + + ctxConn, cancelConn := context.WithTimeout(ctx, c.TimeoutConnect) + defer cancelConn() + + err = client.Connect(ctxConn) + if err != nil { + return nil, nil, fmt.Errorf("mongoDB health check failed on connect: %w", err) + } + + shutdown := func(ctx context.Context) error { + ctxDisc, cancelDisc := context.WithTimeout(ctx, c.TimeoutDisconnect) + defer cancelDisc() + + if err := client.Disconnect(ctxDisc); err != nil { + return fmt.Errorf("mongoDB health check failed on closing connection: %w", err) + } + + return nil + } + + return shutdown, client, nil +} diff --git a/checks/mongo/check_test.go b/checks/mongo/check_test.go index 376c1a6..0aa1fdf 100644 --- a/checks/mongo/check_test.go +++ b/checks/mongo/check_test.go @@ -2,53 +2,51 @@ package mongo import ( "context" - "go.mongodb.org/mongo-driver/mongo" - "go.mongodb.org/mongo-driver/mongo/options" "os" "strings" "testing" - "time" + + "go.mongodb.org/mongo-driver/mongo" + "go.mongodb.org/mongo-driver/mongo/options" "github.com/stretchr/testify/require" ) const mgDSNEnv = "HEALTH_GO_MG_DSN" -func TestNewPingCheck(t *testing.T) { - ctx := context.Background() +func TestNew(t *testing.T) { + check := New(Config{ + DSN: getDSN(t), + }) - t.Run("success check with valid client", func(t *testing.T) { - dsn := getDSN(t) + err := check(context.Background()) + require.NoError(t, err) +} - client, err := mongo.Connect(ctx, options.Client().ApplyURI(dsn)) - require.NoError(t, err) +func TestNew_withClient(t *testing.T) { + ctx := context.Background() - defer func() { - errDisc := client.Disconnect(ctx) - require.NoError(t, errDisc) - }() + client, err := mongo.Connect(ctx, options.Client().ApplyURI(getDSN(t))) + require.NoError(t, err) - check := NewPingCheck(client, 0) + defer func() { + errDisc := client.Disconnect(ctx) + require.NoError(t, errDisc) + }() - err = check(ctx) - require.NoError(t, err) + check := New(Config{ + Client: client, }) - t.Run("fails on nil client", func(t *testing.T) { - check := NewPingCheck(nil, 5*time.Second) - - err := check(ctx) - require.ErrorIs(t, err, errNilClient) - }) + err = check(ctx) + require.NoError(t, err) } -func TestNew(t *testing.T) { - check := New(Config{ - DSN: getDSN(t), - }) +func TestNewWithError(t *testing.T) { + check := New(Config{}) err := check(context.Background()) - require.NoError(t, err) + require.Error(t, err) } func getDSN(t *testing.T) string {