Skip to content

Instantly share code, notes, and snippets.

@rponte
Last active February 7, 2024 02:36
Show Gist options
  • Save rponte/1a31395f58de1cd189daae0a358cec20 to your computer and use it in GitHub Desktop.
Save rponte/1a31395f58de1cd189daae0a358cec20 to your computer and use it in GitHub Desktop.
JPA and Hibernate: Simple and Smart way of using PostgreSQL Advisory Locks with JPQL to prevent Lost Update anomaly
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);
}
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);
}
}
@rponte
Copy link
Author

rponte commented Nov 8, 2023

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

}

@rponte
Copy link
Author

rponte commented Nov 10, 2023

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 😉

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment