| /* |
| * 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; |
| |
| import java.io.File; |
| import java.time.Duration; |
| import java.util.Arrays; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.concurrent.ThreadPoolExecutor; |
| import java.util.concurrent.TimeUnit; |
| import java.util.concurrent.atomic.AtomicLong; |
| import java.util.concurrent.locks.ReentrantReadWriteLock; |
| |
| import com.google.common.collect.RangeMap; |
| import com.google.common.util.concurrent.Uninterruptibles; |
| import io.prometheus.client.Collector; |
| import org.awaitility.Awaitility; |
| import org.junit.jupiter.api.BeforeEach; |
| import org.junit.jupiter.api.Test; |
| import org.junit.jupiter.api.io.TempDir; |
| import org.roaringbitmap.longlong.Roaring64NavigableMap; |
| |
| import org.apache.uniffle.common.PartitionRange; |
| import org.apache.uniffle.common.RemoteStorageInfo; |
| import org.apache.uniffle.common.ShuffleDataResult; |
| import org.apache.uniffle.common.ShufflePartitionedData; |
| import org.apache.uniffle.common.config.ConfigUtils; |
| import org.apache.uniffle.common.rpc.StatusCode; |
| import org.apache.uniffle.common.util.ByteBufUtils; |
| import org.apache.uniffle.common.util.Constants; |
| import org.apache.uniffle.server.DefaultFlushEventHandler; |
| import org.apache.uniffle.server.ShuffleFlushManager; |
| import org.apache.uniffle.server.ShuffleServer; |
| import org.apache.uniffle.server.ShuffleServerConf; |
| import org.apache.uniffle.server.ShuffleServerMetrics; |
| import org.apache.uniffle.server.ShuffleTaskManager; |
| import org.apache.uniffle.server.storage.StorageManager; |
| import org.apache.uniffle.server.storage.StorageManagerFactory; |
| import org.apache.uniffle.storage.util.StorageType; |
| |
| import static org.junit.jupiter.api.Assertions.assertArrayEquals; |
| import static org.junit.jupiter.api.Assertions.assertEquals; |
| import static org.junit.jupiter.api.Assertions.assertFalse; |
| import static org.junit.jupiter.api.Assertions.assertNotNull; |
| import static org.junit.jupiter.api.Assertions.assertNull; |
| import static org.junit.jupiter.api.Assertions.assertTrue; |
| import static org.junit.jupiter.api.Assertions.fail; |
| import static org.mockito.ArgumentMatchers.any; |
| import static org.mockito.Mockito.mock; |
| import static org.mockito.Mockito.times; |
| import static org.mockito.Mockito.verify; |
| import static org.mockito.Mockito.when; |
| |
| public class ShuffleBufferManagerTest extends BufferTestBase { |
| private ShuffleBufferManager shuffleBufferManager; |
| private ShuffleFlushManager mockShuffleFlushManager; |
| private ShuffleServer mockShuffleServer; |
| private ShuffleTaskManager mockShuffleTaskManager; |
| private ShuffleServerConf conf; |
| |
| @BeforeEach |
| public void setUp(@TempDir File tmpDir) { |
| conf = new ShuffleServerConf(); |
| File dataDir = new File(tmpDir, "data"); |
| conf.setString(ShuffleServerConf.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name()); |
| conf.set(ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(dataDir.getAbsolutePath())); |
| conf.set(ShuffleServerConf.SERVER_BUFFER_CAPACITY, 500L); |
| conf.set(ShuffleServerConf.SERVER_MEMORY_SHUFFLE_LOWWATERMARK_PERCENTAGE, 20.0); |
| conf.set(ShuffleServerConf.SERVER_MEMORY_SHUFFLE_HIGHWATERMARK_PERCENTAGE, 80.0); |
| conf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 1024L); |
| mockShuffleFlushManager = mock(ShuffleFlushManager.class); |
| mockShuffleServer = mock(ShuffleServer.class); |
| mockShuffleTaskManager = mock(ShuffleTaskManager.class); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| shuffleBufferManager = new ShuffleBufferManager(conf, mockShuffleFlushManager, false); |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| } |
| |
| @Test |
| public void registerBufferTest() { |
| String appId = "registerBufferTest"; |
| int shuffleId = 1; |
| |
| StatusCode sc = shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| assertEquals(StatusCode.SUCCESS, sc); |
| sc = shuffleBufferManager.registerBuffer(appId, shuffleId, 2, 3); |
| assertEquals(StatusCode.SUCCESS, sc); |
| |
| Map<String, Map<Integer, RangeMap<Integer, ShuffleBuffer>>> bufferPool = |
| shuffleBufferManager.getBufferPool(); |
| |
| assertNotNull(bufferPool.get(appId).get(shuffleId).get(0)); |
| ShuffleBuffer buffer = bufferPool.get(appId).get(shuffleId).get(0); |
| assertEquals(buffer, bufferPool.get(appId).get(shuffleId).get(1)); |
| assertNotNull(bufferPool.get(appId).get(shuffleId).get(2)); |
| assertEquals( |
| bufferPool.get(appId).get(shuffleId).get(2), bufferPool.get(appId).get(shuffleId).get(3)); |
| |
| // register again |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| assertEquals(buffer, bufferPool.get(appId).get(shuffleId).get(0)); |
| } |
| |
| @Test |
| public void getShuffleDataWithExpectedTaskIdsTest() { |
| String appId = "getShuffleDataWithExpectedTaskIdsTest"; |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| shuffleBufferManager.registerBuffer(appId, 1, 0, 1); |
| ShufflePartitionedData spd1 = createData(0, 1, 68); |
| ShufflePartitionedData spd2 = createData(0, 2, 68); |
| ShufflePartitionedData spd3 = createData(0, 1, 68); |
| ShufflePartitionedData spd4 = createData(0, 3, 68); |
| shuffleBufferManager.cacheShuffleData(appId, 1, false, spd1); |
| shuffleBufferManager.cacheShuffleData(appId, 1, false, spd2); |
| shuffleBufferManager.cacheShuffleData(appId, 1, false, spd3); |
| shuffleBufferManager.cacheShuffleData(appId, 1, false, spd4); |
| |
| /** case1: all blocks in cached and read multiple times */ |
| ShuffleDataResult result = |
| shuffleBufferManager.getShuffleData( |
| appId, 1, 0, Constants.INVALID_BLOCK_ID, 60, Roaring64NavigableMap.bitmapOf(1)); |
| assertEquals(1, result.getBufferSegments().size()); |
| assertEquals(0, result.getBufferSegments().get(0).getOffset()); |
| assertEquals(68, result.getBufferSegments().get(0).getLength()); |
| |
| // 2nd read |
| long lastBlockId = result.getBufferSegments().get(0).getBlockId(); |
| result = |
| shuffleBufferManager.getShuffleData( |
| appId, 1, 0, lastBlockId, 60, Roaring64NavigableMap.bitmapOf(1)); |
| assertEquals(1, result.getBufferSegments().size()); |
| assertEquals(0, result.getBufferSegments().get(0).getOffset()); |
| assertEquals(68, result.getBufferSegments().get(0).getLength()); |
| } |
| |
| @Test |
| public void getShuffleDataTest() { |
| String appId = "getShuffleDataTest"; |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| shuffleBufferManager.registerBuffer(appId, 1, 0, 1); |
| shuffleBufferManager.registerBuffer(appId, 2, 0, 1); |
| shuffleBufferManager.registerBuffer(appId, 3, 0, 1); |
| shuffleBufferManager.registerBuffer(appId, 4, 0, 1); |
| ShufflePartitionedData spd1 = createData(0, 68); |
| ShufflePartitionedData spd2 = createData(0, 68); |
| ShufflePartitionedData spd3 = createData(0, 68); |
| ShufflePartitionedData spd4 = createData(0, 68); |
| shuffleBufferManager.cacheShuffleData(appId, 1, false, spd1); |
| shuffleBufferManager.cacheShuffleData(appId, 2, false, spd2); |
| shuffleBufferManager.cacheShuffleData(appId, 2, false, spd3); |
| shuffleBufferManager.cacheShuffleData(appId, 3, false, spd4); |
| // validate buffer, no flush happened |
| Map<String, Map<Integer, RangeMap<Integer, ShuffleBuffer>>> bufferPool = |
| shuffleBufferManager.getBufferPool(); |
| assertEquals(100, bufferPool.get(appId).get(1).get(0).getSize()); |
| assertEquals(200, bufferPool.get(appId).get(2).get(0).getSize()); |
| assertEquals(100, bufferPool.get(appId).get(3).get(0).getSize()); |
| // validate get shuffle data |
| ShuffleDataResult sdr = |
| shuffleBufferManager.getShuffleData(appId, 2, 0, Constants.INVALID_BLOCK_ID, 60); |
| assertArrayEquals(ByteBufUtils.readBytes(spd2.getBlockList()[0].getData()), sdr.getData()); |
| long lastBlockId = spd2.getBlockList()[0].getBlockId(); |
| sdr = shuffleBufferManager.getShuffleData(appId, 2, 0, lastBlockId, 100); |
| assertArrayEquals(ByteBufUtils.readBytes(spd3.getBlockList()[0].getData()), sdr.getData()); |
| // flush happen |
| ShufflePartitionedData spd5 = createData(0, 10); |
| shuffleBufferManager.cacheShuffleData(appId, 4, false, spd5); |
| // according to flush strategy, some buffers should be moved to inFlushMap |
| assertEquals(0, bufferPool.get(appId).get(1).get(0).getBlocks().size()); |
| assertEquals(1, bufferPool.get(appId).get(1).get(0).getInFlushBlockMap().size()); |
| assertEquals(0, bufferPool.get(appId).get(2).get(0).getBlocks().size()); |
| assertEquals(1, bufferPool.get(appId).get(2).get(0).getInFlushBlockMap().size()); |
| assertEquals(0, bufferPool.get(appId).get(3).get(0).getBlocks().size()); |
| assertEquals(1, bufferPool.get(appId).get(3).get(0).getInFlushBlockMap().size()); |
| // keep buffer whose size < low watermark |
| assertEquals(1, bufferPool.get(appId).get(4).get(0).getBlocks().size()); |
| // data in flush buffer now, it also can be got before flush finish |
| sdr = shuffleBufferManager.getShuffleData(appId, 2, 0, Constants.INVALID_BLOCK_ID, 60); |
| assertArrayEquals(ByteBufUtils.readBytes(spd2.getBlockList()[0].getData()), sdr.getData()); |
| lastBlockId = spd2.getBlockList()[0].getBlockId(); |
| sdr = shuffleBufferManager.getShuffleData(appId, 2, 0, lastBlockId, 100); |
| assertArrayEquals(ByteBufUtils.readBytes(spd3.getBlockList()[0].getData()), sdr.getData()); |
| // cache data again, it should cause flush |
| spd1 = createData(0, 10); |
| shuffleBufferManager.cacheShuffleData(appId, 1, false, spd1); |
| assertEquals(1, bufferPool.get(appId).get(1).get(0).getBlocks().size()); |
| // finish flush |
| bufferPool.get(appId).get(1).get(0).getInFlushBlockMap().clear(); |
| bufferPool.get(appId).get(2).get(0).getInFlushBlockMap().clear(); |
| bufferPool.get(appId).get(3).get(0).getInFlushBlockMap().clear(); |
| // empty data return |
| sdr = shuffleBufferManager.getShuffleData(appId, 2, 0, Constants.INVALID_BLOCK_ID, 60); |
| assertEquals(0, sdr.getDataLength()); |
| lastBlockId = spd2.getBlockList()[0].getBlockId(); |
| sdr = shuffleBufferManager.getShuffleData(appId, 2, 0, lastBlockId, 100); |
| assertEquals(0, sdr.getDataLength()); |
| } |
| |
| @Test |
| public void shuffleIdToSizeTest() { |
| String appId1 = "shuffleIdToSizeTest1"; |
| String appId2 = "shuffleIdToSizeTest2"; |
| ReentrantReadWriteLock rwLock1 = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId1)).thenReturn(rwLock1.readLock()); |
| ReentrantReadWriteLock rwLock2 = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId2)).thenReturn(rwLock2.readLock()); |
| shuffleBufferManager.registerBuffer(appId1, 1, 0, 0); |
| shuffleBufferManager.registerBuffer(appId1, 2, 0, 0); |
| shuffleBufferManager.registerBuffer(appId2, 1, 0, 0); |
| shuffleBufferManager.registerBuffer(appId2, 2, 0, 0); |
| ShufflePartitionedData spd1 = createData(0, 67); |
| ShufflePartitionedData spd2 = createData(0, 68); |
| ShufflePartitionedData spd3 = createData(0, 68); |
| ShufflePartitionedData spd4 = createData(0, 68); |
| shuffleBufferManager.cacheShuffleData(appId1, 1, false, spd1); |
| shuffleBufferManager.cacheShuffleData(appId1, 2, false, spd2); |
| shuffleBufferManager.cacheShuffleData(appId1, 2, false, spd3); |
| shuffleBufferManager.cacheShuffleData(appId2, 1, false, spd4); |
| |
| // validate metadata of shuffle size |
| Map<String, Map<Integer, AtomicLong>> shuffleSizeMap = shuffleBufferManager.getShuffleSizeMap(); |
| assertEquals(99, shuffleSizeMap.get(appId1).get(1).get()); |
| assertEquals(200, shuffleSizeMap.get(appId1).get(2).get()); |
| assertEquals(100, shuffleSizeMap.get(appId2).get(1).get()); |
| |
| ShufflePartitionedData spd5 = createData(0, 68); |
| shuffleBufferManager.cacheShuffleData(appId2, 2, false, spd5); |
| // flush happen |
| assertEquals(99, shuffleSizeMap.get(appId1).get(1).get()); |
| assertEquals(0, shuffleSizeMap.get(appId1).get(2).get()); |
| assertEquals(0, shuffleSizeMap.get(appId2).get(1).get()); |
| assertEquals(0, shuffleSizeMap.get(appId2).get(1).get()); |
| shuffleBufferManager.releaseMemory(400, true, false); |
| |
| ShufflePartitionedData spd6 = createData(0, 300); |
| shuffleBufferManager.cacheShuffleData(appId1, 1, false, spd6); |
| // flush happen |
| assertEquals(0, shuffleSizeMap.get(appId1).get(1).get()); |
| shuffleBufferManager.releaseMemory(463, true, false); |
| |
| shuffleBufferManager.cacheShuffleData(appId1, 1, false, spd1); |
| shuffleBufferManager.cacheShuffleData(appId1, 2, false, spd2); |
| shuffleBufferManager.cacheShuffleData(appId2, 1, false, spd4); |
| shuffleBufferManager.removeBuffer(appId1); |
| assertNull(shuffleSizeMap.get(appId1)); |
| assertEquals(100, shuffleSizeMap.get(appId2).get(1).get()); |
| } |
| |
| @Test |
| public void cacheShuffleDataTest() { |
| String appId = "cacheShuffleDataTest"; |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| int startPartitionNum = (int) ShuffleServerMetrics.gaugeTotalPartitionNum.get(); |
| int shuffleId = 1; |
| StatusCode sc = |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 16)); |
| assertEquals(StatusCode.NO_REGISTER, sc); |
| shuffleBufferManager.registerBuffer(appId, shuffleId + 1, 0, 1); |
| assertEquals(startPartitionNum + 1, (int) ShuffleServerMetrics.gaugeTotalPartitionNum.get()); |
| sc = shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 16)); |
| assertEquals(StatusCode.NO_REGISTER, sc); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 100, 101); |
| assertEquals(startPartitionNum + 2, (int) ShuffleServerMetrics.gaugeTotalPartitionNum.get()); |
| sc = shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 16)); |
| assertEquals(StatusCode.NO_REGISTER, sc); |
| |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| assertEquals(startPartitionNum + 3, (int) ShuffleServerMetrics.gaugeTotalPartitionNum.get()); |
| sc = shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 16)); |
| assertEquals(StatusCode.SUCCESS, sc); |
| |
| Map<String, Map<Integer, RangeMap<Integer, ShuffleBuffer>>> bufferPool = |
| shuffleBufferManager.getBufferPool(); |
| ShuffleBuffer buffer = bufferPool.get(appId).get(shuffleId).get(0); |
| assertEquals(48, buffer.getSize()); |
| assertEquals(48, shuffleBufferManager.getUsedMemory()); |
| |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 16)); |
| assertEquals(96, buffer.getSize()); |
| assertEquals(96, shuffleBufferManager.getUsedMemory()); |
| |
| // reach high water lever, flush |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 273)); |
| assertEquals(0, buffer.getSize()); |
| assertEquals(401, shuffleBufferManager.getUsedMemory()); |
| assertEquals(401, shuffleBufferManager.getInFlushSize()); |
| verify(mockShuffleFlushManager, times(1)).addToFlushQueue(any()); |
| |
| // now buffer should be full |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 100)); |
| verify(mockShuffleFlushManager, times(1)).addToFlushQueue(any()); |
| sc = shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 1)); |
| assertEquals(StatusCode.NO_BUFFER, sc); |
| |
| // size won't be reduced which should be processed by flushManager, reset buffer size to 0 |
| shuffleBufferManager.resetSize(); |
| shuffleBufferManager.removeBuffer(appId); |
| assertEquals(startPartitionNum, (int) ShuffleServerMetrics.gaugeTotalPartitionNum.get()); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 0); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 1, 1); |
| shuffleBufferManager.registerBuffer(appId, 2, 0, 0); |
| |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 200)); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(1, 200)); |
| shuffleBufferManager.cacheShuffleData(appId, 2, false, createData(0, 32)); |
| ShuffleBuffer buffer0 = bufferPool.get(appId).get(shuffleId).get(0); |
| ShuffleBuffer buffer1 = bufferPool.get(appId).get(shuffleId).get(1); |
| ShuffleBuffer buffer2 = bufferPool.get(appId).get(2).get(0); |
| assertEquals(0, buffer0.getSize()); |
| assertEquals(0, buffer1.getSize()); |
| assertEquals(64, buffer2.getSize()); |
| assertEquals(528, shuffleBufferManager.getUsedMemory()); |
| assertEquals(464, shuffleBufferManager.getInFlushSize()); |
| verify(mockShuffleFlushManager, times(3)).addToFlushQueue(any()); |
| } |
| |
| @Test |
| public void cacheShuffleDataWithPreAllocationTest() { |
| String appId = "cacheShuffleDataWithPreAllocationTest"; |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| int shuffleId = 1; |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| // pre allocate memory |
| shuffleBufferManager.requireMemory(48, true); |
| assertEquals(48, shuffleBufferManager.getUsedMemory()); |
| assertEquals(48, shuffleBufferManager.getPreAllocatedSize()); |
| // receive data with preAllocation |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, true, createData(0, 16)); |
| assertEquals(48, shuffleBufferManager.getUsedMemory()); |
| assertEquals(48, shuffleBufferManager.getPreAllocatedSize()); |
| // release memory |
| shuffleBufferManager.releaseMemory(48, false, true); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getPreAllocatedSize()); |
| // receive data without preAllocation |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 17)); |
| assertEquals(49, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getPreAllocatedSize()); |
| // release memory |
| shuffleBufferManager.releaseMemory(49, false, false); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getPreAllocatedSize()); |
| |
| // release memory with preAllocation |
| shuffleBufferManager.requireMemory(16, true); |
| shuffleBufferManager.releaseMemory(16, false, true); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getPreAllocatedSize()); |
| |
| // pre allocate all memory |
| shuffleBufferManager.requireMemory(500, true); |
| assertEquals(500, shuffleBufferManager.getUsedMemory()); |
| assertEquals(500, shuffleBufferManager.getPreAllocatedSize()); |
| |
| // no buffer if data without pre allocation |
| StatusCode sc = |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(1, 16)); |
| assertEquals(StatusCode.NO_BUFFER, sc); |
| |
| // actual data size < spillThreshold, won't flush |
| sc = shuffleBufferManager.cacheShuffleData(appId, shuffleId, true, createData(1, 16)); |
| shuffleBufferManager.releasePreAllocatedSize(48); |
| assertEquals(StatusCode.SUCCESS, sc); |
| assertEquals(500, shuffleBufferManager.getUsedMemory()); |
| assertEquals(452, shuffleBufferManager.getPreAllocatedSize()); |
| |
| // actual data size > highWaterMark, flush |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, true, createData(0, 400)); |
| shuffleBufferManager.releasePreAllocatedSize(432); |
| // trigger flush manually |
| shuffleBufferManager.flushIfNecessary(); |
| assertEquals(StatusCode.SUCCESS, sc); |
| assertEquals(500, shuffleBufferManager.getUsedMemory()); |
| assertEquals(20, shuffleBufferManager.getPreAllocatedSize()); |
| verify(mockShuffleFlushManager, times(1)).addToFlushQueue(any()); |
| } |
| |
| @Test |
| public void bufferSizeTest() throws Exception { |
| ShuffleServer mockShuffleServer = mock(ShuffleServer.class); |
| StorageManager storageManager = StorageManagerFactory.getInstance().createStorageManager(conf); |
| ShuffleFlushManager shuffleFlushManager = |
| new ShuffleFlushManager(conf, mockShuffleServer, storageManager); |
| shuffleBufferManager = new ShuffleBufferManager(conf, shuffleFlushManager, false); |
| |
| when(mockShuffleServer.getShuffleFlushManager()).thenReturn(shuffleFlushManager); |
| when(mockShuffleServer.getShuffleBufferManager()).thenReturn(shuffleBufferManager); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mock(ShuffleTaskManager.class)); |
| |
| String appId = "bufferSizeTest"; |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| int shuffleId = 1; |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 2, 3); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 4, 5); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 6, 7); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 8, 9); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 16)); |
| assertEquals(48, shuffleBufferManager.getUsedMemory()); |
| |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 16)); |
| assertEquals(96, shuffleBufferManager.getUsedMemory()); |
| |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 300)); |
| waitForFlush(shuffleFlushManager, appId, shuffleId, 3); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getInFlushSize()); |
| |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 64)); |
| assertEquals(96, shuffleBufferManager.getUsedMemory()); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(2, 64)); |
| assertEquals(192, shuffleBufferManager.getUsedMemory()); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(4, 64)); |
| assertEquals(288, shuffleBufferManager.getUsedMemory()); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(6, 64)); |
| assertEquals(384, shuffleBufferManager.getUsedMemory()); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(8, 64)); |
| waitForFlush(shuffleFlushManager, appId, shuffleId, 5); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getInFlushSize()); |
| |
| shuffleBufferManager.registerBuffer("bufferSizeTest1", shuffleId, 0, 1); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 32)); |
| assertEquals(64, shuffleBufferManager.getUsedMemory()); |
| shuffleBufferManager.cacheShuffleData("bufferSizeTest1", shuffleId, false, createData(0, 32)); |
| assertEquals(128, shuffleBufferManager.getUsedMemory()); |
| assertEquals(2, shuffleBufferManager.getBufferPool().keySet().size()); |
| shuffleBufferManager.removeBuffer(appId); |
| assertEquals(64, shuffleBufferManager.getUsedMemory()); |
| assertEquals(1, shuffleBufferManager.getBufferPool().keySet().size()); |
| } |
| |
| @Test |
| public void flushSingleBufferForHugePartitionTest(@TempDir File tmpDir) throws Exception { |
| ShuffleServerConf shuffleConf = new ShuffleServerConf(); |
| File dataDir = new File(tmpDir, "data"); |
| shuffleConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name()); |
| shuffleConf.set( |
| ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(dataDir.getAbsolutePath())); |
| shuffleConf.set(ShuffleServerConf.SERVER_MEMORY_SHUFFLE_LOWWATERMARK_PERCENTAGE, 20.0); |
| shuffleConf.set(ShuffleServerConf.SERVER_MEMORY_SHUFFLE_HIGHWATERMARK_PERCENTAGE, 80.0); |
| shuffleConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 1024L); |
| shuffleConf.set(ShuffleServerConf.SERVER_BUFFER_CAPACITY, 200L); |
| shuffleConf.set(ShuffleServerConf.HUGE_PARTITION_MEMORY_USAGE_LIMITATION_RATIO, 0.1); |
| shuffleConf.set(ShuffleServerConf.HUGE_PARTITION_SIZE_THRESHOLD, 100L); |
| shuffleConf.set(ShuffleServerConf.SINGLE_BUFFER_FLUSH_ENABLED, false); |
| shuffleConf.setSizeAsBytes(ShuffleServerConf.SINGLE_BUFFER_FLUSH_THRESHOLD, 64L); |
| |
| ShuffleServer mockShuffleServer = mock(ShuffleServer.class); |
| StorageManager storageManager = |
| StorageManagerFactory.getInstance().createStorageManager(shuffleConf); |
| ShuffleFlushManager shuffleFlushManager = |
| new ShuffleFlushManager(shuffleConf, mockShuffleServer, storageManager); |
| shuffleBufferManager = new ShuffleBufferManager(shuffleConf, shuffleFlushManager, false); |
| ShuffleTaskManager shuffleTaskManager = |
| new ShuffleTaskManager( |
| shuffleConf, shuffleFlushManager, shuffleBufferManager, storageManager); |
| |
| when(mockShuffleServer.getShuffleFlushManager()).thenReturn(shuffleFlushManager); |
| when(mockShuffleServer.getShuffleBufferManager()).thenReturn(shuffleBufferManager); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(shuffleTaskManager); |
| |
| String appId = "flushSingleBufferForHugePartitionTest_appId"; |
| int shuffleId = 1; |
| |
| shuffleTaskManager.registerShuffle( |
| appId, shuffleId, Arrays.asList(new PartitionRange(0, 0)), new RemoteStorageInfo(""), ""); |
| |
| // case1: its partition is not huge partition |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 0); |
| ShufflePartitionedData partitionedData = createData(0, 1); |
| shuffleTaskManager.cacheShuffleData(appId, shuffleId, false, partitionedData); |
| shuffleTaskManager.updateCachedBlockIds(appId, shuffleId, 0, partitionedData.getBlockList()); |
| assertEquals(1 + 32, shuffleBufferManager.getUsedMemory()); |
| long usedSize = shuffleTaskManager.getPartitionDataSize(appId, shuffleId, 0); |
| assertEquals(1 + 32, usedSize); |
| assertFalse( |
| shuffleBufferManager.limitHugePartition( |
| appId, shuffleId, 0, shuffleTaskManager.getPartitionDataSize(appId, shuffleId, 0))); |
| |
| // case2: its partition is huge partition, its buffer will be flushed to DISK directly |
| partitionedData = createData(0, 36); |
| shuffleTaskManager.cacheShuffleData(appId, shuffleId, false, partitionedData); |
| shuffleTaskManager.updateCachedBlockIds(appId, shuffleId, 0, partitionedData.getBlockList()); |
| assertEquals(33 + 36 + 32, shuffleBufferManager.getUsedMemory()); |
| assertTrue( |
| shuffleBufferManager.limitHugePartition( |
| appId, shuffleId, 0, shuffleTaskManager.getPartitionDataSize(appId, shuffleId, 0))); |
| partitionedData = createData(0, 1); |
| shuffleTaskManager.cacheShuffleData(appId, shuffleId, false, partitionedData); |
| shuffleTaskManager.updateCachedBlockIds(appId, shuffleId, 0, partitionedData.getBlockList()); |
| waitForFlush(shuffleFlushManager, appId, shuffleId, 3); |
| } |
| |
| @Test |
| public void flushSingleBufferTest(@TempDir File tmpDir) throws Exception { |
| ShuffleServerConf shuffleConf = new ShuffleServerConf(); |
| File dataDir = new File(tmpDir, "data"); |
| shuffleConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name()); |
| shuffleConf.set( |
| ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(dataDir.getAbsolutePath())); |
| shuffleConf.set(ShuffleServerConf.SERVER_BUFFER_CAPACITY, 200L); |
| shuffleConf.set(ShuffleServerConf.SERVER_MEMORY_SHUFFLE_LOWWATERMARK_PERCENTAGE, 20.0); |
| shuffleConf.set(ShuffleServerConf.SERVER_MEMORY_SHUFFLE_HIGHWATERMARK_PERCENTAGE, 80.0); |
| shuffleConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 1024L); |
| shuffleConf.setBoolean(ShuffleServerConf.SINGLE_BUFFER_FLUSH_ENABLED, true); |
| shuffleConf.setSizeAsBytes(ShuffleServerConf.SINGLE_BUFFER_FLUSH_THRESHOLD, 128L); |
| |
| ShuffleServer mockShuffleServer = mock(ShuffleServer.class); |
| StorageManager storageManager = |
| StorageManagerFactory.getInstance().createStorageManager(shuffleConf); |
| ShuffleFlushManager shuffleFlushManager = |
| new ShuffleFlushManager(shuffleConf, mockShuffleServer, storageManager); |
| shuffleBufferManager = new ShuffleBufferManager(shuffleConf, shuffleFlushManager, false); |
| |
| when(mockShuffleServer.getShuffleFlushManager()).thenReturn(shuffleFlushManager); |
| when(mockShuffleServer.getShuffleBufferManager()).thenReturn(shuffleBufferManager); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mock(ShuffleTaskManager.class)); |
| |
| String appId = "bufferSizeTest"; |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| int shuffleId = 1; |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 2, 3); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 64)); |
| assertEquals(96, shuffleBufferManager.getUsedMemory()); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(2, 64)); |
| waitForFlush(shuffleFlushManager, appId, shuffleId, 2); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getInFlushSize()); |
| |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 32)); |
| assertEquals(64, shuffleBufferManager.getUsedMemory()); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(1, 48)); |
| waitForFlush(shuffleFlushManager, appId, shuffleId, 4); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getInFlushSize()); |
| } |
| |
| @Test |
| public void shuffleFlushThreshold() throws Exception { |
| ShuffleServerConf serverConf = new ShuffleServerConf(); |
| serverConf.addAll(conf); |
| serverConf.set(ShuffleServerConf.SERVER_BUFFER_CAPACITY, 180L); |
| serverConf.set(ShuffleServerConf.SERVER_SHUFFLE_FLUSH_THRESHOLD, 64L); |
| |
| StorageManager storageManager = StorageManagerFactory.getInstance().createStorageManager(conf); |
| ShuffleFlushManager shuffleFlushManager = |
| new ShuffleFlushManager(conf, mockShuffleServer, storageManager); |
| shuffleBufferManager = new ShuffleBufferManager(serverConf, shuffleFlushManager, false); |
| |
| String appId = "shuffleFlushTest"; |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| |
| int shuffleId = 0; |
| int smallShuffleId = 1; |
| int smallShuffleIdTwo = 2; |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 2, 3); |
| shuffleBufferManager.registerBuffer(appId, smallShuffleId, 0, 1); |
| shuffleBufferManager.registerBuffer(appId, smallShuffleIdTwo, 0, 1); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 64)); |
| assertEquals(96, shuffleBufferManager.getUsedMemory()); |
| shuffleBufferManager.cacheShuffleData(appId, smallShuffleId, false, createData(0, 31)); |
| assertEquals(96 + 63, shuffleBufferManager.getUsedMemory()); |
| Thread.sleep(100); |
| Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS); |
| // small shuffle id is kept in memory |
| assertEquals(63, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getInFlushSize()); |
| |
| // more data will trigger the flush |
| shuffleBufferManager.cacheShuffleData(appId, smallShuffleId, false, createData(0, 31)); |
| shuffleBufferManager.cacheShuffleData(appId, smallShuffleId, false, createData(0, 31)); |
| assertEquals(63 * 3, shuffleBufferManager.getUsedMemory()); |
| Thread.sleep(100); |
| Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getInFlushSize()); |
| |
| // all the small data in shuffle server, which could be extremely rare |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 22)); |
| shuffleBufferManager.cacheShuffleData(appId, smallShuffleId, false, createData(0, 21)); |
| shuffleBufferManager.cacheShuffleData(appId, smallShuffleIdTwo, false, createData(0, 20)); |
| assertEquals(54 + 53 + 52, shuffleBufferManager.getUsedMemory()); |
| Thread.sleep(100); |
| Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS); |
| assertEquals(52, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getInFlushSize()); |
| } |
| |
| private void waitForFlush( |
| ShuffleFlushManager shuffleFlushManager, String appId, int shuffleId, int expectedBlockNum) |
| throws Exception { |
| int retry = 0; |
| long committedCount = 0; |
| do { |
| committedCount = |
| shuffleFlushManager.getCommittedBlockIds(appId, shuffleId).getLongCardinality(); |
| if (committedCount < expectedBlockNum) { |
| Thread.sleep(500); |
| } |
| retry++; |
| if (retry > 10) { |
| fail("Flush data time out"); |
| } |
| } while (committedCount < expectedBlockNum); |
| |
| // Need to wait for `event.doCleanup` to be executed |
| // to ensure the correctness of subsequent checks of |
| // `shuffleBufferManager.getUsedMemory()` and `shuffleBufferManager.getInFlushSize()`. |
| Awaitility.await() |
| .atMost(Duration.ofSeconds(5)) |
| .until(() -> shuffleBufferManager.getUsedMemory() == 0); |
| } |
| |
| @Test |
| public void bufferManagerInitTest() { |
| ShuffleServerConf serverConf = new ShuffleServerConf(); |
| shuffleBufferManager = new ShuffleBufferManager(serverConf, mockShuffleFlushManager, false); |
| double ratio = ShuffleServerConf.SERVER_BUFFER_CAPACITY_RATIO.defaultValue(); |
| double readRatio = ShuffleServerConf.SERVER_READ_BUFFER_CAPACITY_RATIO.defaultValue(); |
| assertEquals( |
| (long) (Runtime.getRuntime().maxMemory() * ratio), shuffleBufferManager.getCapacity()); |
| assertEquals( |
| (long) (Runtime.getRuntime().maxMemory() * readRatio), |
| shuffleBufferManager.getReadCapacity()); |
| ratio = 0.6; |
| readRatio = 0.1; |
| serverConf.set(ShuffleServerConf.SERVER_BUFFER_CAPACITY_RATIO, ratio); |
| serverConf.set(ShuffleServerConf.SERVER_READ_BUFFER_CAPACITY_RATIO, readRatio); |
| shuffleBufferManager = new ShuffleBufferManager(serverConf, mockShuffleFlushManager, false); |
| assertEquals( |
| (long) (Runtime.getRuntime().maxMemory() * ratio), shuffleBufferManager.getCapacity()); |
| assertEquals( |
| (long) (Runtime.getRuntime().maxMemory() * readRatio), |
| shuffleBufferManager.getReadCapacity()); |
| } |
| |
| @Test |
| public void flushBufferTestWhenNotSelectedStorage(@TempDir File tmpDir) throws Exception { |
| // In this test, rss.server.single.buffer.flush.threshold and |
| // rss.server.flush.cold.storage.threshold.size are 16. |
| // When cacheShuffleData with 64 bytes, will flush to HDFS storage, but we do not register |
| // remote storage. |
| // Then storageManager.selectStorage will return null, we should make sure that when we can not |
| // select a storage, |
| // the resources will not leak. |
| ShuffleServerConf shuffleConf = new ShuffleServerConf(); |
| File dataDir = new File(tmpDir, "data"); |
| shuffleConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name()); |
| shuffleConf.set( |
| ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(dataDir.getAbsolutePath())); |
| shuffleConf.set(ShuffleServerConf.SERVER_BUFFER_CAPACITY, 200L); |
| shuffleConf.set(ShuffleServerConf.SERVER_MEMORY_SHUFFLE_LOWWATERMARK_PERCENTAGE, 20.0); |
| shuffleConf.set(ShuffleServerConf.SERVER_MEMORY_SHUFFLE_HIGHWATERMARK_PERCENTAGE, 80.0); |
| shuffleConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 1024L); |
| shuffleConf.setBoolean(ShuffleServerConf.SINGLE_BUFFER_FLUSH_ENABLED, true); |
| shuffleConf.setSizeAsBytes(ShuffleServerConf.SINGLE_BUFFER_FLUSH_THRESHOLD, 16L); |
| shuffleConf.setSizeAsBytes(ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, 16L); |
| shuffleConf.setString( |
| ShuffleServerConf.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE_HDFS.name()); |
| ShuffleServer mockShuffleServer = mock(ShuffleServer.class); |
| StorageManager storageManager = |
| StorageManagerFactory.getInstance().createStorageManager(shuffleConf); |
| ShuffleFlushManager shuffleFlushManager = |
| new ShuffleFlushManager(shuffleConf, mockShuffleServer, storageManager); |
| shuffleBufferManager = new ShuffleBufferManager(shuffleConf, shuffleFlushManager, false); |
| |
| when(mockShuffleServer.getShuffleFlushManager()).thenReturn(shuffleFlushManager); |
| when(mockShuffleServer.getShuffleBufferManager()).thenReturn(shuffleBufferManager); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mock(ShuffleTaskManager.class)); |
| |
| String appId = "bufferSizeTest"; |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| int shuffleId = 1; |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 2, 3); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(0, 64)); |
| shuffleBufferManager.cacheShuffleData(appId, shuffleId, false, createData(2, 64)); |
| // wait flush event drained |
| Awaitility.await() |
| .atMost(Duration.ofSeconds(5)) |
| .until(() -> shuffleFlushManager.getEventNumInFlush() == 0); |
| // make sure all cleanup tasks are done. |
| DefaultFlushEventHandler flushEventHandler = |
| (DefaultFlushEventHandler) shuffleFlushManager.getEventHandler(); |
| ThreadPoolExecutor executor = |
| ((ThreadPoolExecutor) flushEventHandler.getFallbackThreadPoolExecutor()); |
| executor.shutdown(); |
| assertEquals(0, shuffleBufferManager.getUsedMemory()); |
| assertEquals(0, shuffleBufferManager.getInFlushSize()); |
| } |
| |
| @Test |
| public void blockSizeMetricsTest() { |
| String appId = "blockSizeMetricsTest"; |
| shuffleBufferManager.setShuffleTaskManager(mockShuffleTaskManager); |
| ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock(); |
| when(mockShuffleTaskManager.getAppReadLock(appId)).thenReturn(rwLock.readLock()); |
| when(mockShuffleServer.getShuffleTaskManager()).thenReturn(mockShuffleTaskManager); |
| int shuffleId = 1; |
| shuffleBufferManager.registerBuffer(appId, shuffleId, 0, 1); |
| |
| // cache shuffle block data, and record metrics |
| double[] buckets = |
| ConfigUtils.convertBytesStringToDoubleArray( |
| new ShuffleServerConf() |
| .get(ShuffleServerConf.APP_LEVEL_SHUFFLE_BLOCK_SIZE_METRIC_BUCKETS)); |
| Arrays.stream(buckets) |
| .sorted() |
| .forEach( |
| bucket -> { |
| StatusCode sc = |
| shuffleBufferManager.cacheShuffleData( |
| appId, shuffleId, true, createData(0, (int) bucket)); |
| assertEquals(StatusCode.SUCCESS, sc); |
| }); |
| // check metrics values |
| List<Collector.MetricFamilySamples> samples = |
| ShuffleServerMetrics.appHistogramWriteBlockSize.collect(); |
| assertEquals(samples.size(), 1); |
| int index = 1; |
| Arrays.stream(buckets) |
| .sorted() |
| .forEach( |
| bucket -> { |
| for (Collector.MetricFamilySamples.Sample s : samples.get(0).samples) { |
| if (s.labelValues.contains(bucket)) { |
| assertEquals(s.value, index); |
| } |
| } |
| }); |
| } |
| } |
