fix: memory and mongodb connection leak (#45)

Signed-off-by: Raffael Sahli <raffael.sahli@doodle.com>

Author: raffis

internal/controllers/instance_controller.go

@@ -66,7 +66,7 @@ const (
 )
 
 // MongoDBProvider returns a storage.Database for MongoDB
-func MongoDBProvider(ctx context.Context, instance v1beta1.GrowthbookInstance, username, password string) (storage.Database, error) {
+func MongoDBProvider(ctx context.Context, instance v1beta1.GrowthbookInstance, username, password string) (storage.Disconnector, storage.Database, error) {
 	opts := options.Client().ApplyURI(instance.Spec.MongoDB.URI)
 	if username != "" || password != "" {
 		opts.SetAuth(options.Credential{
@@ -78,22 +78,22 @@ func MongoDBProvider(ctx context.Context, instance v1beta1.GrowthbookInstance, u
 	opts.SetAppName("k8sgrowthbook-controller")
 	mongoClient, err := mongo.NewClient(opts)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	u, err := url.Parse(instance.Spec.MongoDB.URI)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	dbName := strings.TrimLeft(u.Path, "/")
 	db := mongoClient.Database(dbName)
 
 	if err := mongoClient.Connect(ctx); err != nil {
-		return nil, fmt.Errorf("failed connecting to mongodb: %w", err)
+		return nil, nil, fmt.Errorf("failed connecting to mongodb: %w", err)
 	}
 
-	return mongodb.New(db), nil
+	return mongoClient, mongodb.New(db), nil
 }
 
 // GrowthbookInstance reconciles a GrowthbookInstance object
@@ -102,7 +102,7 @@ type GrowthbookInstanceReconciler struct {
 	Log              logr.Logger
 	Scheme           *runtime.Scheme
 	Recorder         record.EventRecorder
-	DatabaseProvider func(ctx context.Context, instance v1beta1.GrowthbookInstance, username, password string) (storage.Database, error)
+	DatabaseProvider func(ctx context.Context, instance v1beta1.GrowthbookInstance, username, password string) (storage.Disconnector, storage.Database, error)
 }
 
 type GrowthbookInstanceReconcilerOptions struct {
@@ -332,11 +332,19 @@ func (r *GrowthbookInstanceReconciler) reconcile(ctx context.Context, instance v
 		}
 	}
 
-	db, err := r.DatabaseProvider(ctx, instance, usr, pw)
+	disconnector, db, err := r.DatabaseProvider(ctx, instance, usr, pw)
 	if err != nil {
 		return instance, err
 	}
 
+	defer func() {
+		ctx, cancel := context.WithTimeout(context.TODO(), time.Second*10)
+		defer cancel()
+		if err := disconnector.Disconnect(ctx); err != nil {
+			logger.Error(err, "failed disconnet mongodb")
+		}
+	}()
+
 	instance.Status.SubResourceCatalog = []v1beta1.ResourceReference{}
 
 	instance, err = r.reconcileUsers(ctx, instance, db, logger)

internal/controllers/instance_controller_test.go

@@ -16,16 +16,16 @@ import (
 )
 
 // MockProvider returns a storage.Database for testing
-func MockProvider(ctx context.Context, instance v1beta1.GrowthbookInstance, username, password string) (storage.Database, error) {
+func MockProvider(ctx context.Context, instance v1beta1.GrowthbookInstance, username, password string) (storage.Disconnector, storage.Database, error) {
 	// See test "reconciling a GrowthbookInstance with a timeout"
 	if instance.Spec.Timeout != nil && instance.Spec.Timeout.Duration == 500*time.Millisecond {
 		<-ctx.Done()
-		return nil, ctx.Err()
+		return nil, nil, ctx.Err()
 	}
 
 	// For testing the controller logic the storage adapter just does nothing and returns no error
 	if instance.Spec.MongoDB.URI == "" {
-		return &growthbook.MockDatabase{
+		return &growthbook.MockDisconnect{}, &growthbook.MockDatabase{
 			FindOne: func(ctx context.Context, filter interface{}, dst interface{}) error {
 				return nil
 			},

internal/growthbook/storage_mock.go

@@ -12,6 +12,13 @@ var (
 	_ storage.Collection = &MockCollection{}
 )
 
+type MockDisconnect struct {
+}
+
+func (d *MockDisconnect) Disconnect(ctx context.Context) error {
+	return nil
+}
+
 type MockDatabase struct {
 	FindOne    func(ctx context.Context, filter interface{}, dst interface{}) error
 	DeleteOne  func(ctx context.Context, filter interface{}) error

internal/storage/storage.go

@@ -2,6 +2,10 @@ package storage
 
 import "context"
 
+type Disconnector interface {
+	Disconnect(ctx context.Context) error
+}
+
 type Database interface {
 	Collection(collName string) Collection
 }