How to write Multi-Threaded Tests in Java

Spending a weekend at home gives you endless possibilities. We, for example, played board games, watched online courses and even rebuild Stonehenge from clay. But we are not multi-core CPUs that can do all these things at the same time (parallelism). Our human attention span is more like a Thread that has to switch from task to task efficiently (concurrency). And clay mixing blocks my mind from pretty much any other activity!

And threads are what this post is about. This post tries to give you an example how to write tests with multiple threads. You could use it to prove that your application is thread-safe. Or like in my case: Ensure that a database lock is working how it's supposed to be.

Setup

Imagine having a resource that can be accessed by virtually everyone but should not be accessed by two at the same time. Say: toilet paper in the super market.

@Entity
@Data
public class ToiletPaper {
    @Id
    @GeneratedValue
    long id;
    boolean available = true;
}

As you see, our ToiletPaper is simply identified by an ID and carries information on whether it is available or not.

It is always good to have a repository full of toilet paper! We want to keep it simple and define two functions that fetch an available pack of toilet paper. The only difference is that one will have a PESSIMISTIC WRITE lock implemented.

@Repository
public interface ToiletPaperRepository extends JpaRepository<ToiletPaper, Long> {

    Optional<ToiletPaper> findTopByAvailableTrue();

    @Lock(LockModeType.PESSIMISTIC_WRITE)
    Optional<ToiletPaper> findFirstByAvailableTrue();
}

This lock will prevent any other process to read, update or delete data once a thread gets hold of it.

Once we grabbed a toilet paper roll, we update its availability to false. Note that this toilet paper instance will not be found by the above queries anymore.

@Transactional
public ToiletPaper grabToiletPaper() {
    return toiletPaperRepository.findTopByAvailableTrue()
            .map(this::updateToiletPaperToUnavailable)
            .orElseThrow(OutOfToiletPaperException::new);
}

private ToiletPaper updateToiletPaperToUnavailable(ToiletPaper toiletPaper) {
    toiletPaper.setAvailable(false);
    return toiletPaperRepository.save(toiletPaper);
}

Keep in mind that there is a similar grabToiletPaper method that will invoke the Lock-implemented version.

A multi-threaded test

Now, let us simulate two real-life supermarket situations:

• Everyone reaches out for toilet paper at the same time.
• If someone accesses toilet paper, it cannot be accessed by someone else.

Each supermarket customer will be represented by a Thread. And the instructions s*he will execute is defined in a class called GreedyRunner or PatientRunner that both will implement the Runnable interface like this:

@AllArgsConstructor
class GreedyWorker implements Runnable {
    private CountDownLatch threadsReadyCounter;
    private CountDownLatch threadsCalledBlocker;
    private CountDownLatch threadsCompletedCounter;

    @Override
    public void run() {
        long threadId = Thread.currentThread().getId();
        threadsReadyCounter.countDown();
        log.info("Thread-{} ready!", threadId);
        try {
            threadsCalledBlocker.await();
            try {
                ToiletPaper toiletPaper = toiletPaperService.grabToiletPaper();
                log.info("Thread-{} got Toilet Paper no. {}!", threadId, toiletPaper.getId());
            } catch (OutOfToiletPaperException ootpe) {
                log.info("No Toilet Paper in stock!");
            }
        } catch (InterruptedException ie) {
            log.info("Thread-{} got interrupted!", threadId);
        } finally {
            threadsCompletedCounter.countDown();
        }

    }
}

You see that we defined three so-called CountDownLatch objects. You can imagine these like Countdowns. You can define a number, for example 3. Calling the await() function makes the thread wait until the CountDownLatch reaches 0. So, if we were to call the countdown() three times, the execution would still be suspended during the first two countdown() calls but continue as soon as the third countdown(). So, what such a worker does is:

• countdown threadsReadyCounter - signalling it is ready to go
• wait until threadsCalledBlocker reaches 0
• grab toilet paper and log whether that was successful
• countdown threadsCompletedCounter - signalling it has completed its task

If all workers execute this logic, we can ensure that all threads will start at the same time. This can show us how the database lock will influence the result. And this is how the multi-threaded test would look like:

