Initial commit for RW lock

keva-dev · Dec 19, 2021 · f22c026 · f22c026
1 parent 44cb670
commit f22c026
Show file tree

Hide file tree

Showing 7 changed files with 200 additions and 162 deletions.
diff --git a/core/src/main/java/dev/keva/core/command/impl/key/manager/ExpirationManager.java b/core/src/main/java/dev/keva/core/command/impl/key/manager/ExpirationManager.java
@@ -8,6 +8,7 @@
 import dev.keva.ioc.annotation.Component;
 import dev.keva.protocol.resp.Command;
 import dev.keva.store.KevaDatabase;
+import dev.keva.store.lock.SpinLock;
 
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -71,14 +72,14 @@ public void executeExpire(byte[] key) {
                 data[0] = "delete".getBytes();
                 data[1] = key;
                 Command command = Command.newInstance(data, false);
-                Lock lock = database.getLock();
-                lock.lock();
+                SpinLock lock = database.getLock();
+                lock.exclusiveLock();
                 try {
                     aof.write(command);
                     database.remove(key);
                     clearExpiration(key);
                 } finally {
-                    lock.unlock();
+                    lock.exclusiveUnlock();
                 }
             } else {
                 database.remove(key);

diff --git a/core/src/main/java/dev/keva/core/command/impl/transaction/manager/TransactionContext.java b/core/src/main/java/dev/keva/core/command/impl/transaction/manager/TransactionContext.java
@@ -2,6 +2,7 @@
 
 import dev.keva.core.command.mapping.CommandMapper;
 import dev.keva.protocol.resp.Command;
+import dev.keva.store.lock.SpinLock;
 import dev.keva.util.hashbytes.BytesKey;
 import dev.keva.util.hashbytes.BytesValue;
 import dev.keva.protocol.resp.reply.MultiBulkReply;
@@ -42,8 +43,8 @@ public void discard() {
         isQueuing = false;
     }
 
-    public Reply<?> exec(ChannelHandlerContext ctx, Lock txLock) throws InterruptedException {
-        txLock.lock();
+    public Reply<?> exec(ChannelHandlerContext ctx, SpinLock txLock) throws InterruptedException {
+        txLock.exclusiveLock();
         try {
             for (val watch : watchMap.entrySet()) {
                 val key = watch.getKey();
@@ -74,7 +75,7 @@ public Reply<?> exec(ChannelHandlerContext ctx, Lock txLock) throws InterruptedE
 
             return new MultiBulkReply(replies);
         } finally {
-            txLock.unlock();
+            txLock.exclusiveUnlock();
         }
     }
 }
diff --git a/core/src/main/java/dev/keva/core/command/mapping/CommandMapper.java b/core/src/main/java/dev/keva/core/command/mapping/CommandMapper.java
@@ -25,11 +25,17 @@
 import java.lang.reflect.InvocationTargetException;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
 @Component
 @Slf4j
 public class CommandMapper {
+
+    private static final Set<String> EXCLUSIVE_COMMANDS = new HashSet<>(Arrays.asList(
+            "exec", "expire", "expireat", "restore", "flushdb"));
+
     @Getter
     private final Map<BytesKey, CommandWrapper> methods = new HashMap<>();
 
@@ -98,25 +104,40 @@ public void init() {
 
                         try {
                             val lock = database.getLock();
-                            lock.lock();
+                            boolean locked = false, exclusive = false, writeToAOF = isAoF && isMutate;
                             try {
-                                if (ctx != null && isAoF && isMutate) {
-                                    try {
-                                        aof.write(command);
-                                    } catch (Exception e) {
-                                        log.error("Error writing to AOF", e);
-                                    }
-                                }
                                 Object[] objects = new Object[types.length];
                                 command.toArguments(objects, types, ctx);
+                                if (ctx != null) {
+                                    locked = true;
+                                    if (isMutate || EXCLUSIVE_COMMANDS.contains(name)) {
+                                        lock.exclusiveLock();
+                                        exclusive = true;
+                                        if (writeToAOF) {
+                                            try {
+                                                aof.write(command);
+                                            } catch (Exception e) {
+                                                log.error("Error writing to AOF", e);
+                                            }
+                                        }
+                                    } else {
+                                        lock.sharedLock();
+                                    }
+                                }
                                 // If not in AOF mode, then recycle(),
                                 // else, the command will be recycled in the AOF dump
                                 if (!kevaConfig.getAof()) {
                                     command.recycle();
                                 }
                                 return (Reply<?>) method.invoke(instance, objects);
                             } finally {
-                                lock.unlock();
+                                if (locked) {
+                                    if (exclusive) {
+                                        lock.exclusiveUnlock();
+                                    } else {
+                                        lock.sharedUnlock();
+                                    }
+                                }
                             }
                         } catch (Exception e) {
                             log.error(e.getMessage(), e);

diff --git a/store/src/main/java/dev/keva/store/KevaDatabase.java b/store/src/main/java/dev/keva/store/KevaDatabase.java
@@ -1,12 +1,13 @@
 package dev.keva.store;
 
+import dev.keva.store.lock.SpinLock;
 import dev.keva.util.hashbytes.BytesKey;
 
 import java.util.AbstractMap;
 import java.util.concurrent.locks.Lock;
 
 public interface KevaDatabase {
-    Lock getLock();
+    SpinLock getLock();
 
     void flush();