SpriCoder commented on code in PR #14710: URL: https://github.com/apache/iotdb/pull/14710#discussion_r1971610209
########## iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryManager.java: ########## @@ -0,0 +1,523 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.memory; + +import org.apache.iotdb.commons.utils.TestOnly; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.LongUnaryOperator; + +public class MemoryManager { + private static final Logger LOGGER = LoggerFactory.getLogger(MemoryManager.class); + + /** The max retry times for memory allocation */ + private static final int MEMORY_ALLOCATE_MAX_RETRIES = 3; + + /** The retry interval for memory allocation */ + private static final long MEMORY_ALLOCATE_RETRY_INTERVAL_IN_MS = 1000; + + /** The min memory size to allocate */ + private static final long MEMORY_ALLOCATE_MIN_SIZE_IN_BYTES = 32; + + /** The name of memory manager */ + private final String name; + + /** Whether memory management is enabled */ + private final boolean enable; + + /** The total memory size in byte of memory manager */ + private long totalMemorySizeInBytes; + + /** The allocated memory size */ + private long allocatedMemorySizeInBytes = 0L; + + /** The parent memory manager */ + private final MemoryManager parentMemoryManager; + + /** The child memory manager */ + private final Map<String, MemoryManager> children = new ConcurrentHashMap<>(); + + /** The allocated memory blocks of this memory manager */ + private final Map<String, IMemoryBlock> allocatedMemoryBlocks = new ConcurrentHashMap<>(); + + @TestOnly + public MemoryManager(long totalMemorySizeInBytes) { + this.name = "Test"; + this.parentMemoryManager = null; + this.totalMemorySizeInBytes = totalMemorySizeInBytes; + this.enable = false; + } + + private MemoryManager( + String name, MemoryManager parentMemoryManager, long totalMemorySizeInBytes) { + this.name = name; + this.parentMemoryManager = parentMemoryManager; + this.totalMemorySizeInBytes = totalMemorySizeInBytes; + this.enable = false; + } + + private MemoryManager( + String name, MemoryManager parentMemoryManager, long totalMemorySizeInBytes, boolean enable) { + this.name = name; + this.parentMemoryManager = parentMemoryManager; + this.totalMemorySizeInBytes = totalMemorySizeInBytes; + this.enable = enable; + } + + // region The Methods Of IMemoryBlock Management + + /** + * Try to force allocate memory block with specified size in bytes + * + * @param name the name of memory block + * @param sizeInBytes the size in bytes of memory block try to allocate + * @param type the type of memory block + * @return the memory block if success, otherwise throw MemoryException + */ + public synchronized IMemoryBlock forceAllocate( + String name, long sizeInBytes, MemoryBlockType type) { + if (!enable) { + return registerMemoryBlock(name, sizeInBytes, type); + } + for (int i = 0; i < MEMORY_ALLOCATE_MAX_RETRIES; i++) { + if (totalMemorySizeInBytes - allocatedMemorySizeInBytes >= sizeInBytes) { + return registerMemoryBlock(name, sizeInBytes, type); + } + + try { + // TODO @spricoder: consider to find more memory in active way + Thread.sleep(MEMORY_ALLOCATE_RETRY_INTERVAL_IN_MS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.warn("forceAllocate: interrupted while waiting for available memory", e); + } + } + + throw new MemoryException( + String.format( + "forceAllocate: failed to allocate memory after %d retries, " + + "total memory size %d bytes, used memory size %d bytes, " + + "requested memory size %d bytes", + MEMORY_ALLOCATE_MAX_RETRIES, + totalMemorySizeInBytes, + allocatedMemorySizeInBytes, + sizeInBytes)); + } + + /** + * Try to force allocate memory block with total memory size in bytes + * + * @param name the name of memory block + * @param memoryBlockType the type of memory block + */ + public synchronized IMemoryBlock forceAllocate(String name, MemoryBlockType memoryBlockType) { + return forceAllocate( + name, totalMemorySizeInBytes - allocatedMemorySizeInBytes, memoryBlockType); + } + + /** + * Try to force allocate memory block with specified size in bytes when memory is sufficient. + * + * @param name the name of memory block + * @param sizeInBytes the size in bytes of memory block try to allocate + * @param maxRatio the used threshold of allocatedMemorySizeInBytes / totalMemorySizeInBytes + * @param memoryBlockType the type of memory block + * @return the memory block if success, otherwise null + */ + public synchronized IMemoryBlock forceAllocateIfSufficient( + String name, long sizeInBytes, float maxRatio, MemoryBlockType memoryBlockType) { + if (maxRatio < 0.0f || maxRatio > 1.0f) { + return null; + } + if (!enable) { + return registerMemoryBlock(name, sizeInBytes, memoryBlockType); + } + if (totalMemorySizeInBytes - allocatedMemorySizeInBytes >= sizeInBytes + && (float) allocatedMemorySizeInBytes / totalMemorySizeInBytes < maxRatio) { + return forceAllocate(name, sizeInBytes, memoryBlockType); + } else { + // TODO @spricoder: consider to find more memory in active way + LOGGER.debug( + "forceAllocateIfSufficient: failed to allocate memory, " + + "total memory size {} bytes, used memory size {} bytes, " + + "requested memory size {} bytes, used threshold {}", + totalMemorySizeInBytes, + allocatedMemorySizeInBytes, + sizeInBytes, + maxRatio); + } + + return null; + } + + /** + * Try to allocate memory block with customAllocateStrategy + * + * @param name the name of memory block + * @param sizeInBytes the size in bytes of memory block try to allocate + * @param customAllocateStrategy the custom allocate strategy when memory is insufficient + * @param type the type of memory block + * @return the memory block if success, otherwise null + */ + public synchronized IMemoryBlock tryAllocate( + String name, + long sizeInBytes, + LongUnaryOperator customAllocateStrategy, + MemoryBlockType type) { + if (!enable) { + return registerMemoryBlock(name, sizeInBytes, type); + } + + if (totalMemorySizeInBytes - allocatedMemorySizeInBytes >= sizeInBytes) { + return registerMemoryBlock(name, sizeInBytes, type); + } + + long sizeToAllocateInBytes = sizeInBytes; + while (sizeToAllocateInBytes > MEMORY_ALLOCATE_MIN_SIZE_IN_BYTES) { + if (totalMemorySizeInBytes - allocatedMemorySizeInBytes >= sizeToAllocateInBytes) { + LOGGER.debug( + "tryAllocate: allocated memory, " + + "total memory size {} bytes, used memory size {} bytes, " + + "original requested memory size {} bytes, " + + "actual requested memory size {} bytes", + totalMemorySizeInBytes, + allocatedMemorySizeInBytes, + sizeInBytes, + sizeToAllocateInBytes); + return registerMemoryBlock(name, sizeToAllocateInBytes, type); + } + + sizeToAllocateInBytes = + Math.max( + customAllocateStrategy.applyAsLong(sizeToAllocateInBytes), + MEMORY_ALLOCATE_MIN_SIZE_IN_BYTES); + } + + // TODO @spricoder: consider to find more memory in active way + LOGGER.warn( + "tryAllocate: failed to allocate memory, " + + "total memory size {} bytes, used memory size {} bytes, " + + "requested memory size {} bytes", + totalMemorySizeInBytes, + allocatedMemorySizeInBytes, + sizeInBytes); + return registerMemoryBlock(name, 0, type); + } + + /** + * Try to register memory block with specified size in bytes + * + * @param name the name of memory block, UNIQUE + * @param sizeInBytes the size in bytes of memory block + * @param type the type of memory block + * @return the memory block + */ + private IMemoryBlock registerMemoryBlock(String name, long sizeInBytes, MemoryBlockType type) { + if (sizeInBytes <= 0) { + LOGGER.warn("forceAllocate {}: sizeInBytes should be positive", name); + } + return allocatedMemoryBlocks.compute( + name, + (blockName, block) -> { + if (block != null) { + LOGGER.warn("register memory block already exists: {}", block); + return block; + } else { + allocatedMemorySizeInBytes += sizeInBytes; + return new AtomicLongMemoryBlock(name, this, sizeInBytes, type); + } + }); + } + + /** + * Release memory block and notify all waiting threads + * + * @param block the memory block to release + */ + public synchronized void release(IMemoryBlock block) { + if (block == null || block.isReleased()) { + return; + } + releaseWithOutNotify(block); + this.notifyAll(); + } + + /** + * Release memory block without notify + * + * @param block the memory block to release + */ + public synchronized void releaseWithOutNotify(IMemoryBlock block) { + if (block == null || block.isReleased()) { + return; + } + + block.markAsReleased(); + allocatedMemorySizeInBytes -= block.getTotalMemorySizeInBytes(); + allocatedMemoryBlocks.remove(block.getName()); + try { + block.close(); + } catch (Exception e) { + LOGGER.error("releaseWithOutNotify: failed to close memory block", e); + } + } + + // endregion + + // region The Methods Of MemoryManager Management + + /** + * Try to create a new memory manager with specified name and total memory size in bytes, then put + * it into children map. NOTICE: if there are same name memory manager, it will return the + * existing one. + * + * @param name the name of memory manager + * @param sizeInBytes the total memory size in bytes of memory manager + * @param enable whether memory management is enabled + * @return the memory manager + */ + public synchronized MemoryManager getOrCreateMemoryManager( + String name, long sizeInBytes, boolean enable) { + if (sizeInBytes <= 0) { + LOGGER.warn("getOrCreateMemoryManager {}: sizeInBytes should be positive", name); + } Review Comment: Change to throw exception here -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
