From 1a1bd11ff730ba8e36ff43f6828adddafb80f0f4 Mon Sep 17 00:00:00 2001
From: Nikita Koksharov <nkoksharov@redisson.pro>
Date: Tue, 10 Sep 2024 15:16:21 +0300
Subject: [PATCH] refactoring

---
 .../org/redisson/QueueTransferService.java    | 28 +++++++++----------
 .../java/org/redisson/QueueTransferTask.java  |  7 +----
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/redisson/src/main/java/org/redisson/QueueTransferService.java b/redisson/src/main/java/org/redisson/QueueTransferService.java
index e17a664de..283ed379f 100644
--- a/redisson/src/main/java/org/redisson/QueueTransferService.java
+++ b/redisson/src/main/java/org/redisson/QueueTransferService.java
@@ -28,27 +28,27 @@ public class QueueTransferService {
     private final Map<String, QueueTransferTask> tasks = new ConcurrentHashMap<>();
     
     public void schedule(String name, QueueTransferTask task) {
-        QueueTransferTask oldTask = tasks.putIfAbsent(name, task);
-        if (oldTask == null) {
-            task.start();
-        } else {
-            oldTask.getLock().execute(() -> {
-                oldTask.incUsage();
-            });
-        }
+        tasks.compute(name, (k, t) -> {
+            if (t == null) {
+                task.start();
+                return task;
+            }
+            t.incUsage();
+            return t;
+        });
     }
     
     public void remove(String name) {
-        QueueTransferTask task = tasks.get(name);
-        if (task == null) {
-            return;
-        }
+        tasks.compute(name, (k, task) -> {
+            if (task == null) {
+                return null;
+            }
 
-        task.getLock().execute(() -> {
             if (task.decUsage() == 0) {
-                tasks.remove(name, task);
                 task.stop();
+                return null;
             }
+            return task;
         });
     }
     
diff --git a/redisson/src/main/java/org/redisson/QueueTransferTask.java b/redisson/src/main/java/org/redisson/QueueTransferTask.java
index 6d3bda023..8c4d57904 100644
--- a/redisson/src/main/java/org/redisson/QueueTransferTask.java
+++ b/redisson/src/main/java/org/redisson/QueueTransferTask.java
@@ -22,7 +22,6 @@ import org.redisson.api.RTopic;
 import org.redisson.api.listener.BaseStatusListener;
 import org.redisson.api.listener.MessageListener;
 import org.redisson.connection.ServiceManager;
-import org.redisson.misc.WrappedLock;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -62,8 +61,7 @@ public abstract class QueueTransferTask {
     private int usage = 1;
     private final AtomicReference<TimeoutTask> lastTimeout = new AtomicReference<TimeoutTask>();
     private final ServiceManager serviceManager;
-    private final WrappedLock lock = new WrappedLock();
-    
+
     public QueueTransferTask(ServiceManager serviceManager) {
         super();
         this.serviceManager = serviceManager;
@@ -157,7 +155,4 @@ public abstract class QueueTransferTask {
         });
     }
 
-    public WrappedLock getLock() {
-        return lock;
-    }
 }