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

AccountRepository.java

package br.com.rponte.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.rponte.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.rponte.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 FailedToAcquireLockForAccountException}
     */
    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);
    }
}

Hibernate 5.x and 6.x

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
                )
        );
    }

    /**
     * Registers our custom SQL function during the application bootstrap
     */
    @Bean
    public HibernatePropertiesCustomizer hibernatePropertiesCustomizer() {
        return hibernateProperties -> {
            hibernateProperties.put(
                    "hibernate.metadata_builder_contributor",
                    AdvisoryLockSqlFunctionsBuilderContributor.class.getName()
            );
        };
    }

}

Hibernate 6.x new versions

On Hibernate 6.x, the MetadataBuilderContributor interface was deprecated and will be removed soon. So, you should use the FunctionContributor instead, as we can see below:

import org.hibernate.boot.model.FunctionContributor;

/**
 * We must register it via {@code /resources/META-INF/services/org.hibernate.boot.model.FunctionContributor}
 * file with its fully qualified class name.
 *
 * For more details, this article can help:
 * https://aregall.tech/hibernate-6-custom-functions#heading-how-to-register-custom-functions-in-hibernate-6
 */
@Configuration
public class AdvisoryLockSqlFunctionsBuilderContributor implements FunctionContributor {

    @Override
    public void contributeFunctions(FunctionContributions functionContributions) {
        functionContributions.getFunctionRegistry()
                .register(
                        "pg_try_advisory_xact_lock",
                        new StandardSQLFunction(
                                "pg_try_advisory_xact_lock",
                                StandardBasicTypes.BOOLEAN
                        )
                );
    }

    /**
     * ⚠️ It does NOT work with Spring Boot 3.2.0
     * https://discourse.hibernate.org/t/migrate-hibernate-5-to-6-with-spring-boot-2-7-x-to-3/7787/2
     */
    @Bean
    public HibernatePropertiesCustomizer hibernatePropertiesCustomizer() {
        return hibernateProperties -> {
            hibernateProperties.put(
                    "hibernate.function_contributor", // 😭😭
                    AdvisoryLockSqlFunctionsBuilderContributor.class.getName()
            );
        };
    }

}

Finally, we must register our custom SQL function. Unfortunately, I've not found a good way to do it programmatically with Spring Boot 3.x as I've done using the HibernatePropertiesCustomizer bean. That's why we must register it via Java Service Locator in the /resources/META-INF/services/org.hibernate.boot.model.FunctionContributor file with its fully qualified class name:

br.com.rponte.nullbank.shared.hibernate.AdvisoryLockSqlFunctionsBuilderContributor

For more details, this article can help: Hibernate 6: custom SQL functions in Spring Data JPA repositories

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 😉

Discussion with Francky Pachot on Twitter about using the pg_try_advisory_xact_lock () function on both the WHERE and SELECT clauses and their possible corner cases.

Francky brought up an issue when using the function on the WHERE clause without an index (or when it's not used):

Ok. There's always a risk when the behavior depends on when the function is evaluated because an execution plan can change.
Maybe it is safer within a CASE statement.
Example:
https://dbfiddle.uk/QUP_t0T7
(that's for WHERE - not sure if the same can happen with SELECT)

Below, we can see the DB Fiddle:

According to Francky, maybe we should use a CASE statement to prevent this issue:

SELECT id, value, pg_try_advisory_xact_lock(id)
  FROM demo
 WHERE id=2 
   AND CASE WHEN id=2 THEN pg_try_advisory_xact_lock(id) END is not null ;