You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@uniffle.apache.org by ro...@apache.org on 2022/07/12 08:33:57 UTC
[incubator-uniffle] branch master updated: [improvement] Remove experimental feature with ShuffleUploader (#51)
This is an automated email from the ASF dual-hosted git repository.
roryqi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git
The following commit(s) were added to refs/heads/master by this push:
new d653cd2 [improvement] Remove experimental feature with ShuffleUploader (#51)
d653cd2 is described below
commit d653cd2e8b425ff2d411425ede9ea1e44ef9cab9
Author: Colin <co...@tencent.com>
AuthorDate: Tue Jul 12 16:33:53 2022 +0800
[improvement] Remove experimental feature with ShuffleUploader (#51)
### What changes were proposed in this pull request?
ShuffleUploader is an experimental feature which is target to merge and upload data from local disk to remote storage if local disk hasn't enough space. It is replaced by `MEMORY_LOCAL_HDFS` and should be removed to avoid maintenance in the future.
### Why are the changes needed?
The experimental feature is redundancy and we have to maintain it during update with new features about storage.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
with exist UTs
---
.../apache/uniffle/client/util/ClientUtils.java | 3 +-
.../apache/uniffle/test/IntegrationTestBase.java | 1 -
.../test/MultiStorageFaultToleranceTest.java | 112 +--
.../org/apache/uniffle/test/MultiStorageTest.java | 815 ---------------------
.../RepartitionWithUploadMultiStorageRssTest.java | 71 --
.../apache/uniffle/test/SparkFallbackReadTest.java | 153 ----
.../uniffle/test/SparkSQLMultiStorageRssTest.java | 87 ---
.../apache/uniffle/server/ShuffleServerConf.java | 67 --
.../org/apache/uniffle/server/ShuffleUploader.java | 493 -------------
.../server/storage/MultiStorageManager.java | 32 +-
.../server/storage/StorageManagerFactory.java | 3 +-
.../uniffle/server/ShuffleFlushManagerTest.java | 1 -
.../uniffle/server/ShuffleServerConfTest.java | 11 -
.../apache/uniffle/server/ShuffleUploaderTest.java | 704 ------------------
.../server/storage/MultiStorageManagerTest.java | 2 +-
.../src/test/resources/confBySizeStringTest.conf | 1 -
server/src/test/resources/confTest.conf | 3 -
.../storage/factory/ShuffleHandlerFactory.java | 47 --
.../factory/ShuffleUploadHandlerFactory.java | 54 --
.../storage/handler/api/ShuffleUploadHandler.java | 41 --
.../handler/impl/HdfsShuffleUploadHandler.java | 268 -------
.../impl/UploadedHdfsClientReadHandler.java | 116 ---
.../UploadedStorageHdfsShuffleReadHandler.java | 102 ---
.../request/CreateShuffleUploadHandlerRequest.java | 157 ----
.../uniffle/storage/util/ShuffleStorageUtils.java | 1 -
.../uniffle/storage/util/ShuffleUploadResult.java | 77 --
.../apache/uniffle/storage/util/StorageType.java | 1 -
.../handler/impl/HdfsShuffleUploadHandlerTest.java | 454 ------------
.../MultiStorageHdfsClientReadHandlerTest.java | 373 ----------
.../MultiStorageHdfsShuffleReadHandlerTest.java | 255 -------
30 files changed, 6 insertions(+), 4499 deletions(-)
diff --git a/client/src/main/java/org/apache/uniffle/client/util/ClientUtils.java b/client/src/main/java/org/apache/uniffle/client/util/ClientUtils.java
index 795369d..26adc79 100644
--- a/client/src/main/java/org/apache/uniffle/client/util/ClientUtils.java
+++ b/client/src/main/java/org/apache/uniffle/client/util/ClientUtils.java
@@ -74,7 +74,6 @@ public class ClientUtils {
return StorageType.MEMORY_HDFS.name().equals(storageType)
|| StorageType.MEMORY_LOCALFILE_HDFS.name().equals(storageType)
|| StorageType.HDFS.name().equals(storageType)
- || StorageType.LOCALFILE_HDFS.name().equals(storageType)
- || StorageType.LOCALFILE_HDFS_2.name().equals(storageType);
+ || StorageType.LOCALFILE_HDFS.name().equals(storageType);
}
}
diff --git a/integration-test/common/src/test/java/org/apache/uniffle/test/IntegrationTestBase.java b/integration-test/common/src/test/java/org/apache/uniffle/test/IntegrationTestBase.java
index cd157f6..0237aa7 100644
--- a/integration-test/common/src/test/java/org/apache/uniffle/test/IntegrationTestBase.java
+++ b/integration-test/common/src/test/java/org/apache/uniffle/test/IntegrationTestBase.java
@@ -107,7 +107,6 @@ abstract public class IntegrationTestBase extends HdfsTestBase {
serverConf.setInteger("rss.jetty.corePool.size", 64);
serverConf.setInteger("rss.rpc.executor.size", 10);
serverConf.setString("rss.server.hadoop.dfs.replication", "2");
- serverConf.setString("rss.server.uploader.base.path", "test");
serverConf.setLong("rss.server.disk.capacity", 10L * 1024L * 1024L * 1024L);
serverConf.setBoolean("rss.server.health.check.enable", false);
return serverConf;
diff --git a/integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageFaultToleranceTest.java b/integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageFaultToleranceTest.java
index df9e3a2..7e83a7b 100644
--- a/integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageFaultToleranceTest.java
+++ b/integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageFaultToleranceTest.java
@@ -46,12 +46,10 @@ import org.apache.uniffle.common.ShuffleBlockInfo;
import org.apache.uniffle.common.ShuffleServerInfo;
import org.apache.uniffle.coordinator.CoordinatorConf;
import org.apache.uniffle.server.ShuffleServerConf;
-import org.apache.uniffle.storage.util.ShuffleStorageUtils;
import org.apache.uniffle.storage.util.StorageType;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assertions.fail;
public class MultiStorageFaultToleranceTest extends ShuffleReadWriteBase {
private ShuffleServerGrpcClient shuffleServerClient;
@@ -65,14 +63,11 @@ public class MultiStorageFaultToleranceTest extends ShuffleReadWriteBase {
shuffleServerConf.setDouble(ShuffleServerConf.CLEANUP_THRESHOLD, 0.0);
shuffleServerConf.setDouble(ShuffleServerConf.HIGH_WATER_MARK_OF_WRITE, 100.0);
shuffleServerConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 100);
- shuffleServerConf.setBoolean(ShuffleServerConf.UPLOADER_ENABLE, true);
shuffleServerConf.setLong(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 30L);
- shuffleServerConf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, REMOTE_STORAGE);
- shuffleServerConf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 1L);
shuffleServerConf.setLong(ShuffleServerConf.SHUFFLE_EXPIRED_TIMEOUT_MS, 5000L);
shuffleServerConf.setLong(ShuffleServerConf.SERVER_APP_EXPIRED_WITHOUT_HEARTBEAT, 60L * 1000L * 60L);
shuffleServerConf.setLong(ShuffleServerConf.SERVER_COMMIT_TIMEOUT, 20L * 1000L);
- shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS_2.name());
+ shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS.name());
shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_BASE_PATH, basePath);
shuffleServerConf.setLong(ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, 400L * 1024L * 1024L);
createAndStartServers(shuffleServerConf, coordinatorConf);
@@ -88,60 +83,6 @@ public class MultiStorageFaultToleranceTest extends ShuffleReadWriteBase {
shuffleServerClient.close();
}
- @Test
- public void hdfsFaultTolerance() {
- try {
- String appId = "app_hdfs_fault_tolerance_data";
- Map<Long, byte[]> expectedData = Maps.newHashMap();
- Map<Integer, List<Integer>> map = Maps.newHashMap();
- map.put(2, Lists.newArrayList(0, 3));
- map.put(3, Lists.newArrayList(3));
- registerShuffle(appId, map);
-
- Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap2 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap3 = Roaring64NavigableMap.bitmapOf();
-
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 2, 0, 1,11, 10 * 1024 * 1024, blockIdBitmap1, expectedData);
-
- List<ShuffleBlockInfo> blocks2 = createShuffleBlockList(
- 3, 3, 2,9, 10 * 1024 * 1024, blockIdBitmap2, expectedData);
-
- List<ShuffleBlockInfo> blocks3 = createShuffleBlockList(
- 2, 3, 2,9, 10 * 1024 * 1024, blockIdBitmap3, expectedData);
-
- assertEquals(0, ShuffleStorageUtils.getStorageIndex(2, appId, 2, 0));
- assertEquals(0, ShuffleStorageUtils.getStorageIndex(2, appId, 3, 3));
- assertEquals(0, ShuffleStorageUtils.getStorageIndex(2, appId, 2, 3));
- assertEquals(1, cluster.getDataNodes().size());
- cluster.stopDataNode(0);
- assertEquals(0, cluster.getDataNodes().size());
-
- sendSinglePartitionToShuffleServer(appId, 2, 0, 1, blocks1);
- boolean isException = false;
- try {
- sendSinglePartitionToShuffleServer(appId, 3, 3,2, blocks2);
- } catch (RuntimeException re) {
- isException = true;
- assertTrue(re.getMessage().contains("Fail to finish"));
- }
- assertTrue(isException);
-
- cluster.startDataNodes(conf, 1, true, HdfsServerConstants.StartupOption.REGULAR,
- null, null, null, false, true);
- assertEquals(1, cluster.getDataNodes().size());
-
- sendSinglePartitionToShuffleServer(appId, 2, 3, 2, blocks3);
-
- validateResult(appId, 2, 0, blockIdBitmap1, Roaring64NavigableMap.bitmapOf(1), expectedData);
- validateResult(appId, 2, 3, blockIdBitmap3, Roaring64NavigableMap.bitmapOf(2), expectedData);
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-
@Test
public void hdfsFallbackTest() throws Exception {
String appId = "fallback_test";
@@ -157,9 +98,6 @@ public class MultiStorageFaultToleranceTest extends ShuffleReadWriteBase {
assertEquals(0, cluster.getDataNodes().size());
sendSinglePartitionToShuffleServer(appId, 0, 0, 0, blocks);
validateResult(appId, 0, 0, blockBitmap, Roaring64NavigableMap.bitmapOf(0), expectedData);
- cluster.startDataNodes(conf, 1, true, HdfsServerConstants.StartupOption.REGULAR,
- null, null, null, false, true);
- assertEquals(1, cluster.getDataNodes().size());
}
private void registerShuffle(String appId, Map<Integer, List<Integer>> registerMap) {
@@ -172,51 +110,6 @@ public class MultiStorageFaultToleranceTest extends ShuffleReadWriteBase {
}
}
- @Test
- public void diskFaultTolerance() {
- String appId = "app_disk_fault_tolerance_data";
- Map<Long, byte[]> expectedData = Maps.newHashMap();
-
- Map<Integer, List<Integer>> map = Maps.newHashMap();
- map.put(2, Lists.newArrayList(1, 3));
- map.put(3, Lists.newArrayList(1));
- registerShuffle(appId, map);
-
- Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap2 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap3 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap4 = Roaring64NavigableMap.bitmapOf();
-
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 2, 1, 1,11, 10 * 1024 * 1024, blockIdBitmap1, expectedData);
-
- List<ShuffleBlockInfo> blocks2 = createShuffleBlockList(
- 3, 1, 2,9, 10 * 1024 * 1024, blockIdBitmap2, expectedData);
-
- List<ShuffleBlockInfo> blocks3 = createShuffleBlockList(
- 2, 3, 2,9, 10 * 1024 * 1024, blockIdBitmap3, expectedData);
-
- List<ShuffleBlockInfo> blocks4 = createShuffleBlockList(
- 2, 1, 1, 11, 10 * 1024 * 1024, blockIdBitmap4, expectedData);
-
- assertEquals(1, ShuffleStorageUtils.getStorageIndex(2, appId, 2, 1));
- assertEquals(1, ShuffleStorageUtils.getStorageIndex(2, appId, 3, 1));
- assertEquals(1, ShuffleStorageUtils.getStorageIndex(2, appId, 2, 3));
- assertEquals(1, ShuffleStorageUtils.getStorageIndex(2, appId, 2, 1));
- try {
- sendSinglePartitionToShuffleServer(appId, 2, 1, 1, blocks1);
- sendSinglePartitionToShuffleServer(appId, 3, 1,2, blocks2);
- sendSinglePartitionToShuffleServer(appId, 2, 3, 2, blocks3);
- sendSinglePartitionToShuffleServer(appId, 2, 1, 1, blocks4);
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- validateResult(appId, 2, 1, blockIdBitmap1, Roaring64NavigableMap.bitmapOf(1), expectedData);
- validateResult(appId, 3, 1, blockIdBitmap2, Roaring64NavigableMap.bitmapOf(2), expectedData);
- validateResult(appId, 2, 3, blockIdBitmap3, Roaring64NavigableMap.bitmapOf(2), expectedData);
- }
-
private void sendSinglePartitionToShuffleServer(String appId, int shuffle, int partition,
long taskAttemptId, List<ShuffleBlockInfo> blocks) {
Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
@@ -239,7 +132,7 @@ public class MultiStorageFaultToleranceTest extends ShuffleReadWriteBase {
protected void validateResult(String appId, int shuffleId, int partitionId, Roaring64NavigableMap blockBitmap,
Roaring64NavigableMap taskBitmap, Map<Long, byte[]> expectedData) {
- ShuffleReadClientImpl readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS_2.name(),
+ ShuffleReadClientImpl readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS.name(),
appId, shuffleId, partitionId, 100, 1, 10, 1000, REMOTE_STORAGE, blockBitmap, taskBitmap,
Lists.newArrayList(new ShuffleServerInfo("test", LOCALHOST, SHUFFLE_SERVER_PORT)), conf, new DefaultIdHelper());
CompressedShuffleBlock csb = readClient.readShuffleBlockData();
@@ -256,7 +149,6 @@ public class MultiStorageFaultToleranceTest extends ShuffleReadWriteBase {
assertTrue(blockBitmap.equals(matched));
}
-
private Set<Long> getExpectBlockIds(List<ShuffleBlockInfo> blocks) {
List<Long> expectBlockIds = Lists.newArrayList();
blocks.forEach(b -> expectBlockIds.add(b.getBlockId()));
diff --git a/integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageTest.java b/integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageTest.java
deleted file mode 100644
index 3eba116..0000000
--- a/integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageTest.java
+++ /dev/null
@@ -1,815 +0,0 @@
-/*
- * 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.test;
-
-import java.io.File;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.TimeUnit;
-
-import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
-import com.google.common.collect.Sets;
-import com.google.common.util.concurrent.Uninterruptibles;
-import org.apache.uniffle.client.util.DefaultIdHelper;
-import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.Test;
-import org.roaringbitmap.longlong.Roaring64NavigableMap;
-
-import org.apache.uniffle.client.impl.ShuffleReadClientImpl;
-import org.apache.uniffle.client.impl.grpc.ShuffleServerGrpcClient;
-import org.apache.uniffle.client.request.RssAppHeartBeatRequest;
-import org.apache.uniffle.client.request.RssFinishShuffleRequest;
-import org.apache.uniffle.client.request.RssGetShuffleResultRequest;
-import org.apache.uniffle.client.request.RssRegisterShuffleRequest;
-import org.apache.uniffle.client.request.RssReportShuffleResultRequest;
-import org.apache.uniffle.client.request.RssSendCommitRequest;
-import org.apache.uniffle.client.request.RssSendShuffleDataRequest;
-import org.apache.uniffle.client.response.CompressedShuffleBlock;
-import org.apache.uniffle.common.BufferSegment;
-import org.apache.uniffle.common.PartitionRange;
-import org.apache.uniffle.common.ShuffleBlockInfo;
-import org.apache.uniffle.common.ShuffleDataResult;
-import org.apache.uniffle.common.ShuffleDataSegment;
-import org.apache.uniffle.common.ShuffleServerInfo;
-import org.apache.uniffle.common.util.ChecksumUtils;
-import org.apache.uniffle.coordinator.CoordinatorConf;
-import org.apache.uniffle.server.ShuffleDataReadEvent;
-import org.apache.uniffle.server.ShuffleServerConf;
-import org.apache.uniffle.storage.common.LocalStorage;
-import org.apache.uniffle.storage.handler.api.ClientReadHandler;
-import org.apache.uniffle.storage.handler.impl.ComposedClientReadHandler;
-import org.apache.uniffle.storage.handler.impl.HdfsClientReadHandler;
-import org.apache.uniffle.storage.handler.impl.LocalFileQuorumClientReadHandler;
-import org.apache.uniffle.storage.handler.impl.UploadedHdfsClientReadHandler;
-import org.apache.uniffle.storage.util.ShuffleStorageUtils;
-import org.apache.uniffle.storage.util.StorageType;
-
-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.assertTrue;
-import static org.junit.jupiter.api.Assertions.fail;
-
-public class MultiStorageTest extends ShuffleReadWriteBase {
- private ShuffleServerGrpcClient shuffleServerClient;
- private static String REMOTE_STORAGE = HDFS_URI + "rss/multi_storage";
-
- @BeforeAll
- public static void setupServers() throws Exception {
- CoordinatorConf coordinatorConf = getCoordinatorConf();
- Map<String, String> dynamicConf = Maps.newHashMap();
- dynamicConf.put(CoordinatorConf.COORDINATOR_REMOTE_STORAGE_PATH.key(), REMOTE_STORAGE);
- addDynamicConf(coordinatorConf, dynamicConf);
- String basePath = generateBasePath();
- ShuffleServerConf shuffleServerConf = getShuffleServerConf();
- shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS_2.name());
- shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_BASE_PATH, basePath);
- shuffleServerConf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, REMOTE_STORAGE);
- shuffleServerConf.setDouble(ShuffleServerConf.CLEANUP_THRESHOLD, 0.0);
- shuffleServerConf.setLong(ShuffleServerConf.CLEANUP_INTERVAL_MS, 1000);
- shuffleServerConf.setDouble(ShuffleServerConf.HIGH_WATER_MARK_OF_WRITE, 100.0);
- shuffleServerConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 100);
- shuffleServerConf.setBoolean(ShuffleServerConf.UPLOADER_ENABLE, true);
- shuffleServerConf.setLong(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 30L);
- shuffleServerConf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 1L);
- shuffleServerConf.setLong(ShuffleServerConf.SHUFFLE_EXPIRED_TIMEOUT_MS, 4000L);
- shuffleServerConf.setLong(ShuffleServerConf.SERVER_APP_EXPIRED_WITHOUT_HEARTBEAT, 60L * 1000L * 60L);
- shuffleServerConf.setLong(ShuffleServerConf.SERVER_COMMIT_TIMEOUT, 20L * 1000L);
- shuffleServerConf.setLong(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 30);
- shuffleServerConf.setLong(ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, 400L * 1024L * 1024L);
- createAndStartServers(shuffleServerConf, coordinatorConf);
- }
-
- @BeforeEach
- public void createClient() {
- shuffleServerClient = new ShuffleServerGrpcClient(LOCALHOST, SHUFFLE_SERVER_PORT);
- }
-
- @AfterEach
- public void closeClient() {
- shuffleServerClient.close();
- }
-
- @Test
- public void readUploadedDataTest() {
- String appId = "ap_read_uploaded_data";
- RssRegisterShuffleRequest rr1 = new RssRegisterShuffleRequest(appId, 0,
- Lists.newArrayList(new PartitionRange(0, 0)), REMOTE_STORAGE);
- RssRegisterShuffleRequest rr2 = new RssRegisterShuffleRequest(appId, 0,
- Lists.newArrayList(new PartitionRange(1, 1)), REMOTE_STORAGE);
- RssRegisterShuffleRequest rr3 = new RssRegisterShuffleRequest(appId, 0,
- Lists.newArrayList(new PartitionRange(2, 2)), REMOTE_STORAGE);
- RssRegisterShuffleRequest rr4 = new RssRegisterShuffleRequest(appId, 0,
- Lists.newArrayList(new PartitionRange(4, 4)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr1);
- shuffleServerClient.registerShuffle(rr2);
- shuffleServerClient.registerShuffle(rr3);
- shuffleServerClient.registerShuffle(rr4);
-
- Map<Long, byte[]> expectedData = Maps.newHashMap();
- Set<Long> expectedBlock1 = Sets.newHashSet();
- Set<Long> expectedBlock2 = Sets.newHashSet();
-
- Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap2 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap3 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap4 = Roaring64NavigableMap.bitmapOf();
-
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 0, 0, 1,3, 25, blockIdBitmap1, expectedData);
- List<ShuffleBlockInfo> blocks2 = createShuffleBlockList(
- 0, 1, 1,5,1024 * 1024, blockIdBitmap2, expectedData);
- List<ShuffleBlockInfo> blocks3 = createShuffleBlockList(
- 0, 2, 2,4, 25, blockIdBitmap3, expectedData);
- List<ShuffleBlockInfo> blocks4 = createShuffleBlockList(
- 0, 4, 3,1, 1024 * 1024, blockIdBitmap4, expectedData);
-
-
- blocks1.forEach(b -> expectedBlock1.add(b.getBlockId()));
- blocks2.forEach(b -> expectedBlock2.add(b.getBlockId()));
-
- Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
- partitionToBlocks.put(0, blocks1);
- partitionToBlocks.put(1, blocks2);
- Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleToBlocks = Maps.newHashMap();
- shuffleToBlocks.put(0, partitionToBlocks);
- RssSendShuffleDataRequest rs1 = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- shuffleServerClient.sendShuffleData(rs1);
-
- RssSendCommitRequest rc1 = new RssSendCommitRequest(appId, 0);
- shuffleServerClient.sendCommit(rc1);
- RssFinishShuffleRequest rf1 = new RssFinishShuffleRequest(appId, 0);
- shuffleServerClient.finishShuffle(rf1);
- Map<Integer, List<Long>> partitionToBlockIds = Maps.newHashMap();
- partitionToBlockIds.put(0, new ArrayList<>(expectedBlock1));
- partitionToBlockIds.put(1, new ArrayList<>(expectedBlock2));
- RssReportShuffleResultRequest rrp1 = new RssReportShuffleResultRequest(
- appId, 0, 1L, partitionToBlockIds, 2);
- shuffleServerClient.reportShuffleResult(rrp1);
-
- LocalStorage item = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, 0));
- assertTrue(item.canWrite());
- assertEquals(3 * 25, item.getNotUploadedSize(appId + "/" + 0));
- item = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, 1));
- assertTrue(item.canWrite());
- assertEquals(5 * 1024 * 1024, item.getNotUploadedSize(appId + "/" + 0));
-
- sendSinglePartitionToShuffleServer(appId, 0,2, 2L, blocks3);
- sendSinglePartitionToShuffleServer(appId, 0, 4, 3L, blocks4);
-
- item = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, 2));
- assertTrue(item.canWrite());
- assertEquals(3 * 25 + 4 * 25, item.getNotUploadedSize(appId + "/" + 0));
-
- item = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, 4));
- assertTrue(item.canWrite());
- assertEquals(5 * 1024 * 1024 + 1024 * 1024, item.getNotUploadedSize(appId + "/" + 0));
-
-
- RssGetShuffleResultRequest rg1 = new RssGetShuffleResultRequest(appId, 0, 0);
- shuffleServerClient.getShuffleResult(rg1);
- RssGetShuffleResultRequest rg2 = new RssGetShuffleResultRequest(appId, 0, 1);
- shuffleServerClient.getShuffleResult(rg2);
- RssGetShuffleResultRequest rg3 = new RssGetShuffleResultRequest(appId, 0, 2);
- shuffleServerClient.getShuffleResult(rg3);
- RssGetShuffleResultRequest rg4 = new RssGetShuffleResultRequest(appId, 0, 4);
- shuffleServerClient.getShuffleResult(rg4);
-
- readShuffleData(shuffleServerClient, appId, 0, 0, 1, 10, 100, 0);
- readShuffleData(shuffleServerClient, appId, 0, 1, 1, 10, 100, 0);
-
-
- wait(appId);
-
- item = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, 0));
- assertTrue(item.canWrite());
- assertEquals(0, item.getNotUploadedSize(appId + "/" + 0));
-
- item = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, 1));
- assertTrue(item.canWrite());
- assertEquals(0, item.getNotUploadedSize(appId + "/" + 0));
-
- ShuffleDataResult result = readShuffleData(shuffleServerClient, appId, 0, 0,
- 1, 10, 1000, 0);
- assertEquals(0, result.getData().length);
-
- ShuffleReadClientImpl readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS_2.name(),
- appId, 0, 0, 100, 1, 10, 1000, REMOTE_STORAGE,
- blockIdBitmap1, Roaring64NavigableMap.bitmapOf(1), Lists.newArrayList(), conf, new DefaultIdHelper());
- validateResult(readClient, expectedData, blockIdBitmap1);
-
- readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS_2.name(),
- appId, 0, 1, 100, 1, 10, 1000, REMOTE_STORAGE,
- blockIdBitmap2, Roaring64NavigableMap.bitmapOf(1), Lists.newArrayList(), conf, new DefaultIdHelper());
- validateResult(readClient, expectedData, blockIdBitmap2);
-
- readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS_2.name(),
- appId, 0, 2, 100, 1, 10, 1000, REMOTE_STORAGE,
- blockIdBitmap3, Roaring64NavigableMap.bitmapOf(2), Lists.newArrayList(), conf, new DefaultIdHelper());
- validateResult(readClient, expectedData, blockIdBitmap3);
-
- readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS_2.name(),
- appId, 0, 4, 100, 1, 10, 1000, REMOTE_STORAGE,
- blockIdBitmap4, Roaring64NavigableMap.bitmapOf(3), Lists.newArrayList(), conf, new DefaultIdHelper());
- validateResult(readClient, expectedData, blockIdBitmap4);
- }
-
- private void wait(String appId) {
- do {
- try {
- RssAppHeartBeatRequest ra = new RssAppHeartBeatRequest(appId, 1000);
- shuffleServerClient.sendHeartBeat(ra);
- boolean uploadFinished = true;
- for (int i = 0; i < 4; i++) {
- LocalStorage localStorage = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, i));
- String path = ShuffleStorageUtils.getFullShuffleDataFolder(localStorage.getBasePath(),
- ShuffleStorageUtils.getShuffleDataPath(appId, 0, i, i));
- File file = new File(path);
- if (file.exists()) {
- uploadFinished = false;
- break;
- }
- }
- if (uploadFinished) {
- break;
- }
- Thread.sleep(1000);
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- } while (true);
- }
-
- @Test
- public void readLocalDataTest() {
- String appId = "app_read_not_uploaded_data";
- Map<Long, byte[]> expectedData = Maps.newHashMap();
- RssRegisterShuffleRequest rr1 = new RssRegisterShuffleRequest(appId, 1,
- Lists.newArrayList(new PartitionRange(0, 0)), REMOTE_STORAGE);
- RssRegisterShuffleRequest rr2 = new RssRegisterShuffleRequest(appId, 1,
- Lists.newArrayList(new PartitionRange(1, 1)), REMOTE_STORAGE);
- RssRegisterShuffleRequest rr3 = new RssRegisterShuffleRequest(appId, 1,
- Lists.newArrayList(new PartitionRange(2, 2)), REMOTE_STORAGE);
- RssRegisterShuffleRequest rr4 = new RssRegisterShuffleRequest(appId, 1,
- Lists.newArrayList(new PartitionRange(3, 3)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr1);
- shuffleServerClient.registerShuffle(rr2);
- shuffleServerClient.registerShuffle(rr3);
- shuffleServerClient.registerShuffle(rr4);
-
- Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap2 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap3 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap4 = Roaring64NavigableMap.bitmapOf();
-
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 1, 0, 1,3, 25, blockIdBitmap1, expectedData);
- List<ShuffleBlockInfo> blocks2 = createShuffleBlockList(
- 1, 1, 2,5,1024 * 1024, blockIdBitmap2, expectedData);
- List<ShuffleBlockInfo> blocks3 = createShuffleBlockList(
- 1, 2, 3,4, 25, blockIdBitmap3, expectedData);
- List<ShuffleBlockInfo> blocks4 = createShuffleBlockList(
- 1, 3, 4,1, 1024 * 1024, blockIdBitmap4, expectedData);
-
- sendSinglePartitionToShuffleServer(appId, 1,0, 1L, blocks1);
- sendSinglePartitionToShuffleServer(appId, 1,1, 2L, blocks2);
- sendSinglePartitionToShuffleServer(appId, 1,2, 3L, blocks3);
- sendSinglePartitionToShuffleServer(appId, 1,3, 4L, blocks4);
-
- RssGetShuffleResultRequest rg1 = new RssGetShuffleResultRequest(appId, 1, 0);
- shuffleServerClient.getShuffleResult(rg1);
- RssGetShuffleResultRequest rg2 = new RssGetShuffleResultRequest(appId, 1, 1);
- shuffleServerClient.getShuffleResult(rg2);
- RssGetShuffleResultRequest rg3 = new RssGetShuffleResultRequest(appId, 1, 2);
- shuffleServerClient.getShuffleResult(rg3);
- RssGetShuffleResultRequest rg4 = new RssGetShuffleResultRequest(appId, 1, 3);
- shuffleServerClient.getShuffleResult(rg4);
-
- Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
- validateResult(appId, 1, 2, expectedData, getExpectBlockIds(blocks3));
- Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
- validateResult(appId, 1, 0, expectedData, getExpectBlockIds(blocks1));
- Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
- validateResult(appId, 1, 1, expectedData, getExpectBlockIds(blocks2));
- Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
- validateResult(appId, 1, 3, expectedData, getExpectBlockIds(blocks4));
- Uninterruptibles.sleepUninterruptibly(20, TimeUnit.SECONDS);
- ShuffleDataResult sdr = readShuffleData(shuffleServerClient, appId, 1, 0,
- 1, 10, 1000, 0);
- assertEquals(0, sdr.getData().length);
- }
-
- @Test
- public void readMixedDataTest() {
- String appId = "app_read_mix_data";
- RssRegisterShuffleRequest rr1 = new RssRegisterShuffleRequest(appId, 0,
- Lists.newArrayList(new PartitionRange(0, 0)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr1);
-
- Map<Long, byte[]> expectedData = Maps.newHashMap();
- Set<Long> expectedBlock1 = Sets.newHashSet();
-
- Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
-
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 0, 0, 1,15, 1024 * 1024, blockIdBitmap1, expectedData);
-
- blocks1.forEach(b -> expectedBlock1.add(b.getBlockId()));
-
- Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
- partitionToBlocks.put(0, blocks1);
- Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleToBlocks = Maps.newHashMap();
- shuffleToBlocks.put(0, partitionToBlocks);
- RssSendShuffleDataRequest rs1 = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- shuffleServerClient.sendShuffleData(rs1);
-
- RssSendCommitRequest rc1 = new RssSendCommitRequest(appId, 0);
- shuffleServerClient.sendCommit(rc1);
- RssFinishShuffleRequest rf1 = new RssFinishShuffleRequest(appId, 0);
- shuffleServerClient.finishShuffle(rf1);
- Map<Integer, List<Long>> partitionToBlockIds = Maps.newHashMap();
- partitionToBlockIds.put(0, new ArrayList<>(expectedBlock1));
- RssReportShuffleResultRequest rrp1 = new RssReportShuffleResultRequest(
- appId, 0, 1L, partitionToBlockIds, 1);
- shuffleServerClient.reportShuffleResult(rrp1);
-
- RssGetShuffleResultRequest rg1 = new RssGetShuffleResultRequest(appId, 0, 0);
- shuffleServerClient.getShuffleResult(rg1);
-
- ShuffleReadClientImpl readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS_2.name(),
- appId, 0, 0, 100, 1, 10, 1000, REMOTE_STORAGE,
- blockIdBitmap1, Roaring64NavigableMap.bitmapOf(1),
- Lists.newArrayList(new ShuffleServerInfo("test", LOCALHOST, SHUFFLE_SERVER_PORT)), conf, new DefaultIdHelper());
-
- CompressedShuffleBlock csb = readClient.readShuffleBlockData();
- Roaring64NavigableMap matched = Roaring64NavigableMap.bitmapOf();
- assertNotNull(csb);
- assertNotNull(csb.getByteBuffer());
- for (Map.Entry<Long, byte[]> entry : expectedData.entrySet()) {
- if (compareByte(entry.getValue(), csb.getByteBuffer())) {
- matched.addLong(entry.getKey());
- }
- }
- wait(appId);
-
- csb = readClient.readShuffleBlockData();
- while (csb != null && csb.getByteBuffer() != null) {
- for (Map.Entry<Long, byte[]> entry : expectedData.entrySet()) {
- if (compareByte(entry.getValue(), csb.getByteBuffer())) {
- matched.addLong(entry.getKey());
- break;
- }
- }
- csb = readClient.readShuffleBlockData();
- }
- assertTrue(blockIdBitmap1.equals(matched));
-
- ShuffleDataResult sdr = readShuffleData(shuffleServerClient, appId, 0, 0,
- 1, 10, 1000, 0);
- assertEquals(0, sdr.getData().length);
-
- List<ShuffleBlockInfo> blocks5 = createShuffleBlockList(
- 0, 0, 1,15, 1024 * 1024, blockIdBitmap1, expectedData);
- partitionToBlocks.clear();
- shuffleToBlocks.clear();
- partitionToBlocks.put(0, blocks5);
- shuffleToBlocks.put(0, partitionToBlocks);
- RssSendShuffleDataRequest rs5 = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- LocalStorage localStorage = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, 0));
- String path = ShuffleStorageUtils.getFullShuffleDataFolder(localStorage.getBasePath(),
- ShuffleStorageUtils.getShuffleDataPath(appId, 0, 0, 0));
- File file = new File(path);
- assertFalse(file.exists());
- try {
- shuffleServerClient.sendShuffleData(rs5);
- shuffleServerClient.sendCommit(rc1);
- shuffleServerClient.finishShuffle(rf1);
- shuffleServerClient.reportShuffleResult(rrp1);
- } catch (Exception e) {
- fail();
- }
- assertFalse(file.exists());
- }
-
- @Test
- public void readDifferentStorageData() {
- String appId = "app_read_diff_data";
- RssRegisterShuffleRequest rr1 = new RssRegisterShuffleRequest(appId, 0,
- Lists.newArrayList(new PartitionRange(0, 0)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr1);
- Map<Long, byte[]> expectedData = Maps.newHashMap();
- Set<Long> expectedBlock1 = Sets.newHashSet();
- Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 0, 0, 1, 40, 10 * 1024 * 1024, blockIdBitmap1, expectedData);
- blocks1.forEach(b -> expectedBlock1.add(b.getBlockId()));
- Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
- partitionToBlocks.put(0, blocks1);
- Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleToBlocks = Maps.newHashMap();
- shuffleToBlocks.put(0, partitionToBlocks);
- RssSendShuffleDataRequest rs1 = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- shuffleServerClient.sendShuffleData(rs1);
- RssSendCommitRequest rc1 = new RssSendCommitRequest(appId, 0);
- shuffleServerClient.sendCommit(rc1);
- RssFinishShuffleRequest rf1 = new RssFinishShuffleRequest(appId, 0);
- shuffleServerClient.finishShuffle(rf1);
- Map<Integer, List<Long>> partitionToBlockIds = Maps.newHashMap();
- partitionToBlockIds.put(0, new ArrayList<>(expectedBlock1));
- RssReportShuffleResultRequest rrp1 = new RssReportShuffleResultRequest(
- appId, 0, 1L, partitionToBlockIds, 1);
- shuffleServerClient.reportShuffleResult(rrp1);
- Set<Long> expectedBlock2 = Sets.newHashSet();
- List<ShuffleBlockInfo> blocks2 = createShuffleBlockList(
- 0, 0, 2, 40, 1024, blockIdBitmap1, expectedData);
- blocks2.forEach(b -> expectedBlock2.add(b.getBlockId()));
- partitionToBlocks.put(0, blocks2);
- RssSendShuffleDataRequest rs2 = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- shuffleServerClient.sendShuffleData(rs2);
- shuffleServerClient.sendCommit(rc1);
- shuffleServerClient.finishShuffle(rf1);
- partitionToBlockIds.put(0, new ArrayList<>(expectedBlock1));
- RssReportShuffleResultRequest rrp2 = new RssReportShuffleResultRequest(
- appId, 0, 2L, partitionToBlockIds, 1);
- shuffleServerClient.reportShuffleResult(rrp2);
- shuffleServerClient.finishShuffle(rf1);
- RssGetShuffleResultRequest rg1 = new RssGetShuffleResultRequest(appId, 0, 0);
- shuffleServerClient.getShuffleResult(rg1);
-
- ShuffleReadClientImpl readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS_2.name(),
- appId, 0, 0, 100, 1, 10, 1000, REMOTE_STORAGE,
- blockIdBitmap1, Roaring64NavigableMap.bitmapOf(1, 2),
- Lists.newArrayList(new ShuffleServerInfo("test", LOCALHOST, SHUFFLE_SERVER_PORT)), conf, new DefaultIdHelper());
-
- CompressedShuffleBlock csb = readClient.readShuffleBlockData();
- Roaring64NavigableMap matched = Roaring64NavigableMap.bitmapOf();
- assertNotNull(csb);
- assertNotNull(csb.getByteBuffer());
- for (Map.Entry<Long, byte[]> entry : expectedData.entrySet()) {
- if (compareByte(entry.getValue(), csb.getByteBuffer())) {
- matched.addLong(entry.getKey());
- }
- }
- wait(appId);
-
- csb = readClient.readShuffleBlockData();
- while (csb != null && csb.getByteBuffer() != null) {
- for (Map.Entry<Long, byte[]> entry : expectedData.entrySet()) {
- if (compareByte(entry.getValue(), csb.getByteBuffer())) {
- matched.addLong(entry.getKey());
- break;
- }
- }
- csb = readClient.readShuffleBlockData();
- }
- assertTrue(blockIdBitmap1.equals(matched));
- }
-
- @Test
- public void readDifferentStorageDataWithFilter() {
- String appId = "app_read_diff_data_filter";
- RssRegisterShuffleRequest rr1 = new RssRegisterShuffleRequest(appId, 0,
- Lists.newArrayList(new PartitionRange(0, 0)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr1);
- Map<Long, byte[]> expectedData = Maps.newHashMap();
- Set<Long> expectedBlock1 = Sets.newHashSet();
- Set<Long> expectedBlock2 = Sets.newHashSet();
- Set<Long> actualBlock1 = Sets.newHashSet();
- Set<Long> actualBlock2 = Sets.newHashSet();
-
- Roaring64NavigableMap blockIdBitmap = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap processBitmap = Roaring64NavigableMap.bitmapOf();
-
- // send large blocks to HDFS
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 0, 0, 1, 40, 10 * 1024 * 1024, blockIdBitmap, expectedData);
- blocks1.forEach(b -> expectedBlock1.add(b.getBlockId()));
- Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
- partitionToBlocks.put(0, blocks1);
- Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleToBlocks = Maps.newHashMap();
- shuffleToBlocks.put(0, partitionToBlocks);
- RssSendShuffleDataRequest rs1 = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- shuffleServerClient.sendShuffleData(rs1);
- RssSendCommitRequest rc1 = new RssSendCommitRequest(appId, 0);
- shuffleServerClient.sendCommit(rc1);
- RssFinishShuffleRequest rf1 = new RssFinishShuffleRequest(appId, 0);
- shuffleServerClient.finishShuffle(rf1);
- Map<Integer, List<Long>> partitionToBlockIds = Maps.newHashMap();
- partitionToBlockIds.put(0, new ArrayList<>(expectedBlock1));
- RssReportShuffleResultRequest rrp1 = new RssReportShuffleResultRequest(
- appId, 0, 1L, partitionToBlockIds, 1);
- shuffleServerClient.reportShuffleResult(rrp1);
-
- // send small blocks to local file
- List<ShuffleBlockInfo> blocks2 = createShuffleBlockList(
- 0, 0, 2, 40, 1024, blockIdBitmap, expectedData);
- blocks2.forEach(b -> expectedBlock2.add(b.getBlockId()));
- partitionToBlocks.put(0, blocks2);
- RssSendShuffleDataRequest rs2 = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- shuffleServerClient.sendShuffleData(rs2);
- shuffleServerClient.sendCommit(rc1);
- shuffleServerClient.finishShuffle(rf1);
- partitionToBlockIds.put(0, new ArrayList<>(expectedBlock1));
- RssReportShuffleResultRequest rrp2 = new RssReportShuffleResultRequest(
- appId, 0, 2L, partitionToBlockIds, 1);
- shuffleServerClient.reportShuffleResult(rrp2);
- shuffleServerClient.finishShuffle(rf1);
- RssGetShuffleResultRequest rg1 = new RssGetShuffleResultRequest(appId, 0, 0);
- shuffleServerClient.getShuffleResult(rg1);
-
- LocalFileQuorumClientReadHandler localFileQuorumClientReadHandler = new LocalFileQuorumClientReadHandler(
- appId, 0, 0, 0, 1, 3,
- 20 * 1024 , blockIdBitmap, processBitmap, Lists.newArrayList(shuffleServerClient));
- HdfsClientReadHandler hdfsClientReadHandler = new HdfsClientReadHandler(
- appId, 0, 0, 0, 1, 3,
- 100 * 1024 * 1024, blockIdBitmap, processBitmap, REMOTE_STORAGE, conf);
- UploadedHdfsClientReadHandler uploadedHdfsClientReadHandler = new UploadedHdfsClientReadHandler(
- appId, 0, 0, 0, 1, 3,
- 20 * 1024, blockIdBitmap, processBitmap, REMOTE_STORAGE, conf);
- ClientReadHandler[] handlers = new ClientReadHandler[3];
- handlers[0] = localFileQuorumClientReadHandler;
- handlers[1] = hdfsClientReadHandler;
- handlers[2] = uploadedHdfsClientReadHandler;
- ComposedClientReadHandler handler = new ComposedClientReadHandler(handlers);
-
- // read the 1-th segment from localfile
- ShuffleDataResult result = null;
- result = handler.readShuffleData();
- result.getBufferSegments().forEach(b -> actualBlock1.add(b.getBlockId()));
- result.getBufferSegments().forEach(b -> processBitmap.add(b.getBlockId()));
-
- // flush the localfile to uploaded HDFS
- wait(appId);
-
- // read the 2-th segment from uploaded HDFS and other segments from HDFS
- while(result != null) {
- result = handler.readShuffleData();
- if (result != null) {
- result.getBufferSegments().forEach(b -> actualBlock2.add(b.getBlockId()));
- result.getBufferSegments().forEach(b -> processBitmap.add(b.getBlockId()));
- }
- }
-
- assertEquals(expectedBlock1.size() + expectedBlock2.size(),
- actualBlock1.size() + actualBlock2.size());
-
- expectedBlock2.forEach(b -> expectedBlock1.add(b));
- actualBlock2.forEach(b -> actualBlock1.add(b));
- assert (expectedBlock1.equals(actualBlock1));
- }
-
- protected void validateResult(
- Map<Long, byte[]> expectedData,
- ShuffleDataResult sdr) {
- byte[] buffer = sdr.getData();
- List<BufferSegment> bufferSegments = sdr.getBufferSegments();
- assertEquals(expectedData.size(), bufferSegments.size());
- for (Map.Entry<Long, byte[]> entry : expectedData.entrySet()) {
- BufferSegment bs = findBufferSegment(entry.getKey(), bufferSegments);
- assertNotNull(bs);
- byte[] data = new byte[bs.getLength()];
- System.arraycopy(buffer, bs.getOffset(), data, 0, bs.getLength());
- }
- }
-
- private BufferSegment findBufferSegment(long blockId, List<BufferSegment> bufferSegments) {
- for (BufferSegment bs : bufferSegments) {
- if (bs.getBlockId() == blockId) {
- return bs;
- }
- }
- return null;
- }
-
- @Test
- public void diskUsageTest() {
- String appId = "app_read_diskusage_data";
- long originSize = shuffleServers.get(0).getShuffleBufferManager().getCapacity();
- Map<Long, byte[]> expectedData = Maps.newHashMap();
-
- RssRegisterShuffleRequest rr1 = new RssRegisterShuffleRequest(appId, 2,
- Lists.newArrayList(new PartitionRange(0, 0)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr1);
-
- RssRegisterShuffleRequest rr2 = new RssRegisterShuffleRequest(appId, 3,
- Lists.newArrayList(new PartitionRange(1, 1)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr2);
-
- RssRegisterShuffleRequest rr3 = new RssRegisterShuffleRequest(appId, 2,
- Lists.newArrayList(new PartitionRange(1, 1)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr3);
-
- Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap2 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap3 = Roaring64NavigableMap.bitmapOf();
-
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 2, 0, 1,30, 10 * 1024 * 1024, blockIdBitmap1, expectedData);
-
- List<ShuffleBlockInfo> blocks2 = createShuffleBlockList(
- 3, 1, 2,9, 10 * 1024 * 1024, blockIdBitmap2, expectedData);
-
- List<ShuffleBlockInfo> blocks3 = createShuffleBlockList(
- 2, 1, 2,9, 10 * 1024 * 1024, blockIdBitmap3, expectedData);
-
- LocalStorage item = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 2, 0));
- item.createMetadataIfNotExist(appId + "/" + 2);
- item.getLock(appId + "/" + 2).readLock().lock();
- sendSinglePartitionToShuffleServer(appId, 2, 0, 1, blocks1);
- assertFalse(item.canWrite());
- assertEquals(30 * 1024 * 1024 * 10, item.getNotUploadedSize(appId + "/" + 2));
- assertEquals(1, item.getNotUploadedPartitions(appId + "/" + 2).getCardinality());
- boolean isException = false;
- try {
- sendSinglePartitionToShuffleServer(appId, 2, 1, 2, blocks3);
- } catch (RuntimeException re) {
- isException = true;
- assertTrue(re.getMessage().contains("Can't finish shuffle process"));
- }
- item.getLock(appId + "/" + 2).readLock().unlock();
- Uninterruptibles.sleepUninterruptibly(6, TimeUnit.SECONDS);
- assertEquals(originSize, shuffleServers.get(0).getShuffleBufferManager().getCapacity());
- assertTrue(isException);
- RssGetShuffleResultRequest rg1 = new RssGetShuffleResultRequest(appId, 2, 0);
- shuffleServerClient.getShuffleResult(rg1);
- validateResult(appId, 2, 0, expectedData, Sets.newHashSet());
- ShuffleReadClientImpl readClient = new ShuffleReadClientImpl(StorageType.LOCALFILE_HDFS_2.name(),
- appId, 2, 0, 100, 1, 10, 1000, REMOTE_STORAGE,
- blockIdBitmap1, Roaring64NavigableMap.bitmapOf(1),
- Lists.newArrayList(new ShuffleServerInfo("test", LOCALHOST, SHUFFLE_SERVER_PORT)), conf, new DefaultIdHelper());
- validateResult(readClient, expectedData, blockIdBitmap1);
- try {
- sendSinglePartitionToShuffleServer(appId, 3, 1,2, blocks2);
- } catch (RuntimeException re) {
- fail();
- }
- RssGetShuffleResultRequest rg2 = new RssGetShuffleResultRequest(appId, 3, 1);
- shuffleServerClient.getShuffleResult(rg2);
- validateResult(appId, 3, 1, expectedData,
- getExpectBlockIds(blocks2));
-
- Uninterruptibles.sleepUninterruptibly(5, TimeUnit.SECONDS);
- }
-
- @Test
- public void removeMetaTest() {
- String appId = "app_read_diskusage_data_without_report";
- Map<Long, byte[]> expectedData = Maps.newHashMap();
- RssRegisterShuffleRequest rr1 = new RssRegisterShuffleRequest(appId, 2,
- Lists.newArrayList(new PartitionRange(0, 0)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr1);
- RssRegisterShuffleRequest rr2 = new RssRegisterShuffleRequest(appId, 3,
- Lists.newArrayList(new PartitionRange(1, 1)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr2);
-
- Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap blockIdBitmap2 = Roaring64NavigableMap.bitmapOf();
-
- List<ShuffleBlockInfo> blocks1 = createShuffleBlockList(
- 2, 0, 1,30, 10 * 1024 * 1024, blockIdBitmap1, expectedData);
- List<ShuffleBlockInfo> blocks2 = createShuffleBlockList(
- 3, 1, 2,9, 10 * 1024 * 1024, blockIdBitmap2, expectedData);
-
- sendSinglePartitionToShuffleServerWithoutReport(appId, 2, 2, 2, blocks1);
- sendSinglePartitionToShuffleServerWithoutReport(appId, 3, 1,2, blocks2);
- shuffleServers.get(0).getShuffleTaskManager().removeResources(appId);
- LocalStorage storage = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 2, 0));
- Uninterruptibles.sleepUninterruptibly(1500, TimeUnit.MILLISECONDS);
- Set<String> keys = storage.getShuffleMetaSet();
- assertTrue(keys.isEmpty());
- storage = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 3, 1));
- keys = storage.getShuffleMetaSet();
- assertTrue(keys.isEmpty());
-
- appId = "app_read_diskusage_data_with_report";
- rr1 = new RssRegisterShuffleRequest(appId, 0,
- Lists.newArrayList(new PartitionRange(0, 0)), REMOTE_STORAGE);
- shuffleServerClient.registerShuffle(rr1);
- blocks1 = createShuffleBlockList(
- 0, 0, 1,30, 10 * 1024, blockIdBitmap1, expectedData);
- sendSinglePartitionToShuffleServer(appId, 0, 0, 2, blocks1);
- shuffleServers.get(0).getShuffleTaskManager().removeResources(appId);
- storage = (LocalStorage) shuffleServers.get(0).getStorageManager()
- .selectStorage(new ShuffleDataReadEvent(appId, 0, 0));
- Uninterruptibles.sleepUninterruptibly(1500, TimeUnit.MILLISECONDS);
- keys = storage.getShuffleMetaSet();
- assertTrue(keys.isEmpty());
- }
-
- private void sendSinglePartitionToShuffleServer(String appId, int shuffle, int partition,
- long taskAttemptId, List<ShuffleBlockInfo> blocks) {
- Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
- Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleToBlocks = Maps.newHashMap();
- Map<Integer, List<Long>> partitionToBlockIds = Maps.newHashMap();
- Set<Long> expectBlockIds = getExpectBlockIds(blocks);
- partitionToBlocks.put(partition, blocks);
- shuffleToBlocks.put(shuffle, partitionToBlocks);
- RssSendShuffleDataRequest rs = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- shuffleServerClient.sendShuffleData(rs);
- RssSendCommitRequest rc = new RssSendCommitRequest(appId, shuffle);
- shuffleServerClient.sendCommit(rc);
- RssFinishShuffleRequest rf = new RssFinishShuffleRequest(appId, shuffle);
- shuffleServerClient.finishShuffle(rf);
- partitionToBlockIds.put(shuffle, new ArrayList<>(expectBlockIds));
- RssReportShuffleResultRequest rrp = new RssReportShuffleResultRequest(
- appId, shuffle, taskAttemptId, partitionToBlockIds, 1);
- shuffleServerClient.reportShuffleResult(rrp);
- }
-
- private void sendSinglePartitionToShuffleServerWithoutReport(String appId, int shuffle, int partition,
- long taskAttemptId, List<ShuffleBlockInfo> blocks) {
- Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
- Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleToBlocks = Maps.newHashMap();
- partitionToBlocks.put(partition, blocks);
- shuffleToBlocks.put(shuffle, partitionToBlocks);
- RssSendShuffleDataRequest rs = new RssSendShuffleDataRequest(appId, 3, 1000, shuffleToBlocks);
- shuffleServerClient.sendShuffleData(rs);
- RssSendCommitRequest rc = new RssSendCommitRequest(appId, shuffle);
- shuffleServerClient.sendCommit(rc);
- RssFinishShuffleRequest rf = new RssFinishShuffleRequest(appId, shuffle);
- shuffleServerClient.finishShuffle(rf);
- }
-
- protected void validateResult(String appId, int shuffleId, int partitionId,
- Map<Long, byte[]> expectedData, Set<Long> expectBlockIds) {
- int matched = 0;
- int index = 0;
- ShuffleDataResult result = null;
- List<ShuffleDataSegment> shuffleDataSegments = readShuffleIndexSegments(
- shuffleServerClient, appId, shuffleId, partitionId, 1, 10, 1000);
- do {
- if (result != null) {
- byte[] buffer = result.getData();
- for (BufferSegment bs : result.getBufferSegments()) {
- if (expectBlockIds.contains(bs.getBlockId())) {
- byte[] data = new byte[bs.getLength()];
- System.arraycopy(buffer, bs.getOffset(), data, 0, bs.getLength());
- assertEquals(bs.getCrc(), ChecksumUtils.getCrc32(data));
- assertTrue(Arrays.equals(data, expectedData.get(bs.getBlockId())));
- matched++;
- }
- assertTrue(expectBlockIds.contains(bs.getBlockId()));
- }
- }
- result = readShuffleData(shuffleServerClient, appId, shuffleId, partitionId,
- 1, 10, index, shuffleDataSegments);
- ++index;
- } while(result != null && result.getData() != null
- && result.getBufferSegments() != null && !result.getBufferSegments().isEmpty());
- assertEquals(expectBlockIds.size(), matched);
- }
-
- protected void validateResult(ShuffleReadClientImpl readClient, Map<Long, byte[]> expectedData,
- Roaring64NavigableMap blockIdBitmap) {
- CompressedShuffleBlock csb = readClient.readShuffleBlockData();
- Roaring64NavigableMap matched = Roaring64NavigableMap.bitmapOf();
- while (csb != null && csb.getByteBuffer() != null) {
- for (Map.Entry<Long, byte[]> entry : expectedData.entrySet()) {
- if (compareByte(entry.getValue(), csb.getByteBuffer())) {
- matched.addLong(entry.getKey());
- break;
- }
- }
- csb = readClient.readShuffleBlockData();
- }
- assertTrue(blockIdBitmap.equals(matched));
- }
-
- private Set<Long> getExpectBlockIds(List<ShuffleBlockInfo> blocks) {
- List<Long> expectBlockIds = Lists.newArrayList();
- blocks.forEach(b -> expectBlockIds.add(b.getBlockId()));
- return Sets.newHashSet(expectBlockIds);
- }
-}
diff --git a/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RepartitionWithUploadMultiStorageRssTest.java b/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RepartitionWithUploadMultiStorageRssTest.java
deleted file mode 100644
index b694f51..0000000
--- a/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RepartitionWithUploadMultiStorageRssTest.java
+++ /dev/null
@@ -1,71 +0,0 @@
-/*
- * 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.test;
-
-import java.io.File;
-import java.util.Map;
-
-import com.google.common.collect.Maps;
-import com.google.common.io.Files;
-import org.apache.spark.SparkConf;
-import org.apache.spark.shuffle.RssSparkConfig;
-import org.junit.jupiter.api.BeforeAll;
-
-import org.apache.uniffle.coordinator.CoordinatorConf;
-import org.apache.uniffle.server.ShuffleServerConf;
-import org.apache.uniffle.storage.util.StorageType;
-
-public class RepartitionWithUploadMultiStorageRssTest extends RepartitionTest {
- @BeforeAll
- public static void setupServers() throws Exception {
- CoordinatorConf coordinatorConf = getCoordinatorConf();
- Map<String, String> dynamicConf = Maps.newHashMap();
- dynamicConf.put(CoordinatorConf.COORDINATOR_REMOTE_STORAGE_PATH.key(), HDFS_URI + "rss/test");
- dynamicConf.put(RssSparkConfig.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS_2.name());
- addDynamicConf(coordinatorConf, dynamicConf);
- createCoordinatorServer(coordinatorConf);
- ShuffleServerConf shuffleServerConf = getShuffleServerConf();
-
- // local storage config
- File tmpDir = Files.createTempDir();
- tmpDir.deleteOnExit();
- File dataDir1 = new File(tmpDir, "data1");
- File dataDir2 = new File(tmpDir, "data2");
- String basePath = dataDir1.getAbsolutePath() + "," + dataDir2.getAbsolutePath();
- shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_BASE_PATH, basePath);
- shuffleServerConf.set(ShuffleServerConf.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS_2.name());
-
- // uploader and remote storage config
- shuffleServerConf.setBoolean("rss.server.uploader.enable", true);
- shuffleServerConf.setLong("rss.server.uploader.combine.threshold.MB", 32);
- shuffleServerConf.setLong("rss.server.uploader.references.speed.mbps", 128);
- shuffleServerConf.setString("rss.server.uploader.remote.storage.type", StorageType.HDFS.name());
- shuffleServerConf.setString("rss.server.uploader.base.path", HDFS_URI + "rss/test");
- shuffleServerConf.setLong("rss.server.uploader.interval.ms", 10);
- shuffleServerConf.setInteger("rss.server.uploader.thread.number", 4);
-
- shuffleServerConf.setLong(ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, 1024L * 1024L);
-
- createShuffleServer(shuffleServerConf);
- startServers();
- }
-
- @Override
- public void updateRssStorage(SparkConf sparkConf) {
- }
-}
diff --git a/integration-test/spark-common/src/test/java/org/apache/uniffle/test/SparkFallbackReadTest.java b/integration-test/spark-common/src/test/java/org/apache/uniffle/test/SparkFallbackReadTest.java
deleted file mode 100644
index 951ce5e..0000000
--- a/integration-test/spark-common/src/test/java/org/apache/uniffle/test/SparkFallbackReadTest.java
+++ /dev/null
@@ -1,153 +0,0 @@
-/*
- * 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.test;
-
-import java.io.File;
-import java.io.FileNotFoundException;
-import java.util.Map;
-import java.util.concurrent.TimeUnit;
-
-import com.google.common.io.Files;
-import com.google.common.util.concurrent.Uninterruptibles;
-import org.apache.hadoop.fs.FileStatus;
-import org.apache.hadoop.fs.Path;
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.shuffle.RssSparkConfig;
-import org.apache.spark.sql.SparkSession;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.Test;
-import scala.Tuple2;
-
-import org.apache.uniffle.coordinator.CoordinatorConf;
-import org.apache.uniffle.server.ShuffleServerConf;
-import org.apache.uniffle.storage.util.StorageType;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assertions.fail;
-
-public class SparkFallbackReadTest extends SparkIntegrationTestBase {
-
- private static File tmpDir;
- @BeforeAll
- public static void setupServers() throws Exception {
- tmpDir = Files.createTempDir();
- File dataDir1 = new File(tmpDir, "data1");
- File dataDir2 = new File(tmpDir, "data2");
- tmpDir.deleteOnExit();
- CoordinatorConf coordinatorConf = getCoordinatorConf();
- createCoordinatorServer(coordinatorConf);
- ShuffleServerConf shuffleServerConf = getShuffleServerConf();
- String basePath = dataDir1.getAbsolutePath() + "," + dataDir2.getAbsolutePath();
- shuffleServerConf.setString("rss.storage.type", StorageType.LOCALFILE_HDFS_2.name());
- shuffleServerConf.setString("rss.storage.basePath", basePath);
- shuffleServerConf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, HDFS_URI + "rss/multi_storage_integration");
- shuffleServerConf.setDouble(ShuffleServerConf.CLEANUP_THRESHOLD, 0.0);
- shuffleServerConf.setDouble(ShuffleServerConf.HIGH_WATER_MARK_OF_WRITE, 100.0);
- shuffleServerConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 100);
- shuffleServerConf.setBoolean(ShuffleServerConf.UPLOADER_ENABLE, true);
- shuffleServerConf.setLong(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 30L);
- shuffleServerConf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 1L);
- shuffleServerConf.setLong(ShuffleServerConf.SHUFFLE_EXPIRED_TIMEOUT_MS, 5000L);
- shuffleServerConf.setLong(ShuffleServerConf.SERVER_APP_EXPIRED_WITHOUT_HEARTBEAT, 5L * 1000L);
- shuffleServerConf.setLong(ShuffleServerConf.SERVER_COMMIT_TIMEOUT, 20L * 1000L);
- shuffleServerConf.setLong(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 15);
- createShuffleServer(shuffleServerConf);
- startServers();
- }
-
- @Test
- public void resultCompareTest() throws Exception {
- run();
- checkShuffleData();
- }
-
- @Override
- Map runTest(SparkSession spark, String fileName) throws Exception {
- Uninterruptibles.sleepUninterruptibly(3, TimeUnit.SECONDS);
- JavaSparkContext jsc = new JavaSparkContext(spark.sparkContext());
- JavaPairRDD<String, Tuple2<Integer, Integer>> javaPairRDD = TestUtils.combineByKeyRDD(TestUtils.getRDD(jsc));
- final long ts = System.currentTimeMillis();
- javaPairRDD.foreach(partition -> {
- long local = System.currentTimeMillis();
- if (partition._1.equals("duck") && local - ts < 10000) {
- Uninterruptibles.sleepUninterruptibly(20, TimeUnit.SECONDS);
- throw new RuntimeException("oops", new IllegalArgumentException("test exception"));
- }
- });
- return javaPairRDD.collectAsMap();
- }
-
- @Override
- public void updateCommonSparkConf(SparkConf sparkConf) {
- sparkConf.setMaster("local[4, 2]");
- }
-
- @Override
- public void updateSparkConfCustomer(SparkConf sparkConf) {
- sparkConf.setMaster("local[4, 2]");
- sparkConf.set(RssSparkConfig.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS_2.name());
- sparkConf.set(RssSparkConfig.RSS_REMOTE_STORAGE_PATH, HDFS_URI + "rss/multi_storage_integration");
- }
-
- private void checkShuffleData() {
- try {
- String hdfsPath = HDFS_URI + "rss/multi_storage_integration";
- File localPath1 = new File(tmpDir, "data1");
- File localPath2 = new File(tmpDir, "data2");
- File[] files1 = localPath1.listFiles();
- File[] files2 = localPath2.listFiles();
- if (files1 != null) {
- for (File file : files1) {
- File[] shuffles = file.listFiles();
- assertEquals(0, shuffles.length);
- }
- }
-
- if (files2 != null) {
- for (File file : files2) {
- File[] shuffles = file.listFiles();
- assertEquals(0, shuffles.length);
- }
- }
- FileStatus[] files = fs.listStatus(new Path(hdfsPath));
- if (files != null) {
- for (FileStatus file : files) {
- FileStatus[] shuffles = fs.listStatus(file.getPath());
- for (FileStatus another : shuffles) {
- assertEquals(1, shuffles.length);
- }
- Uninterruptibles.sleepUninterruptibly(20, TimeUnit.SECONDS);
- boolean isException = false;
- try {
- fs.listStatus(file.getPath());
- } catch (FileNotFoundException fe) {
- isException = true;
- assertTrue(fe.getMessage().contains("multi_storage_integration"));
- }
- assertTrue(isException);
- }
- }
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-}
diff --git a/integration-test/spark-common/src/test/java/org/apache/uniffle/test/SparkSQLMultiStorageRssTest.java b/integration-test/spark-common/src/test/java/org/apache/uniffle/test/SparkSQLMultiStorageRssTest.java
deleted file mode 100644
index 6650bf4..0000000
--- a/integration-test/spark-common/src/test/java/org/apache/uniffle/test/SparkSQLMultiStorageRssTest.java
+++ /dev/null
@@ -1,87 +0,0 @@
-/*
- * 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.test;
-
-import java.io.File;
-import java.util.Map;
-
-import com.google.common.collect.Maps;
-import com.google.common.io.Files;
-import org.apache.spark.SparkConf;
-import org.apache.spark.shuffle.RssSparkConfig;
-import org.junit.jupiter.api.BeforeAll;
-
-import org.apache.uniffle.coordinator.CoordinatorConf;
-import org.apache.uniffle.server.ShuffleServerConf;
-import org.apache.uniffle.storage.util.StorageType;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-
-public class SparkSQLMultiStorageRssTest extends SparkSQLTest {
- private static String basePath;
-
- @BeforeAll
- public static void setupServers() throws Exception {
- CoordinatorConf coordinatorConf = getCoordinatorConf();
- coordinatorConf.setLong("rss.coordinator.app.expired", 5000);
- Map<String, String> dynamicConf = Maps.newHashMap();
- dynamicConf.put(CoordinatorConf.COORDINATOR_REMOTE_STORAGE_PATH.key(), HDFS_URI + "rss/test");
- dynamicConf.put(RssSparkConfig.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS_2.name());
- addDynamicConf(coordinatorConf, dynamicConf);
- createCoordinatorServer(coordinatorConf);
-
- ShuffleServerConf shuffleServerConf = getShuffleServerConf();
- shuffleServerConf.setLong("rss.server.heartbeat.interval", 5000);
- shuffleServerConf.setLong("rss.server.app.expired.withoutHeartbeat", 4000);
-
- // local storage config
- File tmpDir = Files.createTempDir();
- tmpDir.deleteOnExit();
- File dataDir1 = new File(tmpDir, "data1");
- File dataDir2 = new File(tmpDir, "data2");
- basePath = dataDir1.getAbsolutePath() + "," + dataDir2.getAbsolutePath();
- shuffleServerConf.setString("rss.storage.type", StorageType.LOCALFILE_HDFS_2.name());
- shuffleServerConf.setString("rss.storage.basePath", basePath);
-
- // uploader and remote storage config
- shuffleServerConf.setBoolean("rss.server.uploader.enable", true);
- shuffleServerConf.setLong("rss.server.uploader.combine.threshold.MB", 32);
- shuffleServerConf.setLong("rss.server.uploader.references.speed.mbps", 128);
- shuffleServerConf.setString("rss.server.uploader.remote.storage.type", StorageType.HDFS.name());
- shuffleServerConf.setString("rss.server.uploader.base.path", HDFS_URI + "rss/test");
- shuffleServerConf.setLong("rss.server.uploader.interval.ms", 10);
- shuffleServerConf.setInteger("rss.server.uploader.thread.number", 4);
-
- createShuffleServer(shuffleServerConf);
- startServers();
- }
-
- @Override
- public void updateRssStorage(SparkConf sparkConf) {
- }
-
- @Override
- public void checkShuffleData() throws Exception {
- Thread.sleep(12000);
- String[] paths = basePath.split(",");
- for (String path : paths) {
- File f = new File(path);
- assertEquals(0, f.list().length);
- }
- }
-}
diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java b/server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java
index c4131db..b5b0249 100644
--- a/server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java
+++ b/server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java
@@ -174,52 +174,6 @@ public class ShuffleServerConf extends RssBaseConf {
.defaultValue(600L)
.withDescription("If disk cannot be written for timeout seconds, the flush data event will fail");
- public static final ConfigOption<Boolean> UPLOADER_ENABLE = ConfigOptions
- .key("rss.server.uploader.enable")
- .booleanType()
- .defaultValue(false)
- .withDescription("A switch of the uploader");
-
- public static final ConfigOption<Integer> UPLOADER_THREAD_NUM = ConfigOptions
- .key("rss.server.uploader.thread.number")
- .intType()
- .checkValue(ConfigUtils.positiveIntegerValidator2, "uploader thread number must be positive")
- .defaultValue(4)
- .withDescription("The thread number of the uploader");
-
- public static final ConfigOption<Long> UPLOADER_INTERVAL_MS = ConfigOptions
- .key("rss.server.uploader.interval.ms")
- .longType()
- .checkValue(ConfigUtils.positiveLongValidator, "uploader interval must be positive")
- .defaultValue(3000L)
- .withDescription("The interval for the uploader");
-
- public static final ConfigOption<Long> UPLOAD_COMBINE_THRESHOLD_MB = ConfigOptions
- .key("rss.server.uploader.combine.threshold.mb")
- .longType()
- .checkValue(ConfigUtils.positiveLongValidator, "uploader combine threshold must be positive")
- .defaultValue(32L)
- .withDescription("The threshold of the combine mode");
-
- public static final ConfigOption<String> UPLOADER_BASE_PATH = ConfigOptions
- .key("rss.server.uploader.base.path")
- .stringType()
- .noDefaultValue()
- .withDescription("The base path of the uploader");
-
- public static final ConfigOption<String> UPLOAD_STORAGE_TYPE = ConfigOptions
- .key("rss.server.uploader.remote.storage.type")
- .stringType()
- .defaultValue("HDFS")
- .withDescription("The remote storage type of the uploader");
-
- public static final ConfigOption<Long> REFERENCE_UPLOAD_SPEED_MBS = ConfigOptions
- .key("rss.server.uploader.references.speed.mbps")
- .longType()
- .checkValue(ConfigUtils.positiveLongValidator, "uploader reference speed must be positive")
- .defaultValue(8L)
- .withDescription("The speed for the uploader");
-
public static final ConfigOption<Long> DISK_CAPACITY = ConfigOptions
.key("rss.server.disk.capacity")
.longType()
@@ -227,13 +181,6 @@ public class ShuffleServerConf extends RssBaseConf {
.defaultValue(1024L * 1024L * 1024L * 1024L)
.withDescription("Disk capacity that shuffle server can use");
- public static final ConfigOption<Long> CLEANUP_INTERVAL_MS = ConfigOptions
- .key("rss.server.cleanup.interval.ms")
- .longType()
- .checkValue(ConfigUtils.positiveLongValidator, "cleanup interval must be positive")
- .defaultValue(3000L)
- .withDescription("The interval for cleanup");
-
public static final ConfigOption<Long> SHUFFLE_EXPIRED_TIMEOUT_MS = ConfigOptions
.key("rss.server.shuffle.expired.timeout.ms")
.longType()
@@ -242,20 +189,6 @@ public class ShuffleServerConf extends RssBaseConf {
.withDescription("If the shuffle is not read for the long time, and shuffle is uploaded totally,"
+ " , we can delete the shuffle");
- public static final ConfigOption<Long> SHUFFLE_MAX_UPLOAD_SIZE = ConfigOptions
- .key("rss.server.shuffle.max.upload.size")
- .longType()
- .checkValue(ConfigUtils.positiveLongValidator, "max upload size must be positive")
- .defaultValue(1024L * 1024L * 1024L)
- .withDescription("The max value of upload shuffle size");
-
- public static final ConfigOption<Double> SHUFFLE_MAX_FORCE_UPLOAD_TIME_RATIO = ConfigOptions
- .key("rss.server.shuffle.max.force.upload.time.ratio")
- .doubleType()
- .checkValue(ConfigUtils.percentageDoubleValidator, "max upload time radio must between 0.0 and 100.0")
- .defaultValue(95.0)
- .withDescription("The max upload time ratio in force mode");
-
public static final ConfigOption<Long> SERVER_SHUFFLE_INDEX_SIZE_HINT = ConfigOptions
.key("rss.server.index.size.hint")
.longType()
diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleUploader.java b/server/src/main/java/org/apache/uniffle/server/ShuffleUploader.java
deleted file mode 100644
index f1b83b4..0000000
--- a/server/src/main/java/org/apache/uniffle/server/ShuffleUploader.java
+++ /dev/null
@@ -1,493 +0,0 @@
-/*
- * 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;
-
-import java.io.File;
-import java.util.List;
-import java.util.Queue;
-import java.util.Set;
-import java.util.concurrent.Callable;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Sets;
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
-import com.google.common.util.concurrent.Uninterruptibles;
-import org.apache.commons.io.FileUtils;
-import org.apache.commons.lang3.StringUtils;
-import org.apache.commons.lang3.exception.ExceptionUtils;
-import org.apache.hadoop.conf.Configuration;
-import org.roaringbitmap.RoaringBitmap;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.uniffle.common.util.ByteUnit;
-import org.apache.uniffle.storage.common.LocalStorage;
-import org.apache.uniffle.storage.common.ShuffleFileInfo;
-import org.apache.uniffle.storage.factory.ShuffleUploadHandlerFactory;
-import org.apache.uniffle.storage.handler.api.ShuffleUploadHandler;
-import org.apache.uniffle.storage.request.CreateShuffleUploadHandlerRequest;
-import org.apache.uniffle.storage.util.ShuffleStorageUtils;
-import org.apache.uniffle.storage.util.ShuffleUploadResult;
-import org.apache.uniffle.storage.util.StorageType;
-
-/**
- * ShuffleUploader contains force mode and normal mode, which is decided by the remain
- * space of the disk, and it can be retrieved from the metadata. In force mode, shuffle
- * data will be upload to remote storage whether the shuffle is finished or not.
- * In normal mode, shuffle data will be upload only when it is finished and the shuffle
- * files Both mode will leave the uploaded data to be delete by the cleaner.
- *
- * The underlying handle to upload files to remote storage support HDFS at present, but
- * it will support other storage type (eg, COS, OZONE) and will add more optional parameters,
- * so ShuffleUploader use Joshua Bloch builder pattern to construct and validate the parameters.
- *
- */
-public class ShuffleUploader {
-
- private static final Logger LOG = LoggerFactory.getLogger(ShuffleUploader.class);
-
- private final LocalStorage localStorage;
- private final int uploadThreadNum;
- private final long uploadIntervalMS;
- private final long uploadCombineThresholdMB;
- private final long referenceUploadSpeedMBS;
- private final StorageType remoteStorageType;
- private final String hdfsBasePath;
- private final String serverId;
- private final Configuration hadoopConf;
- private final Thread daemonThread;
- private final long maxShuffleSize;
- private final long maxForceUploadExpireTimeS;
- private final double cleanupThreshold;
-
- private final ExecutorService executorService;
- private volatile boolean isStopped;
-
- public ShuffleUploader(Builder builder) {
- this.localStorage = builder.localStorage;
- this.uploadThreadNum = builder.uploadThreadNum;
- this.uploadIntervalMS = builder.uploadIntervalMS;
- this.uploadCombineThresholdMB = builder.uploadCombineThresholdMB;
- this.referenceUploadSpeedMBS = builder.referenceUploadSpeedMBS;
- this.remoteStorageType = builder.remoteStorageType;
- // HDFS related parameters
- this.hdfsBasePath = builder.hdfsBasePath;
- this.serverId = builder.serverId;
- this.hadoopConf = builder.hadoopConf;
- this.maxShuffleSize = builder.maxShuffleSize;
- this.maxForceUploadExpireTimeS = builder.maxForceUploadExpireTimeS;
- this.cleanupThreshold = builder.cleanupThreshold;
-
- Runnable runnable = () -> {
- run();
- };
-
- daemonThread = new ThreadFactoryBuilder()
- .setDaemon(true)
- .setNameFormat(localStorage.getBasePath() + " - ShuffleUploader-%d")
- .build()
- .newThread(runnable);
-
- executorService = Executors.newFixedThreadPool(
- uploadThreadNum,
- new ThreadFactoryBuilder()
- .setDaemon(true)
- .setNameFormat(localStorage.getBasePath() + " ShuffleUploadWorker-%d")
- .build());
- }
-
- public static class Builder {
- private LocalStorage localStorage;
- private int uploadThreadNum;
- private long uploadIntervalMS;
- private long uploadCombineThresholdMB;
- private long referenceUploadSpeedMBS;
- private StorageType remoteStorageType;
- private String hdfsBasePath;
- private String serverId;
- private Configuration hadoopConf;
- private long maxShuffleSize = (long) ByteUnit.MiB.toBytes(256);
- private long maxForceUploadExpireTimeS;
- private double cleanupThreshold;
-
- public Builder() {
- // use HDFS and not force upload by default
- this.remoteStorageType = StorageType.HDFS;
- }
-
- public Builder localStorage(LocalStorage localStorage) {
- this.localStorage = localStorage;
- return this;
- }
-
- public Builder configuration(ShuffleServerConf conf) {
- uploadThreadNum = conf.get(ShuffleServerConf.UPLOADER_THREAD_NUM);
- uploadIntervalMS = conf.get(ShuffleServerConf.UPLOADER_INTERVAL_MS);
- uploadCombineThresholdMB = conf.get(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB);
- referenceUploadSpeedMBS = conf.get(ShuffleServerConf.REFERENCE_UPLOAD_SPEED_MBS);
- cleanupThreshold = conf.get(ShuffleServerConf.CLEANUP_THRESHOLD);
-
- hdfsBasePath = conf.get(ShuffleServerConf.UPLOADER_BASE_PATH);
- if (StringUtils.isEmpty(hdfsBasePath)) {
- throw new IllegalArgumentException("hdfsBasePath couldn't be empty");
- }
-
- remoteStorageType = StorageType.valueOf(conf.getString(ShuffleServerConf.UPLOAD_STORAGE_TYPE));
- if (StorageType.LOCALFILE.equals(remoteStorageType)) {
- throw new IllegalArgumentException("uploadRemoteStorageType couldn't be LOCALFILE or FILE");
- }
- maxShuffleSize = conf.get(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE);
- double maxShuffleForceUploadTimeRatio = conf.get(ShuffleServerConf.SHUFFLE_MAX_FORCE_UPLOAD_TIME_RATIO);
- long pendingEventTimeoutS = conf.get(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC);
- this.maxForceUploadExpireTimeS = (long) (pendingEventTimeoutS * (maxShuffleForceUploadTimeRatio / 100.0));
- this.hadoopConf = conf.getHadoopConf();
- return this;
- }
-
- public Builder serverId(String serverId) {
- this.serverId = serverId;
- return this;
- }
-
- @VisibleForTesting
- Builder maxForceUploadExpireTimeS(long time) {
- this.maxForceUploadExpireTimeS = time;
- return this;
- }
-
- public ShuffleUploader build() throws IllegalArgumentException {
- return new ShuffleUploader(this);
- }
- }
-
- public void run() {
- while (!isStopped) {
- try {
- long start = System.currentTimeMillis();
- upload();
- long uploadTime = System.currentTimeMillis() - start;
- ShuffleServerMetrics.counterTotalUploadTimeS.inc(uploadTime / 1000.0);
-
- if (uploadTime < uploadIntervalMS) {
- Uninterruptibles.sleepUninterruptibly(uploadIntervalMS - uploadTime, TimeUnit.MILLISECONDS);
- }
- } catch (Exception e) {
- LOG.error("{} - upload exception: {}", Thread.currentThread().getName(), ExceptionUtils.getStackTrace(e));
- }
- }
- }
-
- public void start() {
- daemonThread.start();
- }
-
- public void stop() {
- isStopped = true;
- executorService.shutdownNow();
- Uninterruptibles.joinUninterruptibly(daemonThread, 5, TimeUnit.SECONDS);
- }
-
- // upload is a blocked until uploading success or timeout exception
- @VisibleForTesting
- void upload() {
-
- boolean forceUpload = !localStorage.canWrite();
- LOG.debug("Upload force mode is {}, disk size is {}, shuffle keys are {}",
- forceUpload, localStorage.getDiskSize(), localStorage.getShuffleMetaSet());
-
- List<ShuffleFileInfo> shuffleFileInfos = selectShuffleFiles(uploadThreadNum, forceUpload);
- if (shuffleFileInfos == null || shuffleFileInfos.isEmpty()) {
- cleanUploadedShuffle(Sets.newHashSet());
- return;
- }
-
- List<Callable<ShuffleUploadResult>> callableList = Lists.newLinkedList();
- List<String> shuffleKeys = Lists.newArrayList();
- long totalSize = 0;
- long maxSize = 0;
- for (ShuffleFileInfo shuffleFileInfo : shuffleFileInfos) {
- if (!shuffleFileInfo.isValid()) {
- continue;
- }
-
- if (forceUpload) {
- boolean locked = localStorage.lockShuffleExcluded(shuffleFileInfo.getKey());
- if (!locked) {
- continue;
- }
- shuffleKeys.add(shuffleFileInfo.getKey());
- }
- totalSize = totalSize + shuffleFileInfo.getSize();
- maxSize = Math.max(maxSize, shuffleFileInfo.getSize());
- Callable<ShuffleUploadResult> callable = () -> {
- try {
- CreateShuffleUploadHandlerRequest request =
- new CreateShuffleUploadHandlerRequest.Builder()
- .remoteStorageType(remoteStorageType)
- .remoteStorageBasePath(
- ShuffleStorageUtils.getFullShuffleDataFolder(hdfsBasePath, shuffleFileInfo.getKey()))
- .hadoopConf(hadoopConf)
- .hdfsFilePrefix(serverId)
- .combineUpload(shuffleFileInfo.shouldCombine(uploadCombineThresholdMB))
- .build();
-
- ShuffleUploadHandler handler = getHandlerFactory().createShuffleUploadHandler(request);
- ShuffleUploadResult shuffleUploadResult = handler.upload(
- shuffleFileInfo.getDataFiles(),
- shuffleFileInfo.getIndexFiles(),
- shuffleFileInfo.getPartitions());
- if (shuffleUploadResult == null) {
- return null;
- }
- ShuffleServerMetrics.counterTotalUploadSize.inc(shuffleUploadResult.getSize());
- shuffleUploadResult.setShuffleKey(shuffleFileInfo.getKey());
- return shuffleUploadResult;
- } catch (Exception e) {
- LOG.error("Fail to construct upload callable list {}", ExceptionUtils.getStackTrace(e));
- return null;
- }
- };
- callableList.add(callable);
- }
-
- long uploadTimeoutS = calculateUploadTime(maxSize, totalSize, forceUpload);
- LOG.info("Start to upload {} shuffle info maxSize {} totalSize {} and timeout is {} Seconds and disk size is {}",
- callableList.size(), maxSize, totalSize, uploadTimeoutS, localStorage.getDiskSize());
- long startTimeMs = System.currentTimeMillis();
- try {
- Set<String> successUploadShuffles = Sets.newHashSet();
- List<Future<ShuffleUploadResult>> futures =
- executorService.invokeAll(callableList, uploadTimeoutS, TimeUnit.SECONDS);
- for (Future<ShuffleUploadResult> future : futures) {
- if (future.isDone()) {
- ShuffleUploadResult shuffleUploadResult = future.get();
- if (shuffleUploadResult == null) {
- LOG.info("shuffleUploadResult is empty..");
- continue;
- }
- LOG.info("force mode enable {} upload shuffle {} partitions: {}", forceUpload,
- shuffleUploadResult.getShuffleKey(), shuffleUploadResult.getPartitions().size());
- LOG.debug("upload partitions detail: {}", shuffleUploadResult.getPartitions());
- if (forceUpload) {
- deleteForceUploadPartitions(shuffleUploadResult.getShuffleKey(), shuffleUploadResult.getPartitions());
- localStorage.removeShuffle(shuffleUploadResult.getShuffleKey(), shuffleUploadResult.getSize(),
- shuffleUploadResult.getPartitions());
- } else {
- String shuffleKey = shuffleUploadResult.getShuffleKey();
- localStorage.updateUploadedShuffle(shuffleKey, shuffleUploadResult.getSize(),
- shuffleUploadResult.getPartitions());
- if (localStorage.getNotUploadedPartitions(shuffleKey).isEmpty()) {
- successUploadShuffles.add(shuffleKey);
- }
- }
- } else {
- future.cancel(true);
- }
- }
- cleanUploadedShuffle(successUploadShuffles);
- } catch (Exception e) {
- LOG.error(
- "Fail to upload {}, {}",
- shuffleFileInfos.stream().map(ShuffleFileInfo::getKey).collect(Collectors.joining("\n")),
- ExceptionUtils.getStackTrace(e));
- } finally {
- for (String shuffleKey : shuffleKeys) {
- localStorage.unlockShuffleExcluded(shuffleKey);
- }
- LOG.info("{} upload use {}ms and disk size is {}",
- Thread.currentThread().getName(), System.currentTimeMillis() - startTimeMs, localStorage.getDiskSize());
- }
- }
-
- private void deleteForceUploadPartitions(String shuffleKey, List<Integer> partitions) {
- int failDeleteFiles = 0;
- for (int partition : partitions) {
- String filePrefix = ShuffleStorageUtils.generateAbsoluteFilePrefix(
- localStorage.getBasePath(), shuffleKey, partition, serverId);
- String dataFileName = ShuffleStorageUtils.generateDataFileName(filePrefix);
- String indexFileName = ShuffleStorageUtils.generateIndexFileName(filePrefix);
- File dataFile = new File(dataFileName);
- boolean suc = dataFile.delete();
- if (!suc) {
- failDeleteFiles++;
- }
- File indexFile = new File(indexFileName);
- suc = indexFile.delete();
- if (!suc) {
- failDeleteFiles++;
- }
- if (failDeleteFiles > 0) {
- LOG.error("Force upload process delete file fail {} times", failDeleteFiles);
- }
- }
- }
-
- @VisibleForTesting
- long calculateUploadTime(long maxSize, long totalSize, boolean forceUpload) {
- long defaultUploadTimeoutS = 1L;
- long size = 0;
- if (maxSize > totalSize / uploadThreadNum) {
- size = maxSize;
- } else {
- size = totalSize / uploadThreadNum;
- }
- long cur = ByteUnit.BYTE.toMiB(size) / referenceUploadSpeedMBS;
- if (cur <= defaultUploadTimeoutS) {
- cur = defaultUploadTimeoutS * 2;
- } else {
- cur = cur * 2;
- }
-
- // Use a hard upload limitation in force mode
- if (forceUpload) {
- cur = Math.min(cur, maxForceUploadExpireTimeS);
- }
-
- return cur;
- }
-
- @VisibleForTesting
- List<ShuffleFileInfo> selectShuffleFiles(int num, boolean forceUpload) {
- List<ShuffleFileInfo> shuffleFileInfoList = Lists.newLinkedList();
- List<String> shuffleKeys = localStorage.getSortedShuffleKeys(!forceUpload, num);
- if (shuffleKeys.isEmpty()) {
- return Lists.newArrayList();
- }
-
- LOG.info("Get {} candidate shuffles {}", shuffleKeys.size(), shuffleKeys);
- for (String shuffleKey : shuffleKeys) {
- List<Integer> partitions = getNotUploadedPartitions(shuffleKey);
- long sz = getNotUploadedSize(shuffleKey);
- if (partitions.isEmpty() || sz <= 0) {
- LOG.warn("upload shuffle data empty shuffle {} size {} partitions {}", shuffleKey, sz, partitions);
- continue;
- }
- ShuffleFileInfo shuffleFileInfo = new ShuffleFileInfo();
- shuffleFileInfo.setKey(shuffleKey);
- for (int partition : partitions) {
- long size = addPartition(shuffleFileInfo, partition);
- shuffleFileInfo.setSize(shuffleFileInfo.getSize() + size);
- if (shuffleFileInfo.getSize() > maxShuffleSize) {
- shuffleFileInfoList.add(shuffleFileInfo);
-
- // Restrict the max upload segment to uploadThreadNum to make the
- // uploading finish and release the disk space asap in force mode.
- if (forceUpload && shuffleFileInfoList.size() >= uploadThreadNum) {
- return shuffleFileInfoList;
- }
-
- shuffleFileInfo = new ShuffleFileInfo();
- shuffleFileInfo.setKey(shuffleKey);
- }
- }
-
- if (!shuffleFileInfo.isEmpty()) {
- shuffleFileInfoList.add(shuffleFileInfo);
- }
-
- if (forceUpload && shuffleFileInfoList.size() >= uploadThreadNum) {
- return shuffleFileInfoList;
- }
- }
- return shuffleFileInfoList;
- }
-
- private List<Integer> getNotUploadedPartitions(String key) {
- RoaringBitmap bitmap = localStorage.getNotUploadedPartitions(key);
- List<Integer> partitionList = Lists.newArrayList();
- for (int p : bitmap) {
- partitionList.add(p);
- }
- partitionList.sort(Integer::compare);
- return partitionList;
- }
-
- private long getNotUploadedSize(String key) {
- return localStorage.getNotUploadedSize(key);
- }
-
- @VisibleForTesting
- ShuffleUploadHandlerFactory getHandlerFactory() {
- return ShuffleUploadHandlerFactory.getInstance();
- }
-
-
- private long addPartition(ShuffleFileInfo shuffleFileInfo, int partition) {
- String filePrefix = ShuffleStorageUtils.generateAbsoluteFilePrefix(
- localStorage.getBasePath(), shuffleFileInfo.getKey(), partition, serverId);
- String dataFileName = ShuffleStorageUtils.generateDataFileName(filePrefix);
- String indexFileName = ShuffleStorageUtils.generateIndexFileName(filePrefix);
-
- File dataFile = new File(dataFileName);
- if (!dataFile.exists()) {
- LOG.error("{} don't exist!", dataFileName);
- return 0;
- }
-
- File indexFile = new File(indexFileName);
- if (!indexFile.exists()) {
- LOG.error("{} don't exist!", indexFileName);
- return 0;
- }
-
- shuffleFileInfo.getDataFiles().add(dataFile);
- shuffleFileInfo.getIndexFiles().add(indexFile);
- shuffleFileInfo.getPartitions().add(partition);
- return dataFile.length();
- }
-
- @VisibleForTesting
- void cleanUploadedShuffle(Set<String> successUploadShuffles) {
- Queue<String> expiredShuffleKeys = localStorage.getExpiredShuffleKeys();
- if (localStorage.getDiskSize() * 100.0 / localStorage.getCapacity() < cleanupThreshold) {
- return;
- }
-
- for (String key = expiredShuffleKeys.poll(); key != null; key = expiredShuffleKeys.poll()) {
- successUploadShuffles.add(key);
- }
-
- successUploadShuffles.forEach((shuffleKey) -> {
- // If shuffle data is started to read, shuffle data won't be appended. When shuffle is
- // uploaded totally, the partitions which is not uploaded is empty.
- if (localStorage.isShuffleLongTimeNotRead(shuffleKey)) {
- String shufflePath = ShuffleStorageUtils.getFullShuffleDataFolder(localStorage.getBasePath(), shuffleKey);
- long start = System.currentTimeMillis();
- try {
- File baseFolder = new File(shufflePath);
- FileUtils.deleteDirectory(baseFolder);
- LOG.info("Clean shuffle {}", shuffleKey);
- localStorage.removeResources(shuffleKey);
- LOG.info("Delete shuffle data for shuffle [" + shuffleKey + "] with " + shufflePath
- + " cost " + (System.currentTimeMillis() - start) + " ms");
- } catch (Exception e) {
- LOG.warn("Can't delete shuffle data for shuffle [" + shuffleKey + "] with " + shufflePath, e);
- expiredShuffleKeys.add(shuffleKey);
- }
- } else {
- expiredShuffleKeys.add(shuffleKey);
- }
- });
- }
-}
diff --git a/server/src/main/java/org/apache/uniffle/server/storage/MultiStorageManager.java b/server/src/main/java/org/apache/uniffle/server/storage/MultiStorageManager.java
index d3a88a2..633c2ae 100644
--- a/server/src/main/java/org/apache/uniffle/server/storage/MultiStorageManager.java
+++ b/server/src/main/java/org/apache/uniffle/server/storage/MultiStorageManager.java
@@ -18,10 +18,8 @@
package org.apache.uniffle.server.storage;
import java.io.IOException;
-import java.util.List;
import java.util.Set;
-import com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -30,8 +28,6 @@ import org.apache.uniffle.server.Checker;
import org.apache.uniffle.server.ShuffleDataFlushEvent;
import org.apache.uniffle.server.ShuffleDataReadEvent;
import org.apache.uniffle.server.ShuffleServerConf;
-import org.apache.uniffle.server.ShuffleUploader;
-import org.apache.uniffle.storage.common.LocalStorage;
import org.apache.uniffle.storage.common.Storage;
import org.apache.uniffle.storage.handler.api.ShuffleWriteHandler;
import org.apache.uniffle.storage.request.CreateShuffleWriteHandlerRequest;
@@ -42,30 +38,14 @@ public class MultiStorageManager implements StorageManager {
private final StorageManager warmStorageManager;
private final StorageManager coldStorageManager;
- private final List<ShuffleUploader> uploaders = Lists.newArrayList();
- private final boolean uploadShuffleEnable;
private final long flushColdStorageThresholdSize;
private final long fallBackTimes;
- MultiStorageManager(ShuffleServerConf conf, String shuffleServerId) {
+ MultiStorageManager(ShuffleServerConf conf) {
warmStorageManager = new LocalStorageManager(conf);
coldStorageManager = new HdfsStorageManager(conf);
- uploadShuffleEnable = conf.get(ShuffleServerConf.UPLOADER_ENABLE);
fallBackTimes = conf.get(ShuffleServerConf.FALLBACK_MAX_FAIL_TIMES);
flushColdStorageThresholdSize = conf.getSizeAsBytes(ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE);
- if (uploadShuffleEnable) {
- if (!(warmStorageManager instanceof LocalStorageManager)) {
- throw new IllegalArgumentException("Only LOCALFILE type support upload shuffle");
- }
- LocalStorageManager localStorageManager = (LocalStorageManager) warmStorageManager;
- for (LocalStorage storage :localStorageManager.getStorages()) {
- uploaders.add(new ShuffleUploader.Builder()
- .configuration(conf)
- .serverId(shuffleServerId)
- .localStorage(storage)
- .build());
- }
- }
}
@Override
@@ -121,19 +101,9 @@ public class MultiStorageManager implements StorageManager {
}
public void start() {
- if (uploadShuffleEnable) {
- for (ShuffleUploader uploader : uploaders) {
- uploader.start();
- }
- }
}
public void stop() {
- if (uploadShuffleEnable) {
- for (ShuffleUploader uploader : uploaders) {
- uploader.stop();
- }
- }
}
@Override
diff --git a/server/src/main/java/org/apache/uniffle/server/storage/StorageManagerFactory.java b/server/src/main/java/org/apache/uniffle/server/storage/StorageManagerFactory.java
index f549105..b461acc 100644
--- a/server/src/main/java/org/apache/uniffle/server/storage/StorageManagerFactory.java
+++ b/server/src/main/java/org/apache/uniffle/server/storage/StorageManagerFactory.java
@@ -37,9 +37,8 @@ public class StorageManagerFactory {
} else if (StorageType.HDFS.equals(type) || StorageType.MEMORY_HDFS.equals(type)) {
return new HdfsStorageManager(conf);
} else if (StorageType.LOCALFILE_HDFS.equals(type)
- || StorageType.LOCALFILE_HDFS_2.equals(type)
|| StorageType.MEMORY_LOCALFILE_HDFS.equals(type)) {
- return new MultiStorageManager(conf, serverId);
+ return new MultiStorageManager(conf);
} else {
throw new IllegalArgumentException("unknown storageType was found");
}
diff --git a/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerTest.java b/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerTest.java
index 5bcdf05..2da9b4e 100644
--- a/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerTest.java
+++ b/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerTest.java
@@ -384,7 +384,6 @@ public class ShuffleFlushManagerTest extends HdfsTestBase {
shuffleServerConf.set(RssBaseConf.RSS_STORAGE_BASE_PATH, tempDir.getAbsolutePath());
shuffleServerConf.set(ShuffleServerConf.DISK_CAPACITY, 100L);
shuffleServerConf.set(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 5L);
- shuffleServerConf.set(ShuffleServerConf.UPLOADER_BASE_PATH, "test");
StorageManager storageManager =
StorageManagerFactory.getInstance().createStorageManager("shuffleServerId", shuffleServerConf);
ShuffleFlushManager manager =
diff --git a/server/src/test/java/org/apache/uniffle/server/ShuffleServerConfTest.java b/server/src/test/java/org/apache/uniffle/server/ShuffleServerConfTest.java
index c30c9f4..4c49493 100644
--- a/server/src/test/java/org/apache/uniffle/server/ShuffleServerConfTest.java
+++ b/server/src/test/java/org/apache/uniffle/server/ShuffleServerConfTest.java
@@ -43,7 +43,6 @@ public class ShuffleServerConfTest {
assertFalse(shuffleServerConf.loadConfFromFile(null));
assertEquals("GRPC", shuffleServerConf.getString(ShuffleServerConf.RPC_SERVER_TYPE));
assertEquals(256, shuffleServerConf.getInteger(ShuffleServerConf.JETTY_CORE_POOL_SIZE));
- assertFalse(shuffleServerConf.getBoolean(ShuffleServerConf.UPLOADER_ENABLE));
}
@Test
@@ -67,13 +66,7 @@ public class ShuffleServerConfTest {
assertEquals(2, shuffleServerConf.getLong(ShuffleServerConf.SERVER_BUFFER_CAPACITY));
assertEquals("value1", shuffleServerConf.getString("rss.server.hadoop.a.b", ""));
assertEquals("", shuffleServerConf.getString("rss.server.had.a.b", ""));
- assertEquals("COS", shuffleServerConf.getString(ShuffleServerConf.UPLOAD_STORAGE_TYPE));
assertEquals("GRPC", shuffleServerConf.getString(ShuffleServerConf.RPC_SERVER_TYPE));
- assertTrue(shuffleServerConf.getBoolean(ShuffleServerConf.UPLOADER_ENABLE));
- assertEquals(8L, shuffleServerConf.getLong(ShuffleServerConf.REFERENCE_UPLOAD_SPEED_MBS));
- assertEquals(
- 1024L * 1024L * 1024L, shuffleServerConf.getLong(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE));
- assertEquals(13, shuffleServerConf.getInteger(ShuffleServerConf.UPLOADER_THREAD_NUM));
}
@Test
@@ -92,7 +85,6 @@ public class ShuffleServerConfTest {
assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.SERVER_EVENT_SIZE_THRESHOLD_L2), 90 * MB);
assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.SERVER_EVENT_SIZE_THRESHOLD_L3), 120 * MB);
assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.DISK_CAPACITY), 120 * GB);
- assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE), 45 * MB);
// set conf in memory
shuffleServerConf.setSizeAsBytes(ShuffleServerConf.SERVER_BUFFER_CAPACITY, "10KB");
@@ -102,7 +94,6 @@ public class ShuffleServerConfTest {
shuffleServerConf.setSizeAsBytes(ShuffleServerConf.SERVER_EVENT_SIZE_THRESHOLD_L2, "90MB");
shuffleServerConf.setSizeAsBytes(ShuffleServerConf.SERVER_EVENT_SIZE_THRESHOLD_L3, "120MB");
shuffleServerConf.setSizeAsBytes(ShuffleServerConf.DISK_CAPACITY, "120GB");
- shuffleServerConf.setSizeAsBytes(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE, "45MB");
assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.SERVER_BUFFER_CAPACITY), 10 * KB);
assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.SERVER_READ_BUFFER_CAPACITY), 32 * KB);
@@ -111,7 +102,6 @@ public class ShuffleServerConfTest {
assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.SERVER_EVENT_SIZE_THRESHOLD_L2), 90 * MB);
assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.SERVER_EVENT_SIZE_THRESHOLD_L3), 120 * MB);
assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.DISK_CAPACITY), 120 * GB);
- assertEquals(shuffleServerConf.getSizeAsBytes(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE), 45 * MB);
// test the compatibility of old style
assertEquals(shuffleServerConf.getLong(ShuffleServerConf.SERVER_BUFFER_CAPACITY), 10 * KB);
@@ -121,6 +111,5 @@ public class ShuffleServerConfTest {
assertEquals(shuffleServerConf.getLong(ShuffleServerConf.SERVER_EVENT_SIZE_THRESHOLD_L2), 90 * MB);
assertEquals(shuffleServerConf.getLong(ShuffleServerConf.SERVER_EVENT_SIZE_THRESHOLD_L3), 120 * MB);
assertEquals(shuffleServerConf.getLong(ShuffleServerConf.DISK_CAPACITY), 120 * GB);
- assertEquals(shuffleServerConf.getLong(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE), 45 * MB);
}
}
diff --git a/server/src/test/java/org/apache/uniffle/server/ShuffleUploaderTest.java b/server/src/test/java/org/apache/uniffle/server/ShuffleUploaderTest.java
deleted file mode 100644
index 89ba11e..0000000
--- a/server/src/test/java/org/apache/uniffle/server/ShuffleUploaderTest.java
+++ /dev/null
@@ -1,704 +0,0 @@
-/*
- * 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;
-
-import java.io.File;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.OutputStream;
-import java.util.List;
-import java.util.Random;
-import java.util.concurrent.TimeUnit;
-
-import com.google.common.collect.Lists;
-import com.google.common.collect.Sets;
-import com.google.common.util.concurrent.Uninterruptibles;
-import org.junit.jupiter.api.AfterAll;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.io.TempDir;
-import org.roaringbitmap.RoaringBitmap;
-
-import org.apache.uniffle.storage.common.LocalStorage;
-import org.apache.uniffle.storage.common.ShuffleFileInfo;
-import org.apache.uniffle.storage.common.StorageReadMetrics;
-import org.apache.uniffle.storage.factory.ShuffleUploadHandlerFactory;
-import org.apache.uniffle.storage.handler.api.ShuffleUploadHandler;
-import org.apache.uniffle.storage.util.ShuffleUploadResult;
-import org.apache.uniffle.storage.util.StorageType;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-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.spy;
-import static org.mockito.Mockito.when;
-
-
-public class ShuffleUploaderTest {
-
- private static File base;
-
- @BeforeAll
- public static void setUp(@TempDir File tempDir) {
- base = tempDir;
- ShuffleServerMetrics.register();
- }
-
- @AfterAll
- public static void tearDown() {
- ShuffleServerMetrics.clear();
- }
-
- @Test
- public void selectShuffleFiles() {
- try {
- String app1 = "app-1";
- String shuffle1 = "1";
- String shuffleKey1 = String.join("/", app1, shuffle1);
- File partitionDir1 = makePartitionDir(app1, shuffle1, 1);
- File partitionDir2 = makePartitionDir(app1, shuffle1, 2);
- File partitionDir3 = makePartitionDir(app1, shuffle1, 3);
- File partitionDir4 = makePartitionDir(app1, shuffle1, 4);
- File partitionDir5 = makePartitionDir(app1, shuffle1, 5);
-
- File dataFile1 = new File(partitionDir1.getAbsolutePath() + "/127.0.0.1-8080.data");
- File dataFile2 = new File(partitionDir2.getAbsolutePath() + "/127.0.0.1-8080.data");
- File dataFile3 = new File(partitionDir3.getAbsolutePath() + "/127.0.0.1-8080.data");
- File dataFile4 = new File(partitionDir5.getAbsolutePath() + "/127.0.0.1-8080.data");
-
- File indexFile1 = new File(partitionDir1.getAbsolutePath() + "/127.0.0.1-8080.index");
- File indexFile2 = new File(partitionDir2.getAbsolutePath() + "/127.0.0.1-8080.index");
- File indexFile3 = new File(partitionDir3.getAbsolutePath() + "/127.0.0.1-8080.index");
- File indexFile5 = new File(partitionDir4.getAbsolutePath() + "/127.0.0.1-8080.index");
-
- List<File> dataFiles = Lists.newArrayList(dataFile1, dataFile2, dataFile3, dataFile4);
- dataFiles.forEach(f -> writeFile(f, 10));
-
- List<File> indexFiles = Lists.newArrayList(indexFile1, indexFile2, indexFile3, indexFile5);
- indexFiles.forEach(f -> writeFile(f, 10));
-
- ShuffleServerConf conf = new ShuffleServerConf();
- conf.setInteger(ShuffleServerConf.UPLOADER_THREAD_NUM, 2);
- conf.setLong(ShuffleServerConf.UPLOADER_INTERVAL_MS, 3);
- conf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 300);
- conf.setLong(ShuffleServerConf.REFERENCE_UPLOAD_SPEED_MBS, 1);
- conf.setString(ShuffleServerConf.UPLOAD_STORAGE_TYPE, StorageType.HDFS.name());
- conf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, "hdfs://base");
- LocalStorage mockLocalStorage = mock(LocalStorage.class);
- when(mockLocalStorage.getBasePath()).thenReturn(base.getAbsolutePath());
- ShuffleUploader shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("127.0.0.1-8080")
- .configuration(conf)
- .build();
-
- when(mockLocalStorage.getSortedShuffleKeys(true, 4))
- .thenReturn(Lists.newArrayList(shuffleKey1, "zeroPartitionShuffleKey", "zeroSizeShuffleKey"));
- when(mockLocalStorage.getSortedShuffleKeys(true, 1))
- .thenReturn(Lists.newArrayList(shuffleKey1));
- when(mockLocalStorage.getNotUploadedSize("zeroSizeShuffleKey"))
- .thenReturn(10L);
- when(mockLocalStorage.getNotUploadedPartitions("zeroSizeShuffleKey"))
- .thenReturn(RoaringBitmap.bitmapOf());
- when(mockLocalStorage.getNotUploadedSize("zeroSizeShuffleKey"))
- .thenReturn(0L);
- when(mockLocalStorage.getNotUploadedPartitions("zeroSizeShuffleKey"))
- .thenReturn(RoaringBitmap.bitmapOf(1));
- when(mockLocalStorage.getNotUploadedSize(shuffleKey1))
- .thenReturn(30L);
- when(mockLocalStorage.getNotUploadedPartitions(shuffleKey1))
- .thenReturn(RoaringBitmap.bitmapOf(1, 2, 3));
- when(mockLocalStorage.getNotUploadedPartitions("zeroPartitionShuffleKey"))
- .thenReturn(RoaringBitmap.bitmapOf());
- when(mockLocalStorage.getHighWaterMarkOfWrite()).thenReturn(100.0);
- when(mockLocalStorage.getLowWaterMarkOfWrite()).thenReturn(0.0);
- when(mockLocalStorage.getCapacity()).thenReturn(1024L);
-
- List<ShuffleFileInfo> shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, false);
- assertEquals(1, shuffleFileInfos.size());
- ShuffleFileInfo shuffleFileInfo = shuffleFileInfos.get(0);
- assertResult3(dataFiles, indexFiles, shuffleFileInfo);
- conf.setLong(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE, 5);
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("127.0.0.1-8080")
- .configuration(conf)
- .build();
- shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, false);
- assertEquals(3, shuffleFileInfos.size());
- assertResult1(dataFiles, indexFiles, shuffleFileInfos);
- shuffleFileInfos = shuffleUploader.selectShuffleFiles(1, false);
- assertEquals(3, shuffleFileInfos.size());
- assertResult1(dataFiles, indexFiles, shuffleFileInfos);
- conf.setLong(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE, 20);
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("127.0.0.1-8080")
- .configuration(conf)
- .build();
- shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, false);
- assertEquals(1, shuffleFileInfos.size());
- shuffleFileInfo = shuffleFileInfos.get(0);
- assertResult3(dataFiles, indexFiles, shuffleFileInfo);
- conf.setLong(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE, 15);
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("127.0.0.1-8080")
- .configuration(conf)
- .build();
- shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, false);
- assertResult2(dataFiles, indexFiles, shuffleFileInfos);
- conf.setLong(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE, 12);
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("127.0.0.1-8080")
- .configuration(conf)
- .build();
- shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, false);
- assertResult2(dataFiles, indexFiles, shuffleFileInfos);
-
- when(mockLocalStorage.getSortedShuffleKeys(false, 4))
- .thenReturn(Lists.newArrayList(shuffleKey1, "zeroPartitionShuffleKey", "zeroSizeShuffleKey"));
- when(mockLocalStorage.getSortedShuffleKeys(false, 1))
- .thenReturn(Lists.newArrayList(shuffleKey1));
- conf.setLong(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE, 9);
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("127.0.0.1-8080")
- .configuration(conf)
- .build();
- shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, true);
- assertEquals(2, shuffleFileInfos.size());
- assertEquals(2, shuffleFileInfos.size());
- assertEquals(10, shuffleFileInfos.get(0).getSize());
- assertEquals(10, shuffleFileInfos.get(1).getSize());
-
- when(mockLocalStorage.getHighWaterMarkOfWrite()).thenReturn(95.0);
- when(mockLocalStorage.getLowWaterMarkOfWrite()).thenReturn(0.0);
- when(mockLocalStorage.getCapacity()).thenReturn(11L);
- conf.setInteger(ShuffleServerConf.UPLOADER_THREAD_NUM, 1);
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("127.0.0.1-8080")
- .configuration(conf)
- .build();
- shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, true);
- assertEquals(1, shuffleFileInfos.size());
- assertResult4(dataFiles, indexFiles, shuffleFileInfos);
- } catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- }
-
- @Test
- public void selectShuffleFilesRestrictionTest() {
- try {
- final int appNum = 4;
- final int partitionNum = 5;
- final List<String> shuffleKeys = Lists.newArrayList();
- for (int i = 0; i < appNum; ++i) {
- String appId = "app-" + i;
- String shuffleId = "1";
- String shuffleKey = String.join("/", appId, shuffleId);
- shuffleKeys.add(shuffleKey);
-
- for (int j = 0; j < partitionNum; ++j) {
- File partitionDir = makePartitionDir(appId, shuffleId, j);
- File dataFile = new File(partitionDir.getAbsolutePath() + "/127.0.0.1-8080.data");
- File indexFile = new File(partitionDir.getAbsolutePath() + "/127.0.0.1-8080.index");
- writeFile(dataFile, 5 * (i + 1));
- writeFile(indexFile, 5);
- }
- }
-
- LocalStorage mockLocalStorage = mock(LocalStorage.class);
- when(mockLocalStorage.getBasePath()).thenReturn(base.getAbsolutePath());
- when(mockLocalStorage.getSortedShuffleKeys(true, 4))
- .thenReturn(shuffleKeys);
- when(mockLocalStorage.getNotUploadedSize(any()))
- .thenReturn(partitionNum * 10L);
- when(mockLocalStorage.getNotUploadedPartitions(any()))
- .thenReturn(RoaringBitmap.bitmapOf(0, 1, 2, 3, 4));
- when(mockLocalStorage.getHighWaterMarkOfWrite()).thenReturn(100.0);
- when(mockLocalStorage.getLowWaterMarkOfWrite()).thenReturn(0.0);
- when(mockLocalStorage.getCapacity()).thenReturn(1024L);
- when(mockLocalStorage.getSortedShuffleKeys(false, 4))
- .thenReturn(shuffleKeys);
- ShuffleServerConf conf = new ShuffleServerConf();
- conf.setInteger(ShuffleServerConf.UPLOADER_THREAD_NUM, 4);
- conf.setLong(ShuffleServerConf.UPLOADER_INTERVAL_MS, 3);
- conf.setLong(ShuffleServerConf.REFERENCE_UPLOAD_SPEED_MBS, 10);
- conf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 300);
- conf.setLong(ShuffleServerConf.SHUFFLE_MAX_UPLOAD_SIZE, 10L);
- conf.setString(ShuffleServerConf.UPLOAD_STORAGE_TYPE, StorageType.HDFS.name());
- conf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, "hdfs://base");
-
- ShuffleUploader shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("127.0.0.1-8080")
- .configuration(conf)
- .build();
-
- List<ShuffleFileInfo> shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, false);
- assertEquals(15, shuffleFileInfos.size());
- // huge shuffle's segment num limited to thread num
- shuffleFileInfos = shuffleUploader.selectShuffleFiles(4, true);
- assertEquals(4, shuffleFileInfos.size());
-
- } catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- }
-
- private void assertResult1(List<File> dataFiles, List<File> indexFiles, List<ShuffleFileInfo> shuffleFileInfos) {
- for (int i = 0; i < 3; ++i) {
- assertEquals(dataFiles.get(i).getAbsolutePath(), shuffleFileInfos.get(i).getDataFiles().get(0).getAbsolutePath());
- assertEquals(indexFiles.get(i).getAbsolutePath(), shuffleFileInfos.get(i).getIndexFiles().get(0).getAbsolutePath());
- assertEquals(shuffleFileInfos.get(i).getSize(), 10L);
- }
- }
-
- private void assertResult2(List<File> dataFiles, List<File> indexFiles, List<ShuffleFileInfo> shuffleFileInfos) {
- assertEquals(2, shuffleFileInfos.size());
- assertEquals(20, shuffleFileInfos.get(0).getSize());
- assertEquals(10, shuffleFileInfos.get(1).getSize());
- assertEquals(dataFiles.get(0).getAbsolutePath(), shuffleFileInfos.get(0).getDataFiles().get(0).getAbsolutePath());
- assertEquals(indexFiles.get(0).getAbsolutePath(), shuffleFileInfos.get(0).getIndexFiles().get(0).getAbsolutePath());
- assertEquals(dataFiles.get(1).getAbsolutePath(), shuffleFileInfos.get(0).getDataFiles().get(1).getAbsolutePath());
- assertEquals(indexFiles.get(1).getAbsolutePath(), shuffleFileInfos.get(0).getIndexFiles().get(1).getAbsolutePath());
- assertEquals(dataFiles.get(2).getAbsolutePath(), shuffleFileInfos.get(1).getDataFiles().get(0).getAbsolutePath());
- assertEquals(indexFiles.get(2).getAbsolutePath(), shuffleFileInfos.get(1).getIndexFiles().get(0).getAbsolutePath());
- }
-
- private void assertResult3(List<File> dataFiles, List<File> indexFiles, ShuffleFileInfo shuffleFileInfo) {
- for (int i = 0; i < 3; ++i) {
- assertEquals(dataFiles.get(i).getAbsolutePath(), shuffleFileInfo.getDataFiles().get(i).getAbsolutePath());
- assertEquals(indexFiles.get(i).getAbsolutePath(), shuffleFileInfo.getIndexFiles().get(i).getAbsolutePath());
- }
- }
-
- private void assertResult4(List<File> dataFiles, List<File> indexFiles, List<ShuffleFileInfo> shuffleFileInfos) {
- assertEquals(1, shuffleFileInfos.size());
- assertEquals(10, shuffleFileInfos.get(0).getSize());
- assertEquals(dataFiles.get(0).getAbsolutePath(), shuffleFileInfos.get(0).getDataFiles().get(0).getAbsolutePath());
- assertEquals(indexFiles.get(0).getAbsolutePath(), shuffleFileInfos.get(0).getIndexFiles().get(0).getAbsolutePath());
- }
-
- @Test
- public void calculateUploadTimeTest() {
- LocalStorage mockLocalStorage = mock(LocalStorage.class);
- when(mockLocalStorage.getBasePath()).thenReturn(base.getAbsolutePath());
- ShuffleServerConf conf = new ShuffleServerConf();
- conf.setInteger(ShuffleServerConf.UPLOADER_THREAD_NUM, 1);
- conf.setLong(ShuffleServerConf.UPLOADER_INTERVAL_MS, 3);
- conf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 300);
- conf.setLong(ShuffleServerConf.REFERENCE_UPLOAD_SPEED_MBS, 128);
- conf.setString(ShuffleServerConf.UPLOAD_STORAGE_TYPE, StorageType.HDFS.name());
- conf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, "hdfs://base");
- ShuffleUploader shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("prefix")
- .configuration(conf)
- .maxForceUploadExpireTimeS(13)
- .build();
- assertEquals(2, shuffleUploader.calculateUploadTime(0,0, false));
- assertEquals(2, shuffleUploader.calculateUploadTime(0, 128 * 1024, false));
- assertEquals(2, shuffleUploader.calculateUploadTime(0, 128 * 1024 * 1024, false));
- assertEquals(6, shuffleUploader.calculateUploadTime(0,3 * 128 * 1024 * 1024, false));
- assertEquals(12, shuffleUploader.calculateUploadTime(6 * 128 * 1024 * 1024,
- 3 * 128 * 1024 * 1024, false));
- conf.setInteger(ShuffleServerConf.UPLOADER_THREAD_NUM, 2);
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .serverId("prefix")
- .configuration(conf)
- .maxForceUploadExpireTimeS(10)
- .build();
- assertEquals(2, shuffleUploader.calculateUploadTime(0,0, false));
- assertEquals(2, shuffleUploader.calculateUploadTime(0,128 * 1024, false));
- assertEquals(2, shuffleUploader.calculateUploadTime(0,128 * 1024 * 1024, false));
- assertEquals(6, shuffleUploader.calculateUploadTime(0, 6 * 128 * 1024 * 1024, false));
- assertEquals(8, shuffleUploader.calculateUploadTime(4 * 128 * 1024 * 1024,
- 6 * 128 * 1024 * 1024, false));
-
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .configuration(conf)
- .maxForceUploadExpireTimeS(7)
- .build();
- assertEquals(7, shuffleUploader.calculateUploadTime(4 * 128 * 1024 * 1024,
- 6 * 128 * 1024 * 1024, true));
-
- shuffleUploader = new ShuffleUploader.Builder()
- .localStorage(mockLocalStorage)
- .configuration(conf)
- .maxForceUploadExpireTimeS(1)
- .build();
- assertEquals(1, shuffleUploader.calculateUploadTime(0,0, true));
- }
-
- @Test
- public void uploadTest() {
- try {
- ShuffleUploader.Builder builder = new ShuffleUploader.Builder();
- ShuffleServerConf conf = new ShuffleServerConf();
- conf.setInteger(ShuffleServerConf.UPLOADER_THREAD_NUM, 1);
- conf.setLong(ShuffleServerConf.UPLOADER_INTERVAL_MS, 1000);
- conf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 1);
- conf.setLong(ShuffleServerConf.REFERENCE_UPLOAD_SPEED_MBS, 2);
- conf.setString(ShuffleServerConf.UPLOAD_STORAGE_TYPE, StorageType.HDFS.name());
- conf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, "hdfs://test");
- LocalStorage localStorage = LocalStorage.newBuilder()
- .capacity(100)
- .basePath(base.getAbsolutePath())
- .highWaterMarkOfWrite(50)
- .lowWaterMarkOfWrite(45)
- .shuffleExpiredTimeoutMs(1000)
- .build();
- builder
- .localStorage(localStorage)
- .configuration(conf)
- .serverId("test")
- .maxForceUploadExpireTimeS(1);
- ShuffleUploadHandlerFactory mockFactory = mock(ShuffleUploadHandlerFactory.class);
- ShuffleUploadHandler mockHandler = mock(ShuffleUploadHandler.class);
- when(mockFactory.createShuffleUploadHandler(any())).thenReturn(mockHandler);
- ShuffleUploadResult result0 = new ShuffleUploadResult(50, Lists.newArrayList(1, 2));
- ShuffleUploadResult result1 = new ShuffleUploadResult(90, Lists.newArrayList(1, 2, 3));
- ShuffleUploadResult result2 = new ShuffleUploadResult(10, Lists.newArrayList(1, 2));
- ShuffleUploadResult result3 = new ShuffleUploadResult(40, Lists.newArrayList(1, 3, 2, 4));
- when(mockHandler.upload(any(),any(), any())).thenReturn(result0).thenReturn(result1)
- .thenReturn(result2).thenReturn(result3);
-
- ShuffleUploader uploader = spy(builder.build());
- when(uploader.getHandlerFactory()).thenReturn(mockFactory);
- localStorage.createMetadataIfNotExist("key");
- localStorage.updateWrite("key", 70, Lists.newArrayList(1, 2, 3));
- File dir1 = new File(base.getAbsolutePath() + "/key/1-1/");
- dir1.mkdirs();
- File file1d = new File(base.getAbsolutePath() + "/key/1-1/test.data");
- file1d.createNewFile();
- File file1i = new File(base.getAbsolutePath() + "/key/1-1/test.index");
- file1i.createNewFile();
- File dir2 = new File(base.getAbsolutePath() + "/key/2-2/");
- dir2.mkdirs();
- File file2d = new File(base.getAbsolutePath() + "/key/2-2/test.data");
- file2d.createNewFile();
- File file2i = new File(base.getAbsolutePath() + "/key/2-2/test.index");
- file2i.createNewFile();
- File dir3 = new File(base.getAbsolutePath() + "/key/3-3/");
- dir3.mkdirs();
- File file3d = new File(base.getAbsolutePath() + "/key/3-3/test.data");
- file3d.createNewFile();
- File file3i = new File(base.getAbsolutePath() + "/key/3-3/test.index");
- file3i.createNewFile();
- byte[] data = new byte[20];
- new Random().nextBytes(data);
- try (OutputStream out = new FileOutputStream(file1d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- data = new byte[30];
- try (OutputStream out = new FileOutputStream(file2d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- data = new byte[20];
- try (OutputStream out = new FileOutputStream(file3d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- uploader.upload();
- assertEquals(20, localStorage.getNotUploadedSize("key"));
- assertEquals(1, localStorage.getNotUploadedPartitions("key").getCardinality());
- assertTrue(localStorage.getNotUploadedPartitions("key").contains(3));
- assertFalse(file1d.exists());
- assertFalse(file1i.exists());
- assertFalse(file2d.exists());
- assertFalse(file2i.exists());
- assertTrue(file3d.exists());
- assertTrue(file3i.exists());
-
- localStorage.updateWrite("key", 70, Lists.newArrayList(1, 2));
- file1d.createNewFile();
- file1i.createNewFile();
- file2d.createNewFile();
- file2i.createNewFile();
- data = new byte[30];
- new Random().nextBytes(data);
- try (OutputStream out = new FileOutputStream(file1d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- data = new byte[40];
- try (OutputStream out = new FileOutputStream(file2d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- uploader.upload();
- assertEquals(0, localStorage.getNotUploadedSize("key"));
- assertTrue(localStorage.getNotUploadedPartitions("key").isEmpty());
- assertFalse(file1d.exists());
- assertFalse(file1i.exists());
- assertFalse(file2d.exists());
- assertFalse(file2i.exists());
- assertFalse(file3d.exists());
- assertFalse(file3i.exists());
-
- localStorage.updateWrite("key", 30, Lists.newArrayList(1, 2, 3));
- file1d.createNewFile();
- file1i.createNewFile();
- file2d.createNewFile();
- file2i.createNewFile();
- file3i.createNewFile();
- file3d.createNewFile();
- data = new byte[5];
- new Random().nextBytes(data);
- try (OutputStream out = new FileOutputStream(file1d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- data = new byte[5];
- try (OutputStream out = new FileOutputStream(file2d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- data = new byte[20];
- try (OutputStream out = new FileOutputStream(file3d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- uploader.upload();
- assertEquals(30, localStorage.getNotUploadedSize("key"));
- assertEquals(3, localStorage.getNotUploadedPartitions("key").getCardinality());
-
- localStorage.prepareStartRead("key");
- uploader.upload();
- assertEquals(20, localStorage.getNotUploadedSize("key"));
- assertEquals(1, localStorage.getNotUploadedPartitions("key").getCardinality());
- assertTrue(file1d.exists());
- assertTrue(file1i.exists());
- assertTrue(file2d.exists());
- assertTrue(file2i.exists());
- assertTrue(file3d.exists());
- assertTrue(file3i.exists());
-
- localStorage.updateShuffleLastReadTs("key");
- Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
- assertTrue(file1d.exists());
- assertTrue(file1i.exists());
- assertTrue(file2d.exists());
- assertTrue(file2i.exists());
- assertTrue(file3d.exists());
- assertTrue(file3i.exists());
-
- localStorage.updateShuffleLastReadTs("key");
- localStorage.updateWrite("key", 20, Lists.newArrayList(1, 2, 4));
- data = new byte[10];
- new Random().nextBytes(data);
- try (OutputStream out = new FileOutputStream(file1d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- data = new byte[10];
- try (OutputStream out = new FileOutputStream(file2d)) {
- out.write(data);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- uploader.upload();
- assertEquals(0, localStorage.getNotUploadedSize("key"));
- assertTrue(localStorage.getNotUploadedPartitions("key").isEmpty());
- Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
-
- assertFalse(file1d.exists());
- assertFalse(file1i.exists());
- assertFalse(file2d.exists());
- assertFalse(file2i.exists());
- assertFalse(file3d.exists());
- assertFalse(file3i.exists());
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-
- @Test
- public void cleanTest(@TempDir File tempDir) {
- try {
-
- ShuffleServerConf conf = new ShuffleServerConf();
- conf.setInteger(ShuffleServerConf.UPLOADER_THREAD_NUM, 1);
- conf.setLong(ShuffleServerConf.UPLOADER_INTERVAL_MS, 1000);
- conf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 1);
- conf.setLong(ShuffleServerConf.REFERENCE_UPLOAD_SPEED_MBS, 2);
- conf.setString(ShuffleServerConf.UPLOAD_STORAGE_TYPE, StorageType.HDFS.name());
- conf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, "hdfs://test");
-
- LocalStorage localStorage = LocalStorage.newBuilder()
- .basePath(tempDir.getAbsolutePath())
- .cleanupThreshold(50)
- .highWaterMarkOfWrite(100)
- .lowWaterMarkOfWrite(100)
- .capacity(100)
- .cleanIntervalMs(5000)
- .shuffleExpiredTimeoutMs(1)
- .build();
-
- File baseDir = new File(tempDir, "app-1");
- baseDir.mkdir();
- assertTrue(baseDir.exists());
- File dir1 = new File(baseDir, "1");
- File dir2 = new File(baseDir, "2");
- dir1.mkdir();
- dir2.mkdir();
- assertTrue(dir1.exists());
- assertTrue(dir2.exists());
- localStorage.createMetadataIfNotExist("app-1/1");
- localStorage.createMetadataIfNotExist("app-1/2");
- localStorage.updateWrite("app-1/1", 0, Lists.newArrayList());
- localStorage.updateWrite("app-1/2", 0, Lists.newArrayList());
-
- assertTrue(dir1.exists());
- assertTrue(dir2.exists());
- localStorage.updateWrite("app-1/1", 25, Lists.newArrayList(1));
- localStorage.updateWrite("app-1/2", 35, Lists.newArrayList(2));
- assertEquals(60, localStorage.getDiskSize());
-
- assertTrue(dir1.exists());
- assertTrue(dir2.exists());
-
- ShuffleUploader uploader = new ShuffleUploader.Builder()
- .localStorage(localStorage)
- .configuration(conf)
- .serverId("test")
- .maxForceUploadExpireTimeS(1)
- .build();
- uploader.cleanUploadedShuffle(Sets.newHashSet("app-1/1"));
-
- assertTrue(dir1.exists());
- assertTrue(dir2.exists());
- localStorage.updateReadMetrics(new StorageReadMetrics("app-1", 1));
- uploader.cleanUploadedShuffle(Sets.newHashSet("app-1/1"));
- assertTrue(dir1.exists());
- Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
- uploader.cleanUploadedShuffle(Sets.newHashSet("app-1/1"));
-
- assertFalse(dir1.exists());
- assertTrue(dir2.exists());
- localStorage.updateReadMetrics(new StorageReadMetrics("app-1/2", 2));
- uploader.cleanUploadedShuffle(Sets.newHashSet("app-1/2"));
- assertTrue(dir2.exists());
- assertEquals(35, localStorage.getDiskSize());
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-
- @Test
- public void delayCleanTest(@TempDir File tempDir) {
- LocalStorage storage = LocalStorage.newBuilder()
- .basePath(tempDir.getAbsolutePath())
- .cleanupThreshold(0)
- .highWaterMarkOfWrite(100)
- .lowWaterMarkOfWrite(100)
- .capacity(100)
- .cleanIntervalMs(1000)
- .shuffleExpiredTimeoutMs(1)
- .build();
-
- ShuffleServerConf conf = new ShuffleServerConf();
- conf.setInteger(ShuffleServerConf.UPLOADER_THREAD_NUM, 1);
- conf.setLong(ShuffleServerConf.UPLOADER_INTERVAL_MS, 1000);
- conf.setLong(ShuffleServerConf.UPLOAD_COMBINE_THRESHOLD_MB, 1);
- conf.setLong(ShuffleServerConf.REFERENCE_UPLOAD_SPEED_MBS, 2);
- conf.setString(ShuffleServerConf.UPLOAD_STORAGE_TYPE, StorageType.HDFS.name());
- conf.setString(ShuffleServerConf.UPLOADER_BASE_PATH, "hdfs://test");
-
- storage.createMetadataIfNotExist("key1");
- storage.createMetadataIfNotExist("key2");
- storage.createMetadataIfNotExist("key3");
- storage.updateWrite("key1", 100, Lists.newArrayList());
- storage.updateWrite("key2", 50, Lists.newArrayList());
- storage.updateWrite("key3", 95, Lists.newArrayList());
- assertEquals(3, storage.getShuffleMetaSet().size());
- assertTrue(storage.getShuffleMetaSet().contains("key1"));
- assertTrue(storage.getShuffleMetaSet().contains("key2"));
- assertTrue(storage.getShuffleMetaSet().contains("key3"));
- assertEquals(245, storage.getDiskSize());
-
- ShuffleUploader uploader = new ShuffleUploader.Builder()
- .localStorage(storage)
- .configuration(conf)
- .serverId("test")
- .maxForceUploadExpireTimeS(1)
- .build();
-
- storage.updateShuffleLastReadTs("key1");
- storage.updateShuffleLastReadTs("key2");
- storage.updateShuffleLastReadTs("key3");
- Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
- storage.getExpiredShuffleKeys().offer("key1");
- storage.getExpiredShuffleKeys().offer("key2");
- assertEquals(2, storage.getExpiredShuffleKeys().size());
- uploader.cleanUploadedShuffle(Sets.newHashSet());
- assertEquals(1, storage.getShuffleMetaSet().size());
- assertEquals(95, storage.getDiskSize());
- assertTrue(storage.getShuffleMetaSet().contains("key3"));
- assertEquals(0, storage.getExpiredShuffleKeys().size());
-
- storage.getExpiredShuffleKeys().offer("key3");
- assertEquals(1, storage.getExpiredShuffleKeys().size());
- uploader.cleanUploadedShuffle(Sets.newHashSet());
- assertEquals(0, storage.getShuffleMetaSet().size());
- assertEquals(0, storage.getDiskSize());
- assertEquals(0, storage.getExpiredShuffleKeys().size());
- }
-
- private File makePartitionDir(String appId, String shuffleId, int partitionId) {
- File dir = base.toPath().resolve(appId).resolve(shuffleId)
- .resolve(partitionId + "-" + partitionId).toFile();
- dir.mkdirs();
- return dir;
- }
-
- private void writeFile(File f, int size) {
- byte[] data1 = new byte[size];
- new Random().nextBytes(data1);
- try (OutputStream out = new FileOutputStream(f)) {
- out.write(data1);
- } catch (IOException e) {
- fail(e.getMessage());
- }
- }
-}
diff --git a/server/src/test/java/org/apache/uniffle/server/storage/MultiStorageManagerTest.java b/server/src/test/java/org/apache/uniffle/server/storage/MultiStorageManagerTest.java
index b479586..9a944e4 100644
--- a/server/src/test/java/org/apache/uniffle/server/storage/MultiStorageManagerTest.java
+++ b/server/src/test/java/org/apache/uniffle/server/storage/MultiStorageManagerTest.java
@@ -43,7 +43,7 @@ public class MultiStorageManagerTest {
conf.setString(ShuffleServerConf.RSS_STORAGE_BASE_PATH, "test");
conf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 1024L);
conf.setString(ShuffleServerConf.RSS_STORAGE_TYPE, StorageType.MEMORY_LOCALFILE_HDFS.name());
- MultiStorageManager manager = new MultiStorageManager(conf, "shuffleServerId");
+ MultiStorageManager manager = new MultiStorageManager(conf);
manager.registerRemoteStorage(appId, new RemoteStorageInfo(remoteStorage));
List<ShufflePartitionedBlock> blocks = Lists.newArrayList(new ShufflePartitionedBlock(100, 1000, 1, 1, 1L, null));
ShuffleDataFlushEvent event = new ShuffleDataFlushEvent(
diff --git a/server/src/test/resources/confBySizeStringTest.conf b/server/src/test/resources/confBySizeStringTest.conf
index 6751428..dc7e1d5 100644
--- a/server/src/test/resources/confBySizeStringTest.conf
+++ b/server/src/test/resources/confBySizeStringTest.conf
@@ -25,4 +25,3 @@ rss.server.event.size.threshold.l1 45mb
rss.server.event.size.threshold.l2 90mb
rss.server.event.size.threshold.l3 120mb
rss.server.disk.capacity 120gb
-rss.server.shuffle.max.upload.size 45mb
diff --git a/server/src/test/resources/confTest.conf b/server/src/test/resources/confTest.conf
index 65a38aa..69170b4 100644
--- a/server/src/test/resources/confTest.conf
+++ b/server/src/test/resources/confTest.conf
@@ -23,8 +23,5 @@ rss.server.coordinator.port 9527
rss.server.buffer.capacity 2
rss.server.hadoop.a.b value1
rss.server.had.a.b value2
-rss.server.uploader.remote.sTorage.type COS
rss.server.multistorage.enable true
-rss.server.uploader.enable true
-rss.server.uploader.thread.number 13
rss.server.cluster.hadoop.clustere1.
diff --git a/storage/src/main/java/org/apache/uniffle/storage/factory/ShuffleHandlerFactory.java b/storage/src/main/java/org/apache/uniffle/storage/factory/ShuffleHandlerFactory.java
index e2ad7ca..c885cca 100644
--- a/storage/src/main/java/org/apache/uniffle/storage/factory/ShuffleHandlerFactory.java
+++ b/storage/src/main/java/org/apache/uniffle/storage/factory/ShuffleHandlerFactory.java
@@ -32,7 +32,6 @@ import org.apache.uniffle.storage.handler.impl.HdfsShuffleDeleteHandler;
import org.apache.uniffle.storage.handler.impl.LocalFileDeleteHandler;
import org.apache.uniffle.storage.handler.impl.LocalFileQuorumClientReadHandler;
import org.apache.uniffle.storage.handler.impl.MemoryQuorumClientReadHandler;
-import org.apache.uniffle.storage.handler.impl.UploadedHdfsClientReadHandler;
import org.apache.uniffle.storage.request.CreateShuffleDeleteHandlerRequest;
import org.apache.uniffle.storage.request.CreateShuffleReadHandlerRequest;
import org.apache.uniffle.storage.util.StorageType;
@@ -106,52 +105,6 @@ public class ShuffleHandlerFactory {
request.getStorageBasePath(),
request.getHadoopConf());
});
- } else if (StorageType.LOCALFILE_HDFS_2.name().equals(request.getStorageType())) {
- List<ShuffleServerInfo> shuffleServerInfoList = request.getShuffleServerInfoList();
- List<ShuffleServerClient> shuffleServerClients = shuffleServerInfoList.stream().map(
- ssi -> ShuffleServerClientFactory.getInstance().getShuffleServerClient(
- ClientType.GRPC.name(), ssi)).collect(
- Collectors.toList());
-
- return new ComposedClientReadHandler(() -> {
- return new LocalFileQuorumClientReadHandler(
- request.getAppId(),
- request.getShuffleId(),
- request.getPartitionId(),
- request.getIndexReadLimit(),
- request.getPartitionNumPerRange(),
- request.getPartitionNum(),
- request.getReadBufferSize(),
- request.getExpectBlockIds(),
- request.getProcessBlockIds(),
- shuffleServerClients);
- }, () -> {
- return new HdfsClientReadHandler(
- request.getAppId(),
- request.getShuffleId(),
- request.getPartitionId(),
- request.getIndexReadLimit(),
- request.getPartitionNumPerRange(),
- request.getPartitionNum(),
- request.getReadBufferSize(),
- request.getExpectBlockIds(),
- request.getProcessBlockIds(),
- request.getStorageBasePath(),
- request.getHadoopConf());
- }, () -> {
- return new UploadedHdfsClientReadHandler(
- request.getAppId(),
- request.getShuffleId(),
- request.getPartitionId(),
- request.getIndexReadLimit(),
- request.getPartitionNumPerRange(),
- request.getPartitionNum(),
- request.getReadBufferSize(),
- request.getExpectBlockIds(),
- request.getProcessBlockIds(),
- request.getStorageBasePath(),
- request.getHadoopConf());
- });
} else if (StorageType.MEMORY_LOCALFILE.name().equals(request.getStorageType())) {
List<ShuffleServerInfo> shuffleServerInfoList = request.getShuffleServerInfoList();
List<ShuffleServerClient> shuffleServerClients = shuffleServerInfoList.stream().map(
diff --git a/storage/src/main/java/org/apache/uniffle/storage/factory/ShuffleUploadHandlerFactory.java b/storage/src/main/java/org/apache/uniffle/storage/factory/ShuffleUploadHandlerFactory.java
deleted file mode 100644
index 6c62a10..0000000
--- a/storage/src/main/java/org/apache/uniffle/storage/factory/ShuffleUploadHandlerFactory.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * 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.storage.factory;
-
-import java.io.IOException;
-
-import org.apache.uniffle.storage.handler.api.ShuffleUploadHandler;
-import org.apache.uniffle.storage.handler.impl.HdfsShuffleUploadHandler;
-import org.apache.uniffle.storage.request.CreateShuffleUploadHandlerRequest;
-import org.apache.uniffle.storage.util.StorageType;
-
-public class ShuffleUploadHandlerFactory {
-
- public ShuffleUploadHandlerFactory() {
-
- }
-
- private static class LazyHolder {
- static final ShuffleUploadHandlerFactory INSTANCE = new ShuffleUploadHandlerFactory();
- }
-
- public static ShuffleUploadHandlerFactory getInstance() {
- return LazyHolder.INSTANCE;
- }
-
- public ShuffleUploadHandler createShuffleUploadHandler(
- CreateShuffleUploadHandlerRequest request) throws IOException, RuntimeException {
- if (request.getRemoteStorageType() == StorageType.HDFS) {
- return new HdfsShuffleUploadHandler(
- request.getRemoteStorageBasePath(),
- request.getHadoopConf(),
- request.getHdfsFilePrefix(),
- request.getBufferSize(),
- request.getCombineUpload());
- } else {
- throw new RuntimeException("Unsupported remote storage type " + request.getRemoteStorageType().name());
- }
- }
-}
diff --git a/storage/src/main/java/org/apache/uniffle/storage/handler/api/ShuffleUploadHandler.java b/storage/src/main/java/org/apache/uniffle/storage/handler/api/ShuffleUploadHandler.java
deleted file mode 100644
index a6d6e50..0000000
--- a/storage/src/main/java/org/apache/uniffle/storage/handler/api/ShuffleUploadHandler.java
+++ /dev/null
@@ -1,41 +0,0 @@
-/*
- * 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.storage.handler.api;
-
-import java.io.File;
-import java.util.List;
-
-import org.apache.uniffle.storage.util.ShuffleUploadResult;
-
-public interface ShuffleUploadHandler {
-
- /**
- * Upload data files and index files to remote storage, the correctness of relation between
- * items of data files, index files and partitions is guaranteed by caller/user.
- * Normally we best-effort strategy to upload files, so the result may be part of the files.
- *
- * @param dataFiles local data files to be uploaded to remote storage
- * @param indexFiles local index files to be uploaded to remote storage
- * @param partitions partition id of the local data files and index files
- *
- * @return upload result including total successful uploaded file size and partition id list
- *
- */
- ShuffleUploadResult upload(List<File> dataFiles, List<File> indexFiles, List<Integer> partitions);
-
-}
diff --git a/storage/src/main/java/org/apache/uniffle/storage/handler/impl/HdfsShuffleUploadHandler.java b/storage/src/main/java/org/apache/uniffle/storage/handler/impl/HdfsShuffleUploadHandler.java
deleted file mode 100644
index 8d2d12b..0000000
--- a/storage/src/main/java/org/apache/uniffle/storage/handler/impl/HdfsShuffleUploadHandler.java
+++ /dev/null
@@ -1,268 +0,0 @@
-/*
- * 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.storage.handler.impl;
-
-import java.io.File;
-import java.io.IOException;
-import java.util.List;
-import java.util.concurrent.atomic.AtomicLong;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.Lists;
-import org.apache.commons.lang.exception.ExceptionUtils;
-import org.apache.hadoop.conf.Configuration;
-import org.apache.hadoop.fs.FileSystem;
-import org.apache.hadoop.fs.Path;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.uniffle.storage.handler.api.ShuffleUploadHandler;
-import org.apache.uniffle.storage.util.ShuffleStorageUtils;
-import org.apache.uniffle.storage.util.ShuffleUploadResult;
-
-/**
- * Handler to upload local files to hdfs, it has two mode combine and not combine,
- * upload all files to hdfs as one file, if it is combine mode,
- * upload files one by one to remote storage as separate file, if it is not.
- *
- */
-public class HdfsShuffleUploadHandler implements ShuffleUploadHandler {
- private static final Logger LOG = LoggerFactory.getLogger(ShuffleUploadHandler.class);
-
- private final String baseHdfsPath;
- private final Configuration hadoopConf;
- private final String hdfsFilePrefixBase;
- private final FileSystem fileSystem;
- private final int buffSize;
- private final boolean combineUpload;
- private static final AtomicLong sequenceId = new AtomicLong(0);
-
- public HdfsShuffleUploadHandler(
- String baseHdfsPath,
- Configuration hadoopConf,
- String hdfsFilePrefixBase,
- int buffSize,
- boolean combineUpload) throws IOException, IllegalStateException {
- this.baseHdfsPath = baseHdfsPath;
- this.hadoopConf = hadoopConf;
- this.hdfsFilePrefixBase = hdfsFilePrefixBase + "-" + sequenceId.incrementAndGet()
- + "-" + System.currentTimeMillis();
- this.fileSystem = getFileSystem();
- this.buffSize = buffSize;
- this.combineUpload = combineUpload;
- }
-
- // Upload data files and index files to remote storage by reading files one by one and
- // writing to the dfs stream using best-effort strategy to upload files so the result
- // may be part of the files.We do not merge file locally and then upload for there is
- // no obvious speed advantage considering the merge time and potential spindle or starving risk.
- public ShuffleUploadResult upload(
- List<File> dataFiles,
- List<File> indexFiles,
- List<Integer> partitions) {
- ShuffleUploadResult shuffleUploadResult;
- if (combineUpload) {
- shuffleUploadResult = uploadInternal(dataFiles, indexFiles, partitions);
- } else {
- List<ShuffleUploadResult> results = Lists.newLinkedList();
- for (int i = 0; i < partitions.size(); ++i) {
- ShuffleUploadResult cur = uploadInternal(
- dataFiles.subList(i, i + 1),
- indexFiles.subList(i, i + 1),
- partitions.subList(i, i + 1));
- if (cur != null) {
- results.add(cur);
- }
- }
- shuffleUploadResult = ShuffleUploadResult.merge(results);
- }
- return shuffleUploadResult;
- }
-
- private String generateDataFileName(int partition) {
- if (combineUpload) {
- return ShuffleStorageUtils.generateDataFileName(
- "combine/" + hdfsFilePrefixBase + "-" + partition);
- } else {
- return ShuffleStorageUtils.generateDataFileName(
- partition + "/" + hdfsFilePrefixBase + "-" + partition);
- }
- }
-
- private String generateIndexFileName(int partition) {
- if (combineUpload) {
- return ShuffleStorageUtils.generateIndexFileName(
- "combine/" + hdfsFilePrefixBase + "-" + partition);
- } else {
- return ShuffleStorageUtils.generateIndexFileName(
- partition + "/" + hdfsFilePrefixBase + "-" + partition);
- }
- }
-
- // Use best-effort strategy to upload files one by one using the sequence of the partitions,
- // break the upload loop once encounter error and return the upload result.
- @VisibleForTesting
- ShuffleUploadResult uploadInternal(
- List<File> dataFiles,
- List<File> indexFiles,
- List<Integer> partitions) {
- // upload data files
- List<Long> fileSize = Lists.newLinkedList();
-
- String dataFileName = generateDataFileName(partitions.get(0));
- try (HdfsFileWriter writer = new HdfsFileWriter(new Path(baseHdfsPath, dataFileName), hadoopConf)) {
- for (File file : dataFiles) {
- if (!file.exists() || file.length() == 0) {
- LOG.error("Fail to upload data file {}, for it do not exist or length is 0", file.getAbsolutePath());
- break;
- }
-
- try {
- long sz = ShuffleStorageUtils.uploadFile(file, writer, buffSize);
- if (sz == 0) {
- LOG.error("Fail to upload data file {} upload size is 0", file.getAbsolutePath());
- break;
- }
- fileSize.add(sz);
- } catch (IOException e) {
- LOG.error("Fail to upload data file {}, for {}", file.getAbsolutePath(), e.getMessage());
- break;
- }
- }
- } catch (IOException | IllegalStateException e) {
- LOG.error("Fail to create data output stream {}, {}", dataFileName, e.getMessage());
- return null;
- }
-
- // upload index files
- int num = 0;
- String indexFileName = generateIndexFileName(partitions.get(0));
- if (fileSize.size() == 0) {
- LOG.error("No data file uploaded no need to upload index file {}", indexFileName);
- return null;
- }
-
- try (HdfsFileWriter writer = new HdfsFileWriter(new Path(baseHdfsPath, indexFileName), hadoopConf)) {
- try {
- List<Long> sizes = Lists.newArrayList();
- indexFiles.subList(0, fileSize.size()).forEach(f -> sizes.add(f.length()));
- List<Integer> effectivePartitions = partitions.subList(0, fileSize.size());
- writer.writeHeader(effectivePartitions, sizes, fileSize);
- } catch (IOException e) {
- LOG.error("Fail to write header to index output stream {}, {}", indexFileName, ExceptionUtils.getStackTrace(e));
- return null;
- }
-
- for (File file : indexFiles) {
- if (!file.exists() || file.length() == 0) {
- LOG.error("Fail to upload index file {}, for it do not exist or length is 0", file.getAbsolutePath());
- continue;
- }
-
- try {
- ShuffleStorageUtils.uploadFile(file, writer, buffSize);
- ++num;
- } catch (IOException e) {
- LOG.error("Fail to upload index file {}, for {}", file.getAbsolutePath(), ExceptionUtils.getStackTrace(e));
- break;
- }
- }
-
- } catch (IOException | IllegalStateException e) {
- LOG.error("Fail to create index output stream {}, {}", indexFileName, ExceptionUtils.getStackTrace(e));
- return null;
- }
-
- if (num == 0) {
- return null;
- } else {
- long sz = fileSize.subList(0, num).stream().reduce(0L, Long::sum);
- List<Integer> p = partitions.subList(0, num);
- return new ShuffleUploadResult(sz, p);
- }
- }
-
- // The partition files may be more than 10K, which is not suitable to list entire shuffle dir,
- // so we upload files to $BASE/$APPID/$SHUFFLEID/combine if combineUpload is true and upload
- // files to $BASE/$APPID/$SHUFFLEID/ if combineUpload is false.
- // The file num in $SHUFFLEID/combine is limit to disk size of the shuffle server and the client
- // can list the entire dir and find the target files. The file num in $SHUFFLEID/ may be
- // more than 10K, so we create subdir in it (eg, $SHUFFLEID/$PARTITION_ID/) and the client can list
- // the accurate dir and find the target files (eg, $SHUFFLEID/$PARTITION_ID/$SERVER_ID-$TS-$PARTITION_ID.data).
- @VisibleForTesting
- String getOrCreateTargetPath(int partitionId) {
- String targetDir;
- if (combineUpload) {
- targetDir = ShuffleStorageUtils.getFullShuffleDataFolder(baseHdfsPath, "combine");
- } else {
- targetDir = ShuffleStorageUtils.getFullShuffleDataFolder(baseHdfsPath, String.valueOf(partitionId));
- }
-
- if (createDirIfNotExist(targetDir)) {
- return targetDir;
- } else {
- return null;
- }
- }
-
- private boolean createDirIfNotExist(String pathString) {
- try {
- ShuffleStorageUtils.createDirIfNotExist(fileSystem, pathString);
- return true;
- } catch (IOException ioe) {
- LOG.error("Fail to create {}, {} tring to retry", pathString, ioe.getMessage());
- }
-
- try {
- ShuffleStorageUtils.createDirIfNotExist(fileSystem, pathString);
- return true;
- } catch (IOException ioe) {
- LOG.error("Fail to retry to create {}, {}", pathString, ioe.getMessage());
- }
-
- return false;
- }
-
- private FileSystem getFileSystem() throws IOException, IllegalStateException {
- Path path = new Path(baseHdfsPath);
- FileSystem fileSystem = null;
-
- try {
- fileSystem = ShuffleStorageUtils.getFileSystemForPath(path, hadoopConf);
- if (!fileSystem.exists(path)) {
- fileSystem.mkdirs(path);
- }
- } catch (IOException ioe) {
- // if folder exist, ignore the exception
- if (!fileSystem.exists(path)) {
- LOG.error("Can't create shuffle folder:" + baseHdfsPath, ioe);
- throw ioe;
- }
- }
-
- return fileSystem;
- }
-
- public String getHdfsFilePrefixBase() {
- return hdfsFilePrefixBase;
- }
-
- public String getBaseHdfsPath() {
- return baseHdfsPath;
- }
-}
diff --git a/storage/src/main/java/org/apache/uniffle/storage/handler/impl/UploadedHdfsClientReadHandler.java b/storage/src/main/java/org/apache/uniffle/storage/handler/impl/UploadedHdfsClientReadHandler.java
deleted file mode 100644
index 2094f1f..0000000
--- a/storage/src/main/java/org/apache/uniffle/storage/handler/impl/UploadedHdfsClientReadHandler.java
+++ /dev/null
@@ -1,116 +0,0 @@
-/*
- * 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.storage.handler.impl;
-
-import java.io.IOException;
-import java.util.Comparator;
-
-import org.apache.hadoop.conf.Configuration;
-import org.apache.hadoop.fs.FileStatus;
-import org.apache.hadoop.fs.FileSystem;
-import org.apache.hadoop.fs.Path;
-import org.roaringbitmap.longlong.Roaring64NavigableMap;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.uniffle.common.ShuffleDataResult;
-import org.apache.uniffle.common.util.Constants;
-import org.apache.uniffle.storage.util.ShuffleStorageUtils;
-
-public class UploadedHdfsClientReadHandler extends HdfsClientReadHandler {
-
- private static final Logger LOG = LoggerFactory.getLogger(UploadedHdfsClientReadHandler.class);
-
- public UploadedHdfsClientReadHandler(
- String appId,
- int shuffleId,
- int partitionId,
- int indexReadLimit,
- int partitionNumPerRange,
- int partitionNum,
- int readBufferSize,
- Roaring64NavigableMap expectBlockIds,
- Roaring64NavigableMap processBlockIds,
- String storageBasePath,
- Configuration hadoopConf) {
- super(appId,
- shuffleId,
- partitionId,
- indexReadLimit,
- partitionNumPerRange,
- partitionNum,
- readBufferSize,
- expectBlockIds,
- processBlockIds,
- storageBasePath,
- hadoopConf);
- }
-
- @Override
- protected void init(String fullShufflePath) {
- FileSystem fs;
- Path baseFolder = new Path(fullShufflePath);
- try {
- fs = ShuffleStorageUtils.getFileSystemForPath(baseFolder, hadoopConf);
- } catch (IOException ioe) {
- LOG.warn("Can't get FileSystem for {}", baseFolder);
- return;
- }
- FileStatus[] indexFiles;
- String failedGetIndexFileMsg = "Can't list index file in " + baseFolder;
-
- try {
- indexFiles = fs.listStatus(baseFolder,
- file -> file.getName().endsWith(Constants.SHUFFLE_INDEX_FILE_SUFFIX));
- } catch (Exception e) {
- LOG.warn(failedGetIndexFileMsg, e);
- return;
- }
-
- if (indexFiles != null && indexFiles.length != 0) {
- for (FileStatus status : indexFiles) {
- LOG.info("Find index file for shuffleId[" + shuffleId + "], partitionId["
- + partitionId + "] " + status.getPath());
- String fileNamePrefix = getFileNamePrefix(status.getPath().toUri().toString());
- try {
- HdfsShuffleReadHandler handler = new UploadedStorageHdfsShuffleReadHandler(
- appId, shuffleId, partitionId, fileNamePrefix, readBufferSize,
- expectBlockIds, processBlockIds, hadoopConf);
- readHandlers.add(handler);
- } catch (Exception e) {
- LOG.warn("Can't create ShuffleReaderHandler for " + fileNamePrefix, e);
- }
- }
- readHandlers.sort(Comparator.comparing(HdfsShuffleReadHandler::getFilePrefix));
- }
- }
-
- @Override
- public ShuffleDataResult readShuffleData() {
- // init lazily like LocalFileClientRead
- if (readHandlers.isEmpty()) {
- String fullShufflePath = ShuffleStorageUtils.getFullShuffleDataFolder(storageBasePath,
- ShuffleStorageUtils.getUploadShuffleDataPath(appId, shuffleId, partitionId));
- init(fullShufflePath);
- String combinePath = ShuffleStorageUtils.getFullShuffleDataFolder(storageBasePath,
- ShuffleStorageUtils.getCombineDataPath(appId, shuffleId));
- init(combinePath);
- }
- return super.readShuffleData();
- }
-}
diff --git a/storage/src/main/java/org/apache/uniffle/storage/handler/impl/UploadedStorageHdfsShuffleReadHandler.java b/storage/src/main/java/org/apache/uniffle/storage/handler/impl/UploadedStorageHdfsShuffleReadHandler.java
deleted file mode 100644
index 34ef0e9..0000000
--- a/storage/src/main/java/org/apache/uniffle/storage/handler/impl/UploadedStorageHdfsShuffleReadHandler.java
+++ /dev/null
@@ -1,102 +0,0 @@
-/*
- * 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.storage.handler.impl;
-
-import java.io.IOException;
-import java.nio.ByteBuffer;
-import java.util.Arrays;
-
-import org.apache.hadoop.conf.Configuration;
-import org.roaringbitmap.longlong.Roaring64NavigableMap;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.uniffle.common.ShuffleIndexResult;
-
-public class UploadedStorageHdfsShuffleReadHandler extends HdfsShuffleReadHandler {
-
- private static final Logger LOG = LoggerFactory.getLogger(UploadedStorageHdfsShuffleReadHandler.class);
-
- private final int partitionId;
- private long dataFileOffset;
-
- public UploadedStorageHdfsShuffleReadHandler(
- String appId,
- int shuffleId,
- int partitionId,
- String filePrefix,
- int readBufferSize,
- Roaring64NavigableMap expectBlockIds,
- Roaring64NavigableMap processBlockIds,
- Configuration conf) throws IOException {
- super(appId, shuffleId, partitionId, filePrefix, readBufferSize, expectBlockIds, processBlockIds, conf);
- this.partitionId = partitionId;
- }
-
- @Override
- protected byte[] readShuffleData(long offset, int expectedLength) {
- byte[] data = dataReader.read(dataFileOffset + offset, expectedLength);
- if (data.length != expectedLength) {
- LOG.warn("Fail to read expected[{}] data, actual[{}] from file {}.data",
- expectedLength, data.length, filePrefix);
- return new byte[0];
- }
- return data;
- }
-
- @Override
- protected ShuffleIndexResult readShuffleIndex() {
- long start = System.currentTimeMillis();
- try {
- byte[] indexData = indexReader.read();
-
- ByteBuffer byteBuffer = ByteBuffer.wrap(indexData);
- ShuffleIndexHeader shuffleIndexHeader = ShuffleIndexHeader.extractHeader(byteBuffer);
- if (shuffleIndexHeader == null) {
- LOG.error("Fail to read index from {}.index", filePrefix);
- return new ShuffleIndexResult();
- }
-
- int indexFileOffset = shuffleIndexHeader.getHeaderLen();
- int indexPartitionLen = 0;
- long dataFileOffset = 0;
- for (ShuffleIndexHeader.Entry entry : shuffleIndexHeader.getIndexes()) {
- int partitionId = entry.getPartitionId();
- indexPartitionLen = (int) entry.getPartitionIndexLength();
- if (partitionId != this.partitionId) {
- indexFileOffset += entry.getPartitionIndexLength();
- dataFileOffset += entry.getPartitionDataLength();
- continue;
- }
-
- if ((indexFileOffset + indexPartitionLen) > indexData.length) {
- LOG.error("Index of partition {} is invalid, offset = {}, length = {} in {}.index",
- partitionId, indexFileOffset, indexPartitionLen, filePrefix);
- }
-
- LOG.info("Read index files {}.index for {} ms", filePrefix, System.currentTimeMillis() - start);
- this.dataFileOffset = dataFileOffset;
- return new ShuffleIndexResult(
- Arrays.copyOfRange(indexData, indexFileOffset, indexFileOffset + indexPartitionLen));
- }
- } catch (Exception e) {
- LOG.info("Fail to read index files {}.index", filePrefix);
- }
- return new ShuffleIndexResult();
- }
-}
diff --git a/storage/src/main/java/org/apache/uniffle/storage/request/CreateShuffleUploadHandlerRequest.java b/storage/src/main/java/org/apache/uniffle/storage/request/CreateShuffleUploadHandlerRequest.java
deleted file mode 100644
index a8d03ab..0000000
--- a/storage/src/main/java/org/apache/uniffle/storage/request/CreateShuffleUploadHandlerRequest.java
+++ /dev/null
@@ -1,157 +0,0 @@
-/*
- * 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.storage.request;
-
-import org.apache.hadoop.conf.Configuration;
-
-import org.apache.uniffle.storage.util.StorageType;
-
-/**
- * CreateShuffleUploadHandlerRequest is used to hold the parameters to create remote storage for shuffle uploader.
- * For now only HDFS is supported, but COS, OZONE and other storage system will be supported in the future.
- */
-public class CreateShuffleUploadHandlerRequest {
- private final StorageType remoteStorageType;
- private final String shuffleKey;
- private final String remoteStorageBasePath;
- private final String hdfsFilePrefix;
- private final Configuration hadoopConf;
- private final int bufferSize;
- private final boolean combineUpload;
-
- public StorageType getRemoteStorageType() {
- return remoteStorageType;
- }
-
- public String getShuffleKey() {
- return shuffleKey;
- }
-
- public String getRemoteStorageBasePath() {
- return remoteStorageBasePath;
- }
-
- public Configuration getHadoopConf() {
- return hadoopConf;
- }
-
- public String getHdfsFilePrefix() {
- return hdfsFilePrefix;
- }
-
- public int getBufferSize() {
- return bufferSize;
- }
-
- public boolean getCombineUpload() {
- return combineUpload;
- }
-
- public static class Builder {
- private StorageType remoteStorageType;
- private String shuffleKey;
- private String remoteStorageBasePath;
- private String hdfsFilePrefix;
- private Configuration hadoopConf;
- private int bufferSize;
- private boolean combineUpload;
-
- public Builder() {
- // use HDFS by default, we may use COS, OZONE in the future
- this.remoteStorageType = StorageType.HDFS;
- this.bufferSize = 4096;
- this.combineUpload = true;
- }
-
- public Builder remoteStorageType(StorageType remoteStorageType) {
- this.remoteStorageType = remoteStorageType;
- return this;
- }
-
- public Builder shuffleKey(String shuffleKey) {
- this.shuffleKey = shuffleKey;
- return this;
- }
-
- public Builder remoteStorageBasePath(String remoteStorageBasePath) {
- this.remoteStorageBasePath = remoteStorageBasePath;
- return this;
- }
-
- public Builder hdfsFilePrefix(String hdfsFilePrefix) {
- this.hdfsFilePrefix = hdfsFilePrefix;
- return this;
- }
-
- public Builder hadoopConf(Configuration hadoopConf) {
- this.hadoopConf = hadoopConf;
- return this;
- }
-
- public Builder bufferSize(int bufferSize) {
- this.bufferSize = bufferSize;
- return this;
- }
-
- public Builder combineUpload(boolean combineUpload) {
- this.combineUpload = combineUpload;
- return this;
- }
-
- public CreateShuffleUploadHandlerRequest build() throws IllegalArgumentException {
- validate();
- return new CreateShuffleUploadHandlerRequest(this);
- }
-
- private void validate() throws IllegalArgumentException {
- if (remoteStorageType == null) {
- throw new IllegalArgumentException("Remote storage type must be set");
- }
-
- // We only support HDFS at present, so only check HDFS related parameters
- if (remoteStorageType == StorageType.HDFS) {
- if (remoteStorageBasePath == null || remoteStorageBasePath.isEmpty()) {
- throw new IllegalArgumentException("Remote storage path must be set");
- }
-
- if (hdfsFilePrefix == null || hdfsFilePrefix.isEmpty()) {
- throw new IllegalArgumentException("File prefix must be set");
- }
-
- if (hadoopConf == null) {
- throw new IllegalArgumentException("Hadoop conf must be set");
- }
-
- if (bufferSize <= 1024) {
- throw new IllegalArgumentException("Buffer size must be larger than 1K");
- }
- }
- }
-
- }
-
- private CreateShuffleUploadHandlerRequest(Builder builder) {
- this.remoteStorageType = builder.remoteStorageType;
- this.shuffleKey = builder.shuffleKey;
- this.remoteStorageBasePath = builder.remoteStorageBasePath;
- this.hdfsFilePrefix = builder.hdfsFilePrefix;
- this.hadoopConf = builder.hadoopConf;
- this.bufferSize = builder.bufferSize;
- this.combineUpload = builder.combineUpload;
- }
-}
diff --git a/storage/src/main/java/org/apache/uniffle/storage/util/ShuffleStorageUtils.java b/storage/src/main/java/org/apache/uniffle/storage/util/ShuffleStorageUtils.java
index af15a7b..096702a 100644
--- a/storage/src/main/java/org/apache/uniffle/storage/util/ShuffleStorageUtils.java
+++ b/storage/src/main/java/org/apache/uniffle/storage/util/ShuffleStorageUtils.java
@@ -242,7 +242,6 @@ public class ShuffleStorageUtils {
public static boolean containsLocalFile(String storageType) {
return StorageType.LOCALFILE.name().equals(storageType)
|| StorageType.LOCALFILE_HDFS.name().equals(storageType)
- || StorageType.LOCALFILE_HDFS_2.name().equals(storageType)
|| StorageType.MEMORY_LOCALFILE.name().equals(storageType)
|| StorageType.MEMORY_LOCALFILE_HDFS.name().equals(storageType);
}
diff --git a/storage/src/main/java/org/apache/uniffle/storage/util/ShuffleUploadResult.java b/storage/src/main/java/org/apache/uniffle/storage/util/ShuffleUploadResult.java
deleted file mode 100644
index 00162d8..0000000
--- a/storage/src/main/java/org/apache/uniffle/storage/util/ShuffleUploadResult.java
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * 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.storage.util;
-
-import java.util.Collection;
-import java.util.List;
-import java.util.stream.Collectors;
-
-public class ShuffleUploadResult {
- private String shuffleKey;
- private long size;
- private List<Integer> partitions;
-
- public ShuffleUploadResult() {
- this.size = 0;
- partitions = null;
- }
-
- public ShuffleUploadResult(long size, List<Integer> partitions) {
- this.size = size;
- this.partitions = partitions;
- }
-
- public String getShuffleKey() {
- return shuffleKey;
- }
-
- public void setShuffleKey(String shuffleKey) {
- this.shuffleKey = shuffleKey;
- }
-
- public long getSize() {
- return size;
- }
-
- public void setSize(long size) {
- this.size = size;
- }
-
- public List<Integer> getPartitions() {
- return partitions;
- }
-
- public void setPartitions(List<Integer> partitions) {
- this.partitions = partitions;
- }
-
- public static ShuffleUploadResult merge(List<ShuffleUploadResult> results) {
- if (results == null || results.isEmpty()) {
- return null;
- }
-
- long size = results.stream().map(ShuffleUploadResult::getSize).reduce(0L, Long::sum);
- List<Integer> partitions = results
- .stream()
- .map(ShuffleUploadResult::getPartitions)
- .flatMap(Collection::stream)
- .collect(Collectors.toList());
- return new ShuffleUploadResult(size, partitions);
- }
-
-}
diff --git a/storage/src/main/java/org/apache/uniffle/storage/util/StorageType.java b/storage/src/main/java/org/apache/uniffle/storage/util/StorageType.java
index b368292..80c155e 100644
--- a/storage/src/main/java/org/apache/uniffle/storage/util/StorageType.java
+++ b/storage/src/main/java/org/apache/uniffle/storage/util/StorageType.java
@@ -21,7 +21,6 @@ public enum StorageType {
HDFS,
LOCALFILE,
LOCALFILE_HDFS,
- LOCALFILE_HDFS_2,
MEMORY_LOCALFILE,
MEMORY_HDFS,
MEMORY_LOCALFILE_HDFS
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/HdfsShuffleUploadHandlerTest.java b/storage/src/test/java/org/apache/uniffle/storage/handler/impl/HdfsShuffleUploadHandlerTest.java
deleted file mode 100644
index 8d68544..0000000
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/HdfsShuffleUploadHandlerTest.java
+++ /dev/null
@@ -1,454 +0,0 @@
-/*
- * 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.storage.handler.impl;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.fail;
-
-import com.google.common.collect.Lists;
-import org.apache.uniffle.common.util.ChecksumUtils;
-import org.apache.uniffle.storage.HdfsTestBase;
-import org.apache.uniffle.storage.util.ShuffleStorageUtils;
-import org.apache.uniffle.storage.util.ShuffleUploadResult;
-import java.io.EOFException;
-import java.io.File;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.OutputStream;
-import java.util.ArrayList;
-import java.util.Random;
-
-import org.apache.hadoop.fs.FSDataInputStream;
-import org.apache.hadoop.fs.FileStatus;
-import org.apache.hadoop.fs.Path;
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.Test;
-
-public class HdfsShuffleUploadHandlerTest extends HdfsTestBase {
-
- @Test
- public void initTest() throws IOException, IllegalStateException {
- String basePath = HDFS_URI + "test_base";
- Path path = new Path(basePath);
- new HdfsShuffleUploadHandler(basePath, conf, "test", 4096, true);
- assertTrue(fs.isDirectory(path));
- }
-
- @Test
- public void uploadTestCombine() throws IOException, IllegalStateException {
- String basePath = HDFS_URI + "test_base";
- HdfsShuffleUploadHandler handler = new HdfsShuffleUploadHandler(
- basePath, conf, "uploadTestCombine", 4096, true);
-
- File dataFile1 = File.createTempFile("uploadTestCombine1", ".data", baseDir);
- File dataFile2 = File.createTempFile("uploadTestCombine2", ".data", baseDir);
-
- byte[] data1 = new byte[10];
- new Random().nextBytes(data1);
- try (OutputStream out = new FileOutputStream(dataFile1)) {
- out.write(data1);
- }
-
- byte[] data2 = new byte[30];
- new Random().nextBytes(data2);
- try (OutputStream out = new FileOutputStream(dataFile2)) {
- out.write(data2);
- }
-
- File indexFile1 = File.createTempFile("uploadTestCombine1", ".index", baseDir);
- File indexFile2 = File.createTempFile("uploadTestCombine2", ".index", baseDir);
-
- writeAndAssertResult(handler, dataFile1, dataFile2, indexFile1, indexFile2);
-
-
- String id = handler.getHdfsFilePrefixBase().split("-")[1];
- String ts = handler.getHdfsFilePrefixBase().split("-")[2];
- Path dataPath = new Path(
- basePath,"combine/uploadTestCombine-" + id + "-" + ts + "-1.data");
- Path indexPath = new Path(
- basePath,"combine/uploadTestCombine-" + id + "-" + ts + "-1.index");
- assertTrue(fs.isFile(dataPath));
- assertTrue(fs.isFile(indexPath));
-
- FileStatus dataFileStatus = fs.getFileStatus(dataPath);
- FileStatus indexFileStatus = fs.getFileStatus(indexPath);
-
- // check data file and index file length
- assertEquals(40, dataFileStatus.getLen());
- Assertions.assertEquals(ShuffleStorageUtils.getIndexFileHeaderLen(2) + 20, indexFileStatus.getLen());
-
- // check index file header
- try (FSDataInputStream indexStream = fs.open(indexFileStatus.getPath())) {
- assertEquals(2, indexStream.readInt());
- assertEquals(1, indexStream.readInt());
- assertEquals(5L, indexStream.readLong());
- assertEquals(10L, indexStream.readLong());
- assertEquals(2, indexStream.readInt());
- assertEquals(15L, indexStream.readLong());
- assertEquals(30L, indexStream.readLong());
- }
-
- // check data file content
- try (FSDataInputStream dataStream = fs.open(dataFileStatus.getPath())) {
- for (int i = 0; i < data1.length; ++i) {
- assertEquals(data1[i], dataStream.readByte());
- }
-
- for (int i = 0; i < data2.length; ++i) {
- assertEquals(data2[i], dataStream.readByte());
- }
-
- assertThrows(EOFException.class, dataStream::readByte);
- }
- }
-
- private void writeAndAssertResult(
- HdfsShuffleUploadHandler handler,
- File dataFile1,
- File dataFile2,
- File indexFile1,
- File indexFile2) throws IOException {
- writeData(indexFile1, 5);
- writeData(indexFile2, 15);
- uploadAndVerify(
- handler,
- Lists.newArrayList(dataFile1, dataFile2),
- Lists.newArrayList(indexFile1, indexFile2),
- Lists.newArrayList(1, 2),
- 40L,
- Lists.newArrayList(1, 2));
- }
-
- @Test
- public void uploadTestOneByOne() throws IOException, IllegalStateException {
- String basePath = HDFS_URI + "test_base";
- HdfsShuffleUploadHandler handler =
- new HdfsShuffleUploadHandler(basePath, conf, "uploadTestOneByOne", 4096, false);
- File dataFile1 = File.createTempFile("uploadTestOneByOne1", ".data", baseDir);
- File dataFile2 = File.createTempFile("uploadTestOneByOne2", ".data", baseDir);
-
- byte[] data1 = new byte[10];
- new Random().nextBytes(data1);
- try (OutputStream out = new FileOutputStream(dataFile1)) {
- out.write(data1);
- }
-
- byte[] data2 = new byte[30];
- new Random().nextBytes(data2);
- try (OutputStream out = new FileOutputStream(dataFile2)) {
- out.write(data2);
- }
-
- File indexFile1 = File.createTempFile("uploadTestOneByOne1", ".index", baseDir);
- File indexFile2 = File.createTempFile("uploadTestOneByOne2", ".index", baseDir);
-
- writeAndAssertResult(handler, dataFile1, dataFile2, indexFile1, indexFile2);
-
- String id = handler.getHdfsFilePrefixBase().split("-")[1];
- String ts = handler.getHdfsFilePrefixBase().split("-")[2];
- Path dataPath1 = new Path(basePath,"1/uploadTestOneByOne-" + id + "-" + ts + "-1.data");
- Path indexPath1 = new Path(basePath,"1/uploadTestOneByOne-" + id + "-" + ts + "-1.index");
- Path dataPath2 = new Path(basePath,"2/uploadTestOneByOne-" + id + "-" + ts + "-2.data");
- Path indexPath2 = new Path(basePath,"2/uploadTestOneByOne-" + id + "-" + ts + "-2.index");
-
- assertTrue(fs.isFile(dataPath1));
- assertTrue(fs.isFile(indexPath1));
- assertTrue(fs.isFile(dataPath2));
- assertTrue(fs.isFile(indexPath2));
-
- FileStatus dataFileStatus1 = fs.getFileStatus(dataPath1);
- FileStatus indexFileStatus1 = fs.getFileStatus(indexPath1);
- FileStatus dataFileStatus2 = fs.getFileStatus(dataPath2);
- FileStatus indexFileStatus2 = fs.getFileStatus(indexPath2);
-
- // check data file and index file length
- assertEquals(10, dataFileStatus1.getLen());
- assertEquals(ShuffleStorageUtils.getIndexFileHeaderLen(1) + 5, indexFileStatus1.getLen());
- assertEquals(30, dataFileStatus2.getLen());
- assertEquals(ShuffleStorageUtils.getIndexFileHeaderLen(1) + 15, indexFileStatus2.getLen());
-
- // check index file header
- try (FSDataInputStream indexStream = fs.open(indexFileStatus1.getPath())) {
- assertEquals(1, indexStream.readInt());
- assertEquals(1, indexStream.readInt());
- assertEquals(5L, indexStream.readLong());
- }
-
- try (FSDataInputStream indexStream = fs.open(indexFileStatus2.getPath())) {
- assertEquals(1, indexStream.readInt());
- assertEquals(2, indexStream.readInt());
- assertEquals(15L, indexStream.readLong());
- }
-
- // check data file content
- try (FSDataInputStream dataStream = fs.open(dataFileStatus1.getPath())) {
- for (int i = 0; i < data1.length; ++i) {
- assertEquals(data1[i], dataStream.readByte());
- }
- }
-
- try (FSDataInputStream dataStream = fs.open(dataFileStatus2.getPath())) {
- for (int i = 0; i < data2.length; ++i) {
- assertEquals(data2[i], dataStream.readByte());
- }
-
- assertThrows(EOFException.class, dataStream::readByte);
- }
-
- }
-
- @Test
- public void uploadTestCombineBestEffort() throws IOException, IllegalStateException {
- String basePath = HDFS_URI + "test_base";
- HdfsShuffleUploadHandler handler =
- new HdfsShuffleUploadHandler(basePath, conf, "uploadTestCombineBestEffort", 4096, true);
- File dataFile1 = File.createTempFile("uploadTestCombineBestEffort1", ".data", baseDir);
- File dataFile2 = File.createTempFile("uploadTestCombineBestEffort2", ".data", baseDir);
- File dataFile3 = new File(dataFile1.getAbsolutePath() + "null");
-
- byte[] data1 = new byte[10];
- new Random().nextBytes(data1);
- try (OutputStream out = new FileOutputStream(dataFile1)) {
- out.write(data1);
- }
-
- byte[] data2 = new byte[30];
- new Random().nextBytes(data2);
- try (OutputStream out = new FileOutputStream(dataFile2)) {
- out.write(data2);
- }
-
- File indexFile1 = File.createTempFile("uploadTestCombineBestEffort1", ".index", baseDir);
- File indexFile2 = new File(indexFile1.getAbsolutePath() + "null");
- File indexFile3 = new File(indexFile1.getAbsolutePath() + "null");
-
- writeAndAssert2(handler, indexFile1, 5, Lists.newArrayList(dataFile1, dataFile2, dataFile3), Lists.newArrayList(indexFile1, indexFile2, indexFile3), Lists.newArrayList(1, 2, 3), 10L, Lists.newArrayList(1));
-
- String id = handler.getHdfsFilePrefixBase().split("-")[1];
- String ts = handler.getHdfsFilePrefixBase().split("-")[2];
- Path dataPath = new Path(basePath,"combine/uploadTestCombineBestEffort-" + id + "-" + ts + "-1.data");
- Path indexPath = new Path(basePath,"combine/uploadTestCombineBestEffort-" + id + "-" + ts + "-1.index");
-
- assertTrue(fs.isFile(dataPath));
- assertTrue(fs.isFile(indexPath));
-
- FileStatus dataFileStatus = fs.getFileStatus(dataPath);
- FileStatus indexFileStatus = fs.getFileStatus(indexPath);
-
- // check data file and index file length
- assertEquals(40, dataFileStatus.getLen());
- assertEquals(ShuffleStorageUtils.getIndexFileHeaderLen(2) + 5, indexFileStatus.getLen());
-
- // check data file content
- try (FSDataInputStream dataStream = fs.open(dataFileStatus.getPath())) {
- for (int i = 0; i < data1.length; ++i) {
- assertEquals(data1[i], dataStream.readByte());
- }
-
- for (int i = 0; i < data2.length; ++i) {
- assertEquals(data2[i], dataStream.readByte());
- }
- }
-
- // check index file header
- try (FSDataInputStream indexStream = fs.open(indexFileStatus.getPath())) {
- assertEquals(2, indexStream.readInt());
- assertEquals(1, indexStream.readInt());
- assertEquals(5L, indexStream.readLong());
- assertEquals(10L, indexStream.readLong());
- assertEquals(2, indexStream.readInt());
- assertEquals(0L, indexStream.readLong());
- assertEquals(30L, indexStream.readLong());
-
- indexStream.seek(ShuffleStorageUtils.getIndexFileHeaderLen(2) + 5);
- assertThrows(EOFException.class, indexStream::readByte);
-
- }
- }
-
- @Test
- public void uploadTestOneByOneBestEffort() throws IOException, IllegalStateException {
- String basePath = HDFS_URI + "test_base";
- HdfsShuffleUploadHandler handler = new HdfsShuffleUploadHandler(
- basePath, conf, "uploadTestOneByOneBestEffort", 4096, false);
- File dataFile1 = File.createTempFile("uploadTestOneByOneBestEffort1", ".data", baseDir);
- File dataFile2 = new File(dataFile1.getAbsolutePath() + "null");
- File dataFile3 = File.createTempFile("uploadTestOneByOneBestEffort3", ".data", baseDir);
-
- byte[] data1 = new byte[10];
- new Random().nextBytes(data1);
- try (OutputStream out = new FileOutputStream(dataFile1)) {
- out.write(data1);
- }
-
- byte[] data2 = new byte[30];
- new Random().nextBytes(data2);
- try (OutputStream out = new FileOutputStream(dataFile3)) {
- out.write(data2);
- }
-
- File indexFile1 = File.createTempFile("uploadTestOneByOneBestEffort1", ".index", baseDir);
- File indexFile2 = new File(indexFile1.getAbsolutePath() + "null");
- File indexFile3 = new File(indexFile1.getAbsolutePath() + "null");
-
- writeAndAssert2(handler, indexFile1, 5, Lists.newArrayList(dataFile1, dataFile2, dataFile3), Lists.newArrayList(indexFile1, indexFile2, indexFile3), Lists.newArrayList(1, 2, 3), 10L, Lists.newArrayList(1));
-
- String id = handler.getHdfsFilePrefixBase().split("-")[1];
- String ts = handler.getHdfsFilePrefixBase().split("-")[2];
- Path dataPath1 = new Path(basePath,"1/uploadTestOneByOneBestEffort-" + id + "-" + ts + "-1.data");
- Path indexPath1 = new Path(basePath,"1/uploadTestOneByOneBestEffort-" + id + "-" + ts + "-1.index");
- Path dataPath2 = new Path(basePath,"2/uploadTestOneByOneBestEffort-" + id + "-" + ts + "-2.data");
- Path indexPath2 = new Path(basePath,"2/uploadTestOneByOneBestEffort-" + id + "-" + ts + "-2.index");
- Path dataPath3 = new Path(basePath,"3/uploadTestOneByOneBestEffort-" + id + "-" + ts + "-3.data");
- Path indexPath3 = new Path(basePath,"3/uploadTestOneByOneBestEffort-" + id + "-" + ts + "-3.index");
-
- assertTrue(fs.isFile(dataPath1));
- assertTrue(fs.isFile(indexPath1));
- assertTrue(fs.isFile(dataPath2));
- assertFalse(fs.isFile(indexPath2));
- assertTrue(fs.isFile(dataPath3));
- assertTrue(fs.isFile(indexPath3));
-
- FileStatus dataFileStatus1 = fs.getFileStatus(dataPath1);
- FileStatus indexFileStatus1 = fs.getFileStatus(indexPath1);
- FileStatus dataFileStatus2 = fs.getFileStatus(dataPath2);
-
- // check data file and index file length
- assertEquals(10, dataFileStatus1.getLen());
- assertEquals(ShuffleStorageUtils.getIndexFileHeaderLen(1) + 5, indexFileStatus1.getLen());
- assertEquals(0, fs.getFileStatus(dataPath2).getLen());
- assertEquals(30, fs.getFileStatus(dataPath3).getLen());
-
- // check data file content
- try (FSDataInputStream dataStream = fs.open(dataFileStatus1.getPath())) {
- for (int i = 0; i < data1.length; ++i) {
- assertEquals(data1[i], dataStream.readByte());
- }
- }
-
- assertEquals(0, dataFileStatus2.getLen());
-
- // check index file header
- try (FSDataInputStream indexStream = fs.open(indexFileStatus1.getPath())) {
- assertEquals(1, indexStream.readInt());
- assertEquals(1, indexStream.readInt());
- assertEquals(5L, indexStream.readLong());
-
- assertThrows(EOFException.class, () ->
- indexStream.seek(ShuffleStorageUtils.getIndexFileHeaderLen(1) + 5 + 1));
-
- }
-
- }
-
- private void writeAndAssert2(
- HdfsShuffleUploadHandler handler,
- File indexFile,
- int length,
- ArrayList<File> indexFiles,
- ArrayList<File> dataFiles,
- ArrayList<Integer> partitions,
- long successSize,
- ArrayList<Integer> successPartitions) throws IOException {
- writeData(indexFile, length);
- uploadAndVerify(handler, indexFiles, dataFiles, partitions, successSize, successPartitions);
- }
-
- private void uploadAndVerify(
- HdfsShuffleUploadHandler handler,
- ArrayList<File> indexFiles,
- ArrayList<File> dataFiles,
- ArrayList<Integer> partitions,
- long successSize,
- ArrayList<Integer> successPartitions) {
- ShuffleUploadResult ret = handler.upload(
- indexFiles,
- dataFiles,
- partitions);
- assertEquals(successSize, ret.getSize());
- assertEquals(successPartitions, ret.getPartitions());
- }
-
- private void writeData(File file, int length) throws IOException {
- try (OutputStream out = new FileOutputStream(file)) {
- out.write(new byte[length]);
- }
- }
-
- @Test
- public void writeHeaderTest() {
- try {
- String basePath = HDFS_URI + "test_base";
- HdfsShuffleUploadHandler handler = new HdfsShuffleUploadHandler(
- basePath, conf, "uploadTestOneByOneBestEffort", 4096, false);
-
- File indexFile1 = File.createTempFile("writeHeaderTest1", ".index", baseDir);
- File indexFile2 = File.createTempFile("writeHeaderTest2", ".index", baseDir);
- File indexFile3 = File.createTempFile("writeHeaderTest3", ".index", baseDir);
-
- try (OutputStream out = new FileOutputStream(indexFile1)) {
- out.write(new byte[5]);
- }
-
- try (OutputStream out = new FileOutputStream(indexFile2)) {
- out.write(new byte[15]);
- }
-
- try (OutputStream out = new FileOutputStream(indexFile3)) {
- out.write(new byte[25]);
- }
-
- HdfsFileWriter writer = new HdfsFileWriter(new Path(handler.getBaseHdfsPath() + "/writeHeaderTest.index"), conf);
- writer.writeHeader(
- Lists.newArrayList(1, 2, 3),
- Lists.newArrayList(indexFile1.length(), indexFile2.length(), indexFile3.length()),
- Lists.newArrayList(1L, 2L, 3L));
- writer.close();
- byte[] buf = new byte[(int)ShuffleStorageUtils.getIndexFileHeaderLen(3)- ShuffleStorageUtils.getHeaderCrcLen()];
-
- FSDataInputStream readStream1 = fs.open(
- new Path(handler.getBaseHdfsPath() + "/writeHeaderTest.index"));
- assertEquals(3, readStream1.readInt());
- assertEquals(1, readStream1.readInt());
- assertEquals(5, readStream1.readLong());
- assertEquals(1, readStream1.readLong());
- assertEquals(2, readStream1.readInt());
- assertEquals(15, readStream1.readLong());
- assertEquals(2, readStream1.readLong());
- assertEquals(3, readStream1.readInt());
- assertEquals(25, readStream1.readLong());
- assertEquals(3, readStream1.readLong());
-
- FSDataInputStream readStream2 = fs.open(
- new Path(handler.getBaseHdfsPath() + "/writeHeaderTest.index"));
- readStream2.readFully(buf);
- long crc = readStream1.readLong();
- assertEquals(ChecksumUtils.getCrc32(buf), crc);
-
- readStream1.close();
- readStream2.close();
-
- } catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- }
-}
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/MultiStorageHdfsClientReadHandlerTest.java b/storage/src/test/java/org/apache/uniffle/storage/handler/impl/MultiStorageHdfsClientReadHandlerTest.java
deleted file mode 100644
index 8ac9764..0000000
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/MultiStorageHdfsClientReadHandlerTest.java
+++ /dev/null
@@ -1,373 +0,0 @@
-/*
- * 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.storage.handler.impl;
-
-import com.google.common.collect.Lists;
-import com.google.common.collect.Sets;
-import org.apache.uniffle.common.BufferSegment;
-import org.apache.uniffle.common.ShuffleDataResult;
-import org.apache.uniffle.storage.HdfsTestBase;
-import org.apache.uniffle.storage.common.FileBasedShuffleSegment;
-import org.apache.hadoop.fs.Path;
-import org.junit.jupiter.api.Test;
-import org.roaringbitmap.longlong.Roaring64NavigableMap;
-
-import java.io.IOException;
-import java.nio.ByteBuffer;
-import java.util.List;
-import java.util.Random;
-import java.util.Set;
-
-import static org.junit.jupiter.api.Assertions.fail;
-
-public class MultiStorageHdfsClientReadHandlerTest extends HdfsTestBase {
-
- @Test
- public void handlerReadUnCombinedDataTest() {
- try {
- Roaring64NavigableMap expectBlockIds = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap processBlockIds = Roaring64NavigableMap.bitmapOf();
-
- String basePath = HDFS_URI + "test_backup_read";
- Path combinePath = new Path(basePath + "/app1/0/combine");
- fs.mkdirs(combinePath);
- Path partitionPath = new Path(basePath + "/app1/0/1");
- fs.mkdirs(partitionPath);
-
- Path dataPath = new Path(basePath + "/app1/0/1/2.data");
- HdfsFileWriter writer = new HdfsFileWriter(dataPath, conf);
- byte[] data = writeData(writer);
- writer.close();
-
- writeIndexData(basePath, "/app1/0/1/2.index", Lists.newArrayList(), expectBlockIds);
- List<Long> dataSizes;
-
- Path combineDataPath = new Path(basePath + "/app1/0/combine/1.data");
- Path combineIndexPath = new Path(basePath + "/app1/0/combine/1.index");
- HdfsFileWriter combineWriter = new HdfsFileWriter(combineDataPath, conf);
- HdfsFileWriter combineIndexWriter = new HdfsFileWriter(combineIndexPath, conf);
-
- writeData(combineWriter);
- writeData(combineWriter);
-
- writePartition2(combineWriter, combineIndexWriter, expectBlockIds);
- combineIndexWriter.close();
- combineWriter.close();
-
- Set<Long> expects = Sets.newHashSet();
- expects.add(1L);
- List<byte[]> expectData = Lists.newArrayList();
- expectData.add(data);
- compareDataAndIndex("app1", 0, 1, basePath, expectData,
- expectBlockIds, processBlockIds, 2);
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-
- @Test
- public void handlerReadCombinedDataTest() {
- try {
- Roaring64NavigableMap expectBlockIds = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap processBlockIds = Roaring64NavigableMap.bitmapOf();
-
- String basePath = HDFS_URI + "test_backup_read";
- Path combinePath = new Path(basePath + "/app2/0/combine");
- fs.mkdirs(combinePath);
- Path partitionPath = new Path(basePath + "/app2/0/1");
- fs.mkdirs(partitionPath);
- Path dataPath = new Path(basePath + "/app2/0/1/1.data");
-
- HdfsFileWriter writer = new HdfsFileWriter(dataPath, conf);
- writeData(writer);
- writer.close();
- List<Long> dataSizes = Lists.newArrayList();
- writeIndexData(basePath, "/app2/0/1/1.index", dataSizes, expectBlockIds);
-
- Path combineDataPath = new Path(basePath + "/app2/0/combine/1.data");
- Path combineIndexPath = new Path(basePath + "/app2/0/combine/1.index");
- HdfsFileWriter combineWriter = new HdfsFileWriter(combineDataPath, conf);
- HdfsFileWriter combineIndexWriter = new HdfsFileWriter(combineIndexPath, conf);
-
- byte[] data1 = writeData(combineWriter);
-
- writeData(combineWriter);
- writePartition(dataSizes, combineIndexWriter, expectBlockIds);
- combineIndexWriter.close();
- combineWriter.close();
-
- Set<Long> expects = Sets.newHashSet();
- expects.add(2L);
- List<byte[]> expectData = Lists.newArrayList();
- expectData.add(data1);
- compareDataAndIndex("app2", 0, 2, basePath, expectData,
- expectBlockIds, processBlockIds, 2);
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-
-
- @Test
- public void handlerReadTwoKindsDataTest() {
- try {
- Roaring64NavigableMap expectBlockIds = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap processBlockIds = Roaring64NavigableMap.bitmapOf();
-
- String basePath = HDFS_URI + "test_backup_read";
- Path combinePath = new Path(basePath + "/app4/0/combine");
- fs.mkdirs(combinePath);
- Path partitionPath = new Path(basePath + "/app4/0/2");
- fs.mkdirs(partitionPath);
- Path dataPath = new Path(basePath + "/app4/0/2/2.data");
-
- HdfsFileWriter writer = new HdfsFileWriter(dataPath, conf);
- byte[] data = writeData(writer);
- writer.close();
- Path indexPath = new Path(basePath + "/app4/0/2/2.index");
- HdfsFileWriter iWriter = new HdfsFileWriter(indexPath, conf);
- List<Integer> somePartitions = Lists.newArrayList();
- somePartitions.add(2);
- List<Long> someSizes = Lists.newArrayList();
- someSizes.add((long)FileBasedShuffleSegment.SEGMENT_SIZE);
- List<Long> dataSizes = Lists.newArrayList();
- dataSizes.add(256L);
- iWriter.writeHeader(somePartitions, someSizes, dataSizes);
- FileBasedShuffleSegment segment = new FileBasedShuffleSegment(
- 1, 0, 256, 256, 1, 1);
- iWriter.writeIndex(segment);
- iWriter.close();
- expectBlockIds.addLong(1);
- Path combineDataPath = new Path(basePath + "/app4/0/combine/1.data");
- Path combineIndexPath = new Path(basePath + "/app4/0/combine/1.index");
- HdfsFileWriter combineWriter = new HdfsFileWriter(combineDataPath, conf);
- HdfsFileWriter combineIndexWriter = new HdfsFileWriter(combineIndexPath, conf);
-
-
- byte[] data1 = writeData(combineWriter);
- writeData(combineWriter);
- writePartition(dataSizes, combineIndexWriter, expectBlockIds);
- combineIndexWriter.close();
- combineWriter.close();
-
- Set<Long> expects = Sets.newHashSet();
- expects.add(2L);
- List<byte[]> expectData = Lists.newArrayList();
- expectData.add(data);
- expectData.add(data1);
- compareDataAndIndex("app4", 0, 2, basePath, expectData,
- expectBlockIds, processBlockIds, 2);
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-
- @Test
- public void readWithDifferentLimitTest() {
- try {
- List<byte[]> expectData = Lists.newArrayList();
- String basePath = HDFS_URI + "test_limit_read";
- Path partitionPath = new Path(basePath + "/app3/0/combine");
- fs.mkdirs(partitionPath);
- Path dataPath = new Path(basePath + "/app3/0/combine/1.data");
- HdfsFileWriter writer = new HdfsFileWriter(dataPath, conf);
-
- byte[] data = writeData(writer);
- ByteBuffer buffer;
- for (int i = 0; i < 5; i++) {
- new Random().nextBytes(data);
- buffer = ByteBuffer.allocate(data.length);
- buffer.put(data);
- buffer.flip();
- expectData.add(data.clone());
- writer.writeData(buffer);
- }
- buffer = ByteBuffer.allocate(data.length);
- new Random().nextBytes(data);
- buffer.put(data);
- buffer.flip();
- writer.writeData(buffer);
- writer.close();
-
- Roaring64NavigableMap expectBlockIds = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap processBlockIds = Roaring64NavigableMap.bitmapOf();
-
- Path indexPath = new Path(basePath + "/app3/0/combine/1.index");
- HdfsFileWriter iWriter = new HdfsFileWriter(indexPath, conf);
- writePartition3(iWriter, expectBlockIds);
- iWriter.close();
- compareDataAndIndex("app3", 0, 1, basePath, expectData,
- expectBlockIds, processBlockIds,1);
- compareDataAndIndex("app3", 0, 1, basePath, expectData,
- expectBlockIds, processBlockIds,2);
- compareDataAndIndex("app3", 0, 1, basePath, expectData,
- expectBlockIds, processBlockIds,3);
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-
- private void writeIndexData(
- String basePath,
- String file,
- List<Long> dataSizes,
- Roaring64NavigableMap expectBlockIds) throws IOException {
- Path indexPath = new Path(basePath + file);
- HdfsFileWriter iWriter = new HdfsFileWriter(indexPath, conf);
- List<Integer> somePartitions = Lists.newArrayList();
- somePartitions.add(1);
- List<Long> someSizes = Lists.newArrayList();
- someSizes.add((long)FileBasedShuffleSegment.SEGMENT_SIZE);
- dataSizes.clear();
- dataSizes.add(256L);
- iWriter.writeHeader(somePartitions, someSizes, dataSizes);
- FileBasedShuffleSegment segment = new FileBasedShuffleSegment(
- 1, 0, 256, 256, 1, 1);
- iWriter.writeIndex(segment);
- iWriter.close();
- expectBlockIds.addLong(1);
- }
-
- private void writePartition(
- List<Long> dataSizes,
- HdfsFileWriter combineIndexWriter,
- Roaring64NavigableMap expectBlockIds) throws IOException {
- List<Integer> partitions = Lists.newArrayList();
- partitions.add(2);
- partitions.add(3);
- List<Long> sizes = Lists.newArrayList();
- sizes.add((long) FileBasedShuffleSegment.SEGMENT_SIZE);
- sizes.add((long) FileBasedShuffleSegment.SEGMENT_SIZE);
- dataSizes.add(256L);
- dataSizes.add(256L);
- combineIndexWriter.writeHeader(partitions, sizes, dataSizes);
- FileBasedShuffleSegment segment1 = new FileBasedShuffleSegment(
- 2, 0, 256, 256, 2, 1);
- combineIndexWriter.writeIndex(segment1);
- expectBlockIds.addLong(2);
-
- FileBasedShuffleSegment segment2 = new FileBasedShuffleSegment(
- 3, 256, 256, 256, 3, 1);
- combineIndexWriter.writeIndex(segment2);
- expectBlockIds.addLong(3);
- }
-
- private void writePartition2(HdfsFileWriter combineWriter, HdfsFileWriter combineIndexWriter,
- Roaring64NavigableMap expectBlockIds) throws IOException {
- List<Long> dataSizes;
- List<Integer> partitions = Lists.newArrayList();
- partitions.add(2);
- partitions.add(3);
- List<Long> sizes = Lists.newArrayList();
- sizes.add((long)FileBasedShuffleSegment.SEGMENT_SIZE);
- sizes.add((long)FileBasedShuffleSegment.SEGMENT_SIZE);
- dataSizes = Lists.newArrayList();
- dataSizes.add(256L);
- dataSizes.add(256L);
- combineIndexWriter.writeHeader(partitions, sizes, dataSizes);
- FileBasedShuffleSegment segment1 = new FileBasedShuffleSegment(
- 2, 0, 256, 256, 1, 1);
- combineIndexWriter.writeIndex(segment1);
- FileBasedShuffleSegment segment2 = new FileBasedShuffleSegment(
- 3, 256, 256, 256, 1, 1);
- combineIndexWriter.writeIndex(segment2);
- expectBlockIds.addLong(2);
- expectBlockIds.addLong(3);
- }
-
- private void writePartition3(HdfsFileWriter iWriter, Roaring64NavigableMap expectBlockIds) throws IOException {
- List<Integer> somePartitions = Lists.newArrayList();
- somePartitions.add(0);
- somePartitions.add(1);
- somePartitions.add(2);
- List<Long> someSizes = Lists.newArrayList();
- someSizes.add((long)FileBasedShuffleSegment.SEGMENT_SIZE);
- someSizes.add((long)FileBasedShuffleSegment.SEGMENT_SIZE * 5);
- someSizes.add((long)FileBasedShuffleSegment.SEGMENT_SIZE);
- List<Long> dataSizes = Lists.newArrayList();
- dataSizes.add(256L);
- dataSizes.add(256L * 5);
- dataSizes.add(256L);
- iWriter.writeHeader(somePartitions, someSizes, dataSizes);
- FileBasedShuffleSegment segment = new FileBasedShuffleSegment(1, 0, 256, 256, 1, 1L);
- iWriter.writeIndex(segment);
- expectBlockIds.addLong(1);
-
- for (int i = 0; i < 5; i++) {
- segment = new FileBasedShuffleSegment(i + 2, i * 256, 256, 256, 1, 1L);
- iWriter.writeIndex(segment);
- expectBlockIds.addLong(i + 2);
- }
- segment = new FileBasedShuffleSegment(7, 0, 256, 256, 1, 1L);
- iWriter.writeIndex(segment);
- expectBlockIds.addLong(7);
- }
-
- private void compareDataAndIndex(
- String appId,
- int shuffleId,
- int partitionId,
- String basePath,
- List<byte[]> expectedData,
- Roaring64NavigableMap expectBlockIds,
- Roaring64NavigableMap processBlockIds,
- int limit) throws IllegalStateException {
- // read directly and compare
- UploadedHdfsClientReadHandler handler = new UploadedHdfsClientReadHandler(appId,
- shuffleId, partitionId, limit, 1, 3, 1024,
- expectBlockIds, processBlockIds, basePath, conf);
- try {
- List<ByteBuffer> actual = readData(handler);
- compareBytes(expectedData, actual);
- } finally {
- handler.close();
- }
- }
-
- private byte[] writeData(HdfsFileWriter combineWriter) throws IOException {
- byte[] data = new byte[256];
- new Random().nextBytes(data);
- ByteBuffer buffer = ByteBuffer.allocate(data.length);
- buffer.put(data);
- buffer.flip();
- combineWriter.writeData(buffer);
- return data;
- }
-
- private List<ByteBuffer> readData(UploadedHdfsClientReadHandler handler) throws IllegalStateException {
- ShuffleDataResult sdr;
- List<ByteBuffer> result = Lists.newArrayList();
- do {
- sdr = handler.readShuffleData();
- if (sdr == null || sdr.isEmpty()) {
- break;
- }
- List<BufferSegment> bufferSegments = sdr.getBufferSegments();
- for (BufferSegment bs : bufferSegments) {
- byte[] data = new byte[bs.getLength()];
- System.arraycopy(sdr.getData(), bs.getOffset(), data, 0, bs.getLength());
- result.add(ByteBuffer.wrap(data));
- }
- } while(sdr.getData() != null);
- return result;
- }
-}
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/MultiStorageHdfsShuffleReadHandlerTest.java b/storage/src/test/java/org/apache/uniffle/storage/handler/impl/MultiStorageHdfsShuffleReadHandlerTest.java
deleted file mode 100644
index 7302c5d..0000000
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/MultiStorageHdfsShuffleReadHandlerTest.java
+++ /dev/null
@@ -1,255 +0,0 @@
-/*
- * 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.storage.handler.impl;
-
-import static org.junit.jupiter.api.Assertions.assertArrayEquals;
-import static org.junit.jupiter.api.Assertions.assertNull;
-import static org.junit.jupiter.api.Assertions.fail;
-
-import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
-import com.google.common.primitives.Bytes;
-import java.nio.ByteBuffer;
-import java.util.List;
-import java.util.Map;
-import java.util.Random;
-import org.apache.hadoop.fs.Path;
-import org.junit.jupiter.api.Test;
-import org.roaringbitmap.longlong.Roaring64NavigableMap;
-
-import org.apache.uniffle.common.ShuffleDataResult;
-import org.apache.uniffle.common.ShufflePartitionedBlock;
-import org.apache.uniffle.common.util.ChecksumUtils;
-import org.apache.uniffle.storage.HdfsShuffleHandlerTestBase;
-import org.apache.uniffle.storage.common.FileBasedShuffleSegment;
-import org.apache.uniffle.storage.util.ShuffleStorageUtils;
-
-public class MultiStorageHdfsShuffleReadHandlerTest extends HdfsShuffleHandlerTestBase {
-
- @Test
- public void test() {
- try {
- String basePath = HDFS_URI + "test_base";
- int readBufferSize = 7;
- Path testPath = new Path(basePath + "/app1/0/test");
- fs.mkdirs(testPath);
- Path dataPath = new Path(basePath + "/app1/0/test/1-3.data");
- HdfsFileWriter writer = new HdfsFileWriter(dataPath, conf);
- // write data and record meta data
- Map<Integer, Integer> partitionSliceNum = Maps.newHashMap();
- Map<Long, byte[]> expectedBlockIdData = Maps.newHashMap();
- Map<Integer, List<ShufflePartitionedBlock>> expectedBlocks = Maps.newHashMap();
- Map<Integer, List<FileBasedShuffleSegment>> expectedIndexSegments = Maps.newHashMap();
- int partitionNum = 5;
- int dataFileLen = 0;
- int totalIndexSegmentNum = 0;
- List<ShuffleIndexHeader.Entry> indexHeaderEntries = Lists.newArrayList();
- ByteBuffer headerContentBuf = ByteBuffer.allocate(
- (int) ShuffleStorageUtils.getIndexFileHeaderLen(partitionNum) - ShuffleStorageUtils.getHeaderCrcLen());
- headerContentBuf.putInt(partitionNum);
- for (int partitionId = 1; partitionId <= 9; partitionId += 2) {
- int expectTotalBlockNum = new Random().nextInt(37) + 1;
- int blockSize = new Random().nextInt(7) + 1;
- writeTestData(writer, partitionId, expectTotalBlockNum, blockSize, 0,
- expectedBlockIdData, expectedBlocks, expectedIndexSegments, partitionId <= 5);
- int sliceNum = calcExpectedSegmentNum(expectTotalBlockNum, blockSize, readBufferSize);
- partitionSliceNum.put(partitionId, sliceNum);
- List<FileBasedShuffleSegment> segments = expectedIndexSegments.get(partitionId);
- totalIndexSegmentNum += segments.size();
- long indexLen = segments.size() * FileBasedShuffleSegment.SEGMENT_SIZE;
- long dataLen = expectTotalBlockNum * blockSize;
- headerContentBuf.putInt(partitionId);
- headerContentBuf.putLong(indexLen);
- headerContentBuf.putLong(dataLen);
- dataFileLen += dataLen;
- indexHeaderEntries.add(new ShuffleIndexHeader.Entry(partitionId, indexLen, dataLen));
- }
- headerContentBuf.flip();
- long crc = ChecksumUtils.getCrc32(headerContentBuf);
- ShuffleIndexHeader shuffleIndexHeader = new ShuffleIndexHeader(partitionNum, indexHeaderEntries, crc);
- writer.close();
-
- // write index
- Path indexPath = new Path(basePath + "/app1/0/test/1-3.index");
- writer = new HdfsFileWriter(indexPath, conf);
- int indexFileLen =
- shuffleIndexHeader.getHeaderLen() + totalIndexSegmentNum * FileBasedShuffleSegment.SEGMENT_SIZE;
- ByteBuffer byteBuffer = ByteBuffer.allocate(indexFileLen);
- byteBuffer.putInt(shuffleIndexHeader.getPartitionNum());
- for (ShuffleIndexHeader.Entry entry : shuffleIndexHeader.getIndexes()) {
- byteBuffer.putInt(entry.getPartitionId());
- byteBuffer.putLong(entry.getPartitionIndexLength());
- byteBuffer.putLong(entry.getPartitionDataLength());
- }
- byteBuffer.putLong(crc);
- for (int partitionId = 1; partitionId <= 9; partitionId += 2) {
- List<FileBasedShuffleSegment> fileBasedShuffleSegments = expectedIndexSegments.get(partitionId);
- for (FileBasedShuffleSegment segment : fileBasedShuffleSegments) {
- byteBuffer.putLong(segment.getOffset());
- byteBuffer.putInt(segment.getLength());
- byteBuffer.putInt(segment.getUncompressLength());
- if (partitionId <= 5) {
- byteBuffer.putLong(segment.getCrc());
- byteBuffer.putLong(segment.getBlockId());
- byteBuffer.putLong(segment.getTaskAttemptId());
- }
- }
- }
- byteBuffer.flip();
- writer.writeData(byteBuffer);
- writer.close();
-
- for (int partitionId = 1; partitionId <= 9; partitionId += 2) {
- String fileNamePrefix = basePath + "/app1/0/test/1-3";
- Roaring64NavigableMap expectBlockIds = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap processBlockIds = Roaring64NavigableMap.bitmapOf();
- HdfsShuffleReadHandler handler = new UploadedStorageHdfsShuffleReadHandler(
- "app1", 0, partitionId, fileNamePrefix, readBufferSize,
- expectBlockIds, processBlockIds, conf);
-
- int sliceNum = partitionSliceNum.get(partitionId);
- byte[] expectedDataBuf = new byte[0];
- for (ShufflePartitionedBlock spb : expectedBlocks.get(partitionId)) {
- Bytes.concat(expectedDataBuf, spb.getData());
- expectBlockIds.addLong(spb.getBlockId());
- }
-
- if (partitionId <= 5) {
- byte[] actualDataBuf = new byte[0];
- for (int i = 0; i < sliceNum; ++i) {
- ShuffleDataResult shuffleDataResult = handler.readShuffleData();
- checkData(shuffleDataResult, expectedBlockIdData);
- Bytes.concat(actualDataBuf, shuffleDataResult.getData());
- }
- assertNull(handler.readShuffleData());
- assertArrayEquals(expectedDataBuf, actualDataBuf);
- } else {
- assertNull(handler.readShuffleData());
- }
- }
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-
- @Test
- public void testHugeData() {
- try {
- String basePath = HDFS_URI + "testHugeData";
- int readBufferSize = 7;
- Path testPath = new Path(basePath + "/app1/0/test");
- fs.mkdirs(testPath);
- Path dataPath = new Path(basePath + "/app1/0/test/1-3.data");
- HdfsFileWriter writer = new HdfsFileWriter(dataPath, conf);
- // write data and record meta data
- Map<Integer, Integer> partitionSliceNum = Maps.newHashMap();
- Map<Long, byte[]> expectedBlockIdData = Maps.newHashMap();
- Map<Integer, List<ShufflePartitionedBlock>> expectedBlocks = Maps.newHashMap();
- Map<Integer, List<FileBasedShuffleSegment>> expectedIndexSegments = Maps.newHashMap();
- int partitionNum = 3;
- int dataFileLen = 0;
- int totalIndexSegmentNum = 0;
- List<ShuffleIndexHeader.Entry> indexHeaderEntries = Lists.newArrayList();
- ByteBuffer headerContentBuf = ByteBuffer.allocate(
- (int)ShuffleStorageUtils.getIndexFileHeaderLen(partitionNum) - ShuffleStorageUtils.getHeaderCrcLen());
- headerContentBuf.putInt(partitionNum);
- for (int partitionId = 1; partitionId <= 3; partitionId++) {
- int expectTotalBlockNum = 1024 * (new Random().nextInt(3) + 1);
- int blockSize = 1024 * (new Random().nextInt(7) + 11);
- writeTestData(writer, partitionId, expectTotalBlockNum, blockSize, 0,
- expectedBlockIdData, expectedBlocks, expectedIndexSegments, true);
- int sliceNum = calcExpectedSegmentNum(expectTotalBlockNum, blockSize, readBufferSize);
- partitionSliceNum.put(partitionId, sliceNum);
- List<FileBasedShuffleSegment> segments = expectedIndexSegments.get(partitionId);
- totalIndexSegmentNum += segments.size();
- long indexLen = segments.size() * FileBasedShuffleSegment.SEGMENT_SIZE;
- long dataLen = expectTotalBlockNum * blockSize;
- headerContentBuf.putInt(partitionId);
- headerContentBuf.putLong(indexLen);
- headerContentBuf.putLong(dataLen);
- dataFileLen += dataLen;
- indexHeaderEntries.add(new ShuffleIndexHeader.Entry(partitionId, indexLen, dataLen));
- }
- headerContentBuf.flip();
- long crc = ChecksumUtils.getCrc32(headerContentBuf);
- ShuffleIndexHeader shuffleIndexHeader = new ShuffleIndexHeader(partitionNum, indexHeaderEntries, crc);
- writer.close();
-
- // write index
- Path indexPath = new Path(basePath + "/app1/0/test/1-3.index");
- writer = new HdfsFileWriter(indexPath, conf);
- int indexFileLen =
- shuffleIndexHeader.getHeaderLen() + totalIndexSegmentNum * FileBasedShuffleSegment.SEGMENT_SIZE;
- ByteBuffer byteBuffer = ByteBuffer.allocate(indexFileLen);
- byteBuffer.putInt(shuffleIndexHeader.getPartitionNum());
- for (ShuffleIndexHeader.Entry entry : shuffleIndexHeader.getIndexes()) {
- byteBuffer.putInt(entry.getPartitionId());
- byteBuffer.putLong(entry.getPartitionIndexLength());
- byteBuffer.putLong(entry.getPartitionDataLength());
- }
- byteBuffer.putLong(crc);
- for (int partitionId = 1; partitionId <= 3; partitionId++) {
- List<FileBasedShuffleSegment> fileBasedShuffleSegments = expectedIndexSegments.get(partitionId);
- for (FileBasedShuffleSegment segment : fileBasedShuffleSegments) {
- byteBuffer.putLong(segment.getOffset());
- byteBuffer.putInt(segment.getLength());
- byteBuffer.putInt(segment.getUncompressLength());
- byteBuffer.putLong(segment.getCrc());
- byteBuffer.putLong(segment.getBlockId());
- byteBuffer.putLong(segment.getTaskAttemptId());
- }
- }
- byteBuffer.flip();
- writer.writeData(byteBuffer);
- writer.close();
-
- for (int partitionId = 1; partitionId <= 3; partitionId++) {
- String fileNamePrefix = basePath + "/app1/0/test/1-3";
- Roaring64NavigableMap expectBlockIds = Roaring64NavigableMap.bitmapOf();
- Roaring64NavigableMap processBlockIds = Roaring64NavigableMap.bitmapOf();
- HdfsShuffleReadHandler handler = new UploadedStorageHdfsShuffleReadHandler(
- "app1", 0, partitionId, fileNamePrefix, readBufferSize,
- expectBlockIds, processBlockIds, conf);
-
- int sliceNum = partitionSliceNum.get(partitionId);
- byte[] expectedDataBuf = new byte[0];
- for (ShufflePartitionedBlock spb : expectedBlocks.get(partitionId)) {
- Bytes.concat(expectedDataBuf, spb.getData());
- expectBlockIds.addLong(spb.getBlockId());
- }
-
- if (partitionId <= 3) {
- byte[] actualDataBuf = new byte[0];
- for (int i = 0; i < sliceNum; ++i) {
- ShuffleDataResult shuffleDataResult = handler.readShuffleData();
- checkData(shuffleDataResult, expectedBlockIdData);
- Bytes.concat(actualDataBuf, shuffleDataResult.getData());
- }
- assertNull(handler.readShuffleData());
- assertArrayEquals(expectedDataBuf, actualDataBuf);
- } else {
- assertNull(handler.readShuffleData());
- }
- }
- } catch (Exception e) {
- e.printStackTrace();
- fail();
- }
- }
-}