public synchronized MemoryPoolInfo getInfo() { return new MemoryPoolInfo( totalDistributedBytes, reservedDistributedBytes, reservedRevocableDistributedBytes, ImmutableMap.copyOf(queryMemoryReservations), ImmutableMap.copyOf(queryMemoryAllocations), ImmutableMap.copyOf(queryMemoryRevocableReservations)); }
public synchronized MemoryPoolInfo getInfo() { Map<QueryId, List<MemoryAllocation>> memoryAllocations = new HashMap<>(); for (Entry<QueryId, Map<String, Long>> entry : taggedMemoryAllocations.entrySet()) { List<MemoryAllocation> allocations = new ArrayList<>(); if (entry.getValue() != null) { entry.getValue().forEach((tag, allocation) -> allocations.add(new MemoryAllocation(tag, allocation))); } memoryAllocations.put(entry.getKey(), allocations); } return new MemoryPoolInfo(maxBytes, reservedBytes, reservedRevocableBytes, queryMemoryReservations, memoryAllocations, queryMemoryRevocableReservations); }
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()); } }); }
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); } }
if (poolInfo != null) { nodes++; if (poolInfo.getFreeBytes() + poolInfo.getReservedRevocableBytes() <= 0) { blockedNodes++; totalDistributedBytes += poolInfo.getMaxBytes(); reservedDistributedBytes += poolInfo.getReservedBytes(); reservedRevocableDistributedBytes += poolInfo.getReservedRevocableBytes(); for (Map.Entry<QueryId, Long> entry : poolInfo.getQueryMemoryReservations().entrySet()) { queryMemoryReservations.merge(entry.getKey(), entry.getValue(), Long::sum); for (Map.Entry<QueryId, List<MemoryAllocation>> entry : poolInfo.getQueryMemoryAllocations().entrySet()) { queryMemoryAllocations.merge(entry.getKey(), entry.getValue(), this::mergeQueryAllocations); for (Map.Entry<QueryId, Long> entry : poolInfo.getQueryMemoryRevocableReservations().entrySet()) { queryMemoryRevocableReservations.merge(entry.getKey(), entry.getValue(), Long::sum);
private void logQueryKill(QueryId killedQueryId, List<MemoryInfo> nodes) { if (!log.isInfoEnabled()) { return; } StringBuilder nodeDescription = new StringBuilder(); nodeDescription.append("Query Kill Decision: Killed ").append(killedQueryId).append("\n"); for (MemoryInfo node : nodes) { MemoryPoolInfo memoryPoolInfo = node.getPools().get(GENERAL_POOL); if (memoryPoolInfo == null) { continue; } nodeDescription.append("Query Kill Scenario: "); nodeDescription.append("MaxBytes ").append(memoryPoolInfo.getMaxBytes()).append(' '); nodeDescription.append("FreeBytes ").append(memoryPoolInfo.getFreeBytes() + memoryPoolInfo.getReservedRevocableBytes()).append(' '); nodeDescription.append("Queries "); Joiner.on(",").withKeyValueSeparator("=").appendTo(nodeDescription, memoryPoolInfo.getQueryMemoryReservations()); nodeDescription.append('\n'); } log.info(nodeDescription.toString()); }
private List<MemoryAllocation> mergeQueryAllocations(List<MemoryAllocation> left, List<MemoryAllocation> right) { requireNonNull(left, "left is null"); requireNonNull(right, "right is null"); Map<String, MemoryAllocation> mergedAllocations = new HashMap<>(); for (MemoryAllocation allocation : left) { mergedAllocations.put(allocation.getTag(), allocation); } for (MemoryAllocation allocation : right) { mergedAllocations.merge( allocation.getTag(), allocation, (a, b) -> new MemoryAllocation(a.getTag(), a.getAllocation() + b.getAllocation())); } return new ArrayList<>(mergedAllocations.values()); }
@Override public Optional<QueryId> chooseQueryToKill(List<QueryMemoryInfo> runningQueries, List<MemoryInfo> nodes) { Map<QueryId, Long> memoryReservationOnBlockedNodes = new HashMap<>(); for (MemoryInfo node : nodes) { MemoryPoolInfo generalPool = node.getPools().get(GENERAL_POOL); if (generalPool == null) { continue; } if (generalPool.getFreeBytes() + generalPool.getReservedRevocableBytes() > 0) { continue; } Map<QueryId, Long> queryMemoryReservations = generalPool.getQueryMemoryReservations(); queryMemoryReservations.forEach((queryId, memoryReservation) -> { memoryReservationOnBlockedNodes.compute(queryId, (id, oldValue) -> oldValue == null ? memoryReservation : oldValue + memoryReservation); }); } return memoryReservationOnBlockedNodes.entrySet().stream() .max(comparingLong(Map.Entry::getValue)) .map(Map.Entry::getKey); } }
@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(); }
@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); } }
@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 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 } }
@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 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); }
@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(); } }
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()); }
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(); }
@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); }
@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"); }