feat update

Author: KirillKurdyukov

shedlock-ydb/src/main/java/tech/ydb/lock/provider/YdbJDBCLockProvider.java

@@ -31,19 +31,17 @@ public Optional<SimpleLock> lock(LockConfiguration lockConfiguration) {
             try {
                 connection.setAutoCommit(false);
 
-                var selectPS = connection.prepareStatement("SELECT lock_until, locked_by FROM shedlock " +
-                        "WHERE name = ? AND lock_until > CurrentUtcTimestamp() + ?");
+                var selectPS = connection.prepareStatement("SELECT locked_by, lock_until FROM shedlock " +
+                        "WHERE name = ? AND lock_until > CurrentUtcTimestamp()");
 
                 selectPS.setString(1, lockConfiguration.getName());
-                selectPS.setObject(2, lockConfiguration.getLockAtMostFor());
 
-                var haveLeader = false;
                 try (var rs = selectPS.executeQuery()) {
-                    haveLeader = rs.next();
-                }
-
-                if (haveLeader) {
-                    return Optional.empty();
+                    if (rs.next()) {
+                        LOGGER.debug("Instance[{}] acquire lock is failed. Leader is {}, lock_until = {}",
+                                LOCKED_BY, rs.getString(1), rs.getString(2));
+                        return Optional.empty();
+                    }
                 }
 
                 var upsertPS = connection.prepareStatement("" +
@@ -65,31 +63,50 @@ public Optional<SimpleLock> lock(LockConfiguration lockConfiguration) {
                 connection.setAutoCommit(autoCommit);
             }
         } catch (SQLException e) {
-            LOGGER.debug(String.format("Instance[{%s}] acquire lock is failed", LOCKED_BY), e);
+            LOGGER.debug("Instance[{}] acquire lock is failed", LOCKED_BY);
 
             return Optional.empty();
         }
     }
 
     private record YdbJDBCLock(String name, DataSource dataSource) implements SimpleLock {
+        private static final int ATTEMPT_RELEASE_LOCK = 10;
+
         @Override
         public void unlock() {
+            for (int i = 0; i < ATTEMPT_RELEASE_LOCK; i++) {
+                try {
+                    LOGGER.debug("Instance[{}] trying unlock..", LOCKED_BY);
+
+                    doUnlock();
+
+                    return;
+                } catch (SQLException e) {
+                    if (i == ATTEMPT_RELEASE_LOCK - 1) {
+                        throw new RuntimeException(e);
+                    }
+                }
+            }
+        }
+
+        private void doUnlock() throws SQLException {
             try (var connection = dataSource.getConnection()) {
                 var autoCommit = connection.getAutoCommit();
 
                 try {
                     connection.setAutoCommit(true);
                     var ps = connection.prepareStatement(
-                            "UPDATE shedlock SET lock_until = CurrentUtcTimestamp() WHERE name = ?");
-                    ps.setObject(1, name);
+                            "UPDATE shedlock SET lock_until = CurrentUtcTimestamp() WHERE name = ? and locked_by = ?");
+                    ps.setString(1, name);
+                    ps.setString(2, LOCKED_BY);
                     ps.execute();
                 } finally {
                     connection.setAutoCommit(autoCommit);
                 }
             } catch (SQLException e) {
-                LOGGER.error(String.format("Instance[{%s}] release lock is failed", LOCKED_BY), e);
+                LOGGER.debug(String.format("Instance[{%s}] release lock is failed", LOCKED_BY), e);
 
-                throw new RuntimeException(e);
+                throw e;
             }
         }
     }

shedlock-ydb/src/test/java/tech/ydb/lock/provider/YdbLockProviderTest.java

@@ -8,6 +8,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import javax.sql.DataSource;
@@ -16,6 +17,8 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.test.context.DynamicPropertyRegistry;
@@ -27,6 +30,7 @@
  */
 @SpringBootTest(classes = TestApp.class)
 public class YdbLockProviderTest {
+    private static final Logger logger = LoggerFactory.getLogger(YdbLockProviderTest.class);
 
     @RegisterExtension
     private static final YdbHelperExtension ydb = new YdbHelperExtension();
@@ -74,19 +78,20 @@ public void integrationTest() throws ExecutionException, InterruptedException, S
         var atomicInt = new AtomicInteger();
         var locked = new AtomicBoolean();
 
-        for (int i = 0; i < 10; i++) {
+        for (int i = 0; i < 100; i++) {
             lockFutures.add(executorServer.submit(() -> {
                 Optional<SimpleLock> optinal = Optional.empty();
 
                 while (optinal.isEmpty()) {
                     optinal = lockProvider.lock(new LockConfiguration(
-                            Instant.now(), "semaphore", Duration.ofSeconds(10), Duration.ZERO));
+                            Instant.now(), "semaphore", Duration.ofSeconds(100), Duration.ZERO));
 
                     optinal.ifPresent(simpleLock -> {
-                        if (locked.get()) {
+                        if (locked.compareAndExchange(false, true)) {
+                            logger.debug("Failed test! System has two leaders");
+
                             throw new RuntimeException();
                         }
-                        locked.set(true);
 
                         try {
                             Thread.sleep(100);
@@ -96,8 +101,17 @@ public void integrationTest() throws ExecutionException, InterruptedException, S
 
                         atomicInt.addAndGet(50);
                         locked.set(false);
+
+                        logger.info("Leader does UNLOCK!");
+
                         simpleLock.unlock();
                     });
+
+                    try {
+                        Thread.sleep(ThreadLocalRandom.current().nextInt(3_000));
+                    } catch (InterruptedException e) {
+                        throw new RuntimeException(e);
+                    }
                 }
             }));
         }
@@ -106,6 +120,6 @@ public void integrationTest() throws ExecutionException, InterruptedException, S
             future.get();
         }
 
-        Assertions.assertEquals(4950, atomicInt.get());
+        Assertions.assertEquals(5000, atomicInt.get());
     }
 }