Fix negative balance from concurrent transfers with pessimistic row locks

src/main/java/com/example/bankapp/repository/AccountRepository.java

@@ -1,10 +1,18 @@
 package com.example.bankapp.repository;
 
 import com.example.bankapp.model.Account;
+import jakarta.persistence.LockModeType;
 import org.springframework.data.jpa.repository.JpaRepository;
+import org.springframework.data.jpa.repository.Lock;
+import org.springframework.data.jpa.repository.Query;
+import org.springframework.data.repository.query.Param;
 
 import java.util.Optional;
 
 public interface AccountRepository extends JpaRepository<Account, Long> {
     Optional<Account> findByUsername(String username);
+
+    @Lock(LockModeType.PESSIMISTIC_WRITE)
+    @Query("select a from Account a where a.id = :id")
+    Optional<Account> findByIdForUpdate(@Param("id") Long id);
 }

src/main/java/com/example/bankapp/service/AccountService.java

@@ -4,6 +4,8 @@
 import com.example.bankapp.model.Transaction;
 import com.example.bankapp.repository.AccountRepository;
 import com.example.bankapp.repository.TransactionRepository;
+import jakarta.persistence.EntityManager;
+import jakarta.persistence.PersistenceContext;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.security.core.GrantedAuthority;
 import org.springframework.security.core.authority.SimpleGrantedAuthority;
@@ -12,6 +14,7 @@
 import org.springframework.security.core.userdetails.UsernameNotFoundException;
 import org.springframework.security.crypto.password.PasswordEncoder;
 import org.springframework.stereotype.Service;
+import org.springframework.transaction.annotation.Transactional;
 
 import java.math.BigDecimal;
 import java.time.LocalDateTime;
