Merge pull request #3178 from QUBITABHAY/development

Author: Umang01-hash

pkg/gofr/container/datasources.go

@@ -259,8 +259,8 @@ type Mongo interface {
 	UpdateByID(ctx context.Context, collection string, id any, update any) (int64, error)
 
 	// UpdateOne updates a single document in a collection based on a filter.
-	// It returns an error if any.
-	UpdateOne(ctx context.Context, collection string, filter any, update any) error
+	// It returns the number of documents updated and an error if any.
+	UpdateOne(ctx context.Context, collection string, filter any, update any) (int64, error)
 
 	// UpdateMany updates multiple documents in a collection based on a filter.
 	// It returns the number of documents updated and an error if any.

pkg/gofr/container/mock_datasources.go

@@ -255,15 +255,19 @@ func (c *Client) UpdateByID(ctx context.Context, collection string, id, update a
 }
 
 // UpdateOne updates a single document in the specified collection based on the provided filter.
-func (c *Client) UpdateOne(ctx context.Context, collection string, filter, update any) error {
+func (c *Client) UpdateOne(ctx context.Context, collection string, filter, update any) (int64, error) {
 	tracerCtx, span := c.addTrace(ctx, "updateOne", collection)
 
-	_, err := c.Database.Collection(collection).UpdateOne(tracerCtx, filter, update)
+	res, err := c.Database.Collection(collection).UpdateOne(tracerCtx, filter, update)
 
 	defer c.sendOperationStats(&QueryLog{Query: "updateOne", Collection: collection, Filter: filter, Update: update},
 		time.Now(), "updateOne", span)
 
-	return err
+	if err != nil {
+		return 0, err
+	}
+
+	return res.ModifiedCount, nil
 }
 
 // UpdateMany updates multiple documents in the specified collection based on the provided filter.

pkg/gofr/datasource/mongo/mongo.go

@@ -483,7 +483,9 @@ func Test_UpdateCommands(t *testing.T) {
 		mt.AddMockResponses(mtest.CreateSuccessResponse())
 		// Create a document to insert
 
-		err := cl.UpdateOne(context.Background(), mt.Coll.Name(), bson.D{{Key: "name", Value: "test"}}, bson.M{"$set": bson.M{"name": "testing"}})
+		modified, err := cl.UpdateOne(context.Background(), mt.Coll.Name(), bson.D{{Key: "name", Value: "test"}},
+			bson.M{"$set": bson.M{"name": "testing"}})
+		assert.Equal(t, int64(0), modified)
 
 		assert.NoError(t, err)
 	})

pkg/gofr/datasource/mongo/mongo_test.go

@@ -63,7 +63,7 @@ type Mongo interface {
 	DeleteOne(ctx context.Context, collection string, filter any) (int64, error)
 	DeleteMany(ctx context.Context, collection string, filter any) (int64, error)
 	UpdateByID(ctx context.Context, collection string, id any, update any) (int64, error)
-	UpdateOne(ctx context.Context, collection string, filter any, update any) error
+	UpdateOne(ctx context.Context, collection string, filter any, update any) (int64, error)
 	UpdateMany(ctx context.Context, collection string, filter any, update any) (int64, error)
 	Drop(ctx context.Context, collection string) error
 	CreateCollection(ctx context.Context, name string) error

pkg/gofr/migration/interface.go

@@ -2,19 +2,24 @@ package migration
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"strings"
 	"time"
 
 	"gofr.dev/pkg/gofr/container"
 )
 
+var errMongoLockRefreshFailed = errors.New("failed to refresh MongoDB lock: lock lost or stolen")
+
 type mongoDS struct {
 	container.Mongo
 }
 
 type mongoMigrator struct {
 	container.Mongo
 	migrator
+	testInterval time.Duration // Used for testing; if non-zero, overrides defaultRefresh
 }
 
 // apply initializes mongoMigrator using the Mongo interface.
@@ -27,14 +32,30 @@ func (ds mongoDS) apply(m migrator) migrator {
 
 const (
 	mongoMigrationCollection = "gofr_migrations"
+	mongoLockCollection      = "gofr_migration_locks"
+	mongoLockDocumentID      = "gofr_migrations_lock"
 )
 
+func isMongoCollectionExistsError(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	msg := strings.ToLower(err.Error())
+
+	return strings.Contains(msg, "already exists") || strings.Contains(msg, "namespaceexists")
+}
+
 // checkAndCreateMigrationTable initializes a MongoDB collection if it doesn't exist.
 func (mg mongoMigrator) checkAndCreateMigrationTable(c *container.Container) error {
 	err := mg.Mongo.CreateCollection(context.Background(), mongoMigrationCollection)
-	if err != nil {
-		c.Debug("Migration collection might already exist:", err)
 
+	if err != nil && !isMongoCollectionExistsError(err) {
+		return err
+	}
+
+	err = mg.Mongo.CreateCollection(context.Background(), mongoLockCollection)
+	if err != nil && !isMongoCollectionExistsError(err) {
 		return err
 	}
 
@@ -98,11 +119,118 @@ func (mg mongoMigrator) rollback(c *container.Container, data transactionData) {
 	c.Fatalf("Migration %v failed.", data.MigrationNumber)
 }
 
+func (mg mongoMigrator) startRefresh(ctx context.Context, cancel context.CancelFunc, c *container.Container, ownerID string) {
+	interval := defaultRefresh
+	if mg.testInterval > 0 {
+		interval = mg.testInterval
+	}
+
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			now := time.Now()
+
+			filter := map[string]any{
+				"_id":      mongoLockDocumentID,
+				"lockedBy": ownerID,
+			}
+
+			update := map[string]any{
+				"$set": map[string]any{
+					"lockedAt":  now,
+					"expiresAt": now.Add(defaultLockTTL),
+				},
+			}
+
+			modified, err := mg.Mongo.UpdateOne(ctx, mongoLockCollection, filter, update)
+			if err != nil {
+				c.Errorf("failed to refresh mongo lock: %v", err)
+				cancel()
+
+				return
+			}
+
+			if modified == 0 {
+				c.Errorf("%v", errMongoLockRefreshFailed)
+				cancel()
+
+				return
+			}
+
+			c.Debug("Mongo lock refreshed successfully")
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
 func (mg mongoMigrator) lock(ctx context.Context, cancel context.CancelFunc, c *container.Container, ownerID string) error {
-	return mg.migrator.lock(ctx, cancel, c, ownerID)
+	for i := 0; ; i++ {
+		now := time.Now()
+
+		staleFilter := map[string]any{
+			"_id": mongoLockDocumentID,
+			"expiresAt": map[string]any{
+				"$lte": now,
+			},
+		}
+
+		if _, err := mg.Mongo.DeleteOne(ctx, mongoLockCollection, staleFilter); err != nil {
+			c.Errorf("failed to cleanup stale MongoDB lock: %v", err)
+		}
+
+		lockDoc := map[string]any{
+			"_id":       mongoLockDocumentID,
+			"lockedAt":  now,
+			"lockedBy":  ownerID,
+			"expiresAt": now.Add(defaultLockTTL),
+		}
+
+		_, err := mg.Mongo.InsertOne(ctx, mongoLockCollection, lockDoc)
+		if err == nil {
+			c.Debug("Mongo lock acquired successfully")
+
+			go mg.startRefresh(ctx, cancel, c, ownerID)
+
+			return mg.migrator.lock(ctx, cancel, c, ownerID)
+		}
+
+		if !isDuplicateKeyError(err) {
+			c.Errorf("error while acquiring mongodb lock: %v", err)
+
+			return errLockAcquisitionFailed
+		}
+
+		c.Debugf("MongoDB lock already held, retrying in %v... (attempt %d)", defaultRetry, i+1)
+
+		select {
+		case <-time.After(defaultRetry):
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
 }
 
 func (mg mongoMigrator) unlock(c *container.Container, ownerID string) error {
+	deleted, err := mg.Mongo.DeleteOne(context.Background(), mongoLockCollection, map[string]any{
+		"_id":      mongoLockDocumentID,
+		"lockedBy": ownerID,
+	})
+	if err != nil {
+		c.Errorf("unable to release MongoDB lock: %v", err)
+		return errLockReleaseFailed
+	}
+
+	if deleted == 0 {
+		c.Errorf("failed to release MongoDB lock: lock already released or owned by another instance")
+		return errLockReleaseFailed
+	}
+
+	c.Debug("Mongo lock released successfully")
+
 	return mg.migrator.unlock(c, ownerID)
 }