Skip to content

[Refactor] Optimize creating shuffle handlers #259

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
roryqi merged 3 commits into from
Oct 12, 2022
Merged

[Refactor] Optimize creating shuffle handlers #259

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7eb1b5e
[Refactor] Optimize creating shuffle handlers
zuston Oct 10, 2022
af9f553
amend
zuston Oct 11, 2022
271fc33
amend-2
zuston Oct 11, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions client-mr/src/main/java/org/apache/hadoop/mapred/RssMapOutputCollector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,8 +196,6 @@ public void flush() throws IOException, InterruptedException, ClassNotFoundExcep
}

private boolean isMemoryShuffleEnabled(String storageType) {
return StorageType.MEMORY_LOCALFILE.name().equals(storageType)
|| StorageType.MEMORY_HDFS.name().equals(storageType)
|| StorageType.MEMORY_LOCALFILE_HDFS.name().equals(storageType);
return StorageType.withMemory(StorageType.valueOf(storageType));
}
}
4 changes: 1 addition & 3 deletions client-spark/spark2/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,9 +119,7 @@ public RssShuffleWriter(
}

private boolean isMemoryShuffleEnabled(String storageType) {
return StorageType.MEMORY_LOCALFILE.name().equals(storageType)
|| StorageType.MEMORY_HDFS.name().equals(storageType)
|| StorageType.MEMORY_LOCALFILE_HDFS.name().equals(storageType);
return StorageType.withMemory(StorageType.valueOf(storageType));
}

