Merge pull request #3 from mundipagg/fix/mongodb-conn-timeout

Fix mongodb connection timeout config
mundipagg · May 28, 2024 · 126e2b9 · 126e2b9
2 parents e91a033 + bb908d6
commit 126e2b9
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 13 deletions.
diff --git a/checks/mongo/check.go b/checks/mongo/check.go
@@ -16,19 +16,20 @@ import (
 
 var (
 	conn              *mongo.Client
-	ConnectionTimeout = 3 * time.Second
+	connectionTimeout = 3 * time.Second
 	mu                sync.RWMutex
 )
 
 type Config struct {
-	Url         string
-	User        string
-	Password    string
-	Database    string
-	AuthSource  string
-	Timeout     int
-	ForceTLS    bool
-	MaxPoolSize int
+	Url               string
+	User              string
+	Password          string
+	Database          string
+	AuthSource        string
+	Timeout           int
+	ForceTLS          bool
+	MaxPoolSize       int
+	ConnectionTimeout int
 }
 
 type healthCheck struct {
@@ -60,11 +61,12 @@ func (service *healthCheck) executeCheck() error {
 	}
 
 	_, err := service.connectMongo()
+
 	if err != nil {
 		return err
 	}
 
-	return ping()
+	return ping(service.Config)
 }
 
 func (service *healthCheck) connectMongo() (*mongo.Client, error) {
@@ -75,7 +77,7 @@ func (service *healthCheck) connectMongo() (*mongo.Client, error) {
 		return conn, nil
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), ConnectionTimeout)
+	ctx, cancel := context.WithTimeout(context.Background(), getConnectionTimeout(service.Config))
 	defer cancel()
 
 	var err error
@@ -87,14 +89,14 @@ func (service *healthCheck) connectMongo() (*mongo.Client, error) {
 	return conn, nil
 }
 
-func ping() error {
+func ping(config *Config) error {
 	var err error
 	if conn == nil {
 		err = errors.New("Connection is empty")
 		return err
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), ConnectionTimeout)
+	ctx, cancel := context.WithTimeout(context.Background(), getConnectionTimeout(config))
 	defer cancel()
 
 	err = conn.Ping(ctx, readpref.Primary())
@@ -122,6 +124,14 @@ func getClientOptions(config *Config) *options.ClientOptions {
 	return co.ApplyURI(mongoURL).SetAuth(mongoCredential(config))
 }
 
+func getConnectionTimeout(config *Config) time.Duration {
+	if config.ConnectionTimeout > 0 {
+		return time.Duration(config.ConnectionTimeout) * time.Second
+	}
+
+	return connectionTimeout
+}
+
 func mongoCredential(config *Config) options.Credential {
 	user := config.User
 	password := config.Password

diff --git a/checks/mongo/check_test.go b/checks/mongo/check_test.go
@@ -2,6 +2,7 @@ package mongo
 
 import (
 	"testing"
+	"time"
 
 	check "github.com/mundipagg/healthcheck-go/checks"
 	"github.com/stretchr/testify/assert"
@@ -36,3 +37,13 @@ func Test_Execute_WhenConfigurationIsInvalid_ShouldReturnUnhealthy(t *testing.T)
 	assert.Equal(t, result.Status, check.Unhealthy)
 	assert.Contains(t, result.Description, "ERROR:")
 }
+
+func Test_getConnectionTimeout(t *testing.T) {
+	config := newStubMongoConfig()
+
+	assert.Equal(t, getConnectionTimeout(config), 3*time.Second)
+
+	config.ConnectionTimeout = 5
+
+	assert.Equal(t, getConnectionTimeout(config), 5*time.Second)
+}