hyapp-server/services/wallet-service/internal/storage/mysql/account.go

package mysql

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"math"
	"sort"
	"strings"

	"hyapp/pkg/appcode"
	"hyapp/pkg/xerr"
	"hyapp/services/wallet-service/internal/domain/ledger"
)

type walletAccount struct {
	AppCode         string
	UserID          int64
	AssetType       string
	AvailableAmount int64
	FrozenAmount    int64
	Version         int64
	UpdatedAtMS     int64
}

type walletAccountLockRequest struct {
	UserID          int64
	AssetType       string
	CreateIfMissing bool
}

type walletAccountState struct {
	account walletAccount
}

func walletAccountLockKey(userID int64, assetType string) string {
	return fmt.Sprintf("%s:%d", strings.ToUpper(strings.TrimSpace(assetType)), userID)
}

func (r *Repository) lockAccountsOrdered(ctx context.Context, tx *sql.Tx, requests []walletAccountLockRequest, nowMs int64) (map[string]*walletAccountState, error) {
	merged := make(map[string]walletAccountLockRequest, len(requests))
	for _, request := range requests {
		request.AssetType = strings.ToUpper(strings.TrimSpace(request.AssetType))
		if request.UserID <= 0 || strings.TrimSpace(request.AssetType) == "" {
			return nil, xerr.New(xerr.InvalidArgument, "wallet account lock request is invalid")
		}
		if !ledger.ValidAssetType(request.AssetType) {
			return nil, xerr.New(xerr.InvalidArgument, "wallet asset_type is invalid")
		}
		key := walletAccountLockKey(request.UserID, request.AssetType)
		if existing, ok := merged[key]; ok {

			existing.CreateIfMissing = existing.CreateIfMissing || request.CreateIfMissing
			merged[key] = existing
			continue
		}
		merged[key] = request
	}

	ordered := make([]walletAccountLockRequest, 0, len(merged))
	for _, request := range merged {
		ordered = append(ordered, request)
	}
	sort.Slice(ordered, func(i, j int) bool {
		if ordered[i].AssetType == ordered[j].AssetType {
			return ordered[i].UserID < ordered[j].UserID
		}
		return ordered[i].AssetType < ordered[j].AssetType
	})

	states := make(map[string]*walletAccountState, len(ordered))
	for _, request := range ordered {
		account, err := r.lockAccount(ctx, tx, request.UserID, request.AssetType, request.CreateIfMissing, nowMs)
		if err != nil {
			return nil, err
		}
		states[walletAccountLockKey(request.UserID, request.AssetType)] = &walletAccountState{account: account}
	}
	return states, nil
}

func (r *Repository) applyTrackedAccountDelta(ctx context.Context, tx *sql.Tx, state *walletAccountState, availableDelta int64, frozenDelta int64, nowMs int64) (walletAccount, error) {
	if state == nil {
		return walletAccount{}, xerr.New(xerr.Internal, "wallet account state is missing")
	}
	before := state.account
	if err := r.applyAccountDelta(ctx, tx, before, availableDelta, frozenDelta, nowMs); err != nil {
		return walletAccount{}, err
	}
	after := before
	after.AvailableAmount += availableDelta
	after.FrozenAmount += frozenDelta
	after.Version++
	after.UpdatedAtMS = nowMs
	state.account = after
	return after, nil
}

func (r *Repository) lockAccount(ctx context.Context, tx *sql.Tx, userID int64, assetType string, createIfMissing bool, nowMs int64) (walletAccount, error) {
	account, exists, err := r.queryAccountForUpdate(ctx, tx, userID, assetType)
	if err != nil || exists || !createIfMissing {
		if err != nil {
			return walletAccount{}, err
		}
		if !exists {
			return walletAccount{}, xerr.New(xerr.InsufficientBalance, "insufficient balance")
		}
		return account, nil
	}

	if _, err := tx.ExecContext(ctx,
		`INSERT INTO wallet_accounts (app_code, user_id, asset_type, available_amount, frozen_amount, version, created_at_ms, updated_at_ms)
		 VALUES (?, ?, ?, 0, 0, 1, ?, ?)`,
		appcode.FromContext(ctx), userID, assetType, nowMs, nowMs,
	); err != nil {
		return walletAccount{}, err
	}
	return r.queryRequiredAccountForUpdate(ctx, tx, userID, assetType)
}

func (r *Repository) queryAccountForUpdate(ctx context.Context, tx *sql.Tx, userID int64, assetType string) (walletAccount, bool, error) {
	row := tx.QueryRowContext(ctx,
		`SELECT app_code, user_id, asset_type, available_amount, frozen_amount, version, updated_at_ms
		 FROM wallet_accounts
		 WHERE app_code = ? AND user_id = ? AND asset_type = ?
		 FOR UPDATE`,
		appcode.FromContext(ctx),
		userID, assetType,
	)
	var account walletAccount
	if err := row.Scan(&account.AppCode, &account.UserID, &account.AssetType, &account.AvailableAmount, &account.FrozenAmount, &account.Version, &account.UpdatedAtMS); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return walletAccount{}, false, nil
		}
		return walletAccount{}, false, err
	}
	return account, true, nil
}

func (r *Repository) queryRequiredAccountForUpdate(ctx context.Context, tx *sql.Tx, userID int64, assetType string) (walletAccount, error) {
	account, exists, err := r.queryAccountForUpdate(ctx, tx, userID, assetType)
	if err != nil {
		return walletAccount{}, err
	}
	if !exists {
		return walletAccount{}, xerr.New(xerr.Internal, "wallet account creation failed")
	}
	return account, nil
}

func (r *Repository) applyAccountDelta(ctx context.Context, tx *sql.Tx, account walletAccount, availableDelta int64, frozenDelta int64, nowMs int64) error {
	availableAfter := account.AvailableAmount + availableDelta
	frozenAfter := account.FrozenAmount + frozenDelta
	if availableAfter < 0 || frozenAfter < 0 || availableAfter > math.MaxInt64-frozenAfter {

		return xerr.New(xerr.LedgerConflict, "wallet balance delta is invalid")
	}

	result, err := tx.ExecContext(ctx,
		`UPDATE wallet_accounts
		 SET available_amount = ?, frozen_amount = ?, version = version + 1, updated_at_ms = ?
		 WHERE app_code = ? AND user_id = ? AND asset_type = ? AND version = ?`,
		availableAfter,
		frozenAfter,
		nowMs,
		account.AppCode,
		account.UserID,
		account.AssetType,
		account.Version,
	)
	if err != nil {
		return err
	}
	rows, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if rows != 1 {
		return xerr.New(xerr.LedgerConflict, "wallet account version conflict")
	}
	return nil
}