@Override public synchronized void setMemoryPool(MemoryPool newMemoryPool) { // This method first acquires the monitor of this instance. // After that in this method if we acquire the monitors of the // user/revocable memory contexts in the queryMemoryContext instance // (say, by calling queryMemoryContext.getUserMemory()) it's possible // to have a deadlock. Because, the driver threads running the operators // will allocate memory concurrently through the child memory context -> ... -> // root memory context -> this.updateUserMemory() calls, and will acquire // the monitors of the user/revocable memory contexts in the queryMemoryContext instance // first, and then the monitor of this, which may cause deadlocks. // That's why instead of calling methods on queryMemoryContext to get the // user/revocable memory reservations, we call the MemoryPool to get the same // information. requireNonNull(newMemoryPool, "newMemoryPool is null"); if (memoryPool == newMemoryPool) { // Don't unblock our tasks and thrash the pools, if this is a no-op return; } ListenableFuture<?> future = memoryPool.moveQuery(queryId, newMemoryPool); memoryPool = newMemoryPool; future.addListener(() -> { // Unblock all the tasks, if they were waiting for memory, since we're in a new pool. taskContexts.values().forEach(TaskContext::moreMemoryAvailable); }, directExecutor()); }
@Override public synchronized void setMemoryPool(MemoryPool newMemoryPool) { // This method first acquires the monitor of this instance. // After that in this method if we acquire the monitors of the // user/revocable memory contexts in the queryMemoryContext instance // (say, by calling queryMemoryContext.getUserMemory()) it's possible // to have a deadlock. Because, the driver threads running the operators // will allocate memory concurrently through the child memory context -> ... -> // root memory context -> this.updateUserMemory() calls, and will acquire // the monitors of the user/revocable memory contexts in the queryMemoryContext instance // first, and then the monitor of this, which may cause deadlocks. // That's why instead of calling methods on queryMemoryContext to get the // user/revocable memory reservations, we call the MemoryPool to get the same // information. requireNonNull(newMemoryPool, "newMemoryPool is null"); if (memoryPool == newMemoryPool) { // Don't unblock our tasks and thrash the pools, if this is a no-op return; } ListenableFuture<?> future = memoryPool.moveQuery(queryId, newMemoryPool); memoryPool = newMemoryPool; future.addListener(() -> { // Unblock all the tasks, if they were waiting for memory, since we're in a new pool. taskContexts.values().forEach(TaskContext::moreMemoryAvailable); }, directExecutor()); }
@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); }