Add timeout for replication lock query & allow updating metadata with less keys

server/pkg/controller/file.go

@@ -633,7 +633,14 @@ func (c *FileController) validateUpdateMetadataRequest(ctx *gin.Context, req ent
 			}).Error("can't update magic metadata for file which isn't owned by use")
 			return stacktrace.Propagate(ente.ErrPermissionDenied, "")
 		}
-		if existingMetadata != nil && (existingMetadata.Version != updateMMdRequest.MagicMetadata.Version || existingMetadata.Count > updateMMdRequest.MagicMetadata.Count) {
+		oldToNewCountDiff := 0
+		if existingMetadata != nil {
+			oldToNewCountDiff = existingMetadata.Count - updateMMdRequest.MagicMetadata.Count
+		}
+		// Return an error if there is a version mismatch with the previous metadata
+		// or if the new metadata contains an unexpectedly lower number of keys
+		// (oldToNewCountDiff difference is > 2), which may indicate potential data loss due to potentially buggy client.
+		if existingMetadata != nil && (existingMetadata.Version != updateMMdRequest.MagicMetadata.Version || oldToNewCountDiff > 2) {
 			log.WithFields(log.Fields{
 				"existing_count":   existingMetadata.Count,
 				"existing_version": existingMetadata.Version,

server/pkg/controller/replication3.go

@@ -1,6 +1,7 @@
 package controller
 
 import (
+	"context"
 	"database/sql"
 	"encoding/base64"
 	"errors"
@@ -222,7 +223,9 @@ func (c *ReplicationController3) replicate(i int) {
 // objects left to replicate currently.
 func (c *ReplicationController3) tryReplicate() error {
 	// Fetch an object to replicate
-	copies, err := c.ObjectCopiesRepo.GetAndLockUnreplicatedObject()
+	ctxWithTimeout, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	copies, err := c.ObjectCopiesRepo.GetAndLockUnreplicatedObject(ctxWithTimeout)
 	if err != nil {
 		if !errors.Is(err, sql.ErrNoRows) {
 			log.Errorf("Could not fetch an object to replicate: %s", err)

server/pkg/repo/object_copies.go

@@ -22,8 +22,8 @@ type ObjectCopiesRepository struct {
 // to that object to be blocked for 24h before next replication attemp.
 //
 // ObjectCopies is guaranteed to be nil if error is not nil.
-func (repo *ObjectCopiesRepository) GetAndLockUnreplicatedObject() (*ente.ObjectCopies, error) {
-	tx, err := repo.DB.Begin()
+func (repo *ObjectCopiesRepository) GetAndLockUnreplicatedObject(ctx context.Context) (*ente.ObjectCopies, error) {
+	tx, err := repo.DB.BeginTx(ctx, nil)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "")
 	}
@@ -35,7 +35,7 @@ func (repo *ObjectCopiesRepository) GetAndLockUnreplicatedObject() (*ente.Object
 		}
 	}
 
-	row := tx.QueryRow(`
+	row := tx.QueryRowContext(ctx, `
 	SELECT object_key, want_b2, b2, want_wasabi, wasabi, want_scw, scw
 	FROM object_copies
 	WHERE (
@@ -60,7 +60,7 @@ func (repo *ObjectCopiesRepository) GetAndLockUnreplicatedObject() (*ente.Object
 		return nil, stacktrace.Propagate(err, "")
 	}
 
-	err = repo.RegisterReplicationAttempt(tx, r.ObjectKey)
+	err = repo.RegisterReplicationAttempt(tx, ctx, r.ObjectKey)
 	if err != nil {
 		rollback()
 		return nil, stacktrace.Propagate(err, "failed to register replication attempt")
@@ -102,8 +102,8 @@ func (repo *ObjectCopiesRepository) CreateNewWasabiObject(ctx context.Context, t
 
 // RegisterReplicationAttempt sets the last_attempt timestamp so that this row can
 // be skipped over for the next day in case the replication was not succesful.
-func (repo *ObjectCopiesRepository) RegisterReplicationAttempt(tx *sql.Tx, objectKey string) error {
-	_, err := tx.Exec(`
+func (repo *ObjectCopiesRepository) RegisterReplicationAttempt(tx *sql.Tx, ctx context.Context, objectKey string) error {
+	_, err := tx.ExecContext(ctx, `
 	UPDATE object_copies
 	SET last_attempt = now_utc_micro_seconds()
 	WHERE object_key = $1