Using Software Transactional Memory in Java

Software transactional memory is a fantastic tool to handle concurrency, avoiding the common pitfalls of deadlocks, live locks, and lack of composability of atomic operations.

It’s based on the concept of transaction from the world of databases, which is very well known in the software community.

Database transactions follow the ACID principle, which stands for Atomicity, Consistency, Isolation, and Durability, which are essential properties that ensure database transactions are reliable and maintain data integrity.

In the Java world there’s the Multiverse library.

In this article we’ll create code to handle the balance of a bank account, which is a known use case where concurrency can go very wrong.

The full source code can be found on my Github.

First of all, we have to add the multiverse dependency to the POM file:

    <dependency>
        <groupId>org.multiverse</groupId>
        <artifactId>multiverse-core</artifactId>
        <version>0.7.0</version>
    </dependency>

Now we can create the Account class, using the primitives from the Multiverse library. Before that, we create an enum for the Withdrawal process:

public enum WithdrawalResult {
  OK,
  INSUFFICIENT_BALANCE
}

And now the Account class:

import org.multiverse.api.StmUtils;
import org.multiverse.api.references.TxnInteger;

public class Account {
    final private TxnInteger balance;

    public Account(int balance) {
        this.balance = StmUtils.newTxnInteger(balance);
    }

    public Integer balance() {
        return balance.atomicGet();
    }

    public void deposit(int amount) {
        balance.increment(amount);
    }

    public WithdrawalResult withdraw(int amount) {
        int withdrawAmount = -amount;
        return StmUtils.atomic(() -> {
            // Read current balance
            int current = balance.get(); 
            // Check BEFORE changing               
            if (current + withdrawAmount < 0) {                 
                return WithdrawalResult.INSUFFICIENT_BALANCE;
            }

            balance.decrement(amount);
            return WithdrawalResult.OK;
        });
    }

    public void transferTo(Account other, int amount) {
        StmUtils.atomic(() -> {
            WithdrawalResult result = withdraw(amount);
            // if there was balance
            if (result == WithdrawalResult.OK) {
                other.deposit(amount);
            }
        });
    }

}

And now we can test the code. First we check if a valid withdrawal works:

    @Test
    public void shouldBeOkWhenHasFunds() {
        var a = new Account(10);
        var result = a.withdraw(5);

        assertThat(a.balance(), equalTo(5));
        assertThat(result, equalTo(WithdrawalResult.OK));
    }

Next, we check if an invalid withdrawal fails due to insufficient funds.

    @Test
    public void shouldBeInsufficientBalanceWhenNoFunds() {
        var a = new Account(10);
        var result = a.withdraw(11);

        assertThat(a.balance(), equalTo(10));
        assertThat(result, equalTo(WithdrawalResult.INSUFFICIENT_BALANCE));
    }

And now we can check something tastier. Will the account go into an invalid state (negative balance) when 2 threads are competing to perform an withdrawal?

    @Test
    public void shouldNotWithdrawTwiceWhenCompetingThreads() throws InterruptedException {
        var exec = Executors.newFixedThreadPool(2);
        var a = new Account(10);
        var startSignal = new CountDownLatch(1);
        var results = 
          Collections.synchronizedList(new ArrayList<>());

        exec.submit(() -> {
            await(startSignal);
            results.add(a.withdraw(6));
        });
        exec.submit(() -> {
            await(startSignal);
            results.add(a.withdraw(5));
        });

        startSignal.countDown();
        exec.awaitTermination(1, TimeUnit.SECONDS);
        exec.shutdown();

        long okCount = results.stream()
          .filter(r -> r == WithdrawalResult.OK)
          .count();
        long insufficientCount = results.stream()
          .filter(r -> 
             r == WithdrawalResult.INSUFFICIENT_BALANCE
          )
          .count();

        assertThat(okCount, is(1L));
        assertThat(insufficientCount, is(1L));
    }

And now we can check transfers between accounts. A valid transfer should work:

    @Test
    public void shouldBeOkWhenOkTransfer() {
        var a = new Account(10);
        var b = new Account(10);
        a.transferTo(b, 5);

        assertThat(a.balance(), is(5));
        assertThat(b.balance(), is(15));
    }

And even when there are competing threads with multiple transfers at the same time, it should still work:

    @Test
    public void shouldTransferOkWhenCompetingThreads() 
      throws InterruptedException {
        var exec = Executors.newFixedThreadPool(2);
        var a = new Account(10);
        var b = new Account(10);
        var startSignal = new CountDownLatch(1);

        exec.submit(() -> {
            await(startSignal);
            a.transferTo(b, 10);
        });
        exec.submit(() -> {
            await(startSignal);
            b.transferTo(a, 1);
        });

        startSignal.countDown();
        exec.awaitTermination(1, TimeUnit.SECONDS);
        exec.shutdown();


        assertThat(a.balance(), is(1));
        assertThat(b.balance(), is(19));
    }

Leave a comment