fix(dbfs): enforce storage pre allocate before creating file placeholders

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

.gitignore

@@ -10,6 +10,7 @@
 application/statics/assets.zip
 master_v4
 slave_v4
+__debug*
 
 # Test binary, build with `go test -c`
 *.test

inventory/tx.go

@@ -26,6 +26,42 @@ type (
 	TxCtx struct{}
 )
 
+// ReserveStorage atomically adds `size` bytes to user `uid`'s storage inside
+// this transaction, enforcing storage + size <= maxTotal at the database
+// level (when maxTotal > 0). This is the correct API for grabbing the
+// per-user storage quota:
+//
+//   - It MUST be called as the first write of the transaction (or at least
+//     before any other WRITE to files/entities/etc.). That way the users-row
+//     X lock is acquired before any INSERT into files/entities, so
+//     concurrent uploads for the same owner serialize on the users row from
+//     the start and cannot deadlock via foreign-key S locks on the shared
+//     parent folder / storage-policy rows.
+//   - A matching compensating negative diff is appended to storageDiff so
+//     the pre-reservation is netted against the positive diff emitted later
+//     by CreateEntity / fc.Copy / etc. No amount is double-applied.
+//   - On failure (quota exceeded or any other DB error) the caller MUST
+//     Rollback the transaction: any placeholder writes done afterwards
+//     would otherwise be committed without a matching storage reservation.
+//
+// It is safe to call this from an inherited-tx wrapper (returned by
+// InheritTx) — AppendStorageDiff routes to the root tx.
+func (t *Tx) ReserveStorage(ctx context.Context, uc UserClient, uid int, size, maxTotal int64) error {
+	if size <= 0 {
+		return nil
+	}
+
+	txUc, _ := InheritTx(ctx, uc)
+	if err := txUc.ReserveStorage(ctx, uid, size, maxTotal); err != nil {
+		return err
+	}
+
+	// Net out the pre-reservation against any positive diff that
+	// CreateEntity / fc.Copy / ... will append later in the same tx.
+	t.AppendStorageDiff(StorageDiff{uid: -size})
+	return nil
+}
+
 // AppendStorageDiff appends the given storage diff to the transaction.
 func (t *Tx) AppendStorageDiff(diff StorageDiff) {
 	root := t
@@ -99,14 +135,24 @@ func Commit(tx *Tx) error {
 	return err
 }
 
-// CommitWithStorageDiff commits the transaction and applies the storage diff, only if the transaction is not inherited.
+// CommitWithStorageDiff commits the transaction and applies the accumulated
+// storage diff. Only the outermost (non-inherited) transaction performs the
+// commit and the storage mutation.
+//
+// Quota enforcement is done UPFRONT via Tx.ReserveStorage, not here. This
+// function's role is to (a) commit the tx, and (b) apply any residual
+// storage diff produced by CreateEntity / CapEntities / etc. Tx.ReserveStorage
+// pre-emits a matching negative diff, so pre-reserved sizes are netted out
+// against the positive diff CreateEntity appends later. Residual is applied
+// AFTER commit via ApplyStorageDiff (auto-commit, safe to retry on
+// deadlock).
 func CommitWithStorageDiff(ctx context.Context, tx *Tx, l logging.Logger, uc UserClient) error {
-	commited, err := commit(tx)
+	committed, err := commit(tx)
 	if err != nil {
 		return err
 	}
 
-	if !commited {
+	if !committed {
 		return nil
 	}
 

inventory/user.go

@@ -42,6 +42,10 @@ var (
 	ErrorUnknownPasswordType = errors.New("unknown password type")
 	ErrorIncorrectPassword   = errors.New("incorrect password")
 	ErrInsufficientPoints    = errors.New("insufficient points")
+	// ErrInsufficientCapacity is returned by ReserveStorage when the atomic
+	// UPDATE would push the user's storage past their quota. The FS layer
+	// translates this into fs.ErrInsufficientCapacity for API consumers.
+	ErrInsufficientCapacity = errors.New("insufficient storage capacity")
 )
 
 type (
@@ -69,6 +73,15 @@ type (
 		SearchActive(ctx context.Context, limit int, keyword string) ([]*ent.User, error)
 		// ApplyStorageDiff apply storage diff to user.
 		ApplyStorageDiff(ctx context.Context, diffs StorageDiff) error
+		// ReserveStorage atomically adds size bytes to user uid's storage,
+		// enforcing storage + size <= maxTotal at the database level. When
+		// maxTotal <= 0 the guard is skipped (unlimited quota). Returns
+		// ErrInsufficientCapacity when the guard fails (0 rows affected).
+		// Negative size is treated as an unconditional release.
+		ReserveStorage(ctx context.Context, uid int, size int64, maxTotal int64) error
+		// ReleaseStorage unconditionally subtracts size bytes from user uid's
+		// storage. Used to undo a previous ReserveStorage on failure paths.
+		ReleaseStorage(ctx context.Context, uid int, size int64) error
 		// UpdateAvatar updates user avatar.
 		UpdateAvatar(ctx context.Context, u *ent.User, avatar string) (*ent.User, error)
 		// UpdateNickname updates user nickname.
@@ -226,33 +239,88 @@ func (c *userClient) Delete(ctx context.Context, uid int) error {
 func (c *userClient) ApplyStorageDiff(ctx context.Context, diffs StorageDiff) error {
 	ae := serializer.NewAggregateError()
 	for uid, diff := range diffs {
-		// Retry logic for MySQL deadlock (Error 1213)
-		// This is a temporary workaround. TODO: optimize storage mutation
-		maxRetries := 3
-		var lastErr error
-		for attempt := 0; attempt < maxRetries; attempt++ {
-			if err := c.client.User.Update().Where(user.ID(uid)).AddStorage(diff).Exec(ctx); err != nil {
-				lastErr = err
-				// Check if it's a MySQL deadlock error (Error 1213)
-				if strings.Contains(err.Error(), "Error 1213") && attempt < maxRetries-1 {
-					// Wait a bit before retrying with exponential backoff
-					time.Sleep(time.Duration(attempt+1) * 10 * time.Millisecond)
-					continue
-				}
-				ae.Add(fmt.Sprintf("%d", uid), fmt.Errorf("failed to apply storage diff for user %d: %w", uid, err))
-				break
-			}
-			// Success, break out of retry loop
-			lastErr = nil
-			break
+		err := runWithDeadlockRetry(func() error {
+			return c.client.User.Update().Where(user.ID(uid)).AddStorage(diff).Exec(ctx)
+		})
+		if err != nil {
+			ae.Add(fmt.Sprintf("%d", uid), fmt.Errorf("failed to apply storage diff for user %d: %w", uid, err))
 		}
+	}
 
-		if lastErr != nil {
-			ae.Add(fmt.Sprintf("%d", uid), fmt.Errorf("failed to apply storage diff for user %d: %w", uid, lastErr))
+	return ae.Aggregate()
+}
+
+// isMySQLDeadlock reports whether err is a MySQL Error 1213 deadlock.
+func isMySQLDeadlock(err error) bool {
+	return err != nil && strings.Contains(err.Error(), "Error 1213")
+}
+
+// runWithDeadlockRetry runs fn up to maxRetries times, backing off between
+// attempts on MySQL Error 1213 (deadlock) errors. Any other error is returned
+// immediately.
+//
+// CAUTION: Only safe for statements that run in AUTO-COMMIT mode. A deadlock
+// on a statement running inside an explicit BEGIN/COMMIT causes InnoDB to
+// roll back the ENTIRE transaction server-side; re-issuing the same statement
+// against the dead tx handle will not resurrect the transaction — it will
+// either error out or silently apply outside the tx. Use this helper only
+// for ApplyStorageDiff (which issues one auto-commit UPDATE per user).
+func runWithDeadlockRetry(fn func() error) error {
+	const maxRetries = 3
+	var lastErr error
+	for attempt := 0; attempt < maxRetries; attempt++ {
+		err := fn()
+		if err == nil {
+			return nil
+		}
+		lastErr = err
+		if !isMySQLDeadlock(err) || attempt == maxRetries-1 {
+			return err
 		}
+		time.Sleep(time.Duration(attempt+1) * 10 * time.Millisecond)
 	}
+	return lastErr
+}
 
-	return ae.Aggregate()
+// ReserveStorage atomically adds size bytes to user uid's storage while
+// enforcing an optional quota. See UserClient.ReserveStorage.
+//
+// This method deliberately does NOT retry on MySQL deadlocks. Callers that
+// invoke it via a tx-scoped client (see Tx.ReserveStorage) must let a
+// deadlock propagate so the tx is aborted and the placeholder writes are
+// rolled back with it.
+func (c *userClient) ReserveStorage(ctx context.Context, uid int, size, maxTotal int64) error {
+	if size == 0 {
+		return nil
+	}
+	if size < 0 {
+		return c.ReleaseStorage(ctx, uid, -size)
+	}
+
+	q := c.client.User.Update().Where(user.ID(uid))
+	if maxTotal > 0 {
+		// storage + size <= maxTotal  <=>  storage <= maxTotal - size.
+		// When maxTotal < size the predicate is unsatisfiable (returns
+		// zero rows) and we correctly report ErrInsufficientCapacity.
+		q = q.Where(user.StorageLTE(maxTotal - size))
+	}
+	n, err := q.AddStorage(size).Save(ctx)
+	if err != nil {
+		return err
+	}
+	if n == 0 {
+		return ErrInsufficientCapacity
+	}
+	return nil
+}
+
+// ReleaseStorage unconditionally subtracts size bytes from user uid's
+// storage. Like ReserveStorage, it does not retry on MySQL deadlocks.
+func (c *userClient) ReleaseStorage(ctx context.Context, uid int, size int64) error {
+	if size == 0 {
+		return nil
+	}
+	return c.client.User.Update().Where(user.ID(uid)).AddStorage(-size).Exec(ctx)
 }
 
 func (c *userClient) CalculateStorage(ctx context.Context, uid int) (int64, error) {

pkg/filemanager/fs/dbfs/upload.go

@@ -2,6 +2,7 @@ package dbfs
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"math"
 	"path"
@@ -148,8 +149,12 @@ func (f *DBFS) PrepareUpload(ctx context.Context, req *fs.UploadRequest, opts ..
 		return nil, err
 	}
 
-	// Validate available capacity
-	if err := f.validateUserCapacity(ctx, req.Props.Size, ancestor.Owner()); err != nil {
+	owner := ancestor.Owner()
+	capacity, err := f.Capacity(ctx, owner)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get user capacity: %s", err)
+	}
+	if err := f.validateUserCapacityRaw(ctx, req.Props.Size, capacity); err != nil {
 		return nil, err
 	}
 
@@ -176,6 +181,14 @@ func (f *DBFS) PrepareUpload(ctx context.Context, req *fs.UploadRequest, opts ..
 		return nil, serializer.NewError(serializer.CodeDBError, "Failed to start transaction", err)
 	}
 
+	if err := dbTx.ReserveStorage(ctx, f.userClient, owner.ID, req.Props.Size, capacity.Total); err != nil {
+		_ = inventory.Rollback(dbTx)
+		if errors.Is(err, inventory.ErrInsufficientCapacity) {
+			return nil, fs.ErrInsufficientCapacity
+		}
+		return nil, serializer.NewError(serializer.CodeDBError, "Failed to reserve storage capacity", err)
+	}
+
 	if fileExisted {
 		entityType := types.EntityTypeVersion
 		if req.Props.EntityType != nil {