From 44c84932edd8ffa78a81ac75914ae34931a81a52 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Wed, 22 Jul 2026 12:14:34 +0800 Subject: [PATCH 1/2] Fix query memory reservation deregistration race --- .../execution/memory/MemoryPool.java | 138 +++++++++++------- .../execution/memory/MemoryPoolTest.java | 131 +++++++++++++++++ 2 files changed, 216 insertions(+), 53 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 8e1c8d1319cfa..f25511a50a080 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -206,10 +206,13 @@ public int getMemoryReservationSize() { public void registerPlanNodeIdToQueryMemoryMap( String queryId, String fragmentInstanceId, String planNodeId) { synchronized (queryMemoryReservations) { - queryMemoryReservations - .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()) - .computeIfAbsent(fragmentInstanceId, x -> new ConcurrentHashMap<>()) - .putIfAbsent(planNodeId, 0L); + final Map> queryRelatedMemory = + queryMemoryReservations.computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()); + final Map fragmentRelatedMemory = + queryRelatedMemory.computeIfAbsent(fragmentInstanceId, x -> new ConcurrentHashMap<>()); + synchronized (fragmentRelatedMemory) { + fragmentRelatedMemory.putIfAbsent(planNodeId, 0L); + } } } @@ -224,44 +227,45 @@ public void registerPlanNodeIdToQueryMemoryMap( */ public void deRegisterFragmentInstanceFromQueryMemoryMap( String queryId, String fragmentInstanceId, boolean forceDeregister) { - Map> queryRelatedMemory = queryMemoryReservations.get(queryId); - if (queryRelatedMemory != null) { - Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); - boolean hasPotentialMemoryLeak = false; - // fragmentRelatedMemory could be null if the FI has not reserved any memory(For example, - // next() of root operator returns no data) - if (fragmentRelatedMemory != null) { - hasPotentialMemoryLeak = - fragmentRelatedMemory.values().stream().anyMatch(value -> value != 0L); - } - if (!forceDeregister && hasPotentialMemoryLeak) { - // If hasPotentialMemoryLeak is true, it means that LocalSourceChannel/LocalSourceHandles - // have not been closed. - // We should wait for them to be closed. Make sure this method is called again with - // forceDeregister == true, after all LocalSourceChannel/LocalSourceHandles are closed. + List> invalidEntryList = null; + synchronized (queryMemoryReservations) { + final Map> queryRelatedMemory = + queryMemoryReservations.get(queryId); + if (queryRelatedMemory == null) { return; } - synchronized (queryMemoryReservations) { - queryRelatedMemory.remove(fragmentInstanceId); - if (queryRelatedMemory.isEmpty()) { - queryMemoryReservations.remove(queryId); + final Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + if (fragmentRelatedMemory != null) { + synchronized (fragmentRelatedMemory) { + invalidEntryList = + fragmentRelatedMemory.entrySet().stream() + .filter(entry -> entry.getValue() != 0L) + .collect(Collectors.toList()); + if (!forceDeregister && !invalidEntryList.isEmpty()) { + // The LocalSourceChannels/LocalSourceHandles have not been closed yet. The caller will + // invoke this method again with forceDeregister == true after closing them. + return; + } + queryRelatedMemory.remove(fragmentInstanceId, fragmentRelatedMemory); } + } else { + // fragmentRelatedMemory could be null if the FI has not reserved any memory (for example, + // next() of root operator returns no data). + queryRelatedMemory.remove(fragmentInstanceId); } - if (hasPotentialMemoryLeak) { - // hasPotentialMemoryLeak means that fragmentRelatedMemory is not null - List> invalidEntryList = - fragmentRelatedMemory.entrySet().stream() - .filter(entry -> entry.getValue() != 0L) - .collect(Collectors.toList()); - throw new MemoryLeakException( - String.format( - DataNodeQueryMessages - .QUERY_EXCEPTION_PLANNODE_RELATED_MEMORY_IS_NOT_ZERO_WHEN_TRYING_TO_DEREGISTER_E01109C5, - queryId, - fragmentInstanceId, - invalidEntryList)); + if (queryRelatedMemory.isEmpty()) { + queryMemoryReservations.remove(queryId); } } + if (invalidEntryList != null && !invalidEntryList.isEmpty()) { + throw new MemoryLeakException( + String.format( + DataNodeQueryMessages + .QUERY_EXCEPTION_PLANNODE_RELATED_MEMORY_IS_NOT_ZERO_WHEN_TRYING_TO_DEREGISTER_E01109C5, + queryId, + fragmentInstanceId, + invalidEntryList)); + } } /** @@ -388,10 +392,17 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l Validate.isTrue(bytes > 0L); try { - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .computeIfPresent( + while (true) { + final Map> queryRelatedMemory = + queryMemoryReservations.get(queryId); + final Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + synchronized (fragmentRelatedMemory) { + // The fragment map may have been removed while this thread was waiting for its lock. + if (queryMemoryReservations.get(queryId) != queryRelatedMemory + || queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) { + continue; + } + fragmentRelatedMemory.computeIfPresent( planNodeId, (k, reservedMemory) -> { if (reservedMemory < bytes) { @@ -400,6 +411,9 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l } return reservedMemory - bytes; }); + break; + } + } } catch (NullPointerException e) { throw new IllegalArgumentException( DataNodeQueryMessages @@ -457,22 +471,40 @@ public boolean tryReserve( String planNodeId, long bytesToReserve, long maxBytesCanReserve) { - long tryUsedBytes = memoryBlock.forceAllocateWithoutLimitation(bytesToReserve); - long queryRemainingBytes = - maxBytesCanReserve - - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .merge(planNodeId, bytesToReserve, Long::sum); - return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() && queryRemainingBytes >= 0; + while (true) { + final Map> queryRelatedMemory = + queryMemoryReservations.get(queryId); + final Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + synchronized (fragmentRelatedMemory) { + // Serialize reservation changes with fragment deregistration while allowing reservations + // for other fragments to proceed concurrently. + if (queryMemoryReservations.get(queryId) != queryRelatedMemory + || queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) { + continue; + } + final long tryUsedBytes = memoryBlock.forceAllocateWithoutLimitation(bytesToReserve); + final long queryRemainingBytes = + maxBytesCanReserve - fragmentRelatedMemory.merge(planNodeId, bytesToReserve, Long::sum); + return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() && queryRemainingBytes >= 0; + } + } } private void rollbackReserve( String queryId, String fragmentInstanceId, String planNodeId, long bytesToReserve) { - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .merge(planNodeId, -bytesToReserve, Long::sum); - memoryBlock.release(bytesToReserve); + while (true) { + final Map> queryRelatedMemory = + queryMemoryReservations.get(queryId); + final Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + synchronized (fragmentRelatedMemory) { + if (queryMemoryReservations.get(queryId) != queryRelatedMemory + || queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) { + continue; + } + fragmentRelatedMemory.merge(planNodeId, -bytesToReserve, Long::sum); + memoryBlock.release(bytesToReserve); + return; + } + } } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index 8614d39f20a54..e6298e9230a6e 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -28,6 +28,21 @@ import org.junit.Test; import org.mockito.Mockito; +import java.lang.reflect.Field; +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.function.BiFunction; + public class MemoryPoolTest { MemoryPool pool; @@ -390,4 +405,120 @@ public void testReserveWithPriorityShowQueriesWithinAvailable() { Assert.assertTrue(r.getFuture().isDone()); Assert.assertEquals(256L, pool.getReservedBytes()); } + + @SuppressWarnings("unchecked") + @Test(timeout = 15000L) + public void testReserveIsAtomicWithFragmentDeregistration() throws Exception { + Field reservationsField = MemoryPool.class.getDeclaredField("queryMemoryReservations"); + reservationsField.setAccessible(true); + Map>> reservations = + (Map>>) reservationsField.get(pool); + + BlockingMemoryMap blockingMemoryMap = new BlockingMemoryMap(); + blockingMemoryMap.put(PLAN_NODE_ID, 0L); + reservations.get(QUERY_ID).put(FRAGMENT_INSTANCE_ID, blockingMemoryMap); + + ExecutorService executor = Executors.newFixedThreadPool(2); + Future reservation = null; + Future deregistration = null; + try { + reservation = + executor.submit( + () -> + pool.tryReserveForTest( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256L, Long.MAX_VALUE)); + Assert.assertTrue(blockingMemoryMap.reservationStarted.await(5, TimeUnit.SECONDS)); + + deregistration = + executor.submit( + () -> + pool.deRegisterFragmentInstanceFromQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, false)); + + if (blockingMemoryMap.deregistrationSnapshotTaken.await(1, TimeUnit.SECONDS)) { + // Without synchronization, deregistration can snapshot zero and remove the fragment while + // the reservation update is in progress. + blockingMemoryMap.allowDeregistrationSnapshot.countDown(); + deregistration.get(5, TimeUnit.SECONDS); + blockingMemoryMap.allowReservation.countDown(); + } else { + // With the fix, deregistration waits for the reservation update and then observes 256. + blockingMemoryMap.allowReservation.countDown(); + Assert.assertTrue(reservation.get(5, TimeUnit.SECONDS)); + Assert.assertTrue(blockingMemoryMap.deregistrationSnapshotTaken.await(5, TimeUnit.SECONDS)); + blockingMemoryMap.allowDeregistrationSnapshot.countDown(); + } + + Assert.assertTrue(reservation.get(5, TimeUnit.SECONDS)); + deregistration.get(5, TimeUnit.SECONDS); + Assert.assertEquals(256L, pool.getQueryMemoryReservedBytes(QUERY_ID)); + Assert.assertEquals(256L, pool.getReservedBytes()); + + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256L); + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, false); + Assert.assertEquals(0, pool.getQueryMemoryReservationSize()); + Assert.assertEquals(0L, pool.getReservedBytes()); + } finally { + blockingMemoryMap.allowReservation.countDown(); + blockingMemoryMap.allowDeregistrationSnapshot.countDown(); + if (reservation != null) { + reservation.cancel(true); + } + if (deregistration != null) { + deregistration.cancel(true); + } + executor.shutdownNow(); + } + } + + private static class BlockingMemoryMap extends ConcurrentHashMap { + + private final CountDownLatch reservationStarted = new CountDownLatch(1); + private final CountDownLatch allowReservation = new CountDownLatch(1); + private final CountDownLatch deregistrationSnapshotTaken = new CountDownLatch(1); + private final CountDownLatch allowDeregistrationSnapshot = new CountDownLatch(1); + + @Override + public Long merge( + String key, + Long value, + BiFunction remappingFunction) { + reservationStarted.countDown(); + await(allowReservation); + return super.merge(key, value, remappingFunction); + } + + @Override + public Collection values() { + Collection snapshot = new ArrayList<>(super.values()); + waitForDeregistration(); + return snapshot; + } + + @Override + public Set> entrySet() { + Set> snapshot = new HashSet<>(); + for (Map.Entry entry : super.entrySet()) { + snapshot.add(new AbstractMap.SimpleImmutableEntry<>(entry)); + } + waitForDeregistration(); + return snapshot; + } + + private void waitForDeregistration() { + deregistrationSnapshotTaken.countDown(); + await(allowDeregistrationSnapshot); + } + + private static void await(CountDownLatch latch) { + try { + if (!latch.await(5, TimeUnit.SECONDS)) { + throw new AssertionError(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + } } From 7af40afe6a102128d3bb3c190e75f4355cfd9d38 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Wed, 22 Jul 2026 18:47:47 +0800 Subject: [PATCH 2/2] Handle deregistration winning reservation race --- .../iotdb/db/i18n/DataNodeQueryMessages.java | 4 +- .../iotdb/db/i18n/DataNodeQueryMessages.java | 4 +- .../execution/memory/MemoryPool.java | 73 ++++++++++++------- .../execution/memory/MemoryPoolTest.java | 66 +++++++++++++++++ 4 files changed, 117 insertions(+), 30 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index b540ebb25404c..10f3c69d5e5f7 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -2748,8 +2748,8 @@ private DataNodeQueryMessages() {} public static final String QUERY_EXCEPTION_QUERY_IS_ABORTED_SINCE_IT_REQUESTS_MORE_MEMORY_THAN_CAN_D77C2921 = "Query is aborted since it requests more memory than can be allocated, bytesToReserve: %sB, " + "maxBytesCanReserve: %sB"; - public static final String QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2 = - "RelatedMemoryReserved can't be null when freeing memory"; + public static final String EXCEPTION_NO_MEMORY_RESERVATION_IS_REGISTERED_FOR_QUERY_ARG_AND_FRAGMENT_INSTANCE_ARG_06016520 = + "No memory reservation is registered for query %s and fragment instance %s."; public static final String QUERY_EXCEPTION_INTERRUPTED_BY_92FAED2D = "Interrupted By"; public static final String QUERY_EXCEPTION_DRIVER_WAS_INTERRUPTED_737358E4 = "Driver was interrupted"; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 0e73bd1dd967e..bf99a9850e9c7 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -3269,8 +3269,8 @@ private DataNodeQueryMessages() {} "查询已中止，因为其请求的内存超过可分配上限，bytesToReserve：%sB，maxBytesCanReserve：%sB"; - public static final String QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2 = - "释放内存时 RelatedMemoryReserved 不能为 null"; + public static final String EXCEPTION_NO_MEMORY_RESERVATION_IS_REGISTERED_FOR_QUERY_ARG_AND_FRAGMENT_INSTANCE_ARG_06016520 = + "查询 %s、FragmentInstance %s 未注册内存预留。"; public static final String QUERY_EXCEPTION_INTERRUPTED_BY_92FAED2D = "被中断，原因："; public static final String QUERY_EXCEPTION_DRIVER_WAS_INTERRUPTED_737358E4 = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index f25511a50a080..cbc83b9836b48 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -391,34 +391,33 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l Validate.notNull(queryId, DataNodeQueryMessages.EXCEPTION_QUERYID_CAN_NOT_BE_NULL_DOT_16639DBE); Validate.isTrue(bytes > 0L); - try { - while (true) { - final Map> queryRelatedMemory = - queryMemoryReservations.get(queryId); - final Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); - synchronized (fragmentRelatedMemory) { - // The fragment map may have been removed while this thread was waiting for its lock. - if (queryMemoryReservations.get(queryId) != queryRelatedMemory - || queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) { - continue; - } - fragmentRelatedMemory.computeIfPresent( - planNodeId, - (k, reservedMemory) -> { - if (reservedMemory < bytes) { - throw new IllegalArgumentException( - DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); - } - return reservedMemory - bytes; - }); - break; + while (true) { + final Map> queryRelatedMemory = + queryMemoryReservations.get(queryId); + if (queryRelatedMemory == null) { + throw memoryReservationNotRegistered(queryId, fragmentInstanceId); + } + final Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + if (fragmentRelatedMemory == null) { + throw memoryReservationNotRegistered(queryId, fragmentInstanceId); + } + synchronized (fragmentRelatedMemory) { + // The fragment map may have been removed while this thread was waiting for its lock. + if (queryMemoryReservations.get(queryId) != queryRelatedMemory + || queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) { + continue; } + fragmentRelatedMemory.computeIfPresent( + planNodeId, + (k, reservedMemory) -> { + if (reservedMemory < bytes) { + throw new IllegalArgumentException( + DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); + } + return reservedMemory - bytes; + }); + break; } - } catch (NullPointerException e) { - throw new IllegalArgumentException( - DataNodeQueryMessages - .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2, - e); } memoryBlock.release(bytes); @@ -474,7 +473,13 @@ public boolean tryReserve( while (true) { final Map> queryRelatedMemory = queryMemoryReservations.get(queryId); + if (queryRelatedMemory == null) { + throw memoryReservationNotRegistered(queryId, fragmentInstanceId); + } final Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + if (fragmentRelatedMemory == null) { + throw memoryReservationNotRegistered(queryId, fragmentInstanceId); + } synchronized (fragmentRelatedMemory) { // Serialize reservation changes with fragment deregistration while allowing reservations // for other fragments to proceed concurrently. @@ -495,7 +500,13 @@ private void rollbackReserve( while (true) { final Map> queryRelatedMemory = queryMemoryReservations.get(queryId); + if (queryRelatedMemory == null) { + throw memoryReservationNotRegistered(queryId, fragmentInstanceId); + } final Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + if (fragmentRelatedMemory == null) { + throw memoryReservationNotRegistered(queryId, fragmentInstanceId); + } synchronized (fragmentRelatedMemory) { if (queryMemoryReservations.get(queryId) != queryRelatedMemory || queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) { @@ -507,4 +518,14 @@ private void rollbackReserve( } } } + + private IllegalArgumentException memoryReservationNotRegistered( + String queryId, String fragmentInstanceId) { + return new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_NO_MEMORY_RESERVATION_IS_REGISTERED_FOR_QUERY_ARG_AND_FRAGMENT_INSTANCE_ARG_06016520, + queryId, + fragmentInstanceId)); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index e6298e9230a6e..5ca7ec10e5bb8 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -37,6 +37,7 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -471,6 +472,50 @@ public void testReserveIsAtomicWithFragmentDeregistration() throws Exception { } } + @SuppressWarnings("unchecked") + @Test(timeout = 15000L) + public void testDeregistrationWinsBeforeReserveAcquiresFragmentLock() throws Exception { + Field reservationsField = MemoryPool.class.getDeclaredField("queryMemoryReservations"); + reservationsField.setAccessible(true); + Map>> reservations = + (Map>>) reservationsField.get(pool); + + BlockingQueryMap blockingQueryMap = new BlockingQueryMap(); + blockingQueryMap.putAll(reservations.get(QUERY_ID)); + reservations.put(QUERY_ID, blockingQueryMap); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future reservation = null; + try { + reservation = + executor.submit( + () -> { + blockingQueryMap.blockCurrentThreadOnGet(); + return pool.tryReserveForTest( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256L, Long.MAX_VALUE); + }); + Assert.assertTrue(blockingQueryMap.fragmentMapRead.await(5, TimeUnit.SECONDS)); + + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, false); + Assert.assertEquals(0, pool.getQueryMemoryReservationSize()); + blockingQueryMap.allowGetToReturn.countDown(); + + try { + reservation.get(5, TimeUnit.SECONDS); + Assert.fail("Expected the reservation to fail after deregistration"); + } catch (ExecutionException e) { + Assert.assertTrue(e.getCause() instanceof IllegalArgumentException); + } + Assert.assertEquals(0L, pool.getReservedBytes()); + } finally { + blockingQueryMap.allowGetToReturn.countDown(); + if (reservation != null) { + reservation.cancel(true); + } + executor.shutdownNow(); + } + } + private static class BlockingMemoryMap extends ConcurrentHashMap { private final CountDownLatch reservationStarted = new CountDownLatch(1); @@ -521,4 +566,25 @@ private static void await(CountDownLatch latch) { } } } + + private static class BlockingQueryMap extends ConcurrentHashMap> { + + private final CountDownLatch fragmentMapRead = new CountDownLatch(1); + private final CountDownLatch allowGetToReturn = new CountDownLatch(1); + private volatile Thread blockedThread; + + private void blockCurrentThreadOnGet() { + blockedThread = Thread.currentThread(); + } + + @Override + public Map get(Object key) { + Map result = super.get(key); + if (Thread.currentThread() == blockedThread) { + fragmentMapRead.countDown(); + BlockingMemoryMap.await(allowGetToReturn); + } + return result; + } + } }