private static int completeGroupQueries(Set<MockQueryExecution> groupQueries) { int groupRan = 0; for (Iterator<MockQueryExecution> iterator = groupQueries.iterator(); iterator.hasNext(); ) { MockQueryExecution query = iterator.next(); if (query.getState() == RUNNING) { query.complete(); iterator.remove(); groupRan++; } } return groupRan; }
@Override public void start() { state = RUNNING; fireStateChange(); }
@Override public BasicQueryInfo getBasicQueryInfo() { return new BasicQueryInfo(getQueryInfo()); }
@Test(timeOut = 10_000) public void testMemoryLimit() { RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor()); root.setSoftMemoryLimit(new DataSize(1, BYTE)); root.setMaxQueuedQueries(4); root.setHardConcurrencyLimit(3); MockQueryExecution query1 = new MockQueryExecution(2); root.run(query1); // Process the group to refresh stats root.processQueuedQueries(); assertEquals(query1.getState(), RUNNING); MockQueryExecution query2 = new MockQueryExecution(0); root.run(query2); assertEquals(query2.getState(), QUEUED); MockQueryExecution query3 = new MockQueryExecution(0); root.run(query3); assertEquals(query3.getState(), QUEUED); query1.complete(); root.processQueuedQueries(); assertEquals(query2.getState(), RUNNING); assertEquals(query3.getState(), RUNNING); }
@Test(timeOut = 10_000) public void testQueueFull() { RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(1); root.setHardConcurrencyLimit(1); MockQueryExecution query1 = new MockQueryExecution(0); root.run(query1); assertEquals(query1.getState(), RUNNING); MockQueryExecution query2 = new MockQueryExecution(0); root.run(query2); assertEquals(query2.getState(), QUEUED); MockQueryExecution query3 = new MockQueryExecution(0); root.run(query3); assertEquals(query3.getState(), FAILED); assertEquals(query3.getThrowable().getMessage(), "Too many queued queries for \"root\""); }
group2.setMaxQueuedQueries(4); group2.setHardConcurrencyLimit(2); MockQueryExecution query1a = new MockQueryExecution(0); group1.run(query1a); assertEquals(query1a.getState(), RUNNING); MockQueryExecution query1b = new MockQueryExecution(0); group1.run(query1b); assertEquals(query1b.getState(), QUEUED); MockQueryExecution query1c = new MockQueryExecution(0); group1.run(query1c); assertEquals(query1c.getState(), QUEUED); MockQueryExecution query2a = new MockQueryExecution(0); group2.run(query2a); assertEquals(query2a.getState(), QUEUED);
private static Set<MockQueryExecution> fillGroupTo(InternalResourceGroup group, Set<MockQueryExecution> existingQueries, int count, boolean queryPriority) { int existingCount = existingQueries.size(); Set<MockQueryExecution> queries = new HashSet<>(existingQueries); for (int i = 0; i < count - existingCount; i++) { MockQueryExecution query = new MockQueryExecution(0, group.getId().toString().replace(".", "") + Integer.toString(i), queryPriority ? i + 1 : 1); queries.add(query); group.run(query); } return queries; }
@Override public boolean isDone() { return getState().isDone(); }
@Test public void testSubgroupMemoryLimit() { RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor()); root.setSoftMemoryLimit(new DataSize(10, BYTE)); root.setMaxQueuedQueries(4); root.setHardConcurrencyLimit(3); InternalResourceGroup subgroup = root.getOrCreateSubGroup("subgroup"); subgroup.setSoftMemoryLimit(new DataSize(1, BYTE)); subgroup.setMaxQueuedQueries(4); subgroup.setHardConcurrencyLimit(3); MockQueryExecution query1 = new MockQueryExecution(2); subgroup.run(query1); // Process the group to refresh stats root.processQueuedQueries(); assertEquals(query1.getState(), RUNNING); MockQueryExecution query2 = new MockQueryExecution(0); subgroup.run(query2); assertEquals(query2.getState(), QUEUED); MockQueryExecution query3 = new MockQueryExecution(0); subgroup.run(query3); assertEquals(query3.getState(), QUEUED); query1.complete(); root.processQueuedQueries(); assertEquals(query2.getState(), RUNNING); assertEquals(query3.getState(), RUNNING); }
@Setup public void setup() { root = new RootInternalResourceGroup("root", (group, export) -> {}, executor); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(queries); root.setHardConcurrencyLimit(queries); InternalResourceGroup group = root; for (int i = 0; i < children; i++) { group = root.getOrCreateSubGroup(String.valueOf(i)); group.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); group.setMaxQueuedQueries(queries); group.setHardConcurrencyLimit(queries); } for (int i = 0; i < queries; i++) { group.run(new MockQueryExecution(10)); } }
@Test(timeOut = 10_000) public void testSoftCpuLimit() { RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor()); root.setSoftMemoryLimit(new DataSize(1, BYTE)); root.setSoftCpuLimit(new Duration(1, SECONDS)); root.setHardCpuLimit(new Duration(2, SECONDS)); root.setCpuQuotaGenerationMillisPerSecond(2000); root.setMaxQueuedQueries(1); root.setHardConcurrencyLimit(2); MockQueryExecution query1 = new MockQueryExecution(1, "query_id", 1, new Duration(1, SECONDS)); root.run(query1); assertEquals(query1.getState(), RUNNING); MockQueryExecution query2 = new MockQueryExecution(0); root.run(query2); assertEquals(query2.getState(), RUNNING); MockQueryExecution query3 = new MockQueryExecution(0); root.run(query3); assertEquals(query3.getState(), QUEUED); query1.complete(); root.processQueuedQueries(); assertEquals(query2.getState(), RUNNING); assertEquals(query3.getState(), QUEUED); root.generateCpuQuota(2); root.processQueuedQueries(); assertEquals(query2.getState(), RUNNING); assertEquals(query3.getState(), RUNNING); }
@Override public Optional<ErrorCode> getErrorCode() { return Optional.ofNullable(getQueryInfo().getFailureInfo()).map(ExecutionFailureInfo::getErrorCode); }
@Override public void cancelQuery() { state = FAILED; fireStateChange(); }
group2.setMaxQueuedQueries(4); group2.setHardConcurrencyLimit(2); MockQueryExecution query1a = new MockQueryExecution(0); group1.run(query1a); assertEquals(query1a.getState(), RUNNING); MockQueryExecution query1b = new MockQueryExecution(0); group1.run(query1b); assertEquals(query1b.getState(), QUEUED); MockQueryExecution query1c = new MockQueryExecution(0); group1.run(query1c); assertEquals(query1c.getState(), QUEUED); MockQueryExecution query2a = new MockQueryExecution(0); group2.run(query2a); assertEquals(query2a.getState(), QUEUED); query1a.complete(); root.processQueuedQueries(); assertEquals(query1b.getState(), RUNNING); assertEquals(query1c.getState(), QUEUED); assertEquals(query2a.getState(), QUEUED); query1b.complete(); root.processQueuedQueries(); assertEquals(query2a.getState(), RUNNING); assertEquals(query1c.getState(), QUEUED);
@Override public Optional<DateTime> getEndTime() { return Optional.ofNullable(getQueryInfo().getQueryStats().getEndTime()); }
@Override public void fail(Throwable cause) { state = FAILED; failureCause = cause; fireStateChange(); }
@Test(timeOut = 10_000) public void testHardCpuLimit() { RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor()); root.setSoftMemoryLimit(new DataSize(1, BYTE)); root.setHardCpuLimit(new Duration(1, SECONDS)); root.setCpuQuotaGenerationMillisPerSecond(2000); root.setMaxQueuedQueries(1); root.setHardConcurrencyLimit(1); MockQueryExecution query1 = new MockQueryExecution(1, "query_id", 1, new Duration(2, SECONDS)); root.run(query1); assertEquals(query1.getState(), RUNNING); MockQueryExecution query2 = new MockQueryExecution(0); root.run(query2); assertEquals(query2.getState(), QUEUED); query1.complete(); root.processQueuedQueries(); assertEquals(query2.getState(), QUEUED); root.generateCpuQuota(2); root.processQueuedQueries(); assertEquals(query2.getState(), RUNNING); }
@Override public Optional<DateTime> getExecutionStartTime() { return Optional.ofNullable(getQueryInfo().getQueryStats().getExecutionStartTime()); }