@@ -31,6 +34,9 @@ public class AccountService implements UserDetailsService {
     @Autowired
     private TransactionRepository transactionRepository;
 
+    @PersistenceContext
+    private EntityManager entityManager;
+
     public Account findAccountByUsername(String username) {
         return accountRepository.findByUsername(username).orElseThrow(() -> new RuntimeException("Account not found"));
     }
@@ -48,6 +54,7 @@ public Account registerAccount(String username, String password) {
     }
 
 
+    @Transactional
     public void deposit(Account account, BigDecimal amount) {
         account.setBalance(account.getBalance().add(amount));
         accountRepository.save(account);
@@ -61,18 +68,22 @@ public void deposit(Account account, BigDecimal amount) {
         transactionRepository.save(transaction);
     }
 
+    @Transactional
     public void withdraw(Account account, BigDecimal amount) {
-        if (account.getBalance().compareTo(amount) < 0) {
+        Account lockedAccount = accountRepository.findByIdForUpdate(account.getId())
+                .orElseThrow(() -> new RuntimeException("Account not found"));
+        if (lockedAccount.getBalance().compareTo(amount) < 0) {
             throw new RuntimeException("Insufficient funds");
         }
-        account.setBalance(account.getBalance().subtract(amount));
-        accountRepository.save(account);
+        lockedAccount.setBalance(lockedAccount.getBalance().subtract(amount));
+        account.setBalance(lockedAccount.getBalance());
+        accountRepository.save(lockedAccount);
 
         Transaction transaction = new Transaction(
                 amount,
                 "Withdrawal",
                 LocalDateTime.now(),
-                account
+                lockedAccount
         );
         transactionRepository.save(transaction);
     }
@@ -100,36 +111,57 @@ public Collection<? extends GrantedAuthority> authorities() {
         return Arrays.asList(new SimpleGrantedAuthority("USER"));
     }
 
+    @Transactional
     public void transferAmount(Account fromAccount, String toUsername, BigDecimal amount) {
-        if (fromAccount.getBalance().compareTo(amount) < 0) {
-            throw new RuntimeException("Insufficient funds");
-        }
-
         Account toAccount = accountRepository.findByUsername(toUsername)
                 .orElseThrow(() -> new RuntimeException("Recipient account not found"));
+        Long toAccountId = toAccount.getId();
+        if (fromAccount.getId().equals(toAccountId)) {
+            throw new RuntimeException("Cannot transfer to the same account");
+        }
+        entityManager.detach(toAccount);
+
+        Account lockedFrom;
+        Account lockedTo;
+        if (fromAccount.getId().compareTo(toAccountId) < 0) {
+            lockedFrom = accountRepository.findByIdForUpdate(fromAccount.getId())
+                    .orElseThrow(() -> new RuntimeException("Account not found"));
+            lockedTo = accountRepository.findByIdForUpdate(toAccountId)
+                    .orElseThrow(() -> new RuntimeException("Recipient account not found"));
+        } else {
+            lockedTo = accountRepository.findByIdForUpdate(toAccountId)
+                    .orElseThrow(() -> new RuntimeException("Recipient account not found"));
+            lockedFrom = accountRepository.findByIdForUpdate(fromAccount.getId())
+                    .orElseThrow(() -> new RuntimeException("Account not found"));
+        }
+
+        if (lockedFrom.getBalance().compareTo(amount) < 0) {
+            throw new RuntimeException("Insufficient funds");
+        }
 
         // Deduct from sender's account
-        fromAccount.setBalance(fromAccount.getBalance().subtract(amount));
-        accountRepository.save(fromAccount);
+        lockedFrom.setBalance(lockedFrom.getBalance().subtract(amount));
+        fromAccount.setBalance(lockedFrom.getBalance());
+        accountRepository.save(lockedFrom);
 
         // Add to recipient's account
-        toAccount.setBalance(toAccount.getBalance().add(amount));
-        accountRepository.save(toAccount);
+        lockedTo.setBalance(lockedTo.getBalance().add(amount));
+        accountRepository.save(lockedTo);
 
         // Create transaction records for both accounts
         Transaction debitTransaction = new Transaction(
                 amount,
-                "Transfer Out to " + toAccount.getUsername(),
+                "Transfer Out to " + lockedTo.getUsername(),
                 LocalDateTime.now(),
-                fromAccount
+                lockedFrom
         );
         transactionRepository.save(debitTransaction);
 
         Transaction creditTransaction = new Transaction(
                 amount,
-                "Transfer In from " + fromAccount.getUsername(),
+                "Transfer In from " + lockedFrom.getUsername(),
                 LocalDateTime.now(),
-                toAccount
+                lockedTo
         );
         transactionRepository.save(creditTransaction);
     }

src/test/java/com/example/bankapp/service/AccountServiceConcurrencyTest.java

@@ -0,0 +1,146 @@
+package com.example.bankapp.service;
+
+import com.example.bankapp.model.Account;
+import com.example.bankapp.repository.AccountRepository;
+import com.example.bankapp.repository.TransactionRepository;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+
+import java.math.BigDecimal;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+@SpringBootTest
+class AccountServiceConcurrencyTest {
+
+    @Autowired
+    private AccountService accountService;
+
+    @Autowired
+    private AccountRepository accountRepository;
+
+    @Autowired
+    private TransactionRepository transactionRepository;
+
+    @BeforeEach
+    void setUp() {
+        transactionRepository.deleteAll();
+        accountRepository.deleteAll();
+
+        Account alice = accountService.registerAccount("alice", "password");
+        alice.setBalance(new BigDecimal("100"));
+        accountRepository.save(alice);
+
+        Account bob = accountService.registerAccount("bob", "password");
+        bob.setBalance(BigDecimal.ZERO);
+        accountRepository.save(bob);
+    }
+
+    @Test
+    void concurrentTransfersMustNotOverdraw() throws InterruptedException {
+        int threadCount = 10;
+        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+        CountDownLatch startGate = new CountDownLatch(1);
+        CountDownLatch finished = new CountDownLatch(threadCount);
+        AtomicInteger successes = new AtomicInteger();
+        AtomicInteger failures = new AtomicInteger();
+
+        for (int i = 0; i < threadCount; i++) {
+            executor.submit(() -> {
+                try {
+                    startGate.await();
+                    Account from = accountService.findAccountByUsername("alice");
+                    accountService.transferAmount(from, "bob", new BigDecimal("100"));
+                    successes.incrementAndGet();
+                } catch (RuntimeException e) {
+                    failures.incrementAndGet();
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    failures.incrementAndGet();
+                } finally {
+                    finished.countDown();
+                }
+            });
+        }
+
+        startGate.countDown();
+        assertThat(finished.await(30, TimeUnit.SECONDS))
+                .as("all concurrent transfer tasks should finish within 30 seconds")
+                .isTrue();
+        executor.shutdown();
+
+        Account alice = accountRepository.findByUsername("alice").orElseThrow();
+        Account bob = accountRepository.findByUsername("bob").orElseThrow();
+
+        assertThat(alice.getBalance().compareTo(BigDecimal.ZERO))
+                .as("alice must not have a negative balance after concurrent transfers")
+                .isGreaterThanOrEqualTo(0);
+        assertThat(alice.getBalance().add(bob.getBalance()))
+                .as("alice and bob balances must conserve the initial total")
+                .isEqualByComparingTo("100");
+        assertThat(successes.get())
+                .as("exactly one transfer should succeed from alice's 100 balance")
+                .isEqualTo(1);
+        assertThat(failures.get())
+                .as("all other transfers should fail after alice's funds are exhausted")
+                .isEqualTo(threadCount - 1);
+    }
+
+    @Test
+    void concurrentTransfersInOppositeDirectionsDoNotDeadlock() throws InterruptedException {
+        Account bob = accountRepository.findByUsername("bob").orElseThrow();
+        bob.setBalance(new BigDecimal("100"));
+        accountRepository.save(bob);
+
+        int threadCount = 10;
+        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+        CountDownLatch startGate = new CountDownLatch(1);
+        CountDownLatch finished = new CountDownLatch(threadCount);
+        ConcurrentLinkedQueue<Throwable> unexpectedErrors = new ConcurrentLinkedQueue<>();
+
+        for (int i = 0; i < threadCount; i++) {
+            boolean aliceToBob = i < threadCount / 2;
+            executor.submit(() -> {
+                try {
+                    startGate.await();
+                    String fromUsername = aliceToBob ? "alice" : "bob";
+                    Account from = accountService.findAccountByUsername(fromUsername);
+                    accountService.transferAmount(from, aliceToBob ? "bob" : "alice", new BigDecimal("10"));
+                } catch (RuntimeException e) {
+                    if (!"Insufficient funds".equals(e.getMessage())) {
+                        unexpectedErrors.add(e);
+                    }
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    unexpectedErrors.add(e);
+                } finally {
+                    finished.countDown();
+                }
+            });
+        }
+
+        startGate.countDown();
+        assertThat(finished.await(30, TimeUnit.SECONDS))
+                .as("all opposite-direction transfer tasks should finish without deadlocking")
+                .isTrue();
+        executor.shutdown();
+
+        Account alice = accountRepository.findByUsername("alice").orElseThrow();
+        bob = accountRepository.findByUsername("bob").orElseThrow();
+
+        assertThat(unexpectedErrors)
+                .as("opposite-direction transfers must not produce unexpected exceptions")
+                .isEmpty();
+        assertThat(alice.getBalance().add(bob.getBalance()))
+                .as("alice and bob balances must conserve the initial total")
+                .isEqualByComparingTo("200");
+    }
+}