advancedxy commented on code in PR #2492: URL: https://github.com/apache/uniffle/pull/2492#discussion_r2137594099
########## server/src/main/java/org/apache/uniffle/server/buffer/lab/ChunkCreator.java: ########## @@ -0,0 +1,316 @@ +/* + * 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.uniffle.server.buffer.lab; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.LongAdder; + +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.apache.hadoop.util.StringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Does the management of LAB chunk creations. A monotonically incrementing id is associated with + * every chunk + */ +public class ChunkCreator { + private static final Logger LOG = LoggerFactory.getLogger(ChunkCreator.class); + // monotonically increasing chunkid. Starts at 1. + private final AtomicInteger chunkID = new AtomicInteger(1); + + // mapping from chunk IDs to chunks + private final Map<Integer, Chunk> chunkIdMap = new ConcurrentHashMap<Integer, Chunk>(); + static ChunkCreator instance; + private final int maxAlloc; + private ChunkPool chunksPool; + private final int chunkSize; + + ChunkCreator(int chunkSize, long bufferCapacity, int maxAlloc) { + this.chunkSize = chunkSize; + this.maxAlloc = maxAlloc; + initializePools(chunkSize, bufferCapacity); + } + + private void initializePools(int chunkSize, long bufferCapacity) { + this.chunksPool = initializePool(bufferCapacity, chunkSize); + } + + /** + * Initializes the instance of ChunkCreator + * + * @param chunkSize the chunkSize + * @param bufferCapacity the buffer capacity + * @return singleton ChunkCreator + */ + public static synchronized void initialize(int chunkSize, long bufferCapacity, int maxAlloc) { + if (instance != null) { + return; + } + instance = new ChunkCreator(chunkSize, bufferCapacity, maxAlloc); + } + + public static ChunkCreator getInstance() { + return instance; + } + + /** + * Creates and inits a chunk with specific index type and type. + * + * @return the chunk that was initialized + */ + Chunk getChunk() { + return getChunk(chunksPool.getChunkSize()); + } + + /** + * Creates and inits a chunk. + * + * @return the chunk that was initialized + * @param size the size of the chunk to be allocated, in bytes + */ + Chunk getChunk(int size) { + Chunk chunk = null; + ChunkPool pool = null; + + // if the size is suitable for one of the pools + if (chunksPool != null && size == chunksPool.getChunkSize()) { + pool = chunksPool; + } + + if (pool != null) { + chunk = pool.getChunk(); + if (chunk == null) { + LOG.warn( + "The chunk pool is full. Reached maxCount= " + + pool.getMaxCount() + + ". Creating chunk outside of the pool."); + } + } + + if (chunk == null) { + chunk = createChunk(false, size); + } + chunk.init(); + return chunk; + } + + /** + * Creates the chunk + * + * @param pool indicates if the chunks have to be created which will be used by the Pool + * @param size the size of the chunk to be allocated, in bytes + * @return the chunk + */ + private Chunk createChunk(boolean pool, int size) { + Chunk chunk; + int id = chunkID.getAndIncrement(); + assert id > 0; Review Comment: I don't think so .. chunkID should always greater than zero? Even if it's possible to overflow, we should use Preconditions.check instead of assert, which will be no-op if assertion if not enabled. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
