synchronized ListenableFuture<?> moveQuery(QueryId queryId, MemoryPool targetMemoryPool) { long originalReserved = getQueryMemoryReservation(queryId); long originalRevocableReserved = getQueryRevocableMemoryReservation(queryId); // Get the tags before we call free() as that would remove the tags and we will lose the tags. Map<String, Long> taggedAllocations = taggedMemoryAllocations.remove(queryId); ListenableFuture<?> future = targetMemoryPool.reserve(queryId, MOVE_QUERY_TAG, originalReserved); free(queryId, MOVE_QUERY_TAG, originalReserved); targetMemoryPool.reserveRevocable(queryId, originalRevocableReserved); freeRevocable(queryId, originalRevocableReserved); targetMemoryPool.taggedMemoryAllocations.put(queryId, taggedAllocations); return future; }
private void requestMemoryRevoking(MemoryPool memoryPool, Collection<SqlTask> sqlTasks) { long remainingBytesToRevoke = (long) (-memoryPool.getFreeBytes() + (memoryPool.getMaxBytes() * (1.0 - memoryRevokingTarget))); remainingBytesToRevoke -= getMemoryAlreadyBeingRevoked(sqlTasks, memoryPool); requestRevoking(memoryPool, sqlTasks, remainingBytesToRevoke); }
@VisibleForTesting void registerPoolListeners() { memoryPools.forEach(memoryPool -> memoryPool.addListener(memoryPoolListener)); }
private boolean memoryRevokingNeeded(MemoryPool memoryPool) { return memoryPool.getReservedRevocableBytes() > 0 && memoryPool.getFreeBytes() <= memoryPool.getMaxBytes() * (1.0 - memoryRevokingThreshold); }
private synchronized ListenableFuture<?> updateSystemMemory(String allocationTag, long delta) { if (delta >= 0) { systemMemoryPool.reserve(queryId, allocationTag, delta); // Since various operators and the output buffers now support blocking when the system pool is full // we return NOT_BLOCKED to prevent them from blocking, which is the legacy behavior. return NOT_BLOCKED; } systemMemoryPool.free(queryId, allocationTag, -delta); return NOT_BLOCKED; }
for (TestingPrestoServer server : queryRunner.getServers()) { for (MemoryPool pool : server.getLocalMemoryManager().getPools()) { assertTrue(pool.tryReserve(fakeQueryId, "test", pool.getMaxBytes())); assertTrue(reserved.isPresent()); reserved.get().free(fakeQueryId, "test", reserved.get().getMaxBytes()); assertTrue(reserved.get().getFreeBytes() > 0);
@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 void configureLegacyMemoryPools(NodeMemoryConfig config, ReservedSystemMemoryConfig systemMemoryConfig, long availableMemory) { checkArgument(systemMemoryConfig.getReservedSystemMemory().toBytes() < availableMemory, "Reserved memory %s is greater than available heap %s", systemMemoryConfig.getReservedSystemMemory(), new DataSize(availableMemory, BYTE)); maxMemory = new DataSize(availableMemory - systemMemoryConfig.getReservedSystemMemory().toBytes(), BYTE); ImmutableMap.Builder<MemoryPoolId, MemoryPool> builder = ImmutableMap.builder(); checkArgument(config.getMaxQueryMemoryPerNode().toBytes() <= maxMemory.toBytes(), format("%s set to %s, but only %s of useable heap available", QUERY_MAX_MEMORY_PER_NODE_CONFIG, config.getMaxQueryMemoryPerNode(), maxMemory)); long generalPoolSize = maxMemory.toBytes(); if (config.isReservedPoolEnabled()) { builder.put(RESERVED_POOL, new MemoryPool(RESERVED_POOL, config.getMaxQueryMemoryPerNode())); generalPoolSize -= config.getMaxQueryMemoryPerNode().toBytes(); } builder.put(GENERAL_POOL, new MemoryPool(GENERAL_POOL, new DataSize(generalPoolSize, BYTE))); builder.put(SYSTEM_POOL, new MemoryPool(SYSTEM_POOL, systemMemoryConfig.getReservedSystemMemory())); this.pools = builder.build(); }
@Test public void testBlockingOnRevocableMemoryFreeUser() { setupConsumeRevocableMemory(ONE_BYTE, 10); assertTrue(userPool.tryReserve(fakeQueryId, "test", TEN_MEGABYTES_WITHOUT_TWO_BYTES.toBytes())); // we expect 2 iterations as we have 2 bytes remaining in memory pool and we allocate 1 byte per page assertEquals(runDriversUntilBlocked(waitingForRevocableSystemMemory()), 2); assertTrue(userPool.getFreeBytes() <= 0, String.format("Expected empty pool but got [%d]", userPool.getFreeBytes())); // lets free 5 bytes userPool.free(fakeQueryId, "test", 5); assertEquals(runDriversUntilBlocked(waitingForRevocableSystemMemory()), 5); assertTrue(userPool.getFreeBytes() <= 0, String.format("Expected empty pool but got [%d]", userPool.getFreeBytes())); // 3 more bytes is enough for driver to finish userPool.free(fakeQueryId, "test", 3); assertDriversProgress(waitingForRevocableSystemMemory()); assertEquals(userPool.getFreeBytes(), 10); }
MemoryPool reservedPool = new MemoryPool(RESERVED_POOL, new DataSize(10, BYTE)); long secondQueryMemory = reservedPool.getMaxBytes() - 1; if (useReservedPool) { assertTrue(reservedPool.reserve(secondQuery, "test", secondQueryMemory).isDone()); new DataSize(10, BYTE), new DataSize(20, BYTE), new MemoryPool(GENERAL_POOL, new DataSize(10, BYTE)), new TestingGcMonitor(), localQueryRunner.getExecutor(), reservedPool.free(secondQuery, "test", secondQueryMemory);
MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)); QueryContext queryContext = new QueryContext( queryId, long reservedMemoryInBytes = memoryPool.getFreeBytes() - additionalMemoryInBytes; memoryPool.reserve(queryId, "test", reservedMemoryInBytes); memoryPool.free(queryId, "test", reservedMemoryInBytes); memoryPool.free(queryId, "test", reservedMemoryInBytes); memoryPool.free(queryId, "test", reservedMemoryInBytes);
@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); }
/** * Try to reserve the given number of bytes. Return value indicates whether the caller may use the requested memory. */ public boolean tryReserve(QueryId queryId, String allocationTag, long bytes) { checkArgument(bytes >= 0, "bytes is negative"); synchronized (this) { if (getFreeBytes() - bytes < 0) { return false; } reservedBytes += bytes; if (bytes != 0) { queryMemoryReservations.merge(queryId, bytes, Long::sum); updateTaggedMemoryAllocations(queryId, allocationTag, bytes); } } onMemoryReserved(); return true; }
@Test public void testNotifyListenerOnMemoryReserved() { setupConsumeRevocableMemory(ONE_BYTE, 10); AtomicReference<MemoryPool> notifiedPool = new AtomicReference<>(); AtomicLong notifiedBytes = new AtomicLong(); userPool.addListener(MemoryPoolListener.onMemoryReserved(pool -> { notifiedPool.set(pool); notifiedBytes.set(pool.getReservedBytes()); })); userPool.reserve(fakeQueryId, "test", 3); assertEquals(notifiedPool.get(), userPool); assertEquals(notifiedBytes.get(), 3L); }
@Test public void testBlockingOnRevocableMemoryFreeViaRevoke() { RevocableMemoryOperator revocableMemoryOperator = setupConsumeRevocableMemory(ONE_BYTE, 5); assertTrue(userPool.tryReserve(fakeQueryId, "test", TEN_MEGABYTES_WITHOUT_TWO_BYTES.toBytes())); // we expect 2 iterations as we have 2 bytes remaining in memory pool and we allocate 1 byte per page assertEquals(runDriversUntilBlocked(waitingForRevocableSystemMemory()), 2); revocableMemoryOperator.getOperatorContext().requestMemoryRevoking(); // 2 more iterations assertEquals(runDriversUntilBlocked(waitingForRevocableSystemMemory()), 2); revocableMemoryOperator.getOperatorContext().requestMemoryRevoking(); // 3 more bytes is enough for driver to finish assertDriversProgress(waitingForRevocableSystemMemory()); assertEquals(userPool.getFreeBytes(), 2); }
private synchronized ListenableFuture<?> updateSystemMemory(String allocationTag, long delta) { // We call memoryPool.getQueryMemoryReservation(queryId) instead of calling queryMemoryContext.getUserMemory() to // calculate the total memory size. // // Calling the latter can result in a deadlock: // * A thread doing a user allocation will acquire locks in this order: // 1. monitor of queryMemoryContext.userAggregateMemoryContext // 2. monitor of this (QueryContext) // * The current thread doing a system allocation will acquire locks in this order: // 1. monitor of this (QueryContext) // 2. monitor of queryMemoryContext.userAggregateMemoryContext // Deadlock is possible for concurrent user and system allocations when updateSystemMemory()/updateUserMemory // calls queryMemoryContext.getUserMemory()/queryMemoryContext.getSystemMemory(), respectively. For concurrent // allocations of the same type (e.g., tryUpdateUserMemory/updateUserMemory) it is not possible as they share // the same RootAggregatedMemoryContext instance, and one of the threads will be blocked on the monitor of that // RootAggregatedMemoryContext instance even before calling the QueryContext methods (the monitors of // RootAggregatedMemoryContext instance and this will be acquired in the same order). long totalMemory = memoryPool.getQueryMemoryReservation(queryId); if (delta >= 0) { enforceTotalMemoryLimit(totalMemory, delta, maxTotalMemory); return memoryPool.reserve(queryId, allocationTag, delta); } memoryPool.free(queryId, allocationTag, -delta); return NOT_BLOCKED; }
@Override public synchronized String toString() { return toStringHelper(this) .add("id", id) .add("maxBytes", maxBytes) .add("freeBytes", getFreeBytes()) .add("reservedBytes", reservedBytes) .add("reservedRevocableBytes", reservedRevocableBytes) .add("future", future) .toString(); }
@Test public void testTrySetZeroBytesFullPool() { LocalMemoryContext localMemoryContext = operatorContext.localUserMemoryContext(); // fill up the pool memoryPool.reserve(new QueryId("test_query"), "test", memoryPool.getFreeBytes()); // try to reserve 0 bytes in the full pool assertTrue(localMemoryContext.trySetBytes(localMemoryContext.getBytes())); }
memoryPool.free(QUERY_ID, "test", memoryPool.getReservedBytes());
public ListenableFuture<?> reserveRevocable(QueryId queryId, long bytes) { checkArgument(bytes >= 0, "bytes is negative"); ListenableFuture<?> result; synchronized (this) { if (bytes != 0) { queryMemoryRevocableReservations.merge(queryId, bytes, Long::sum); } reservedRevocableBytes += bytes; if (getFreeBytes() <= 0) { if (future == null) { future = NonCancellableMemoryFuture.create(); } checkState(!future.isDone(), "future is already completed"); result = future; } else { result = NOT_BLOCKED; } } onMemoryReserved(); return result; }