public TaskContext build() { MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), memoryPoolSize); SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(maxSpillSize); DefaultQueryContext queryContext = new DefaultQueryContext( queryId, queryMaxMemory, queryMaxTotalMemory, memoryPool, GC_MONITOR, notificationExecutor, yieldExecutor, queryMaxSpillSize, spillSpaceTracker); return createTaskContext(queryContext, session, taskStateMachine); } }
@GET @Path("{poolId}") public Response getMemoryInfo(@PathParam("poolId") String poolId) { if (GENERAL_POOL.getId().equals(poolId)) { return toSuccessfulResponse(memoryManager.getGeneralPool().getInfo()); } if (SYSTEM_POOL.getId().equals(poolId) && memoryManager.getSystemPool().isPresent()) { return toSuccessfulResponse(memoryManager.getSystemPool().get().getInfo()); } if (RESERVED_POOL.getId().equals(poolId) && memoryManager.getReservedPool().isPresent()) { return toSuccessfulResponse(memoryManager.getReservedPool().get().getInfo()); } return Response.status(NOT_FOUND).build(); }
@Override public Optional<QueryId> chooseQueryToKill(List<QueryMemoryInfo> runningQueries, List<MemoryInfo> nodes) { QueryId biggestQuery = null; long maxMemory = 0; for (QueryMemoryInfo query : runningQueries) { long bytesUsed = query.getMemoryReservation(); if (bytesUsed > maxMemory && GENERAL_POOL.equals(query.getMemoryPoolId())) { biggestQuery = query.getQueryId(); maxMemory = bytesUsed; } } return Optional.ofNullable(biggestQuery); } }
TEST_SESSION.toSessionRepresentation(), RUNNING, new MemoryPoolId("reserved"), false, URI.create("1"), assertEquals(basicInfo.getMemoryPool().getId(), "reserved"); assertEquals(basicInfo.isScheduled(), false); assertEquals(basicInfo.getQuery(), "SELECT 4");
@PreDestroy public synchronized void destroy() throws IOException { try (Closer closer = Closer.create()) { for (ClusterMemoryPool pool : pools.values()) { closer.register(() -> exporter.unexport(generatedNameOf(ClusterMemoryPool.class, pool.getId().toString()))); } closer.register(listenerExecutor::shutdownNow); } }
private Map<MemoryPoolId, ClusterMemoryPool> createClusterMemoryPools(boolean systemPoolEnabled, boolean reservedPoolEnabled) { Set<MemoryPoolId> memoryPools = new HashSet<>(); memoryPools.add(GENERAL_POOL); if (systemPoolEnabled) { memoryPools.add(SYSTEM_POOL); } if (reservedPoolEnabled) { memoryPools.add(RESERVED_POOL); } ImmutableMap.Builder<MemoryPoolId, ClusterMemoryPool> builder = ImmutableMap.builder(); for (MemoryPoolId poolId : memoryPools) { ClusterMemoryPool pool = new ClusterMemoryPool(poolId); builder.put(poolId, pool); try { exporter.export(generatedNameOf(ClusterMemoryPool.class, poolId.toString()), pool); } catch (JmxException e) { log.error(e, "Error exporting memory pool %s", poolId); } } return builder.build(); }
protected AbstractResourceConfigurationManager(ClusterMemoryPoolManager memoryPoolManager) { memoryPoolManager.addChangeListener(new MemoryPoolId("general"), poolInfo -> { Map<ResourceGroup, DataSize> memoryLimits = new HashMap<>(); synchronized (generalPoolMemoryFraction) { for (Map.Entry<ResourceGroup, Double> entry : generalPoolMemoryFraction.entrySet()) { double bytes = poolInfo.getMaxBytes() * entry.getValue(); // setSoftMemoryLimit() acquires a lock on the root group of its tree, which could cause a deadlock if done while holding the "generalPoolMemoryFraction" lock memoryLimits.put(entry.getKey(), new DataSize(bytes, BYTE)); } generalPoolBytes = poolInfo.getMaxBytes(); } for (Map.Entry<ResourceGroup, DataSize> entry : memoryLimits.entrySet()) { entry.getKey().setSoftMemoryLimit(entry.getValue()); } }); }
@PreDestroy public synchronized void destroy() { for (MemoryPool pool : pools) { String objectName = ObjectNames.builder(MemoryPool.class, pool.getId().toString()).build(); try { exporter.unexport(objectName); } catch (JmxException e) { // ignored } } pools.clear(); } }
@Override public synchronized void updateMemoryPoolAssignments(MemoryPoolAssignmentsRequest assignments) { if (coordinatorId != null && coordinatorId.equals(assignments.getCoordinatorId()) && assignments.getVersion() <= currentMemoryPoolAssignmentVersion) { return; } currentMemoryPoolAssignmentVersion = assignments.getVersion(); if (coordinatorId != null && !coordinatorId.equals(assignments.getCoordinatorId())) { log.warn("Switching coordinator affinity from " + coordinatorId + " to " + assignments.getCoordinatorId()); } coordinatorId = assignments.getCoordinatorId(); for (MemoryPoolAssignment assignment : assignments.getAssignments()) { if (assignment.getPoolId().equals(GENERAL_POOL)) { queryContexts.getUnchecked(assignment.getQueryId()).setMemoryPool(localMemoryManager.getGeneralPool()); } else if (assignment.getPoolId().equals(RESERVED_POOL)) { MemoryPool reservedPool = localMemoryManager.getReservedPool() .orElseThrow(() -> new IllegalArgumentException(format("Cannot move %s to the reserved pool as the reserved pool is not enabled", assignment.getQueryId()))); queryContexts.getUnchecked(assignment.getQueryId()).setMemoryPool(reservedPool); } else { new IllegalArgumentException(format("Cannot move %s to %s as the target memory pool id is invalid", assignment.getQueryId(), assignment.getPoolId())); } } }
@Test public void testMoveQuery() { QueryId testQuery = new QueryId("test_query"); MemoryPool pool1 = new MemoryPool(new MemoryPoolId("test"), new DataSize(1000, BYTE)); MemoryPool pool2 = new MemoryPool(new MemoryPoolId("test"), new DataSize(1000, BYTE)); pool1.reserve(testQuery, "test_tag", 10); Map<String, Long> allocations = pool1.getTaggedMemoryAllocations().get(testQuery); assertEquals(allocations, ImmutableMap.of("test_tag", 10L)); pool1.moveQuery(testQuery, pool2); assertNull(pool1.getTaggedMemoryAllocations().get(testQuery)); allocations = pool2.getTaggedMemoryAllocations().get(testQuery); assertEquals(allocations, ImmutableMap.of("test_tag", 10L)); assertEquals(pool1.getFreeBytes(), 1000); assertEquals(pool2.getFreeBytes(), 990); pool2.free(testQuery, "test", 10); assertEquals(pool2.getFreeBytes(), 1000); }
private synchronized void addPool(MemoryPool pool) { try { String objectName = ObjectNames.builder(MemoryPool.class, pool.getId().toString()).build(); exporter.export(objectName, pool); pools.add(pool); } catch (JmxException e) { // ignored } }
@Test public void testTaggedAllocations() { QueryId testQuery = new QueryId("test_query"); MemoryPool testPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1000, BYTE)); testPool.reserve(testQuery, "test_tag", 10); Map<String, Long> allocations = testPool.getTaggedMemoryAllocations().get(testQuery); assertEquals(allocations, ImmutableMap.of("test_tag", 10L)); // free 5 bytes for test_tag testPool.free(testQuery, "test_tag", 5); assertEquals(allocations, ImmutableMap.of("test_tag", 5L)); testPool.reserve(testQuery, "test_tag2", 20); assertEquals(allocations, ImmutableMap.of("test_tag", 5L, "test_tag2", 20L)); // free the remaining 5 bytes for test_tag testPool.free(testQuery, "test_tag", 5); assertEquals(allocations, ImmutableMap.of("test_tag2", 20L)); // free all for test_tag2 testPool.free(testQuery, "test_tag2", 20); assertEquals(testPool.getTaggedMemoryAllocations().size(), 0); }
.setSystemProperty("optimizer.optimize-hash-generation", "true") .build(); MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)); SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE));
session.toSessionRepresentation(), state, new MemoryPoolId("test"), !state.isDone(), URI.create("http://test"),
TEST_SESSION.toSessionRepresentation(), state, new MemoryPoolId("reserved"), true, URI.create("1"),
private TaskContext newTestingTaskContext(ScheduledExecutorService taskNotificationExecutor, ScheduledExecutorService driverYieldExecutor, TaskStateMachine taskStateMachine) { DefaultQueryContext queryContext = new DefaultQueryContext( new QueryId("queryid"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), taskNotificationExecutor, driverYieldExecutor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE))); return queryContext.addTaskContext(taskStateMachine, TEST_SESSION, false, false, OptionalInt.empty()); }
@BeforeMethod public void setUpTest() { memoryPool = new MemoryPool(new MemoryPoolId("test"), memoryPoolSize); queryContext = new DefaultQueryContext( new QueryId("test_query"), queryMaxMemory, queryMaxTotalMemory, memoryPool, new TestingGcMonitor(), notificationExecutor, yieldExecutor, queryMaxSpillSize, spillSpaceTracker); taskContext = queryContext.addTaskContext( new TaskStateMachine(new TaskId("query", 0, 0), notificationExecutor), testSessionBuilder().build(), true, true, OptionalInt.empty()); pipelineContext = taskContext.addPipelineContext(0, true, true, false); driverContext = pipelineContext.addDriverContext(); operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("a"), "test-operator"); }
@Test public void testCountAlreadyRevokedMemoryWithinAPool() throws Exception { // Given SqlTask sqlTask1 = newSqlTask(); MemoryPool anotherMemoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(10, BYTE)); sqlTask1.getQueryContext().setMemoryPool(anotherMemoryPool); OperatorContext operatorContext1 = createContexts(sqlTask1); SqlTask sqlTask2 = newSqlTask(); OperatorContext operatorContext2 = createContexts(sqlTask2); List<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2); MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(asList(memoryPool, anotherMemoryPool), () -> tasks, executor, 1.0, 1.0); allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2); /* * sqlTask1 fills its pool */ operatorContext1.localRevocableMemoryContext().setBytes(12); requestMemoryRevoking(scheduler); assertMemoryRevokingRequestedFor(operatorContext1); /* * When sqlTask2 fills its pool */ operatorContext2.localRevocableMemoryContext().setBytes(12); requestMemoryRevoking(scheduler); /* * Then sqlTask2 should be asked to revoke its memory too */ assertMemoryRevokingRequestedFor(operatorContext1, operatorContext2); }
MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)); SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE)); DefaultQueryContext queryContext = new DefaultQueryContext(taskId.getQueryId(),
public SqlTask createInitialTask() { TaskId taskId = new TaskId("query", 0, nextTaskId.incrementAndGet()); URI location = URI.create("fake://task/" + taskId); DefaultQueryContext queryContext = new DefaultQueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), taskNotificationExecutor, driverYieldExecutor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE))); queryContext.addTaskContext(new TaskStateMachine(taskId, taskNotificationExecutor), testSessionBuilder().build(), false, false, OptionalInt.empty()); return createSqlTask( taskId, location, "fake", queryContext, sqlTaskExecutionFactory, taskNotificationExecutor, Functions.identity(), new DataSize(32, MEGABYTE), new CounterStat()); } }