JPA and Hibernate: Simple and Smart way of using PostgreSQL Advisory Locks with JPQL to prevent Lost Update anomaly

AccountRepository.java

package br.com.stackspot.nullbank.withdrawal;

import org.hibernate.LockOptions;
import org.springframework.data.jpa.repository.*;
import org.springframework.stereotype.Repository;

import javax.persistence.LockModeType;
import javax.persistence.QueryHint;
import javax.transaction.Transactional;
import java.util.Optional;

@Repository
public interface AccountRepository extends JpaRepository<Account, Long> {

    /**
     * Loads the entity even when a lock is not acquired
     */ 
    @Transactional
    @Query(value = """
                   select new br.com.stackspot.nullbank.withdrawal.LockableAccount(
                          c
                         ,pg_try_advisory_xact_lock(c.id)
                         )
                     from Account c
                    where c.id = :accountId
                      and pg_try_advisory_xact_lock(c.id) is not null
                   """
    )
    public Optional<LockableAccount> findByIdWithPessimisticAdvisoryLocking(Long accountId);

}

PessimisticLockingWithAdvisoryLockInQueryATMService.java

package br.com.stackspot.nullbank.withdrawal;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.retry.annotation.Backoff;
import org.springframework.retry.annotation.Retryable;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Service
public class PessimisticLockingWithAdvisoryLockInQueryATMService {

    @Autowired
    private AccountRepository repository;
    @Autowired
    private TransactionRepository transactionRepository;

    @Retryable(
            value = FailedToAcquireLockForAccountException.class,
            maxAttempts = 3,
            backoff = @Backoff(delay = 100, random = true, multiplier = 2.0)
    )
    @Transactional
    public void withdraw(Long accountId, double amount) {

        // We load the entity even if a lock is not acquired
        LockableAccount lockedAccount = repository.findByIdWithPessimisticAdvisoryLocking(accountId).orElseThrow(() -> {
            throw new IllegalStateException("account does not exist: " + accountId);
        });

        // But the business logic is executed only if the lock was acquired for the account
        Account account = lockedAccount
                .getAccountIfLockedOrElseThrow();

        double newBalance = (account.getBalance() - amount);
        if (newBalance < 0) {
            throw new IllegalStateException("there's not enough balance");
        }

        account.setBalance(newBalance);
        repository.save(account);

        transactionRepository
                .save(new Transaction(account, amount, "withdraw"));
    }

}

/**
 * Represents an account that may be locked or not
 */
class LockableAccount {
    private Account account;
    private boolean locked;

    public LockableAccount(Account account, boolean locked) {
        this.account = account;
        this.locked = locked;
    }

    /**
     * Returns the actual account if it was locked or else throws an {@code AccountNotFoundOrLockNotAcquiredException}
     */
    public Account getAccountIfLockedOrElseThrow() {
        if (!locked) {
            throw new FailedToAcquireLockForAccountException("Account already locked by another user");
        }
        return account;
    }
    
    public boolean isLocked() {
        return locked;
    }
}

class FailedToAcquireLockForAccountException extends RuntimeException {

    public FailedToAcquireLockForAccountException(String message) {
        super(message);
    }
}

If for any reason Hibernate can not execute the JPQL informing the pg_try_advisory_xact_lock() function does not exist, you can register this function via MetadataBuilderContributor:

import org.hibernate.boot.MetadataBuilder;
import org.hibernate.boot.spi.MetadataBuilderContributor;
import org.hibernate.dialect.function.StandardSQLFunction;
import org.hibernate.type.StandardBasicTypes;
import org.springframework.boot.autoconfigure.orm.jpa.HibernatePropertiesCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * It seems like it's not necessary on Hibernate v5.6.10. 
 * I mean, hibernate can understand this function correctly in JPQL queries.
 */
@Configuration
public class AdvisoryLockSqlFunctionsBuilderContributor implements MetadataBuilderContributor {

    @Override
    public void contribute(MetadataBuilder metadataBuilder) {
        metadataBuilder.applySqlFunction(
                "pg_try_advisory_xact_lock",
                new StandardSQLFunction(
                        "pg_try_advisory_xact_lock",
                        StandardBasicTypes.BOOLEAN
                )
        );
    }

    @Bean
    public HibernatePropertiesCustomizer hibernatePropertiesCustomizer() {
        return hibernateProperties -> {
            hibernateProperties.put(
                    "hibernate.metadata_builder_contributor",
                    AdvisoryLockSqlFunctionsBuilderContributor.class.getName()
            );
        };
    }

}

The previous code acquired a lock with the pg_try_advisory_xact_lock() function using the account ID as a key (just a number), which is brittle since this ID number could be anything. Ideally, we should give some context for the key, something like "account:42" or "product:1234" for example.

Since the pg_try_advisory_xact_lock() function does not support text as argument, we must generate a int/bigint hash for any text informed as key, like "account:42". We could generate a hash for text using the Postgres built-in function pg_catalog.hashtextextended(text, seed) as below:

@Repository
public interface AccountRepository extends JpaRepository<Account, Long> {

    @Transactional
    @Query(value = """
                   select c
                     from Account c
                    where c.id = :accountId
                      and pg_try_advisory_xact_lock(
                                pg_catalog.hashtextextended('account', c.id)
                          ) is true
                   """
    )
    public Optional<Account> findByIdWithPessimisticAdvisoryLocking(Long accountId);

}

Here are other examples of using this function:

select 
    id,
    pg_catalog.hashtextextended('account', id)      as hash_key,
    pg_catalog.hashtextextended('account' || id, 0) as hash_key_with_fixed_seed,
    pg_catalog.hashtextextended('account' || id, (pg_backend_pid() * random())::int) as hash_key_with_random_seed
  from pg_catalog.generate_series(1, 10) as id
  ;

And the query result on Postgres v14.5:

id|hash_key            |hash_key_with_fixed_seed|hash_key_with_random_seed|
--+--------------------+------------------------+-------------------------+
 1| 4789936018558587555|     1611349570782493084|      3918960573832768574|
 2|-9008359969644079684|    -1031716354870497645|     -7807373867093230547|
 3|  815680793056945129|    -8779373722393434404|     -1138859681579334640|
 4|-4958287634170775584|       83050773856503466|      8954598227837031082|
 5| 2494949060428457094|     7761777444195815969|      9039206469815446715|
 6|  989566638122565597|     7929002698029429524|     -3138848592652980239|
 7|-1505647370073910025|     9181563818372204761|     -8062445020047817409|
 8| 2831380095553820023|      151707059225141879|     -5656928723741106830|
 9|  814742393734240117|    -3045715879349183706|      2233070228472412733|
10| 6704697263975292804|     6266128781436705289|     -6153728191813298609|

‼️ Just be aware of that hash_key_with_random_seed may not be useful for this scenario, since it changes the key value every time it runs 😉