From 827bb23038d16e7e616c62e9c3df5c6d562ab9d9 Mon Sep 17 00:00:00 2001
From: Andrea Cavalli <andrea@cavallium.it>
Date: Fri, 22 Jan 2021 04:22:58 +0100
Subject: [PATCH] Bugfix

---
 .../warp/commonutils/locks/TransitLock.java   | 20 ++++++++++++---
 .../commonutils/locks/TransitLockTest.java    | 25 +++++++++++--------
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/src/main/java/org/warp/commonutils/locks/TransitLock.java b/src/main/java/org/warp/commonutils/locks/TransitLock.java
index d7f322c..2322000 100644
--- a/src/main/java/org/warp/commonutils/locks/TransitLock.java
+++ b/src/main/java/org/warp/commonutils/locks/TransitLock.java
@@ -1,25 +1,37 @@
 package org.warp.commonutils.locks;
 
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicInteger;
 
 public class TransitLock {
 
 	private final FlexibleCountDownLatch unfinishedTransits = new FlexibleCountDownLatch(0);
 	private final Semaphore permits = new Semaphore(Integer.MAX_VALUE);
+	private final AtomicInteger allowedTransits = new AtomicInteger(1);
 	private final Object synchronization = new Object();
 
 	public TransitLock() {
 	}
 
 	public void allowTransit() {
-		var permitsToRelease = Math.max(0, Integer.MAX_VALUE - (permits.availablePermits() + 10000));
-		permits.release(permitsToRelease);
+		var allowedTransits = this.allowedTransits.incrementAndGet();
+		if (allowedTransits == 1) {
+			var permitsToRelease = Math.max(0, Integer.MAX_VALUE - (permits.availablePermits() + 10000));
+			permits.release(permitsToRelease);
+		} else if (allowedTransits > 1) {
+			throw new IllegalStateException("Already allowed transit!");
+		}
 	}
 
 	public void disallowTransit() {
 		synchronized (synchronization) {
-			unfinishedTransits.await();
-			permits.drainPermits();
+			var allowedTransits = this.allowedTransits.decrementAndGet();
+			if (allowedTransits == 0) {
+				unfinishedTransits.await();
+				permits.drainPermits();
+			} else if (allowedTransits > 0) {
+				throw new IllegalStateException("Transits are still allowed for an unknown reason!");
+			}
 		}
 	}
 
diff --git a/src/test/java/org/warp/commonutils/locks/TransitLockTest.java b/src/test/java/org/warp/commonutils/locks/TransitLockTest.java
index 45c577c..2d50f5f 100644
--- a/src/test/java/org/warp/commonutils/locks/TransitLockTest.java
+++ b/src/test/java/org/warp/commonutils/locks/TransitLockTest.java
@@ -33,17 +33,20 @@ public class TransitLockTest {
 	@Test
 	public void testAllowTransit() {
 		var lock = new TransitLock();
+		lock.disallowTransit();
 		lock.allowTransit();
 		lock.transit();
 	}
 
 	@Test
 	public void testMultiAllowTransit() {
-		var lock = new TransitLock();
-		lock.allowTransit();
-		lock.allowTransit();
-		lock.allowTransit();
-		lock.transit();
+		Assertions.assertThrows(IllegalStateException.class, () -> {
+			var lock = new TransitLock();
+			lock.allowTransit();
+			lock.allowTransit();
+			lock.allowTransit();
+			lock.transit();
+		});
 	}
 
 	@Test
@@ -54,6 +57,7 @@ public class TransitLockTest {
 		lock.disallowTransit();
 		lock.disallowTransit();
 		lock.allowTransit();
+		lock.allowTransit();
 		lock.transit();
 	}
 
@@ -86,13 +90,13 @@ public class TransitLockTest {
 		var lock = new TransitLock();
 		AtomicInteger alreadyRunningTransits = new AtomicInteger();
 
-		var pool = Executors.newFixedThreadPool(100);
-		AtomicReference<AssertionError> failure = new AtomicReference<>();
+		var pool = Executors.newFixedThreadPool(500);
+		AtomicReference<Exception> failure = new AtomicReference<>();
 		for (int i = 0; i < 100; i++) {
 			int iF = i;
 			pool.submit(() -> {
 				try {
-					for (int j = 0; j < 100; j++) {
+					for (int j = 0; j < 500; j++) {
 						if (iF % 2 == 0) {
 							lock.startTransit();
 							alreadyRunningTransits.getAndIncrement();
@@ -104,15 +108,14 @@ public class TransitLockTest {
 							lock.allowTransit();
 						}
 					}
-				} catch (AssertionError e) {
+				} catch (Exception e) {
 					e.printStackTrace();
 					failure.set(e);
 				}
 			});
 		}
-		lock.allowTransit();
 		pool.shutdown();
-		if (failure.get() != null) throw failure.get();
+		if (failure.get() != null) throw new AssertionError(failure.get());
 		Assertions.assertDoesNotThrow(() -> pool.awaitTermination(10, TimeUnit.SECONDS));
 		Assertions.assertTrue(pool.isTerminated());
 	}