feat: added leader info in semaphore data

Author: KirillKurdyukov

shedlock-ydb/pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>tech.ydb.dialects</groupId>
     <artifactId>shedlock-ydb</artifactId>
-    <version>0.1.0</version>
+    <version>0.2.0</version>
 
     <packaging>jar</packaging>
 

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

@@ -1,15 +1,20 @@
 package tech.ydb.lock.provider;
 
+import java.nio.charset.StandardCharsets;
 import java.sql.SQLException;
 import java.util.Optional;
+import javax.annotation.PostConstruct;
 import javax.annotation.PreDestroy;
 import net.javacrumbs.shedlock.core.LockConfiguration;
 import net.javacrumbs.shedlock.core.LockProvider;
 import net.javacrumbs.shedlock.core.SimpleLock;
+import net.javacrumbs.shedlock.support.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import tech.ydb.coordination.CoordinationClient;
+import tech.ydb.coordination.CoordinationSession;
 import tech.ydb.coordination.SemaphoreLease;
+import tech.ydb.core.Result;
 import tech.ydb.jdbc.YdbConnection;
 
 /**
@@ -19,21 +24,38 @@ public class YdbCoordinationServiceLockProvider implements LockProvider {
     private static final Logger logger = LoggerFactory.getLogger(YdbCoordinationServiceLockProvider.class);
     private static final String YDB_LOCK_NODE_NAME = "shared-lock-ydb";
     private static final int ATTEMPT_CREATE_NODE = 10;
+    private static final String INSTANCE_INFO =
+            "{Hostname=" + Utils.getHostname() + ", " + "Current PID=" + ProcessHandle.current().pid() + "}";
+    private static final byte[] INSTANCE_INFO_BYTES = INSTANCE_INFO.getBytes(StandardCharsets.UTF_8);
 
     private final YdbConnection ydbConnection;
     private final CoordinationClient coordinationClient;
 
+    private volatile CoordinationSession coordinationSession;
+
     public YdbCoordinationServiceLockProvider(YdbConnection ydbConnection) {
         this.ydbConnection = ydbConnection;
         this.coordinationClient = CoordinationClient.newClient(ydbConnection.getCtx().getGrpcTransport());
     }
 
+    @PostConstruct
     public void init() {
         for (int i = 0; i < ATTEMPT_CREATE_NODE; i++) {
             var status = coordinationClient.createNode(YDB_LOCK_NODE_NAME).join();
 
             if (status.isSuccess()) {
-                return;
+                coordinationSession = coordinationClient.createSession(YDB_LOCK_NODE_NAME);
+
+                var statusCS = coordinationSession.connect().join();
+
+                if (statusCS.isSuccess()) {
+                    logger.info("Created coordination node session [{}]", coordinationSession);
+
+                    return;
+                }
+                if (i == ATTEMPT_CREATE_NODE - 1) {
+                    statusCS.expectSuccess("Failed creating coordination node session");
+                }
             }
 
             if (i == ATTEMPT_CREATE_NODE - 1) {
@@ -44,35 +66,41 @@ public void init() {
 
     @Override
     public Optional<SimpleLock> lock(LockConfiguration lockConfiguration) {
-        var coordinationSession = coordinationClient.createSession(YDB_LOCK_NODE_NAME);
-
-        coordinationSession.connect().join()
-                .expectSuccess("Failed creating coordination node session");
+        logger.info("Instance[{}] is trying to become a leader...", INSTANCE_INFO);
 
-        logger.debug("Created coordination node session");
-
-        var semaphoreLease = coordinationSession.acquireEphemeralSemaphore(lockConfiguration.getName(), true,
-                lockConfiguration.getLockAtMostFor()).join();
+        Result<SemaphoreLease> semaphoreLease = coordinationSession.acquireEphemeralSemaphore(
+                lockConfiguration.getName(),
+                true,
+                INSTANCE_INFO_BYTES,
+                lockConfiguration.getLockAtMostFor()
+        ).join();
 
         if (semaphoreLease.isSuccess()) {
-            logger.debug("Semaphore acquired");
+            logger.info("Instance[{}] acquired semaphore[SemaphoreName={}]", INSTANCE_INFO,
+                    semaphoreLease.getValue().getSemaphoreName());
 
             return Optional.of(new YdbSimpleLock(semaphoreLease.getValue()));
         } else {
-            logger.debug("Semaphore is not acquired");
+            logger.info("Instance[{}] did not acquire semaphore", INSTANCE_INFO);
+
             return Optional.empty();
         }
     }
 
     private record YdbSimpleLock(SemaphoreLease semaphoreLease) implements SimpleLock {
         @Override
         public void unlock() {
+            logger.info("Instance[{}] released semaphore[SemaphoreName={}]", INSTANCE_INFO, semaphoreLease.getSemaphoreName());
+
             semaphoreLease.release().join();
         }
     }
 
     @PreDestroy
     private void close() throws SQLException {
+        // closing coordination session
+        coordinationSession.close();
+
         ydbConnection.close();
     }
 }

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

@@ -2,11 +2,9 @@
 
 import java.time.Duration;
 import java.time.Instant;
-import java.util.ArrayList;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import net.javacrumbs.shedlock.core.LockConfiguration;
@@ -55,11 +53,9 @@ public void integrationTest() throws ExecutionException, InterruptedException {
         var executorServer = Executors.newFixedThreadPool(10);
         var atomicInt = new AtomicInteger();
         var locked = new AtomicBoolean();
-        var futures = new ArrayList<Future<?>>();
 
-        for (int i = 0; i < 100; i++) {
-            final var ii = i;
-            futures.add(executorServer.submit(() -> {
+        for (int i = 0; i < 10; i++) {
+            executorServer.submit(() -> {
                 Optional<SimpleLock> optinal = Optional.empty();
 
                 while (optinal.isEmpty()) {
@@ -78,18 +74,14 @@ public void integrationTest() throws ExecutionException, InterruptedException {
                             throw new RuntimeException(e);
                         }
 
-                        atomicInt.addAndGet(ii);
+                        atomicInt.addAndGet(50);
                         locked.set(false);
                         simpleLock.unlock();
                     });
                 }
-            }));
+            }).get();
         }
 
-        for (Future<?> future : futures) {
-            future.get();
-        }
-
-        Assertions.assertEquals(4950, atomicInt.get());
+        Assertions.assertEquals(500, atomicInt.get());
     }
 }