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 8e1c8d1319cfa..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 @@ -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)); + } } /** @@ -387,24 +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 { - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .computeIfPresent( - planNodeId, - (k, reservedMemory) -> { - if (reservedMemory < bytes) { - throw new IllegalArgumentException( - DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); - } - return reservedMemory - bytes; - }); - } catch (NullPointerException e) { - throw new IllegalArgumentException( - DataNodeQueryMessages - .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2, - e); + 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; + } } memoryBlock.release(bytes); @@ -457,22 +470,62 @@ 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); + 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. + 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); + 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) { + continue; + } + fragmentRelatedMemory.merge(planNodeId, -bytesToReserve, Long::sum); + memoryBlock.release(bytesToReserve); + return; + } + } + } + + 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 8614d39f20a54..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 @@ -28,6 +28,22 @@ 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.ExecutionException; +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 +406,185 @@ 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(); + } + } + + @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); + 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); + } + } + } + + 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; + } + } }