feat: support reusing an existing mongo client to prevent connection churn

checks/mongo/check.go

@@ -2,6 +2,7 @@ package mongo
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
@@ -18,9 +19,15 @@ const (
 
 // 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
@@ -30,7 +37,7 @@ type Config struct {
 }
 
 // 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
 func New(config Config) func(ctx context.Context) error {
 	if config.TimeoutConnect == 0 {
@@ -46,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))
-		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)
+		shutdown, client, err := initClient(ctx, config)
 		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
 			}
 		}()
 
@@ -83,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
+}

checks/mongo/check_test.go

@@ -6,6 +6,9 @@ import (
 	"strings"
 	"testing"
 
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+
 	"github.com/stretchr/testify/require"
 )
 
@@ -20,6 +23,32 @@ func TestNew(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestNew_withClient(t *testing.T) {
+	ctx := context.Background()
+
+	client, err := mongo.Connect(ctx, options.Client().ApplyURI(getDSN(t)))
+	require.NoError(t, err)
+
+	defer func() {
+		errDisc := client.Disconnect(ctx)
+		require.NoError(t, errDisc)
+	}()
+
+	check := New(Config{
+		Client: client,
+	})
+
+	err = check(ctx)
+	require.NoError(t, err)
+}
+
+func TestNewWithError(t *testing.T) {
+	check := New(Config{})
+
+	err := check(context.Background())
+	require.Error(t, err)
+}
+
 func getDSN(t *testing.T) string {
 	t.Helper()