[AMORO-3794] Wrap task operation with process lock (#3798)

* wrap task operation with process lock

* Fix some checkstyle errors

* Add a lock timeout while polling task

Author: zhoujinsong

amoro-ams/src/main/java/org/apache/amoro/server/DefaultOptimizingService.java

@@ -32,7 +32,6 @@
 import org.apache.amoro.exception.IllegalTaskStateException;
 import org.apache.amoro.exception.ObjectNotExistsException;
 import org.apache.amoro.exception.PluginRetryAuthException;
-import org.apache.amoro.exception.TaskNotFoundException;
 import org.apache.amoro.resource.ResourceGroup;
 import org.apache.amoro.server.catalog.CatalogManager;
 import org.apache.amoro.server.optimizing.OptimizingProcess;
@@ -54,7 +53,6 @@
 import org.apache.amoro.shade.guava32.com.google.common.base.Preconditions;
 import org.apache.amoro.shade.guava32.com.google.common.collect.Sets;
 import org.apache.amoro.shade.guava32.com.google.common.util.concurrent.ThreadFactoryBuilder;
-import org.apache.amoro.shade.thrift.org.apache.thrift.TException;
 import org.apache.commons.lang3.StringUtils;
 import org.jetbrains.annotations.NotNull;
 import org.slf4j.Logger;
@@ -214,33 +212,21 @@ private OptimizerInstance getAuthenticatedOptimizer(String authToken) {
   @Override
   public OptimizingTask pollTask(String authToken, int threadId) {
     LOG.debug("Optimizer {} (threadId {}) try polling task", authToken, threadId);
+    OptimizerThread optimizerThread = getAuthenticatedOptimizer(authToken).getThread(threadId);
     OptimizingQueue queue = getQueueByToken(authToken);
-    return Optional.ofNullable(queue.pollTask(pollingTimeout, breakQuotaLimit))
-        .map(task -> extractOptimizingTask(task, authToken, threadId, queue))
-        .orElse(null);
-  }
-
-  private OptimizingTask extractOptimizingTask(
-      TaskRuntime<?> task, String authToken, int threadId, OptimizingQueue queue) {
-    try {
-      OptimizerThread optimizerThread = getAuthenticatedOptimizer(authToken).getThread(threadId);
-      task.schedule(optimizerThread);
+    TaskRuntime<?> task = queue.pollTask(optimizerThread, pollingTimeout, breakQuotaLimit);
+    if (task != null) {
       LOG.info("OptimizerThread {} polled task {}", optimizerThread, task.getTaskId());
       return task.extractProtocolTask();
-    } catch (Throwable throwable) {
-      LOG.error("Schedule task {} failed, put it to retry queue", task.getTaskId(), throwable);
-      queue.retryTask(task);
-      return null;
     }
+    return null;
   }
 
   @Override
   public void ackTask(String authToken, int threadId, OptimizingTaskId taskId) {
     LOG.info("Ack task {} by optimizer {} (threadId {})", taskId, authToken, threadId);
     OptimizingQueue queue = getQueueByToken(authToken);
-    Optional.ofNullable(queue.getTask(taskId))
-        .orElseThrow(() -> new TaskNotFoundException(taskId))
-        .ack(getAuthenticatedOptimizer(authToken).getThread(threadId));
+    queue.ackTask(taskId, getAuthenticatedOptimizer(authToken).getThread(threadId));
   }
 
   @Override
@@ -254,9 +240,7 @@ public void completeTask(String authToken, OptimizingTaskResult taskResult) {
     OptimizingQueue queue = getQueueByToken(authToken);
     OptimizerThread thread =
         getAuthenticatedOptimizer(authToken).getThread(taskResult.getThreadId());
-    Optional.ofNullable(queue.getTask(taskResult.getTaskId()))
-        .orElseThrow(() -> new TaskNotFoundException(taskResult.getTaskId()))
-        .complete(thread, taskResult);
+    queue.completeTask(thread, taskResult);
   }
 
   @Override
@@ -284,7 +268,7 @@ public String authenticate(OptimizerRegisterInfo registerInfo) {
   }
 
   @Override
-  public boolean cancelProcess(long processId) throws TException {
+  public boolean cancelProcess(long processId) {
     TableProcessMeta processMeta =
         getAs(TableProcessMapper.class, m -> m.getProcessMeta(processId));
     if (processMeta == null) {

amoro-ams/src/main/java/org/apache/amoro/server/optimizing/OptimizingQueue.java

@@ -23,8 +23,10 @@
 import org.apache.amoro.ServerTableIdentifier;
 import org.apache.amoro.api.BlockableOperation;
 import org.apache.amoro.api.OptimizingTaskId;
+import org.apache.amoro.api.OptimizingTaskResult;
 import org.apache.amoro.exception.OptimizingClosedException;
 import org.apache.amoro.exception.PersistenceException;
+import org.apache.amoro.exception.TaskNotFoundException;
 import org.apache.amoro.optimizing.MetricsSummary;
 import org.apache.amoro.optimizing.OptimizingType;
 import org.apache.amoro.optimizing.RewriteFilesInput;
@@ -44,6 +46,7 @@
 import org.apache.amoro.server.persistence.mapper.TableProcessMapper;
 import org.apache.amoro.server.process.TableProcessMeta;
 import org.apache.amoro.server.resource.OptimizerInstance;
+import org.apache.amoro.server.resource.OptimizerThread;
 import org.apache.amoro.server.resource.QuotaProvider;
 import org.apache.amoro.server.table.DefaultTableRuntime;
 import org.apache.amoro.server.table.blocker.TableBlocker;
@@ -90,7 +93,6 @@ public class OptimizingQueue extends PersistentBase {
 
   private final QuotaProvider quotaProvider;
   private final Queue<TableOptimizingProcess> tableQueue = new LinkedTransferQueue<>();
-  private final Queue<TaskRuntime<?>> retryTaskQueue = new LinkedTransferQueue<>();
   private final SchedulingPolicy scheduler;
   private final CatalogManager catalogManager;
   private final Executor planExecutor;
@@ -200,24 +202,23 @@ public void releaseTable(DefaultTableRuntime tableRuntime) {
 
   private void clearProcess(OptimizingProcess optimizingProcess) {
     tableQueue.removeIf(process -> process.getProcessId() == optimizingProcess.getProcessId());
-    retryTaskQueue.removeIf(
-        taskRuntime -> taskRuntime.getTaskId().getProcessId() == optimizingProcess.getProcessId());
   }
 
-  public TaskRuntime<?> pollTask(long maxWaitTime, boolean breakQuotaLimit) {
+  public TaskRuntime<?> pollTask(
+      OptimizerThread thread, long maxWaitTime, boolean breakQuotaLimit) {
     long deadline = calculateDeadline(maxWaitTime);
-    TaskRuntime<?> task = fetchTask();
+    TaskRuntime<?> task = fetchScheduledTask(thread, true);
     while (task == null && waitTask(deadline)) {
-      task = fetchTask();
+      task = fetchScheduledTask(thread, true);
     }
     if (task == null && breakQuotaLimit && planningTables.isEmpty()) {
-      task = fetchScheduledTask(false);
+      task = fetchScheduledTask(thread, false);
     }
     return task;
   }
 
-  public TaskRuntime<?> pollTask(long maxWaitTime) {
-    return pollTask(maxWaitTime, false);
+  public TaskRuntime<?> pollTask(OptimizerThread thread, long maxWaitTime) {
+    return pollTask(thread, maxWaitTime, true);
   }
 
   private long calculateDeadline(long maxWaitTime) {
@@ -240,14 +241,9 @@ private boolean waitTask(long waitDeadline) {
     }
   }
 
-  private TaskRuntime<?> fetchTask() {
-    TaskRuntime<?> task = retryTaskQueue.poll();
-    return task != null ? task : fetchScheduledTask(true);
-  }
-
-  private TaskRuntime<?> fetchScheduledTask(boolean needQuotaChecking) {
+  private TaskRuntime<?> fetchScheduledTask(OptimizerThread thread, boolean needQuotaChecking) {
     return tableQueue.stream()
-        .map(process -> process.poll(needQuotaChecking))
+        .map(process -> process.poll(thread, needQuotaChecking))
         .filter(Objects::nonNull)
         .findFirst()
         .orElse(null);
@@ -352,12 +348,12 @@ private TableOptimizingProcess planInternal(DefaultTableRuntime tableRuntime) {
     }
   }
 
-  public TaskRuntime<?> getTask(OptimizingTaskId taskId) {
-    return tableQueue.stream()
-        .filter(p -> p.getProcessId() == taskId.getProcessId())
-        .findFirst()
-        .map(p -> p.getTaskMap().get(taskId))
-        .orElse(null);
+  public void ackTask(OptimizingTaskId taskId, OptimizerThread thread) {
+    findProcess(taskId).ackTask(taskId, thread);
+  }
+
+  public void completeTask(OptimizerThread thread, OptimizingTaskResult result) {
+    findProcess(result.getTaskId()).completeTask(thread, result);
   }
 
   public List<TaskRuntime<?>> collectTasks() {
@@ -374,8 +370,7 @@ public List<TaskRuntime<?>> collectTasks(Predicate<TaskRuntime<?>> predicate) {
   }
 
   public void retryTask(TaskRuntime<?> taskRuntime) {
-    taskRuntime.reset();
-    retryTaskQueue.offer(taskRuntime);
+    findProcess(taskRuntime.getTaskId()).resetTask((TaskRuntime<RewriteStageTask>) taskRuntime);
   }
 
   public ResourceGroup getOptimizerGroup() {
@@ -402,6 +397,13 @@ public void dispose() {
     this.metrics.unregister();
   }
 
+  private TableOptimizingProcess findProcess(OptimizingTaskId taskId) {
+    return tableQueue.stream()
+        .filter(p -> p.getProcessId() == taskId.getProcessId())
+        .findFirst()
+        .orElseThrow(() -> new TaskNotFoundException(taskId));
+  }
+
   private double getAvailableCore() {
     // the available core should be at least 1
     return Math.max(quotaProvider.getTotalQuota(optimizerGroup.getName()), 1);
@@ -437,26 +439,32 @@ private class TableOptimizingProcess implements OptimizingProcess {
     private Map<String, Long> toSequence = Maps.newHashMap();
     private boolean hasCommitted = false;
 
-    public TaskRuntime<?> poll(boolean needQuotaChecking) {
-      if (lock.tryLock()) {
-        try {
-          TaskRuntime<?> task = null;
-          if (status != ProcessStatus.KILLED && status != ProcessStatus.FAILED) {
-            int actualQuota = getActualQuota();
-            int quotaLimit = getQuotaLimit();
-            if (!needQuotaChecking || actualQuota < quotaLimit) {
-              task = taskQueue.poll();
+    public TaskRuntime<?> poll(OptimizerThread thread, boolean needQuotaChecking) {
+      try {
+        // Wait 10ms here for some light operation like poll/ack
+        if (lock.tryLock(10, TimeUnit.MILLISECONDS)) {
+          try {
+            TaskRuntime<?> task = null;
+            if (status != ProcessStatus.KILLED && status != ProcessStatus.FAILED) {
+              int actualQuota = getActualQuota();
+              int quotaLimit = getQuotaLimit();
+              if (!needQuotaChecking || actualQuota < quotaLimit) {
+                task = taskQueue.poll();
+              }
             }
+            if (task != null) {
+              optimizingTasksMap
+                  .computeIfAbsent(tableRuntime.getTableIdentifier(), k -> new AtomicInteger(0))
+                  .incrementAndGet();
+              task.schedule(thread);
+            }
+            return task;
+          } finally {
+            lock.unlock();
           }
-          if (task != null) {
-            optimizingTasksMap
-                .computeIfAbsent(tableRuntime.getTableIdentifier(), k -> new AtomicInteger(0))
-                .incrementAndGet();
-          }
-          return task;
-        } finally {
-          lock.unlock();
         }
+      } catch (InterruptedException e) {
+        // ignore it.
       }
       return null;
     }
@@ -543,6 +551,34 @@ public void close(boolean needCommit) {
       }
     }
 
+    private void ackTask(OptimizingTaskId taskId, OptimizerThread thread) {
+      TaskRuntime<?> taskRuntime = getTaskRuntime(taskId);
+      lock.lock();
+      try {
+        taskRuntime.ack(thread);
+      } finally {
+        lock.unlock();
+      }
+    }
+
+    private void completeTask(OptimizerThread thread, OptimizingTaskResult result) {
+      TaskRuntime<?> taskRuntime = getTaskRuntime(result.getTaskId());
+      lock.lock();
+      try {
+        taskRuntime.complete(thread, result);
+      } finally {
+        lock.unlock();
+      }
+    }
+
+    private TaskRuntime<?> getTaskRuntime(OptimizingTaskId taskId) {
+      TaskRuntime<?> taskRuntime = getTaskMap().get(taskId);
+      if (taskRuntime == null) {
+        throw new TaskNotFoundException(taskId);
+      }
+      return taskRuntime;
+    }
+
     private void acceptResult(TaskRuntime<?> taskRuntime) {
       lock.lock();
       try {
@@ -612,6 +648,16 @@ private void acceptResult(TaskRuntime<?> taskRuntime) {
       }
     }
 
+    private void resetTask(TaskRuntime<RewriteStageTask> taskRuntime) {
+      lock.lock();
+      try {
+        taskRuntime.reset();
+        taskQueue.add(taskRuntime);
+      } finally {
+        lock.unlock();
+      }
+    }
+
     @Override
     public boolean isClosed() {
       return status == ProcessStatus.KILLED;

amoro-ams/src/main/java/org/apache/amoro/server/optimizing/TaskRuntime.java

@@ -81,7 +81,7 @@ public SimpleFuture getCompletedFuture() {
     return future;
   }
 
-  public void complete(OptimizerThread thread, OptimizingTaskResult result) {
+  void complete(OptimizerThread thread, OptimizingTaskResult result) {
     invokeConsistency(
         () -> {
           validThread(thread);
@@ -121,7 +121,7 @@ void reset() {
         });
   }
 
-  public void schedule(OptimizerThread thread) {
+  void schedule(OptimizerThread thread) {
     invokeConsistency(
         () -> {
           statusMachine.accept(Status.SCHEDULED);
@@ -132,7 +132,7 @@ public void schedule(OptimizerThread thread) {
         });
   }
 
-  public void ack(OptimizerThread thread) {
+  void ack(OptimizerThread thread) {
     invokeConsistency(
         () -> {
           validThread(thread);

amoro-ams/src/test/java/org/apache/amoro/server/dashboard/utils/TestOptimizingUtil.java

@@ -93,13 +93,12 @@ public void testCalculateQuotaOccupy() {
     DefaultTableRuntime tableRuntime = initTableWithFiles();
     OptimizingQueue queue = buildOptimizingGroupService(tableRuntime);
     Assert.assertEquals(0, queue.collectTasks().size());
-    TaskRuntime task = queue.pollTask(MAX_POLLING_TIME);
-    task.schedule(optimizerThread);
-    task.ack(optimizerThread);
+    TaskRuntime<?> task = queue.pollTask(optimizerThread, MAX_POLLING_TIME);
+    queue.ackTask(task.getTaskId(), optimizerThread);
     Assert.assertEquals(
         1, queue.collectTasks(t -> t.getStatus() == TaskRuntime.Status.ACKED).size());
     Assert.assertNotNull(task);
-    task.complete(
+    queue.completeTask(
         optimizerThread,
         buildOptimizingTaskResult(task.getTaskId(), optimizerThread.getThreadId()));
     Assert.assertEquals(TaskRuntime.Status.SUCCESS, task.getStatus());