/**
Expand Down
4 changes: 1 addition & 3 deletions client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,9 +122,7 @@ public RssShuffleWriter(
}

private boolean isMemoryShuffleEnabled(String storageType) {
return StorageType.MEMORY_LOCALFILE.name().equals(storageType)
|| StorageType.MEMORY_HDFS.name().equals(storageType)
|| StorageType.MEMORY_LOCALFILE_HDFS.name().equals(storageType);
return StorageType.withMemory(StorageType.valueOf(storageType));
}

@Override
Expand Down
218 changes: 84 additions & 134 deletions storage/src/main/java/org/apache/uniffle/storage/factory/ShuffleHandlerFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,16 @@

package org.apache.uniffle.storage.factory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.stream.Collectors;

import org.apache.uniffle.client.api.ShuffleServerClient;
import org.apache.uniffle.client.factory.ShuffleServerClientFactory;
import org.apache.uniffle.client.util.ClientType;
import org.apache.uniffle.common.ShuffleServerInfo;
import org.apache.uniffle.common.exception.RssException;
import org.apache.uniffle.storage.handler.api.ClientReadHandler;
import org.apache.uniffle.storage.handler.api.ShuffleDeleteHandler;
import org.apache.uniffle.storage.handler.impl.ComposedClientReadHandler;
Expand Down Expand Up @@ -51,142 +54,89 @@ public static synchronized ShuffleHandlerFactory getInstance() {
}

public ClientReadHandler createShuffleReadHandler(CreateShuffleReadHandlerRequest request) {
if (StorageType.HDFS.name().equals(request.getStorageType())) {
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());
} else if (StorageType.LOCALFILE.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 LocalFileQuorumClientReadHandler(request.getAppId(), request.getShuffleId(), request.getPartitionId(),
request.getIndexReadLimit(), request.getPartitionNumPerRange(), request.getPartitionNum(),
request.getReadBufferSize(), request.getExpectBlockIds(), request.getProcessBlockIds(),
shuffleServerClients);
} else if (StorageType.LOCALFILE_HDFS.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());
});
} else if (StorageType.MEMORY_LOCALFILE.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());
ClientReadHandler memoryClientReadHandler = new MemoryQuorumClientReadHandler(
request.getAppId(),
request.getShuffleId(),
request.getPartitionId(),
request.getReadBufferSize(),
shuffleServerClients);
ClientReadHandler localClientReadHandler = new LocalFileQuorumClientReadHandler(request.getAppId(),
request.getShuffleId(), request.getPartitionId(), request.getIndexReadLimit(),
request.getPartitionNumPerRange(), request.getPartitionNum(),
request.getReadBufferSize(), request.getExpectBlockIds(), request.getProcessBlockIds(),
shuffleServerClients);
return new ComposedClientReadHandler(memoryClientReadHandler, localClientReadHandler);
} else if (StorageType.MEMORY_HDFS.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 MemoryQuorumClientReadHandler(
request.getAppId(),
request.getShuffleId(),
request.getPartitionId(),
request.getReadBufferSize(),
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());
});
} else if (StorageType.MEMORY_LOCALFILE_HDFS.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 MemoryQuorumClientReadHandler(
request.getAppId(),
request.getShuffleId(),
request.getPartitionId(),
request.getReadBufferSize(),
shuffleServerClients);
}, () -> {
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());
});
} else {
String storageType = request.getStorageType();
StorageType type = StorageType.valueOf(storageType);

if (StorageType.MEMORY == type) {
throw new UnsupportedOperationException(
"Doesn't support storage type for client read handler:" + request.getStorageType());
"Doesn't support storage type for client read handler:" + storageType);
}

if (StorageType.HDFS == type) {
return getHdfsClientReadHandler(request);
}
if (StorageType.LOCALFILE == type) {
return getLocalfileClientReaderHandler(request);
}

List<ClientReadHandler> handlers = new ArrayList<>();
if (StorageType.withMemory(type)) {
handlers.add(
getMemoryClientReadHandler(request)
);
}
if (StorageType.withLocalfile(type)) {
handlers.add(
getLocalfileClientReaderHandler(request)
);
}
if (StorageType.withHDFS(type)) {
handlers.add(
getHdfsClientReadHandler(request)
);
}
if (handlers.isEmpty()) {
throw new RssException("This should not happen due to the unknown storage type: " + storageType);
}

Callable<ClientReadHandler>[] callables =
handlers
.stream()
.map(x -> (Callable<ClientReadHandler>) () -> x)
.collect(Collectors.toList())
.toArray(new Callable[handlers.size()]);
return new ComposedClientReadHandler(callables);
}

private ClientReadHandler getMemoryClientReadHandler(CreateShuffleReadHandlerRequest request) {
List<ShuffleServerInfo> shuffleServerInfoList = request.getShuffleServerInfoList();
List<ShuffleServerClient> shuffleServerClients = shuffleServerInfoList.stream().map(
ssi -> ShuffleServerClientFactory.getInstance().getShuffleServerClient(
ClientType.GRPC.name(), ssi)).collect(
Collectors.toList());
ClientReadHandler memoryClientReadHandler = new MemoryQuorumClientReadHandler(
request.getAppId(),
request.getShuffleId(),
request.getPartitionId(),
request.getReadBufferSize(),
shuffleServerClients);
return memoryClientReadHandler;
}

private ClientReadHandler getLocalfileClientReaderHandler(CreateShuffleReadHandlerRequest request) {
List<ShuffleServerInfo> shuffleServerInfoList = request.getShuffleServerInfoList();
List<ShuffleServerClient> shuffleServerClients = shuffleServerInfoList.stream().map(
ssi -> ShuffleServerClientFactory.getInstance().getShuffleServerClient(ClientType.GRPC.name(), ssi)).collect(
Collectors.toList());
return new LocalFileQuorumClientReadHandler(request.getAppId(), request.getShuffleId(), request.getPartitionId(),
request.getIndexReadLimit(), request.getPartitionNumPerRange(), request.getPartitionNum(),
request.getReadBufferSize(), request.getExpectBlockIds(), request.getProcessBlockIds(),
shuffleServerClients);
}

private ClientReadHandler getHdfsClientReadHandler(CreateShuffleReadHandlerRequest request) {
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());
}

public ShuffleDeleteHandler createShuffleDeleteHandler(CreateShuffleDeleteHandlerRequest request) {
Expand Down
35 changes: 29 additions & 6 deletions storage/src/main/java/org/apache/uniffle/storage/util/StorageType.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,33 @@
package org.apache.uniffle.storage.util;

public enum StorageType {
HDFS,
LOCALFILE,
LOCALFILE_HDFS,
MEMORY_LOCALFILE,
MEMORY_HDFS,
MEMORY_LOCALFILE_HDFS
MEMORY(1),
LOCALFILE(2),
HDFS(4),
LOCALFILE_HDFS(6),
MEMORY_LOCALFILE(3),
MEMORY_HDFS(5),
MEMORY_LOCALFILE_HDFS(7);
Copy link
Contributor

@frankliee frankliee Oct 11, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why this order？

Copy link
Member Author

@zuston zuston Oct 11, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just follow the original order


private int val;

StorageType(int val) {
this.val = val;
}

private int getVal() {
return val;
}

public static boolean withMemory(StorageType storageType) {
return (storageType.getVal() & MEMORY.getVal()) != 0;
}

public static boolean withLocalfile(StorageType storageType) {
return (storageType.getVal() & LOCALFILE.getVal()) != 0;
}

public static boolean withHDFS(StorageType storageType) {
return (storageType.getVal() & HDFS.getVal()) != 0;
}
}
68 changes: 68 additions & 0 deletions storage/src/test/java/org/apache/uniffle/storage/util/StorageTypeTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
/*
* 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 org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class StorageTypeTest {

@Test
public void commonTest() {
String type = "HDFS";
assertEquals(StorageType.valueOf(type), StorageType.HDFS);

StorageType storageType = StorageType.MEMORY;
assertTrue(StorageType.withMemory(storageType));
assertFalse(StorageType.withLocalfile(storageType));
assertFalse(StorageType.withHDFS(storageType));

storageType = StorageType.LOCALFILE;
assertFalse(StorageType.withMemory(storageType));
assertTrue(StorageType.withLocalfile(storageType));
assertFalse(StorageType.withHDFS(storageType));

storageType = StorageType.HDFS;
assertFalse(StorageType.withMemory(storageType));
assertFalse(StorageType.withLocalfile(storageType));
assertTrue(StorageType.withHDFS(storageType));

storageType = StorageType.MEMORY_HDFS;
assertTrue(StorageType.withMemory(storageType));
assertFalse(StorageType.withLocalfile(storageType));
assertTrue(StorageType.withHDFS(storageType));

storageType = StorageType.MEMORY_LOCALFILE;
assertTrue(StorageType.withMemory(storageType));
assertTrue(StorageType.withLocalfile(storageType));
assertFalse(StorageType.withHDFS(storageType));

storageType = StorageType.MEMORY_LOCALFILE_HDFS;
assertTrue(StorageType.withMemory(storageType));
assertTrue(StorageType.withLocalfile(storageType));
assertTrue(StorageType.withHDFS(storageType));

storageType = StorageType.LOCALFILE_HDFS;
assertFalse(StorageType.withMemory(storageType));
assertTrue(StorageType.withLocalfile(storageType));
assertTrue(StorageType.withHDFS(storageType));
}
}