This is an automated email from the ASF dual-hosted git repository. lancelly pushed a commit to branch priorityForShowQueries in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 50759db43292b1af060ebd49fd2538dff813b549 Author: lancelly <[email protected]> AuthorDate: Sat Oct 14 13:18:57 2023 +0800 set show query to highest priority --- .../db/queryengine/execution/driver/Driver.java | 12 ++ .../db/queryengine/execution/driver/IDriver.java | 4 + .../fragment/FragmentInstanceManager.java | 5 + .../execution/schedule/DriverScheduler.java | 5 +- .../multilevelqueue/MultilevelPriorityQueue.java | 21 ++- .../execution/schedule/task/DriverTask.java | 18 ++- .../queryengine/plan/execution/QueryExecution.java | 14 ++ .../plan/planner/plan/FragmentInstance.java | 10 ++ .../queryengine/plan/statement/StatementType.java | 1 + .../plan/statement/sys/ShowQueriesStatement.java | 3 +- .../schedule/DefaultDriverSchedulerTest.java | 28 ++-- .../DriverTaskTimeoutSentinelThreadTest.java | 19 +-- .../queue/MultilevelPriorityQueueTest.java | 157 +++++++++++++++++++++ 13 files changed, 263 insertions(+), 34 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/Driver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/Driver.java index 86cefabeb6f..45b897d9c60 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/Driver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/Driver.java @@ -71,6 +71,8 @@ public abstract class Driver implements IDriver { protected final DriverLock exclusiveLock = new DriverLock(); + private boolean isHighestPriority; + protected enum State { ALIVE, NEED_DESTRUCTION, @@ -171,6 +173,16 @@ public abstract class Driver implements IDriver { this.driverContext.setDriverTaskID(driverTaskId); } + @Override + public boolean isHighestPriority() { + return isHighestPriority; + } + + @Override + public void setHighestPriority(boolean isHighestPriority) { + this.isHighestPriority = isHighestPriority; + } + @Override public void close() { // mark the service for destruction diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/IDriver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/IDriver.java index 92b29158bcb..b424e3645ea 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/IDriver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/IDriver.java @@ -79,4 +79,8 @@ public interface IDriver { ISink getSink(); DriverContext getDriverContext(); + + boolean isHighestPriority(); + + void setHighestPriority(boolean isHighestPriority); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java index 97d6be05df8..4c8323b705b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java @@ -171,6 +171,11 @@ public class FragmentInstanceManager { List<IDriver> drivers = new ArrayList<>(); driverFactories.forEach(factory -> drivers.add(factory.createDriver())); + // For ShowQueries related instances, isHighestPriority == true + if (instance.isHighestPriority()) { + drivers.forEach(driver -> driver.setHighestPriority(true)); + } + context.initializeNumOfDrivers(drivers.size()); // get the sink of last driver ISink sink = drivers.get(drivers.size() - 1).getSink(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java index 0c33c2b1141..9228adbd77f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java @@ -196,7 +196,8 @@ public class DriverScheduler implements IDriverScheduler, IService { timeOut > 0 ? timeOut : QUERY_TIMEOUT_MS, DriverTaskStatus.READY, driverTaskHandle, - driver.getEstimatedMemorySize()))); + driver.getEstimatedMemorySize(), + driver.isHighestPriority()))); List<DriverTask> submittedTasks = new ArrayList<>(); for (DriverTask task : tasks) { @@ -430,7 +431,7 @@ public class DriverScheduler implements IDriverScheduler, IService { } @TestOnly - IndexedBlockingQueue<DriverTask> getReadyQueue() { + public IndexedBlockingQueue<DriverTask> getReadyQueue() { return readyQueue; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/multilevelqueue/MultilevelPriorityQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/multilevelqueue/MultilevelPriorityQueue.java index 6094410d74b..4932f06420a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/multilevelqueue/MultilevelPriorityQueue.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/multilevelqueue/MultilevelPriorityQueue.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.execution.schedule.queue.multilevelqueue; +import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.queryengine.execution.schedule.queue.IndexedBlockingReserveQueue; import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTask; @@ -44,7 +45,10 @@ public class MultilevelPriorityQueue extends IndexedBlockingReserveQueue<DriverT private final PriorityQueue<DriverTask>[] levelWaitingSplits; /** - This queue is independent of the other priority queues and has the highest priority. It is used to assign the highest execution priority to tasks like "ShowQuery," without considering cumulative execution time. */ + * This queue is independent of the other priority queues and has the highest priority. It is used + * to assign the highest execution priority to tasks like "ShowQuery," without considering + * cumulative execution time. + */ private final PriorityQueue<DriverTask> highestPriorityLevelQueue; /** @@ -69,7 +73,8 @@ public class MultilevelPriorityQueue extends IndexedBlockingReserveQueue<DriverT this.levelScheduledTime = new AtomicLong[LEVEL_THRESHOLD_SECONDS.length]; this.levelMinScheduledTime = new AtomicLong[LEVEL_THRESHOLD_SECONDS.length]; this.levelWaitingSplits = new PriorityQueue[LEVEL_THRESHOLD_SECONDS.length]; - this.highestPriorityLevelQueue = new PriorityQueue<>(); + this.highestPriorityLevelQueue = + new PriorityQueue<>(new DriverTask.SchedulePriorityComparator()); for (int level = 0; level < LEVEL_THRESHOLD_SECONDS.length; level++) { levelScheduledTime[level] = new AtomicLong(); levelMinScheduledTime[level] = new AtomicLong(-1); @@ -91,8 +96,9 @@ public class MultilevelPriorityQueue extends IndexedBlockingReserveQueue<DriverT @Override public void pushToQueue(DriverTask task) { checkArgument(task != null, "DriverTask to be pushed is null"); - // Push tasks with the highest priority(Currently, only ShowQuery related tasks) into highestPriorityLevelQueue directly. - if(task.isHighestPriority()){ + // Push tasks with the highest priority(Currently, only ShowQuery related tasks) into + // highestPriorityLevelQueue directly. + if (task.isHighestPriority()) { highestPriorityLevelQueue.offer(task); return; } @@ -113,7 +119,7 @@ public class MultilevelPriorityQueue extends IndexedBlockingReserveQueue<DriverT protected DriverTask pollFirst() { // Always choose tasks in the highestPriorityLevelQueue first. - if(!highestPriorityLevelQueue.isEmpty()){ + if (!highestPriorityLevelQueue.isEmpty()) { return highestPriorityLevelQueue.poll(); } @@ -302,4 +308,9 @@ public class MultilevelPriorityQueue extends IndexedBlockingReserveQueue<DriverT return LEVEL_THRESHOLD_SECONDS.length - 1; } + + @TestOnly + public PriorityQueue<DriverTask> getHighestPriorityLevelQueue() { + return highestPriorityLevelQueue; + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java index f3ed7d29b69..acbbcde6fab 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java @@ -50,7 +50,7 @@ public class DriverTask implements IDIndexedAccessible { private final long ddl; private final Lock lock; - private final boolean isHighestPriority = true; + private final boolean isHighestPriority; private String abortCause; @@ -64,7 +64,7 @@ public class DriverTask implements IDIndexedAccessible { /** Initialize a dummy instance for queryHolder. */ public DriverTask() { - this(new StubFragmentInstance(), 0L, null, null, 0); + this(new StubFragmentInstance(), 0L, null, null, 0, false); } public DriverTask( @@ -72,7 +72,8 @@ public class DriverTask implements IDIndexedAccessible { long timeoutMs, DriverTaskStatus status, DriverTaskHandle driverTaskHandle, - long estimatedMemorySize) { + long estimatedMemorySize, + boolean isHighestPriority) { this.driver = driver; this.setStatus(status); this.ddl = System.currentTimeMillis() + timeoutMs; @@ -80,6 +81,7 @@ public class DriverTask implements IDIndexedAccessible { this.driverTaskHandle = driverTaskHandle; this.priority = new AtomicReference<>(new Priority(0, 0)); this.estimatedMemorySize = estimatedMemorySize; + this.isHighestPriority = isHighestPriority; } @Override @@ -282,5 +284,15 @@ public class DriverTask implements IDIndexedAccessible { public DriverContext getDriverContext() { return null; } + + @Override + public boolean isHighestPriority() { + return false; + } + + @Override + public void setHighestPriority(boolean isHighestPriority) { + // do nothing + } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 14765cfe2aa..547756f7d87 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -26,6 +26,7 @@ import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; import org.apache.iotdb.commons.conf.IoTDBConstant; import org.apache.iotdb.commons.exception.IoTDBException; import org.apache.iotdb.commons.service.metric.PerformanceOverviewMetrics; +import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.query.KilledByOthersException; @@ -60,6 +61,7 @@ import org.apache.iotdb.db.queryengine.plan.scheduler.ClusterScheduler; import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler; import org.apache.iotdb.db.queryengine.plan.scheduler.load.LoadTsFileScheduler; import org.apache.iotdb.db.queryengine.plan.statement.Statement; +import org.apache.iotdb.db.queryengine.plan.statement.StatementType; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertBaseStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsStatement; @@ -367,6 +369,13 @@ public class QueryExecution implements IQueryExecution { QUERY_PLAN_COST_METRIC_SET.recordPlanCost( DISTRIBUTION_PLANNER, System.nanoTime() - startTime); } + + // if is this Statement is ShowQueryStatement, set its instances to the highest priority, so + // that the sub-tasks of the ShowQueries instances could be executed first. + if (StatementType.SHOW_QUERIES.equals(rawStatement.getType())) { + distributedPlan.getInstances().forEach(instance -> instance.setHighestPriority(true)); + } + if (isQuery() && logger.isDebugEnabled()) { logger.debug( "distribution plan done. Fragment instance count is {}, details is: \n {}", @@ -723,6 +732,11 @@ public class QueryExecution implements IQueryExecution { return distributedPlan; } + @TestOnly + public void setLogicalPlan(LogicalQueryPlan logicalPlan) { + this.logicalPlan = logicalPlan; + } + public LogicalQueryPlan getLogicalPlan() { return logicalPlan; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java index 2a59e9f0158..9a211c8384a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java @@ -72,6 +72,8 @@ public class FragmentInstance implements IConsensusRequest { // The num of all FI on the dispatched DataNode in this query private int dataNodeFINum; + private boolean isHighestPriority; + // We can add some more params for a specific FragmentInstance // So that we can make different FragmentInstance owns different data range. @@ -154,6 +156,14 @@ public class FragmentInstance implements IConsensusRequest { return isRoot; } + public boolean isHighestPriority() { + return isHighestPriority; + } + + public void setHighestPriority(boolean highestPriority) { + isHighestPriority = highestPriority; + } + public void setTimeFilter(Filter timeFilter) { this.timeFilter = timeFilter; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java index b818f763e7e..b673703649a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java @@ -84,6 +84,7 @@ public enum StatementType { SHOW, SHOW_MERGE_STATUS, + SHOW_QUERIES, CREATE_INDEX, DROP_INDEX, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/ShowQueriesStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/ShowQueriesStatement.java index 8c434050c34..fb16d1ec00c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/ShowQueriesStatement.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/ShowQueriesStatement.java @@ -22,6 +22,7 @@ package org.apache.iotdb.db.queryengine.plan.statement.sys; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.auth.entity.PrivilegeType; import org.apache.iotdb.db.auth.AuthorityChecker; +import org.apache.iotdb.db.queryengine.plan.statement.StatementType; import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor; import org.apache.iotdb.db.queryengine.plan.statement.component.OrderByComponent; import org.apache.iotdb.db.queryengine.plan.statement.component.OrderByKey; @@ -47,7 +48,7 @@ public class ShowQueriesStatement extends ShowStatement { private ZoneId zoneId; public ShowQueriesStatement() { - // do nothing + this.statementType = StatementType.SHOW_QUERIES; } @Override diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DefaultDriverSchedulerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DefaultDriverSchedulerTest.java index abda6791de5..3a7be5f4dcf 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DefaultDriverSchedulerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DefaultDriverSchedulerTest.java @@ -79,7 +79,7 @@ public class DefaultDriverSchedulerTest { DriverTaskStatus.RUNNING, }; for (DriverTaskStatus status : invalidStates) { - DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0, false); manager.getBlockedTasks().add(testTask); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); @@ -96,7 +96,7 @@ public class DefaultDriverSchedulerTest { clear(); } DriverTask testTask = - new DriverTask(mockDriver, 100L, DriverTaskStatus.BLOCKED, driverTaskHandle, 0); + new DriverTask(mockDriver, 100L, DriverTaskStatus.BLOCKED, driverTaskHandle, 0, false); manager.getBlockedTasks().add(testTask); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); @@ -140,7 +140,7 @@ public class DefaultDriverSchedulerTest { DriverTaskStatus.RUNNING, }; for (DriverTaskStatus status : invalidStates) { - DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0, false); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); Map<FragmentInstanceId, Set<DriverTask>> fragmentRelatedTask = new ConcurrentHashMap<>(); @@ -156,7 +156,7 @@ public class DefaultDriverSchedulerTest { clear(); } DriverTask testTask = - new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, driverTaskHandle, 0); + new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, driverTaskHandle, 0, false); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); Map<FragmentInstanceId, Set<DriverTask>> fragmentRelatedTask = new ConcurrentHashMap<>(); @@ -198,7 +198,7 @@ public class DefaultDriverSchedulerTest { DriverTaskStatus.READY, }; for (DriverTaskStatus status : invalidStates) { - DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0, false); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); Map<FragmentInstanceId, Set<DriverTask>> fragmentRelatedTask = new ConcurrentHashMap<>(); @@ -214,7 +214,7 @@ public class DefaultDriverSchedulerTest { clear(); } DriverTask testTask = - new DriverTask(mockDriver, 100L, DriverTaskStatus.RUNNING, driverTaskHandle, 0); + new DriverTask(mockDriver, 100L, DriverTaskStatus.RUNNING, driverTaskHandle, 0, false); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); Map<FragmentInstanceId, Set<DriverTask>> fragmentRelatedTask = new ConcurrentHashMap<>(); @@ -260,7 +260,7 @@ public class DefaultDriverSchedulerTest { DriverTaskStatus.READY, }; for (DriverTaskStatus status : invalidStates) { - DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0, false); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); Map<FragmentInstanceId, Set<DriverTask>> fragmentRelatedTask = new ConcurrentHashMap<>(); @@ -276,7 +276,7 @@ public class DefaultDriverSchedulerTest { clear(); } DriverTask testTask = - new DriverTask(mockDriver, 100L, DriverTaskStatus.RUNNING, driverTaskHandle, 0); + new DriverTask(mockDriver, 100L, DriverTaskStatus.RUNNING, driverTaskHandle, 0, false); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); Map<FragmentInstanceId, Set<DriverTask>> fragmentRelatedTask = new ConcurrentHashMap<>(); @@ -322,7 +322,7 @@ public class DefaultDriverSchedulerTest { DriverTaskStatus.READY, }; for (DriverTaskStatus status : invalidStates) { - DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, status, driverTaskHandle, 0, false); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); Map<FragmentInstanceId, Set<DriverTask>> fragmentRelatedTask = new ConcurrentHashMap<>(); @@ -338,7 +338,7 @@ public class DefaultDriverSchedulerTest { clear(); } DriverTask testTask = - new DriverTask(mockDriver, 100L, DriverTaskStatus.RUNNING, driverTaskHandle, 0); + new DriverTask(mockDriver, 100L, DriverTaskStatus.RUNNING, driverTaskHandle, 0, false); Set<DriverTask> taskSet = new HashSet<>(); taskSet.add(testTask); Map<FragmentInstanceId, Set<DriverTask>> fragmentRelatedTask = new ConcurrentHashMap<>(); @@ -387,9 +387,9 @@ public class DefaultDriverSchedulerTest { DriverTaskStatus.FINISHED, DriverTaskStatus.ABORTED, }; for (DriverTaskStatus status : invalidStates) { - DriverTask testTask1 = new DriverTask(mockDriver1, 100L, status, driverTaskHandle, 0); + DriverTask testTask1 = new DriverTask(mockDriver1, 100L, status, driverTaskHandle, 0, false); DriverTask testTask2 = - new DriverTask(mockDriver2, 100L, DriverTaskStatus.BLOCKED, driverTaskHandle, 0); + new DriverTask(mockDriver2, 100L, DriverTaskStatus.BLOCKED, driverTaskHandle, 0, false); Set<DriverTask> taskSet1 = new HashSet<>(); taskSet1.add(testTask1); @@ -428,10 +428,10 @@ public class DefaultDriverSchedulerTest { Mockito.reset(mockDriver2); Mockito.when(mockDriver2.getDriverTaskId()).thenReturn(driverTaskId2); - DriverTask testTask1 = new DriverTask(mockDriver1, 100L, status, driverTaskHandle, 0); + DriverTask testTask1 = new DriverTask(mockDriver1, 100L, status, driverTaskHandle, 0, false); DriverTask testTask2 = - new DriverTask(mockDriver2, 100L, DriverTaskStatus.BLOCKED, driverTaskHandle, 0); + new DriverTask(mockDriver2, 100L, DriverTaskStatus.BLOCKED, driverTaskHandle, 0, false); Set<DriverTask> taskSet1 = new HashSet<>(); taskSet1.add(testTask1); Set<DriverTask> taskSet2 = new HashSet<>(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverTaskTimeoutSentinelThreadTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverTaskTimeoutSentinelThreadTest.java index 98366813ced..85575b136a1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverTaskTimeoutSentinelThreadTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverTaskTimeoutSentinelThreadTest.java @@ -73,28 +73,29 @@ public class DriverTaskTimeoutSentinelThreadTest { "0", new ThreadGroup("timeout-test"), taskQueue, mockScheduler, producer); // FINISHED status test - DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.FINISHED, null, 0); + DriverTask testTask = + new DriverTask(mockDriver, 100L, DriverTaskStatus.FINISHED, null, 0, false); executor.execute(testTask); Assert.assertEquals(DriverTaskStatus.FINISHED, testTask.getStatus()); Mockito.verify(mockDriver, Mockito.never()).processFor(Mockito.any()); Mockito.verify(mockDriver, Mockito.never()).failed(Mockito.any()); // ABORTED status test - testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.ABORTED, null, 0); + testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.ABORTED, null, 0, false); executor.execute(testTask); Assert.assertEquals(DriverTaskStatus.ABORTED, testTask.getStatus()); Mockito.verify(mockDriver, Mockito.never()).processFor(Mockito.any()); Mockito.verify(mockDriver, Mockito.never()).failed(Mockito.any()); // RUNNING status test - testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.RUNNING, null, 0); + testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.RUNNING, null, 0, false); executor.execute(testTask); Assert.assertEquals(DriverTaskStatus.RUNNING, testTask.getStatus()); Mockito.verify(mockDriver, Mockito.never()).processFor(Mockito.any()); Mockito.verify(mockDriver, Mockito.never()).failed(Mockito.any()); // BLOCKED status test - testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.BLOCKED, null, 0); + testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.BLOCKED, null, 0, false); executor.execute(testTask); Assert.assertEquals(DriverTaskStatus.BLOCKED, testTask.getStatus()); Mockito.verify(mockDriver, Mockito.never()).processFor(Mockito.any()); @@ -134,7 +135,7 @@ public class DriverTaskTimeoutSentinelThreadTest { AbstractDriverThread executor = new DriverTaskThread( "0", new ThreadGroup("timeout-test"), taskQueue, mockScheduler, producer); - DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0, false); executor.execute(testTask); Mockito.verify(mockDriver, Mockito.times(1)).processFor(Mockito.any()); Assert.assertEquals( @@ -175,7 +176,7 @@ public class DriverTaskTimeoutSentinelThreadTest { AbstractDriverThread executor = new DriverTaskThread( "0", new ThreadGroup("timeout-test"), taskQueue, mockScheduler, producer); - DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0, false); executor.execute(testTask); Mockito.verify(mockDriver, Mockito.times(1)).processFor(Mockito.any()); Assert.assertNull(testTask.getAbortCause()); @@ -225,7 +226,7 @@ public class DriverTaskTimeoutSentinelThreadTest { AbstractDriverThread executor = new DriverTaskThread( "0", new ThreadGroup("timeout-test"), taskQueue, mockScheduler, producer); - DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0, false); executor.execute(testTask); Mockito.verify(mockDriver, Mockito.times(1)).processFor(Mockito.any()); Assert.assertNull(testTask.getAbortCause()); @@ -276,7 +277,7 @@ public class DriverTaskTimeoutSentinelThreadTest { AbstractDriverThread executor = new DriverTaskThread( "0", new ThreadGroup("timeout-test"), taskQueue, mockScheduler, producer); - DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0, false); executor.execute(testTask); Mockito.verify(mockDriver, Mockito.times(1)).processFor(Mockito.any()); Assert.assertNull(testTask.getAbortCause()); @@ -317,7 +318,7 @@ public class DriverTaskTimeoutSentinelThreadTest { executor.close(); throw new RuntimeException("mock exception"); }); - DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0); + DriverTask testTask = new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, null, 0, false); taskQueue.push(testTask); executor.run(); // Here we use run() instead of start() to execute the task in the same thread Mockito.verify(mockDriver, Mockito.times(1)).processFor(Mockito.any()); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java new file mode 100644 index 00000000000..7a4934e7136 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java @@ -0,0 +1,157 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.execution.schedule.queue; + +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.driver.IDriver; +import org.apache.iotdb.db.queryengine.execution.schedule.DriverScheduler; +import org.apache.iotdb.db.queryengine.execution.schedule.queue.multilevelqueue.DriverTaskHandle; +import org.apache.iotdb.db.queryengine.execution.schedule.queue.multilevelqueue.MultilevelPriorityQueue; +import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTask; +import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTaskId; +import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTaskStatus; + +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import java.util.ArrayList; +import java.util.List; +import java.util.OptionalInt; + +public class MultilevelPriorityQueueTest { + @Test + public void testPollBlocked() throws InterruptedException { + IndexedBlockingQueue<DriverTask> queue = new MultilevelPriorityQueue(2, 1000, new DriverTask()); + List<DriverTask> res = new ArrayList<>(); + Thread t1 = + new Thread( + () -> { + try { + DriverTask e = queue.poll(); + res.add(e); + } catch (InterruptedException e) { + e.printStackTrace(); + Assert.fail(); + } + }); + t1.start(); + Thread.sleep(100); + Assert.assertEquals(Thread.State.WAITING, t1.getState()); + DriverTask e2 = mockDriverTask(mockDriverTaskId(), false); + queue.push(e2); + Thread.sleep(100); + Assert.assertEquals(Thread.State.TERMINATED, t1.getState()); + Assert.assertEquals(1, res.size()); + Assert.assertEquals(e2.getDriverTaskId().toString(), res.get(0).getDriverTaskId().toString()); + } + + @Test + public void testPushExceedCapacity() { + IndexedBlockingQueue<DriverTask> queue = new MultilevelPriorityQueue(2, 1, new DriverTask()); + DriverTask e1 = mockDriverTask(mockDriverTaskId(), false); + queue.push(e1); + DriverTask e2 = mockDriverTask(mockDriverTaskId(), false); + try { + queue.push(e2); + Assert.fail(); + } catch (IllegalStateException e) { + // ignore; + } + } + + @Test + public void testPushAndPoll() throws InterruptedException { + IndexedBlockingQueue<DriverTask> queue = new MultilevelPriorityQueue(2, 1000, new DriverTask()); + DriverTask e1 = mockDriverTask(mockDriverTaskId(), false); + queue.push(e1); + Assert.assertEquals(1, queue.size()); + DriverTask e2 = + mockDriverTask( + new DriverTaskId( + new FragmentInstanceId(new PlanFragmentId(new QueryId("test"), 0), "inst-1"), 0), + false); + queue.push(e2); + Assert.assertEquals(2, queue.size()); + Assert.assertEquals(e1.getDriverTaskId().toString(), queue.poll().getDriverTaskId().toString()); + Assert.assertEquals(1, queue.size()); + Assert.assertEquals(e2.getDriverTaskId().toString(), queue.poll().getDriverTaskId().toString()); + Assert.assertEquals(0, queue.size()); + } + + @Test + public void testClear() { + IndexedBlockingQueue<DriverTask> queue = new MultilevelPriorityQueue(2, 1000, new DriverTask()); + DriverTask e1 = mockDriverTask(mockDriverTaskId(), false); + queue.push(e1); + DriverTask e2 = + mockDriverTask( + new DriverTaskId( + new FragmentInstanceId(new PlanFragmentId(new QueryId("test"), 0), "inst-1"), 0), + false); + queue.push(e2); + Assert.assertEquals(2, queue.size()); + queue.clear(); + Assert.assertEquals(0, queue.size()); + } + + @Test + public void testPushAndPollWithHighestLevelPriority() throws InterruptedException { + MultilevelPriorityQueue queue = new MultilevelPriorityQueue(2, 1000, new DriverTask()); + DriverTask e1 = mockDriverTask(mockDriverTaskId(), true); + queue.push(e1); + Assert.assertEquals(1, queue.size()); + Assert.assertEquals(1, queue.getHighestPriorityLevelQueue().size()); + DriverTask e2 = + mockDriverTask( + new DriverTaskId( + new FragmentInstanceId(new PlanFragmentId(new QueryId("test"), 0), "inst-1"), 0), + false); + queue.push(e2); + Assert.assertEquals(2, queue.size()); + Assert.assertEquals(e1.getDriverTaskId().toString(), queue.poll().getDriverTaskId().toString()); + Assert.assertEquals(1, queue.size()); + Assert.assertEquals(0, queue.getHighestPriorityLevelQueue().size()); + Assert.assertEquals(e2.getDriverTaskId().toString(), queue.poll().getDriverTaskId().toString()); + Assert.assertEquals(0, queue.size()); + } + + private DriverTask mockDriverTask(DriverTaskId driverTaskID, boolean isHighestPriority) { + DriverScheduler manager = DriverScheduler.getInstance(); + IDriver mockDriver = Mockito.mock(IDriver.class); + DriverTaskHandle driverTaskHandle = + new DriverTaskHandle( + 1, + (MultilevelPriorityQueue) manager.getReadyQueue(), + OptionalInt.of(Integer.MAX_VALUE)); + Mockito.when(mockDriver.getDriverTaskId()).thenReturn(driverTaskID); + return new DriverTask( + mockDriver, 100L, DriverTaskStatus.READY, driverTaskHandle, 0, isHighestPriority); + } + + private DriverTaskId mockDriverTaskId() { + QueryId queryId = new QueryId("test"); + FragmentInstanceId instanceId = + new FragmentInstanceId(new PlanFragmentId(queryId, 0), "inst-0"); + return new DriverTaskId(instanceId, 0); + } +}