@Test
void multiThreadedGrabToiletPaper_NoLock() throws InterruptedException {
    CountDownLatch readyCounter = new CountDownLatch(NUMBER_OF_THREADS);
    CountDownLatch callBlocker = new CountDownLatch(1);
    CountDownLatch completeCounter = new CountDownLatch(NUMBER_OF_THREADS);

    List<Thread> workers = Stream
            .generate(() -> new Thread(new GreedyWorker(readyCounter, callBlocker, completeCounter)))
            .limit(NUMBER_OF_THREADS)
            .collect(Collectors.toList());

    workers.forEach(Thread::start);

    readyCounter.await();
    log.info("Open the Toilet Paper Hunt!");
    callBlocker.countDown();
    completeCounter.await();
    log.info("Hunt ended!");
}

First, we initialize the CountDownLatches. We want to ensure all workers start at the same time, so we define a countdown starting at NUMBER_OF_THREADS. We create the same amount of workers and start them up. As we saw earlier, each one will countdown on the readyCounter and the test will continue execution once every worker is ready.

Now, we countdown the callBlocker which in turn makes each worker start their toilet paper grabbing logic! And wait until the mayham is over. This is what happens:

5 GreedyWorkers - 3 ToiletPaper in database - No Lock on database:

[           main] c.s.toiletpaperrush.ToiletPaperIT: Open the Toilet Paper Hunt!
[       Thread-3] c.s.toiletpaperrush.ToiletPaperIT: Thread-30 ready!
[       Thread-1] c.s.toiletpaperrush.ToiletPaperIT: Thread-28 ready!
[       Thread-5] c.s.toiletpaperrush.ToiletPaperIT: Thread-32 ready!
[       Thread-2] c.s.toiletpaperrush.ToiletPaperIT: Thread-29 ready!
[       Thread-4] c.s.toiletpaperrush.ToiletPaperIT: Thread-31 ready!
[       Thread-2] c.s.toiletpaperrush.ToiletPaperIT: Thread-29 got Toilet Paper no. 1!
[       Thread-3] c.s.toiletpaperrush.ToiletPaperIT: Thread-30 got Toilet Paper no. 1!
[       Thread-5] c.s.toiletpaperrush.ToiletPaperIT: Thread-32 got Toilet Paper no. 1!
[       Thread-4] c.s.toiletpaperrush.ToiletPaperIT: Thread-31 got Toilet Paper no. 1!
[       Thread-1] c.s.toiletpaperrush.ToiletPaperIT: Thread-28 got Toilet Paper no. 1!
[           main] c.s.toiletpaperrush.ToiletPaperIT: Hunt ended!

5 PatientWorkers - 3 ToiletPaper in database - Lock on database:

[           main] c.s.toiletpaperrush.ToiletPaperIT: Open the Toilet Paper Hunt!
[       Thread-8] c.s.toiletpaperrush.ToiletPaperIT: Thread-35 ready!
[       Thread-6] c.s.toiletpaperrush.ToiletPaperIT: Thread-33 ready!
[       Thread-9] c.s.toiletpaperrush.ToiletPaperIT: Thread-36 ready!
[      Thread-10] c.s.toiletpaperrush.ToiletPaperIT: Thread-37 ready!
[       Thread-7] c.s.toiletpaperrush.ToiletPaperIT: Thread-34 ready!
[       Thread-8] c.s.toiletpaperrush.ToiletPaperIT: Thread-35 got Toilet Paper no. 4!
[      Thread-10] c.s.toiletpaperrush.ToiletPaperIT: Thread-37 got Toilet Paper no. 5!
[       Thread-9] c.s.toiletpaperrush.ToiletPaperIT: Thread-36 got Toilet Paper no. 6!
[       Thread-6] c.s.toiletpaperrush.ToiletPaperIT: No Toilet Paper in stock!
[       Thread-7] c.s.toiletpaperrush.ToiletPaperIT: No Toilet Paper in stock!
[           main] c.s.toiletpaperrush.ToiletPaperIT: Hunt ended!

We see that the greedy ones got grasp on the same toilet paper roll although there were still others left. Patient workers - the case where only one thread could access one ToiletPaper at a time - grabbed toilet paper until the database was empty. But it was ensured that no two threads could access the same entity simultaneously.

I hope that you got some insights on how to use Runnables and CountDownLatches to write a multi-threaded test. For everyone interested, I was inspired by this example. And maybe you will remember this the next time you go on your toilet paper scavenger hunt in these trying times! 😉