From 27ec2c3a400771d5037187444d3e69b6588e54b3 Mon Sep 17 00:00:00 2001
From: Andrea Cavalli <andrea@warp.ovh>
Date: Mon, 29 Jun 2020 12:07:29 +0200
Subject: [PATCH] Update FutureLockUtils.java

---
 .../concurrency/future/FutureLockUtils.java   | 44 +++++++++++++++----
 1 file changed, 36 insertions(+), 8 deletions(-)

diff --git a/src/main/java/org/warp/commonutils/concurrency/future/FutureLockUtils.java b/src/main/java/org/warp/commonutils/concurrency/future/FutureLockUtils.java
index 09e56be..6a9b25d 100644
--- a/src/main/java/org/warp/commonutils/concurrency/future/FutureLockUtils.java
+++ b/src/main/java/org/warp/commonutils/concurrency/future/FutureLockUtils.java
@@ -17,11 +17,18 @@ public class FutureLockUtils {
 		} else {
 			lockValue = 0;
 		}
-		return r.get().whenComplete((x, y) -> {
+		try {
+			return r.get().whenComplete((x, y) -> {
+				if (lock != null) {
+					lock.unlockRead(lockValue);
+				}
+			});
+		} catch (Throwable ex) {
 			if (lock != null) {
 				lock.unlockRead(lockValue);
 			}
-		});
+			throw ex;
+		}
 	}
 
 	public static <T> CompletableFuture<T> writeLock(@Nullable StampedLock lock, @NotNull Supplier<CompletableFuture<T>> r) {
@@ -31,11 +38,18 @@ public class FutureLockUtils {
 		} else {
 			lockValue = 0;
 		}
-		return r.get().whenComplete((x, y) -> {
+		try {
+			return r.get().whenComplete((x, y) -> {
+				if (lock != null) {
+					lock.unlockWrite(lockValue);
+				}
+			});
+		} catch (Throwable ex) {
 			if (lock != null) {
 				lock.unlockWrite(lockValue);
 			}
-		});
+			throw ex;
+		}
 	}
 
 	public static <T> CompletableFuture<T> readLockIO(@Nullable StampedLock lock, @NotNull IOSupplier<CompletableFuture<T>> r) throws IOException {
@@ -45,11 +59,18 @@ public class FutureLockUtils {
 		} else {
 			lockValue = 0;
 		}
-		return r.get().whenComplete((x, y) -> {
+		try {
+			return r.get().whenComplete((x, y) -> {
+				if (lock != null) {
+					lock.unlockRead(lockValue);
+				}
+			});
+		} catch (Throwable ex) {
 			if (lock != null) {
 				lock.unlockRead(lockValue);
 			}
-		});
+			throw ex;
+		}
 	}
 
 	public static <T> CompletableFuture<T> writeLockIO(@Nullable StampedLock lock, @NotNull IOSupplier<CompletableFuture<T>> r) throws IOException {
@@ -59,10 +80,17 @@ public class FutureLockUtils {
 		} else {
 			lockValue = 0;
 		}
-		return r.get().whenComplete((x, y) -> {
+		try {
+			return r.get().whenComplete((x, y) -> {
+				if (lock != null) {
+					lock.unlockWrite(lockValue);
+				}
+			});
+		} catch (Throwable ex) {
 			if (lock != null) {
 				lock.unlockWrite(lockValue);
 			}
-		});
+			throw ex;
+		}
 	}
 }