Fix procedure lock wait scheduling

Author: Caideyipi

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/Procedure.java

@@ -235,6 +235,16 @@ protected void releaseLock(Env env) {
     // no op
   }
 
+  /**
+   * Called after an execution attempt returns {@link ProcedureLockState#LOCK_EVENT_WAIT}. Override
+   * it to put the procedure into the corresponding lock wait queue.
+   *
+   * @param env env
+   */
+  protected void waitForLock(Env env) {
+    // no op
+  }
+
   /**
    * Used to keep procedure lock even when the procedure is yielded or suspended.
    *

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/ProcedureExecutor.java

@@ -23,7 +23,6 @@
 import org.apache.iotdb.commons.utils.RetryUtils;
 import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.confignode.i18n.ProcedureMessages;
-import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv;
 import org.apache.iotdb.confignode.procedure.exception.ProcedureException;
 import org.apache.iotdb.confignode.procedure.scheduler.ProcedureScheduler;
 import org.apache.iotdb.confignode.procedure.scheduler.SimpleProcedureScheduler;
@@ -42,7 +41,6 @@
 import java.util.Deque;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -397,9 +395,9 @@ private void executeProcedure(Procedure<Env> proc) {
     } finally {
       // Only after procedure has completed execution can it be allowed to be rescheduled to prevent
       // data races
-      if (Objects.equals(lockState, ProcedureLockState.LOCK_EVENT_WAIT)) {
+      if (lockState == ProcedureLockState.LOCK_EVENT_WAIT) {
         LOG.info(ProcedureMessages.PROCEDUREID_WAIT_FOR_LOCK, proc.getProcId());
-        ((ConfigNodeProcedureEnv) this.environment).getNodeLock().waitProcedure(proc);
+        proc.waitForLock(this.environment);
       }
     }
   }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/node/AbstractNodeProcedure.java

@@ -56,6 +56,18 @@ protected ProcedureLockState acquireLock(ConfigNodeProcedureEnv configNodeProced
     }
   }
 
+  @Override
+  protected void waitForLock(ConfigNodeProcedureEnv configNodeProcedureEnv) {
+    configNodeProcedureEnv.getSchedulerLock().lock();
+    try {
+      configNodeProcedureEnv
+          .getNodeLock()
+          .waitProcedure(this, configNodeProcedureEnv.getScheduler());
+    } finally {
+      configNodeProcedureEnv.getSchedulerLock().unlock();
+    }
+  }
+
   @Override
   protected void releaseLock(ConfigNodeProcedureEnv configNodeProcedureEnv) {
     configNodeProcedureEnv.getSchedulerLock().lock();

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/scheduler/LockQueue.java

@@ -45,7 +45,11 @@ public boolean releaseLock(Procedure<?> procedure) {
     return true;
   }
 
-  public void waitProcedure(Procedure<?> procedure) {
+  public void waitProcedure(Procedure<?> procedure, ProcedureScheduler procedureScheduler) {
+    if (lockOwnerProcedure == null) {
+      procedureScheduler.addFront(procedure);
+      return;
+    }
     deque.addLast(procedure);
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/scheduler/SimpleProcedureScheduler.java

@@ -22,6 +22,7 @@
 import org.apache.iotdb.confignode.procedure.Procedure;
 
 import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.locks.ReentrantLock;
 
 /** Simple scheduler for procedures */
 public class SimpleProcedureScheduler extends AbstractProcedureScheduler {
@@ -48,6 +49,7 @@ public void clear() {
     schedLock();
     try {
       runnables.clear();
+      waitings.clear();
     } finally {
       schedUnlock();
     }
@@ -68,12 +70,37 @@ public int queueSize() {
     return runnables.size();
   }
 
-  public void addWaiting(Procedure proc) {
-    waitings.add(proc);
+  public void waitProcedure(Procedure proc, ReentrantLock lock) {
+    boolean signal = false;
+    schedLock();
+    try {
+      if (lock.isLocked()) {
+        waitings.add(proc);
+      } else {
+        runnables.addFirst(proc);
+        signal = true;
+      }
+    } finally {
+      schedUnlock();
+    }
+    if (signal) {
+      signalAll();
+    }
   }
 
-  public void releaseWaiting() {
-    runnables.addAll(waitings);
-    waitings.clear();
+  public void releaseWaiting(ReentrantLock lock) {
+    boolean signal;
+    schedLock();
+    try {
+      lock.unlock();
+      signal = !waitings.isEmpty();
+      runnables.addAll(waitings);
+      waitings.clear();
+    } finally {
+      schedUnlock();
+    }
+    if (signal) {
+      signalAll();
+    }
   }
 }

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/entity/SimpleLockProcedure.java

@@ -53,18 +53,21 @@ protected ProcedureLockState acquireLock(TestProcEnv testProcEnv) {
 
       return ProcedureLockState.LOCK_ACQUIRED;
     }
-    SimpleProcedureScheduler scheduler = (SimpleProcedureScheduler) testProcEnv.getScheduler();
-    scheduler.addWaiting(this);
     System.out.println(procName + " wait for lock.");
     return ProcedureLockState.LOCK_EVENT_WAIT;
   }
 
+  @Override
+  protected void waitForLock(TestProcEnv testProcEnv) {
+    SimpleProcedureScheduler scheduler = (SimpleProcedureScheduler) testProcEnv.getScheduler();
+    scheduler.waitProcedure(this, testProcEnv.getEnvLock());
+  }
+
   @Override
   protected void releaseLock(TestProcEnv testProcEnv) {
     System.out.println(procName + " release lock.");
-    testProcEnv.getEnvLock().unlock();
     SimpleProcedureScheduler scheduler = (SimpleProcedureScheduler) testProcEnv.getScheduler();
-    scheduler.releaseWaiting();
+    scheduler.releaseWaiting(testProcEnv.getEnvLock());
   }
 
   @Override