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 2026/05/12 02:05:10 UTC
(uniffle) branch master updated: [#1805] feat(spark): Support Spark 4.0.2 (#2748)
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/uniffle.git
The following commit(s) were added to refs/heads/master by this push:
new c18774dd7 [#1805] feat(spark): Support Spark 4.0.2 (#2748)
c18774dd7 is described below
commit c18774dd7e60312ed5066487c390f3dd6c149396
Author: YangJie <ya...@baidu.com>
AuthorDate: Tue May 12 10:05:04 2026 +0800
[#1805] feat(spark): Support Spark 4.0.2 (#2748)
* [#1805] feat(spark): Support Spark 4.0.2
Add Spark 4.0.2 support with the following changes:
- New client-spark/spark4 and client-spark/spark4-shaded modules
- New integration-test/spark4 module
- spark4 Maven profile with Scala 2.13, Java 17, Hadoop 3.4.1
- CI workflow entries for unit and integration tests
Key compatibility fixes for Hadoop 3.4.1:
- Replace removed NodeHealthScriptRunner with inline logic
- Add ImpersonationProvider.authorize(UGI, InetAddress) for Hadoop 3.4
- Handle wrapped BindException in JettyServer for Jetty 9.4.53
- Add commons-logging explicit test dependency
- Skip Kerberos tests on JDK 17+ via @DisabledIfSystemProperty
Closes #1805
* [#1805] refactor(spark): Drop SparkVersionUtils inheritance in Spark4VersionUtils
Utility classes with only static members should not extend another static-only utility; static methods are not polymorphic, so the inheritance is misleading. Make Spark4VersionUtils standalone, qualify MAJOR_VERSION/MINOR_VERSION via SparkVersionUtils, and add a private constructor to prevent instantiation. Drop the duplicated isSpark2/3/4 assertions from testSpark4Version since testSparkVersion already covers them through SparkVersionUtils directly.
---
.github/workflows/parallel.yml | 59 +-
.github/workflows/sequential.yml | 4 +
.../apache/spark/shuffle/RssSparkShuffleUtils.java | 3 +-
.../apache/spark/shuffle/SparkVersionUtils.java | 4 +
client-spark/spark4-shaded/pom.xml | 342 ++++++
client-spark/spark4/pom.xml | 111 ++
.../spark/shuffle/DelegationRssShuffleManager.java | 251 +++++
.../org/apache/spark/shuffle/FunctionUtils.java | 41 +-
.../org/apache/spark/shuffle/RssShuffleDataIo.java | 35 +-
.../spark/shuffle/RssShuffleDriverComponents.java | 36 +-
.../apache/spark/shuffle/RssShuffleManager.java | 646 +++++++++++
.../apache/spark/shuffle/Spark4VersionUtils.java} | 40 +-
.../spark/shuffle/reader/RssShuffleReader.java | 508 +++++++++
.../shuffle/writer/PartitionLengthStatistic.java | 52 +
.../spark/shuffle/writer/RssShuffleWriter.java | 911 ++++++++++++++++
.../shuffle/DelegationRssShuffleManagerTest.java | 237 ++++
.../apache/spark/shuffle/FunctionUtilsTests.java | 31 +-
.../spark/shuffle/RssShuffleManagerTest.java | 300 ++++++
.../spark/shuffle/RssShuffleManagerTestBase.java | 70 ++
.../spark/shuffle/RssSpark4ShuffleUtilsTest.java | 62 ++
.../spark/shuffle/SparkVersionUtilsTest.java | 30 +-
.../java/org/apache/spark/shuffle/TestUtils.java | 51 +
.../spark/shuffle/reader/RssShuffleReaderTest.java | 165 +++
.../writer/PartitionLengthStatisticTest.java | 67 ++
.../spark/shuffle/writer/RssShuffleWriterTest.java | 1134 ++++++++++++++++++++
.../org/apache/uniffle/common/web/JettyServer.java | 18 +-
.../apache/uniffle/common/KerberizedHadoop.java | 11 +
.../filesystem/HadoopFilesystemProviderTest.java | 5 +
.../common/security/HadoopSecurityContextTest.java | 5 +
.../security/SecurityContextFactoryTest.java | 5 +
integration-test/common/pom.xml | 2 +-
...ccessCandidatesCheckerKerberizedHadoopTest.java | 5 +
...amicClientConfServiceKerberlizedHadoopTest.java | 5 +
.../ShuffleServerWithKerberizedHadoopTest.java | 5 +
integration-test/spark-common/pom.xml | 2 +-
integration-test/{common => spark4}/pom.xml | 145 +--
.../apache/uniffle/test/AQERepartitionTest.java | 104 ++
.../apache/uniffle/test/MapSideCombineTest.java | 148 +++
.../spark4/src/test/resources/log4j2.xml | 33 +
pom.xml | 140 ++-
.../apache/uniffle/server/HealthScriptChecker.java | 19 +-
.../server/KerberizedShuffleTaskManagerTest.java | 5 +
.../ShuffleFlushManagerOnKerberizedHadoopTest.java | 5 +
.../KerberizedHadoopClientReadHandlerTest.java | 5 +
.../KerberizedHadoopShuffleReadHandlerTest.java | 5 +
.../ShuffleKerberizedHadoopStorageUtilsTest.java | 5 +
46 files changed, 5623 insertions(+), 244 deletions(-)
diff --git a/.github/workflows/parallel.yml b/.github/workflows/parallel.yml
index bc671c789..c36b69119 100644
--- a/.github/workflows/parallel.yml
+++ b/.github/workflows/parallel.yml
@@ -53,23 +53,40 @@ jobs:
runs-on: ubuntu-latest
strategy:
matrix:
- profile:
+ include:
# Pay attention! Don't use profile with character ':'!
- - spark2
- - spark2.3
- - spark3.0
- - spark3
- - spark3.2
- - spark3.2.0
- - spark3.3
- - spark3.4
- - spark3.5
- - spark3.5-scala2.13
- - mr-hadoop2.8
- - mr-hadoop3.2
- - tez
- - tez-hadoop3.2
- - dashboard
+ - profile: spark2
+ java-version: 8
+ - profile: spark2.3
+ java-version: 8
+ - profile: spark3.0
+ java-version: 8
+ - profile: spark3
+ java-version: 8
+ - profile: spark3.2
+ java-version: 8
+ - profile: spark3.2.0
+ java-version: 8
+ - profile: spark3.3
+ java-version: 8
+ - profile: spark3.4
+ java-version: 8
+ - profile: spark3.5
+ java-version: 8
+ - profile: spark3.5-scala2.13
+ java-version: 8
+ - profile: spark4
+ java-version: 17
+ - profile: mr-hadoop2.8
+ java-version: 8
+ - profile: mr-hadoop3.2
+ java-version: 8
+ - profile: tez
+ java-version: 8
+ - profile: tez-hadoop3.2
+ java-version: 8
+ - profile: dashboard
+ java-version: 8
fail-fast: false
name: -P${{ matrix.profile }}
steps:
@@ -77,19 +94,19 @@ jobs:
run: sudo hostname "action-host" | sudo echo "127.0.0.1 action-host" | sudo tee -a /etc/hosts
- name: Checkout project
uses: actions/checkout@v3
- - name: Set up JDK ${{ inputs.java-version }}
+ - name: Set up JDK ${{ matrix.java-version }}
uses: actions/setup-java@v3
with:
- java-version: ${{ inputs.java-version }}
+ java-version: ${{ matrix.java-version }}
distribution: ${{ inputs.jdk-distro }}
- name: Cache local Maven repository
uses: actions/cache@v3
with:
path: ~/.m2/repository
- key: mvn-${{ inputs.java-version }}-package-${{ matrix.profile }}-${{ hashFiles('**/pom.xml') }}
+ key: mvn-${{ matrix.java-version }}-package-${{ matrix.profile }}-${{ hashFiles('**/pom.xml') }}
restore-keys: |
- mvn-${{ inputs.java-version }}-package-${{ matrix.profile }}-
- mvn-${{ inputs.java-version }}-package-
+ mvn-${{ matrix.java-version }}-package-${{ matrix.profile }}-
+ mvn-${{ matrix.java-version }}-package-
- name: Execute `./mvnw ${{ inputs.maven-args }} -P${{ matrix.profile }}`
run: |
PROFILES="${{ matrix.profile }}"
diff --git a/.github/workflows/sequential.yml b/.github/workflows/sequential.yml
index c483ee82f..f82abe439 100644
--- a/.github/workflows/sequential.yml
+++ b/.github/workflows/sequential.yml
@@ -68,6 +68,10 @@ jobs:
key: mvn-${{ inputs.java-version }}-${{ inputs.cache-key }}-${{ hashFiles('**/pom.xml') }}
restore-keys: |
mvn-${{ inputs.java-version }}-${{ inputs.cache-key }}-
+ - name: Execute `./mvnw ${{ inputs.maven-args }} -Pspark4`
+ if: inputs.java-version == '17'
+ run: ./mvnw -B -fae ${{ inputs.maven-args }} -Pspark4 | tee -a /tmp/maven.log;
+ shell: bash
- name: Execute `./mvnw ${{ inputs.maven-args }} -Pspark3`
run: ./mvnw -B -fae ${{ inputs.maven-args }} -Pspark3 | tee -a /tmp/maven.log;
shell: bash
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java
index 56ec29caa..3d986203f 100644
--- a/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java
@@ -363,7 +363,8 @@ public class RssSparkShuffleUtils {
// FetchFailedException.
// Therefore, the stage re-computation feature is only enabled for Spark versions larger than or
// equal to 2.3.
- return SparkVersionUtils.isSpark3()
+ return SparkVersionUtils.isSpark4()
+ || SparkVersionUtils.isSpark3()
|| (SparkVersionUtils.isSpark2() && SparkVersionUtils.MINOR_VERSION >= 3);
}
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java
index a55f0aa7d..f987d7034 100644
--- a/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java
@@ -48,6 +48,10 @@ public class SparkVersionUtils {
return MAJOR_VERSION == 3;
}
+ public static boolean isSpark4() {
+ return MAJOR_VERSION == 4;
+ }
+
public static boolean isSpark320() {
return SPARK_VERSION.matches("^3.2.0([^\\d].*)?$");
}
diff --git a/client-spark/spark4-shaded/pom.xml b/client-spark/spark4-shaded/pom.xml
new file mode 100644
index 000000000..14cd5c206
--- /dev/null
+++ b/client-spark/spark4-shaded/pom.xml
@@ -0,0 +1,342 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ ~ 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.
+ -->
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+ xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+ <parent>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>uniffle-parent</artifactId>
+ <version>0.11.0-SNAPSHOT</version>
+ <relativePath>../../pom.xml</relativePath>
+ </parent>
+
+ <artifactId>rss-client-spark4-shaded</artifactId>
+ <packaging>jar</packaging>
+ <name>Apache Uniffle Shaded Client (Spark 4)</name>
+
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-client-spark4</artifactId>
+ <version>${project.version}</version>
+ <!-- use the lz4 from spark env -->
+ <exclusions>
+ <exclusion>
+ <groupId>org.lz4</groupId>
+ <artifactId>lz4-java</artifactId>
+ </exclusion>
+ <exclusion>
+ <artifactId>log4j-slf4j-impl</artifactId>
+ <groupId>org.apache.logging.log4j</groupId>
+ </exclusion>
+ <exclusion>
+ <artifactId>slf4j-api</artifactId>
+ <groupId>org.slf4j</groupId>
+ </exclusion>
+ <exclusion>
+ <groupId>com.esotericsoftware</groupId>
+ <artifactId>kryo-shaded</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+ <!-- log4j2 -->
+ <dependency>
+ <groupId>org.slf4j</groupId>
+ <artifactId>slf4j-api</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Use log4j-slf4j2-impl for Spark 4 (SLF4J 2.x) -->
+ <dependency>
+ <groupId>org.apache.logging.log4j</groupId>
+ <artifactId>log4j-slf4j2-impl</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.logging.log4j</groupId>
+ <artifactId>log4j-api</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.logging.log4j</groupId>
+ <artifactId>log4j-core</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <!-- end -->
+ </dependencies>
+
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-shade-plugin</artifactId>
+ <executions>
+ <execution>
+ <id>rss-client-spark4-shaded-jar</id>
+ <phase>package</phase>
+ <goals>
+ <goal>shade</goal>
+ </goals>
+ <configuration>
+ <artifactSet>
+ <excludes>
+ <exclude>org.apache.logging.log4j:log4j-slf4j-impl</exclude>
+ <exclude>org.apache.logging.log4j:log4j-slf4j2-impl</exclude>
+ </excludes>
+ </artifactSet>
+ <finalName>${project.artifactId}-${project.version}</finalName>
+ <transformers>
+ <transformer
+ implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
+ <transformer
+ implementation="org.apache.maven.plugins.shade.resource.ApacheLicenseResourceTransformer"/>
+ <transformer
+ implementation="org.apache.maven.plugins.shade.resource.ApacheNoticeResourceTransformer"/>
+ </transformers>
+ <filters>
+ <filter>
+ <artifact>*:*</artifact>
+ <excludes>
+ <exclude>module-info.class</exclude>
+ <exclude>META-INF/*.SF</exclude>
+ <exclude>META-INF/*.DSA</exclude>
+ <exclude>META-INF/*.RSA</exclude>
+ <exclude>**/*.proto</exclude>
+ <exclude>**/*.css</exclude>
+ <exclude>**/*.html</exclude>
+ </excludes>
+ </filter>
+ </filters>
+ <relocations>
+ <relocation>
+ <pattern>com.google</pattern>
+ <shadedPattern>${rss.shade.packageName}.com.google</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>google.protobuf</pattern>
+ <shadedPattern>${rss.shade.packageName}.google.protobuf</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>com.fasterxml.jackson.core</pattern>
+ <shadedPattern>${rss.shade.packageName}.jackson.core</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>com.fasterxml.jackson.databind</pattern>
+ <shadedPattern>${rss.shade.packageName}.jackson.databind</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>com.fasterxml.jackson.annotation</pattern>
+ <shadedPattern>${rss.shade.packageName}.jackson.annotation</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>io.netty</pattern>
+ <shadedPattern>${rss.shade.packageName}.io.netty</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>io.grpc</pattern>
+ <shadedPattern>${rss.shade.packageName}.io.grpc</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>io.perfmark</pattern>
+ <shadedPattern>${rss.shade.packageName}.io.perfmark</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>org.roaringbitmap</pattern>
+ <shadedPattern>${rss.shade.packageName}.org.roaringbitmap</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>org.apache</pattern>
+ <shadedPattern>${rss.shade.packageName}.org.apache</shadedPattern>
+ <excludes>
+ <!-- Exclude the provided dependencies -->
+ <!-- Exclude the packages belonging to uniffle -->
+ <exclude>org/apache/uniffle/**/*</exclude>
+ <!-- Exclude the packages belonging to hadoop -->
+ <exclude>org/apache/hadoop/*</exclude>
+ <exclude>org/apache/hadoop/**/*</exclude>
+ <!-- Exclude the logging packages-->
+ <exclude>**/pom.xml</exclude>
+ <exclude>org/apache/commons/logging/*</exclude>
+ <exclude>org/apache/commons/logging/**/*</exclude>
+ <exclude>org/apache/log4j/*</exclude>
+ <exclude>org/apache/log4j/**/*</exclude>
+ <exclude>org/apache/logging/**/*</exclude>
+ <!-- Exclude spark packages -->
+ <exclude>org/apache/spark/**/*</exclude>
+ <!-- Exclude commons-io packages -->
+ <exclude>org/apache/commons/io/**/*</exclude>
+ <!-- Exclude the packages belonging to collections 3 -->
+ <exclude>org/apache/commons/collections/**/*</exclude>
+ </excludes>
+ </relocation>
+ <relocation>
+ <pattern>org.codehaus</pattern>
+ <shadedPattern>${rss.shade.packageName}.codehaus</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>android.annotation</pattern>
+ <shadedPattern>${rss.shade.packageName}.android.annotation</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>javassist</pattern>
+ <shadedPattern>${rss.shade.packageName}.javassist</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>io.prometheus.client</pattern>
+ <shadedPattern>${rss.shade.packageName}.io.prometheus.client</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>com.sun.activation</pattern>
+ <shadedPattern>${rss.shade.packageName}.com.sun.activation</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>org.checkerframework</pattern>
+ <shadedPattern>${rss.shade.packageName}.org.checkerframework</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>org.eclipse</pattern>
+ <shadedPattern>${rss.shade.packageName}.org.eclipse</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>picocli</pattern>
+ <shadedPattern>${rss.shade.packageName}.picocli</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>META-INF/license</pattern>
+ <shadedPattern>META-INF/license-tmp</shadedPattern>
+ </relocation>
+ </relocations>
+ </configuration>
+ </execution>
+ </executions>
+ </plugin>
+ <!-- shade the native netty libs as well -->
+ <plugin>
+ <groupId>org.codehaus.mojo</groupId>
+ <artifactId>build-helper-maven-plugin</artifactId>
+ <executions>
+ <execution>
+ <id>regex-property</id>
+ <goals>
+ <goal>regex-property</goal>
+ </goals>
+ <configuration>
+ <name>rss.shade.native.packageName</name>
+ <value>${rss.shade.packageName}</value>
+ <regex>\.</regex>
+ <replacement>_</replacement>
+ <failIfNoMatch>true</failIfNoMatch>
+ </configuration>
+ </execution>
+ </executions>
+ </plugin>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-antrun-plugin</artifactId>
+ <executions>
+ <execution>
+ <id>unpack</id>
+ <phase>package</phase>
+ <goals>
+ <goal>run</goal>
+ </goals>
+ <configuration>
+ <target>
+ <echo message="Shade netty native libraries to ${rss.shade.native.packageName}"/>
+ <unzip src="${project.build.directory}/${project.artifactId}-${project.version}.jar"
+ dest="${project.build.directory}/unpacked/"/>
+ <echo message="renaming native epoll library"></echo>
+ <move includeemptydirs="false"
+ todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_transport_native_epoll_x86_64.so"
+ to="lib${rss.shade.native.packageName}_netty_transport_native_epoll_x86_64.so"
+ type="glob"></mapper>
+ </move>
+ <move includeemptydirs="false"
+ todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_transport_native_epoll_aarch_64.so"
+ to="lib${rss.shade.native.packageName}_netty_transport_native_epoll_aarch_64.so"
+ type="glob"></mapper>
+ </move>
+ <move includeemptydirs="false"
+ todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_transport_native_epoll_riscv64.so"
+ to="lib${rss.shade.native.packageName}_netty_transport_native_epoll_riscv64.so"
+ type="glob"></mapper>
+ </move>
+ <move includeemptydirs="false"
+ todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_transport_native_kqueue_x86_64.jnilib"
+ to="lib${rss.shade.native.packageName}_netty_transport_native_kqueue_x86_64.jnilib"
+ type="glob"></mapper>
+ </move>
+ <move includeemptydirs="false"
+ todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_resolver_dns_native_macos_aarch_64.jnilib"
+ to="lib${rss.shade.native.packageName}_netty_resolver_dns_native_macos_aarch_64.jnilib"
+ type="glob"></mapper>
+ </move>
+ <move includeemptydirs="false"
+ todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_resolver_dns_native_macos_x86_64.jnilib"
+ to="lib${rss.shade.native.packageName}_netty_resolver_dns_native_macos_x86_64.jnilib"
+ type="glob"></mapper>
+ </move>
+ <move includeemptydirs="false"
+ todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_transport_native_kqueue_aarch_64.jnilib"
+ to="lib${rss.shade.native.packageName}_netty_transport_native_kqueue_aarch_64.jnilib"
+ type="glob"></mapper>
+ </move>
+ <!-- Since it may be compiled on case-insensitive operating system, it may be executed multiple times.
+ Here the license directory, licenses directory, and LICENSE file are processed with aliases. -->
+ <move failonerror="false" quiet="true"
+ todir="${project.build.directory}/unpacked/META-INF/license-tmps">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/license-tmp">
+ <include name="LICENSE*.txt"/>
+ </fileset>
+ </move>
+ <delete failonerror="false" quiet="true"
+ dir="${project.build.directory}/unpacked/META-INF/license-tmp"/>
+ <move failonerror="false" quiet="true"
+ file="${project.build.directory}/unpacked/META-INF/license-tmps"
+ tofile="${project.build.directory}/unpacked/META-INF/licenses"/>
+ <move failonerror="false" quiet="true"
+ file="${project.build.directory}/unpacked/META-INF/LICENSE-binary"
+ tofile="${project.build.directory}/unpacked/META-INF/LICENSE"/>
+ <echo message="repackaging netty jar"></echo>
+ <jar destfile="${project.build.directory}/${project.artifactId}-${project.version}.jar"
+ basedir="${project.build.directory}/unpacked"/>
+ <exec executable="bash" failonerror="true" dir="${project.build.directory}/">
+ <arg line="${project.basedir}/../../dev/scripts/checkshade.sh ${project.build.directory}/${project.artifactId}-${project.version}.jar"/>
+ </exec>
+ </target>
+ </configuration>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ </build>
+</project>
diff --git a/client-spark/spark4/pom.xml b/client-spark/spark4/pom.xml
new file mode 100644
index 000000000..8ff70fc50
--- /dev/null
+++ b/client-spark/spark4/pom.xml
@@ -0,0 +1,111 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ ~ 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.
+ -->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+ xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+ <parent>
+ <artifactId>uniffle-parent</artifactId>
+ <groupId>org.apache.uniffle</groupId>
+ <version>0.11.0-SNAPSHOT</version>
+ <relativePath>../../pom.xml</relativePath>
+ </parent>
+
+ <artifactId>rss-client-spark4</artifactId>
+ <packaging>jar</packaging>
+ <name>Apache Uniffle Client (Spark 4)</name>
+
+ <dependencies>
+ <dependency>
+ <groupId>com.google.protobuf</groupId>
+ <artifactId>protobuf-java</artifactId>
+ <version>${protobuf.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-core_${scala.binary.version}</artifactId>
+ <version>${spark.version}</version>
+ <scope>provided</scope>
+ <exclusions>
+ <exclusion>
+ <groupId>org.slf4j</groupId>
+ <artifactId>slf4j-log4j12</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-client</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-client-spark-common</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>io.netty</groupId>
+ <artifactId>netty-all</artifactId>
+ </dependency>
+ <dependency>
+ <groupId>org.mockito</groupId>
+ <artifactId>mockito-inline</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-client-spark-common</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>shuffle-storage</artifactId>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-minicluster</artifactId>
+ <scope>test</scope>
+ <exclusions>
+ <exclusion>
+ <groupId>org.slf4j</groupId>
+ <artifactId>slf4j-log4j12</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-catalyst_${scala.binary.version}</artifactId>
+ <version>${spark.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ </dependencies>
+
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-surefire-plugin</artifactId>
+ </plugin>
+ </plugins>
+ </build>
+</project>
diff --git a/client-spark/spark4/src/main/java/org/apache/spark/shuffle/DelegationRssShuffleManager.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/DelegationRssShuffleManager.java
new file mode 100644
index 000000000..33828e1f8
--- /dev/null
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/DelegationRssShuffleManager.java
@@ -0,0 +1,251 @@
+/*
+ * 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.spark.shuffle;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import com.google.common.collect.Maps;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.spark.ShuffleDependency;
+import org.apache.spark.SparkConf;
+import org.apache.spark.TaskContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.CoordinatorClient;
+import org.apache.uniffle.client.request.RssAccessClusterRequest;
+import org.apache.uniffle.client.response.RssAccessClusterResponse;
+import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.common.config.RssConf;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.rpc.StatusCode;
+import org.apache.uniffle.common.util.Constants;
+
+import static org.apache.uniffle.common.util.Constants.ACCESS_INFO_REQUIRED_SHUFFLE_NODES_NUM;
+
+public class DelegationRssShuffleManager implements ShuffleManager {
+
+ private static final Logger LOG = LoggerFactory.getLogger(DelegationRssShuffleManager.class);
+
+ private final ShuffleManager delegate;
+ private final int accessTimeoutMs;
+ private final SparkConf sparkConf;
+ private String user;
+ private String uuid;
+
+ public DelegationRssShuffleManager(SparkConf sparkConf, boolean isDriver) throws Exception {
+ LOG.info(
+ "Uniffle {} version: {}", this.getClass().getName(), Constants.VERSION_AND_REVISION_SHORT);
+ this.sparkConf = sparkConf;
+ accessTimeoutMs = sparkConf.get(RssSparkConfig.RSS_ACCESS_TIMEOUT_MS);
+ if (isDriver) {
+ try (CoordinatorClient coordinatorClient =
+ RssSparkShuffleUtils.createCoordinatorClientsForAccessCluster(sparkConf)) {
+ delegate = createShuffleManagerInDriver(coordinatorClient);
+ }
+ } else {
+ delegate = createShuffleManagerInExecutor();
+ }
+
+ if (delegate == null) {
+ throw new RssException("Fail to create shuffle manager!");
+ }
+ }
+
+ private ShuffleManager createShuffleManagerInDriver(CoordinatorClient coordinatorClient)
+ throws RssException {
+ ShuffleManager shuffleManager;
+ user = "user";
+ try {
+ user = UserGroupInformation.getCurrentUser().getShortUserName();
+ } catch (Exception e) {
+ LOG.error("Error on getting user from ugi.", e);
+ }
+ boolean canAccess = tryAccessCluster(coordinatorClient);
+ if (uuid == null || "".equals(uuid)) {
+ uuid = String.valueOf(System.currentTimeMillis());
+ }
+ if (canAccess) {
+ try {
+ sparkConf.set("spark.rss.quota.user", user);
+ sparkConf.set("spark.rss.quota.uuid", uuid);
+ shuffleManager = new RssShuffleManager(sparkConf, true);
+ sparkConf.set(RssSparkConfig.RSS_ENABLED.key(), "true");
+ sparkConf.set("spark.shuffle.manager", RssShuffleManager.class.getCanonicalName());
+ LOG.info("Use RssShuffleManager");
+ return shuffleManager;
+ } catch (Exception exception) {
+ LOG.warn(
+ "Fail to create RssShuffleManager, fallback to SortShuffleManager {}",
+ exception.getMessage());
+ }
+ }
+
+ try {
+ shuffleManager =
+ RssSparkShuffleUtils.loadShuffleManager(
+ Constants.SORT_SHUFFLE_MANAGER_NAME, sparkConf, true);
+ sparkConf.set(RssSparkConfig.RSS_ENABLED.key(), "false");
+ sparkConf.set("spark.shuffle.manager", "sort");
+ if (sparkConf.getBoolean(Constants.SPARK_DYNAMIC_ENABLED, false)) {
+ sparkConf.set("spark.shuffle.service.enabled", "true");
+ LOG.info("Auto enable shuffle service while dynamic allocation is enabled.");
+ }
+ LOG.info("Use SortShuffleManager");
+ } catch (Exception e) {
+ throw new RssException(e.getMessage(), e);
+ }
+
+ return shuffleManager;
+ }
+
+ private boolean tryAccessCluster(CoordinatorClient coordinatorClient) {
+ String accessId = DelegationRssShuffleManagerUtils.acquireAccessId(sparkConf);
+ if (accessId == null) {
+ LOG.warn("Access id key is null");
+ return false;
+ }
+ long retryInterval = sparkConf.get(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_INTERVAL_MS);
+ int retryTimes = sparkConf.get(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_TIMES);
+
+ int assignmentShuffleNodesNum =
+ sparkConf.get(RssSparkConfig.RSS_CLIENT_ASSIGNMENT_SHUFFLE_SERVER_NUMBER);
+ Map<String, String> extraProperties = Maps.newHashMap();
+ extraProperties.put(
+ ACCESS_INFO_REQUIRED_SHUFFLE_NODES_NUM, String.valueOf(assignmentShuffleNodesNum));
+
+ RssConf rssConf = RssSparkConfig.toRssConf(sparkConf);
+ List<String> excludeProperties =
+ rssConf.get(RssClientConf.RSS_CLIENT_REPORT_EXCLUDE_PROPERTIES);
+ List<String> includeProperties =
+ rssConf.get(RssClientConf.RSS_CLIENT_REPORT_INCLUDE_PROPERTIES);
+ rssConf.getAll().stream()
+ .filter(entry -> includeProperties == null || includeProperties.contains(entry.getKey()))
+ .filter(entry -> !excludeProperties.contains(entry.getKey()))
+ .forEach(entry -> extraProperties.put(entry.getKey(), (String) entry.getValue()));
+
+ Set<String> assignmentTags = RssSparkShuffleUtils.getAssignmentTags(sparkConf);
+ try {
+ if (coordinatorClient != null) {
+ RssAccessClusterResponse response =
+ coordinatorClient.accessCluster(
+ new RssAccessClusterRequest(
+ accessId,
+ assignmentTags,
+ accessTimeoutMs,
+ extraProperties,
+ user,
+ retryInterval,
+ retryTimes));
+ if (response.getStatusCode() == StatusCode.SUCCESS) {
+ LOG.warn("Success to access cluster using {}", accessId);
+ uuid = response.getUuid();
+ return true;
+ } else if (response.getStatusCode() == StatusCode.ACCESS_DENIED) {
+ throw new RssException(
+ "Request to access cluster is denied using "
+ + accessId
+ + " for "
+ + response.getMessage());
+ } else {
+ throw new RssException("Fail to reach cluster for " + response.getMessage());
+ }
+ }
+ } catch (Throwable e) {
+ LOG.warn("Fail to access cluster using {} for ", accessId, e);
+ }
+
+ return false;
+ }
+
+ private ShuffleManager createShuffleManagerInExecutor() throws RssException {
+ ShuffleManager shuffleManager;
+ boolean useRSS = sparkConf.get(RssSparkConfig.RSS_ENABLED);
+ if (useRSS) {
+ // Executors always use the manager chosen at startup without fallback
+ shuffleManager = new RssShuffleManager(sparkConf, false);
+ LOG.info("Use RssShuffleManager");
+ } else {
+ try {
+ shuffleManager =
+ RssSparkShuffleUtils.loadShuffleManager(
+ Constants.SORT_SHUFFLE_MANAGER_NAME, sparkConf, false);
+ LOG.info("Use SortShuffleManager");
+ } catch (Exception e) {
+ throw new RssException(e.getMessage(), e);
+ }
+ }
+ return shuffleManager;
+ }
+
+ public ShuffleManager getDelegate() {
+ return delegate;
+ }
+
+ @Override
+ public <K, V, C> ShuffleHandle registerShuffle(
+ int shuffleId, ShuffleDependency<K, V, C> dependency) {
+ return delegate.registerShuffle(shuffleId, dependency);
+ }
+
+ @Override
+ public <K, V> ShuffleWriter<K, V> getWriter(
+ ShuffleHandle handle, long mapId, TaskContext context, ShuffleWriteMetricsReporter metrics) {
+ return delegate.getWriter(handle, mapId, context, metrics);
+ }
+
+ @Override
+ public <K, C> ShuffleReader<K, C> getReader(
+ ShuffleHandle handle,
+ int startPartition,
+ int endPartition,
+ TaskContext context,
+ ShuffleReadMetricsReporter metrics) {
+ return delegate.getReader(handle, startPartition, endPartition, context, metrics);
+ }
+
+ @Override
+ public <K, C> ShuffleReader<K, C> getReader(
+ ShuffleHandle handle,
+ int startMapIndex,
+ int endMapIndex,
+ int startPartition,
+ int endPartition,
+ TaskContext context,
+ ShuffleReadMetricsReporter metrics) {
+ return delegate.getReader(
+ handle, startMapIndex, endMapIndex, startPartition, endPartition, context, metrics);
+ }
+
+ @Override
+ public boolean unregisterShuffle(int shuffleId) {
+ return delegate.unregisterShuffle(shuffleId);
+ }
+
+ @Override
+ public void stop() {
+ delegate.stop();
+ }
+
+ @Override
+ public ShuffleBlockResolver shuffleBlockResolver() {
+ return delegate.shuffleBlockResolver();
+ }
+}
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/FunctionUtils.java
similarity index 51%
copy from storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
copy to client-spark/spark4/src/main/java/org/apache/spark/shuffle/FunctionUtils.java
index d84bdfead..d72ac4473 100644
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/FunctionUtils.java
@@ -15,30 +15,29 @@
* limitations under the License.
*/
-package org.apache.uniffle.storage.handler.impl;
+package org.apache.spark.shuffle;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
+import scala.Function0;
-import org.apache.uniffle.common.KerberizedHadoopBase;
+public class FunctionUtils {
-public class KerberizedHadoopClientReadHandlerTest extends KerberizedHadoopBase {
+ /**
+ * Returns a Function0 that delegates to {@code f} at most once; subsequent calls return the
+ * cached result. Thread-safe: concurrent callers block until the first invocation completes.
+ */
+ public static <T> Function0<T> once(Function0<T> f) {
+ return new Function0<T>() {
+ private boolean computed = false;
+ private T value;
- @BeforeAll
- public static void beforeAll() throws Exception {
- testRunner = KerberizedHadoopClientReadHandlerTest.class;
- KerberizedHadoopBase.init();
- }
-
- @BeforeEach
- public void beforeEach() throws Exception {
- initHadoopSecurityContext();
- }
-
- @Test
- public void test() throws Exception {
- HadoopClientReadHandlerTest.createAndRunCases(
- kerberizedHadoop.getSchemeAndAuthorityPrefix(), kerberizedHadoop.getConf(), "alex");
+ @Override
+ public synchronized T apply() {
+ if (!computed) {
+ value = f.apply();
+ computed = true;
+ }
+ return value;
+ }
+ };
}
}
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleDataIo.java
similarity index 50%
copy from storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
copy to client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleDataIo.java
index d84bdfead..edbe25c9f 100644
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleDataIo.java
@@ -15,30 +15,29 @@
* limitations under the License.
*/
-package org.apache.uniffle.storage.handler.impl;
+package org.apache.spark.shuffle;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
+import org.apache.spark.SparkConf;
+import org.apache.spark.shuffle.api.ShuffleDataIO;
+import org.apache.spark.shuffle.api.ShuffleDriverComponents;
+import org.apache.spark.shuffle.api.ShuffleExecutorComponents;
+import org.apache.spark.shuffle.sort.io.LocalDiskShuffleExecutorComponents;
-import org.apache.uniffle.common.KerberizedHadoopBase;
+public class RssShuffleDataIo implements ShuffleDataIO {
+ private final SparkConf sparkConf;
-public class KerberizedHadoopClientReadHandlerTest extends KerberizedHadoopBase {
-
- @BeforeAll
- public static void beforeAll() throws Exception {
- testRunner = KerberizedHadoopClientReadHandlerTest.class;
- KerberizedHadoopBase.init();
+ public RssShuffleDataIo(SparkConf sparkConf) {
+ this.sparkConf = sparkConf;
}
- @BeforeEach
- public void beforeEach() throws Exception {
- initHadoopSecurityContext();
+ /** Compatible with SortShuffleManager when DelegationRssShuffleManager fallback */
+ @Override
+ public ShuffleExecutorComponents executor() {
+ return new LocalDiskShuffleExecutorComponents(sparkConf);
}
- @Test
- public void test() throws Exception {
- HadoopClientReadHandlerTest.createAndRunCases(
- kerberizedHadoop.getSchemeAndAuthorityPrefix(), kerberizedHadoop.getConf(), "alex");
+ @Override
+ public ShuffleDriverComponents driver() {
+ return new RssShuffleDriverComponents(sparkConf);
}
}
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleDriverComponents.java
similarity index 52%
copy from storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
copy to client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleDriverComponents.java
index d84bdfead..a69ffa96e 100644
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleDriverComponents.java
@@ -15,30 +15,28 @@
* limitations under the License.
*/
-package org.apache.uniffle.storage.handler.impl;
+package org.apache.spark.shuffle;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
+import org.apache.spark.SparkConf;
+import org.apache.spark.shuffle.sort.io.LocalDiskShuffleDriverComponents;
-import org.apache.uniffle.common.KerberizedHadoopBase;
+public class RssShuffleDriverComponents extends LocalDiskShuffleDriverComponents {
-public class KerberizedHadoopClientReadHandlerTest extends KerberizedHadoopBase {
+ private final SparkConf sparkConf;
- @BeforeAll
- public static void beforeAll() throws Exception {
- testRunner = KerberizedHadoopClientReadHandlerTest.class;
- KerberizedHadoopBase.init();
+ public RssShuffleDriverComponents(SparkConf sparkConf) {
+ this.sparkConf = sparkConf;
}
- @BeforeEach
- public void beforeEach() throws Exception {
- initHadoopSecurityContext();
- }
-
- @Test
- public void test() throws Exception {
- HadoopClientReadHandlerTest.createAndRunCases(
- kerberizedHadoop.getSchemeAndAuthorityPrefix(), kerberizedHadoop.getConf(), "alex");
+ /**
+ * This method is called after DelegationRssShuffleManager initialize, so
+ * RssSparkConfig.RSS_ENABLED must be already set
+ */
+ @Override
+ public boolean supportsReliableStorage() {
+ return sparkConf.get(RssSparkConfig.RSS_ENABLED)
+ || RssShuffleManager.class
+ .getCanonicalName()
+ .equals(sparkConf.get("spark.shuffle.manager"));
}
}
diff --git a/client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java
new file mode 100644
index 000000000..1d62fc9b9
--- /dev/null
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java
@@ -0,0 +1,646 @@
+/*
+ * 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.spark.shuffle;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+import scala.Tuple2;
+import scala.Tuple3;
+import scala.collection.Iterator;
+import scala.collection.Seq;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Sets;
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.spark.ShuffleDependency;
+import org.apache.spark.SparkConf;
+import org.apache.spark.SparkEnv;
+import org.apache.spark.TaskContext;
+import org.apache.spark.broadcast.Broadcast;
+import org.apache.spark.executor.ShuffleReadMetrics;
+import org.apache.spark.executor.ShuffleWriteMetrics;
+import org.apache.spark.rdd.DeterministicLevel;
+import org.apache.spark.shuffle.events.ShuffleAssignmentInfoEvent;
+import org.apache.spark.shuffle.handle.MutableShuffleHandleInfo;
+import org.apache.spark.shuffle.handle.ShuffleHandleInfo;
+import org.apache.spark.shuffle.handle.SimpleShuffleHandleInfo;
+import org.apache.spark.shuffle.handle.StageAttemptShuffleHandleInfo;
+import org.apache.spark.shuffle.reader.RssShuffleReader;
+import org.apache.spark.shuffle.writer.DataPusher;
+import org.apache.spark.shuffle.writer.RssShuffleWriter;
+import org.apache.spark.sql.internal.SQLConf;
+import org.apache.spark.storage.BlockId;
+import org.apache.spark.storage.BlockManagerId;
+import org.roaringbitmap.longlong.Roaring64NavigableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.PartitionDataReplicaRequirementTracking;
+import org.apache.uniffle.client.api.ShuffleResult;
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+import org.apache.uniffle.client.impl.FailedBlockSendTracker;
+import org.apache.uniffle.client.impl.MergedPartitionStats;
+import org.apache.uniffle.client.util.ClientUtils;
+import org.apache.uniffle.client.util.RssClientConfig;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.common.config.RssConf;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.exception.RssFetchFailedException;
+import org.apache.uniffle.common.util.RssUtils;
+import org.apache.uniffle.shuffle.RssShuffleClientFactory;
+import org.apache.uniffle.shuffle.ShuffleWriteTaskStats;
+import org.apache.uniffle.shuffle.manager.RssShuffleManagerBase;
+
+public class RssShuffleManager extends RssShuffleManagerBase {
+ private static final Logger LOG = LoggerFactory.getLogger(RssShuffleManager.class);
+
+ public RssShuffleManager(SparkConf conf, boolean isDriver) {
+ super(conf, isDriver);
+ this.dataDistributionType = getDataDistributionType(sparkConf);
+ if (isIntegrityValidationEnabled(rssConf)) {
+ LOG.info("shuffle integrity validation has been enabled.");
+ }
+ }
+
+ @VisibleForTesting
+ RssShuffleManager(
+ SparkConf conf,
+ boolean isDriver,
+ DataPusher dataPusher,
+ Map<String, Set<Long>> taskToSuccessBlockIds,
+ Map<String, FailedBlockSendTracker> taskToFailedBlockSendTracker) {
+ super(conf, isDriver, dataPusher, taskToSuccessBlockIds, taskToFailedBlockSendTracker);
+ }
+
+ // Called on the driver to assign shuffle servers via the coordinator and broadcast the handle to executors.
+ @Override
+ public <K, V, C> ShuffleHandle registerShuffle(
+ int shuffleId, ShuffleDependency<K, V, C> dependency) {
+
+ // fail fast if number of partitions is not supported by block id layout
+ if (dependency.partitioner().numPartitions() > blockIdLayout.maxNumPartitions) {
+ throw new RssException(
+ "Cannot register shuffle with "
+ + dependency.partitioner().numPartitions()
+ + " partitions because the configured block id layout supports at most "
+ + blockIdLayout.maxNumPartitions
+ + " partitions.");
+ }
+
+ // RSS requires serialized object relocation support; JavaSerializer does not provide this.
+ if (!SparkEnv.get().serializer().supportsRelocationOfSerializedObjects()) {
+ throw new IllegalArgumentException(
+ "Can't use serialized shuffle for shuffleId: "
+ + shuffleId
+ + ", because the"
+ + " serializer: "
+ + SparkEnv.get().serializer().getClass().getName()
+ + " does not support object "
+ + "relocation.");
+ }
+
+ if (id.get() == null) {
+ id.compareAndSet(null, SparkEnv.get().conf().getAppId() + "_" + uuid);
+ appId = id.get();
+ if (dataPusher != null) {
+ dataPusher.setRssAppId(id.get());
+ }
+ }
+ LOG.info("Generate application id used in rss: {}", id.get());
+ // If stage retry is enabled, the Deterministic status of the ShuffleId needs to be recorded.
+ if (rssStageRetryEnabled) {
+ shuffleIdMappingManager.recordShuffleIdDeterminate(
+ shuffleId,
+ dependency.rdd().getOutputDeterministicLevel() != DeterministicLevel.INDETERMINATE());
+ }
+
+ if (dependency.partitioner().numPartitions() == 0) {
+ shuffleIdToPartitionNum.putIfAbsent(shuffleId, 0);
+ shuffleIdToNumMapTasks.computeIfAbsent(
+ shuffleId, key -> dependency.rdd().partitions().length);
+ LOG.info(
+ "RegisterShuffle with ShuffleId[{}], partitionNum is 0, "
+ + "return the empty RssShuffleHandle directly",
+ shuffleId);
+ Broadcast<SimpleShuffleHandleInfo> hdlInfoBd =
+ RssSparkShuffleUtils.broadcastShuffleHdlInfo(
+ RssSparkShuffleUtils.getActiveSparkContext(),
+ shuffleId,
+ Collections.emptyMap(),
+ RemoteStorageInfo.EMPTY_REMOTE_STORAGE);
+ return new RssShuffleHandle<>(
+ shuffleId, id.get(), dependency.rdd().getNumPartitions(), dependency, hdlInfoBd);
+ }
+
+ String storageType = sparkConf.get(RssSparkConfig.RSS_STORAGE_TYPE.key());
+ RemoteStorageInfo defaultRemoteStorage = getDefaultRemoteStorageInfo(sparkConf);
+ RemoteStorageInfo remoteStorage =
+ ClientUtils.fetchRemoteStorage(
+ id.get(), defaultRemoteStorage, dynamicConfEnabled, storageType, shuffleWriteClient);
+
+ Set<String> assignmentTags = RssSparkShuffleUtils.getAssignmentTags(sparkConf);
+ ClientUtils.validateClientType(clientType);
+ assignmentTags.add(clientType);
+
+ int requiredShuffleServerNumber =
+ RssSparkShuffleUtils.getRequiredShuffleServerNumber(sparkConf);
+ int estimateTaskConcurrency = RssSparkShuffleUtils.estimateTaskConcurrency(sparkConf);
+
+ Map<Integer, List<ShuffleServerInfo>> partitionToServers =
+ requestShuffleAssignment(
+ shuffleId,
+ dependency.partitioner().numPartitions(),
+ 1,
+ requiredShuffleServerNumber,
+ estimateTaskConcurrency,
+ rssStageResubmitManager.getServerIdBlackList());
+ startHeartbeat();
+ shuffleIdToPartitionNum.computeIfAbsent(
+ shuffleId, key -> dependency.partitioner().numPartitions());
+ shuffleIdToNumMapTasks.computeIfAbsent(shuffleId, key -> dependency.rdd().partitions().length);
+ if (shuffleManagerRpcServiceEnabled && rssStageRetryForWriteFailureEnabled) {
+ ShuffleHandleInfo shuffleHandleInfo =
+ new MutableShuffleHandleInfo(
+ shuffleId, partitionToServers, remoteStorage, partitionSplitMode);
+ StageAttemptShuffleHandleInfo handleInfo =
+ new StageAttemptShuffleHandleInfo(shuffleId, remoteStorage, shuffleHandleInfo);
+ shuffleHandleInfoManager.register(shuffleId, handleInfo);
+ } else if (shuffleManagerRpcServiceEnabled && partitionReassignEnabled) {
+ ShuffleHandleInfo shuffleHandleInfo =
+ new MutableShuffleHandleInfo(
+ shuffleId, partitionToServers, remoteStorage, partitionSplitMode);
+ shuffleHandleInfoManager.register(shuffleId, shuffleHandleInfo);
+ }
+ Broadcast<SimpleShuffleHandleInfo> hdlInfoBd =
+ RssSparkShuffleUtils.broadcastShuffleHdlInfo(
+ RssSparkShuffleUtils.getActiveSparkContext(),
+ shuffleId,
+ partitionToServers,
+ remoteStorage);
+ LOG.info(
+ "RegisterShuffle with ShuffleId[{}], uniffleShuffleId[{}], partitionNum[{}], shuffleServerForResult: {}",
+ shuffleId,
+ shuffleId,
+ partitionToServers.size(),
+ partitionToServers);
+
+ // Post assignment event
+ RssSparkShuffleUtils.getActiveSparkContext()
+ .listenerBus()
+ .post(
+ new ShuffleAssignmentInfoEvent(
+ shuffleId,
+ new ArrayList<>(
+ partitionToServers.values().stream()
+ .flatMap(x -> x.stream())
+ .map(x -> x.getId())
+ .collect(Collectors.toSet()))));
+
+ return new RssShuffleHandle<>(
+ shuffleId, id.get(), dependency.rdd().getNumPartitions(), dependency, hdlInfoBd);
+ }
+
+ @Override
+ public <K, V> ShuffleWriter<K, V> getWriter(
+ ShuffleHandle handle, long mapId, TaskContext context, ShuffleWriteMetricsReporter metrics) {
+ if (!(handle instanceof RssShuffleHandle)) {
+ throw new RssException("Unexpected ShuffleHandle:" + handle.getClass().getName());
+ }
+ RssShuffleHandle<K, V, ?> rssHandle = (RssShuffleHandle<K, V, ?>) handle;
+ setPusherAppId(rssHandle);
+ int shuffleId = rssHandle.getShuffleId();
+ ShuffleWriteMetrics writeMetrics;
+ if (metrics != null) {
+ writeMetrics = new WriteMetrics(metrics);
+ } else {
+ writeMetrics = context.taskMetrics().shuffleWriteMetrics();
+ }
+ String taskId = "" + context.taskAttemptId() + "_" + context.attemptNumber();
+ return new RssShuffleWriter<>(
+ rssHandle.getAppId(),
+ shuffleId,
+ taskId,
+ getTaskAttemptIdForBlockId(context.partitionId(), context.attemptNumber()),
+ writeMetrics,
+ this,
+ sparkConf,
+ shuffleWriteClient,
+ managerClientSupplier,
+ rssHandle,
+ this::markFailedTask,
+ context);
+ }
+
+ public void setPusherAppId(RssShuffleHandle rssShuffleHandle) {
+ // TODO: Decouple appId initialization from the first ShuffleHandle access
+ if (id.get() == null) {
+ id.compareAndSet(null, rssShuffleHandle.getAppId());
+ if (dataPusher != null) {
+ dataPusher.setRssAppId(id.get());
+ }
+ }
+ }
+
+ @Override
+ public <K, C> ShuffleReader<K, C> getReader(
+ ShuffleHandle handle,
+ int startPartition,
+ int endPartition,
+ TaskContext context,
+ ShuffleReadMetricsReporter metrics) {
+ return getReader(handle, 0, Integer.MAX_VALUE, startPartition, endPartition, context, metrics);
+ }
+
+ @Override
+ public <K, C> ShuffleReader<K, C> getReader(
+ ShuffleHandle handle,
+ int startMapIndex,
+ int endMapIndex,
+ int startPartition,
+ int endPartition,
+ TaskContext context,
+ ShuffleReadMetricsReporter metrics) {
+ long start = System.currentTimeMillis();
+ Pair<Roaring64NavigableMap, Long> info =
+ getExpectedTasksAndRecordsForReader(
+ handle.shuffleId(), startPartition, endPartition, startMapIndex, endMapIndex);
+ Roaring64NavigableMap taskIdBitmap = info.getLeft();
+ long expectedRecordsRead = info.getRight();
+ return getReaderImpl(
+ handle,
+ startMapIndex,
+ endMapIndex,
+ startPartition,
+ endPartition,
+ context,
+ metrics,
+ taskIdBitmap,
+ expectedRecordsRead,
+ System.currentTimeMillis() - start);
+ }
+
+ public <K, C> ShuffleReader<K, C> getReaderImpl(
+ ShuffleHandle handle,
+ int startMapIndex,
+ int endMapIndex,
+ int startPartition,
+ int endPartition,
+ TaskContext context,
+ ShuffleReadMetricsReporter metrics,
+ Roaring64NavigableMap taskIdBitmap,
+ long expectedRecordsRead,
+ long taskIdRetrievedMillis) {
+ if (!(handle instanceof RssShuffleHandle)) {
+ throw new RssException("Unexpected ShuffleHandle:" + handle.getClass().getName());
+ }
+ RssShuffleHandle<K, ?, C> rssShuffleHandle = (RssShuffleHandle<K, ?, C>) handle;
+ final int partitionNum = rssShuffleHandle.getDependency().partitioner().numPartitions();
+ int shuffleId = rssShuffleHandle.getShuffleId();
+
+ ShuffleHandleInfo shuffleHandleInfo;
+ Supplier<ShuffleHandleInfo> func =
+ () ->
+ getShuffleHandleInfo(
+ context.stageId(), context.stageAttemptNumber(), rssShuffleHandle, false);
+ if (readShuffleHandleCacheEnabled) {
+ shuffleHandleInfo = super.getOrFetchShuffleHandle(shuffleId, func);
+ } else {
+ shuffleHandleInfo = func.get();
+ }
+
+ Map<ShuffleServerInfo, Set<Integer>> serverToPartitions =
+ getPartitionDataServers(shuffleHandleInfo, startPartition, endPartition);
+ long start = System.currentTimeMillis();
+ ShuffleResult shuffleResult =
+ getShuffleResultForMultiPart(
+ clientType,
+ serverToPartitions,
+ rssShuffleHandle.getAppId(),
+ shuffleId,
+ context.stageAttemptNumber(),
+ shuffleHandleInfo.createPartitionReplicaTracking());
+ Roaring64NavigableMap blockIdBitmap = shuffleResult.getBlockIds();
+
+ if (isIntegrityValidationServerManagementEnabled(rssConf)) {
+ MergedPartitionStats mergedPartitionStats = shuffleResult.getMergedPartitionStats();
+ if (mergedPartitionStats != null) {
+ long records = mergedPartitionStats.getExpectedRecordNumberByTaskIds(taskIdBitmap);
+ if (records > 0) {
+ expectedRecordsRead = records;
+ }
+ }
+ }
+
+ LOG.info(
+ "Retrieved {} upstream task ids in {} ms and {} block IDs from {} shuffle-servers in {} ms for shuffleId[{}], partitionId[{},{}]",
+ taskIdBitmap.getLongCardinality(),
+ taskIdRetrievedMillis,
+ blockIdBitmap.getLongCardinality(),
+ serverToPartitions.size(),
+ System.currentTimeMillis() - start,
+ handle.shuffleId(),
+ startPartition,
+ endPartition);
+
+ ShuffleReadMetrics readMetrics;
+ if (metrics != null) {
+ readMetrics = new ReadMetrics(metrics);
+ } else {
+ readMetrics = context.taskMetrics().shuffleReadMetrics();
+ }
+
+ final RemoteStorageInfo shuffleRemoteStorageInfo = rssShuffleHandle.getRemoteStorage();
+ LOG.debug("Shuffle reader using remote storage {}", shuffleRemoteStorageInfo);
+ final String shuffleRemoteStoragePath = shuffleRemoteStorageInfo.getPath();
+ Configuration readerHadoopConf =
+ RssSparkShuffleUtils.getRemoteStorageHadoopConf(sparkConf, shuffleRemoteStorageInfo);
+
+ return new RssShuffleReader<K, C>(
+ startPartition,
+ endPartition,
+ startMapIndex,
+ endMapIndex,
+ context,
+ rssShuffleHandle,
+ shuffleRemoteStoragePath,
+ readerHadoopConf,
+ partitionNum,
+ RssUtils.generatePartitionToBitmap(
+ blockIdBitmap, startPartition, endPartition, blockIdLayout),
+ taskIdBitmap,
+ readMetrics,
+ managerClientSupplier,
+ RssSparkConfig.toRssConf(sparkConf),
+ dataDistributionType,
+ shuffleHandleInfo.getAllPartitionServersForReader(),
+ expectedRecordsRead);
+ }
+
+ private Map<ShuffleServerInfo, Set<Integer>> getPartitionDataServers(
+ ShuffleHandleInfo shuffleHandleInfo, int startPartition, int endPartition) {
+ Map<Integer, List<ShuffleServerInfo>> allPartitionToServers =
+ shuffleHandleInfo.getAllPartitionServersForReader();
+ Map<Integer, List<ShuffleServerInfo>> requirePartitionToServers =
+ allPartitionToServers.entrySet().stream()
+ .filter(x -> x.getKey() >= startPartition && x.getKey() < endPartition)
+ .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+ Map<ShuffleServerInfo, Set<Integer>> serverToPartitions =
+ RssUtils.generateServerToPartitions(requirePartitionToServers);
+ return serverToPartitions;
+ }
+
+ public static boolean isIntegrityValidationEnabled(RssConf rssConf) {
+ assert rssConf != null;
+ // disable integrity validation when the multi replicas is enabled.
+ if (rssConf.getInteger(
+ RssClientConfig.RSS_DATA_REPLICA, RssClientConfig.RSS_DATA_REPLICA_DEFAULT_VALUE)
+ > 1) {
+ return false;
+ }
+ // Spark 4.x always supports integrity validation
+ return rssConf.get(RssSparkConfig.RSS_CLIENT_INTEGRITY_VALIDATION_ENABLED);
+ }
+
+ public static boolean isIntegrityValidationServerManagementEnabled(RssConf rssConf) {
+ if (!isIntegrityValidationEnabled(rssConf)) {
+ return false;
+ }
+ return rssConf.get(RssSparkConfig.RSS_DATA_INTEGRITY_VALIDATION_SERVER_MANAGEMENT_ENABLED);
+ }
+
+ public static boolean isIntegrityValidationClientManagementEnabled(RssConf rssConf) {
+ if (!isIntegrityValidationEnabled(rssConf)) {
+ return false;
+ }
+ return !isIntegrityValidationServerManagementEnabled(rssConf);
+ }
+
+ public static boolean isIntegrationValidationFailureAnalysisEnabled(RssConf rssConf) {
+ // TODO: Support validation failure analysis with server-side management mode
+ if (isIntegrityValidationServerManagementEnabled(rssConf)) {
+ return false;
+ }
+ return rssConf.get(RssSparkConfig.RSS_DATA_INTEGRATION_VALIDATION_ANALYSIS_ENABLED);
+ }
+
+ private static Iterator<BlockManagerId> getUpstreamBlockManagerIdsForShuffleReader(
+ int shuffleId, int startPartition, int endPartition, int startMapIndex, int endMapIndex) {
+ Iterator<Tuple2<BlockManagerId, Seq<Tuple3<BlockId, Object, Object>>>> mapStatusIter =
+ SparkEnv.get()
+ .mapOutputTracker()
+ .getMapSizesByExecutorId(shuffleId, startMapIndex, endMapIndex, startPartition, endPartition);
+ final Iterator<Tuple2<BlockManagerId, Seq<Tuple3<BlockId, Object, Object>>>> immutableIter =
+ mapStatusIter;
+ Iterator<BlockManagerId> iter =
+ new Iterator<BlockManagerId>() {
+ @Override
+ public boolean hasNext() {
+ return immutableIter.hasNext();
+ }
+
+ @Override
+ public BlockManagerId next() {
+ return immutableIter.next()._1();
+ }
+ };
+ return iter;
+ }
+
+ public static Map<Long, ShuffleWriteTaskStats> getUpstreamWriteTaskStats(
+ RssConf rssConf,
+ int shuffleId,
+ int startPartition,
+ int endPartition,
+ int startMapIndex,
+ int endMapIndex) {
+ if (isIntegrityValidationServerManagementEnabled(rssConf)) {
+ return Collections.emptyMap();
+ }
+ Iterator<BlockManagerId> iter =
+ getUpstreamBlockManagerIdsForShuffleReader(
+ shuffleId, startPartition, endPartition, startMapIndex, endMapIndex);
+ Map<Long, ShuffleWriteTaskStats> upstreamStats = new HashMap<>();
+ while (iter.hasNext()) {
+ BlockManagerId blockManagerId = iter.next();
+ ShuffleWriteTaskStats shuffleWriteTaskStats =
+ ShuffleWriteTaskStats.decode(rssConf, blockManagerId.topologyInfo().get());
+ upstreamStats.put(shuffleWriteTaskStats.getTaskAttemptId(), shuffleWriteTaskStats);
+ }
+ return upstreamStats;
+ }
+
+ private Pair<Roaring64NavigableMap, Long> getExpectedTasksAndRecordsForReader(
+ int shuffleId, int startPartition, int endPartition, int startMapIndex, int endMapIndex) {
+ Iterator<BlockManagerId> iter =
+ getUpstreamBlockManagerIdsForShuffleReader(
+ shuffleId, startPartition, endPartition, startMapIndex, endMapIndex);
+ Roaring64NavigableMap taskIdBitmap = Roaring64NavigableMap.bitmapOf();
+ long expectedRecordsRead = 0;
+ while (iter.hasNext()) {
+ BlockManagerId blockManagerId = iter.next();
+ if (!blockManagerId.topologyInfo().isDefined()) {
+ throw new RssException("Can't get expected taskAttemptId");
+ }
+
+ String raw = blockManagerId.topologyInfo().get();
+ if (isIntegrityValidationClientManagementEnabled(rssConf)) {
+ ShuffleWriteTaskStats shuffleWriteTaskStats = ShuffleWriteTaskStats.decode(rssConf, raw);
+ taskIdBitmap.add(shuffleWriteTaskStats.getTaskAttemptId());
+ for (int i = startPartition; i < endPartition; i++) {
+ expectedRecordsRead += shuffleWriteTaskStats.getRecordsWritten(i);
+ }
+ } else {
+ taskIdBitmap.add(Long.parseLong(raw));
+ }
+ }
+ return Pair.of(taskIdBitmap, expectedRecordsRead);
+ }
+
+ @Override
+ public ShuffleBlockResolver shuffleBlockResolver() {
+ throw new RssException("RssShuffleManager.shuffleBlockResolver is not implemented");
+ }
+
+ @Override
+ protected ShuffleWriteClient createShuffleWriteClient() {
+ int unregisterThreadPoolSize =
+ sparkConf.get(RssSparkConfig.RSS_CLIENT_UNREGISTER_THREAD_POOL_SIZE);
+ int unregisterTimeoutSec = sparkConf.get(RssSparkConfig.RSS_CLIENT_UNREGISTER_TIMEOUT_SEC);
+ int unregisterRequestTimeoutSec =
+ sparkConf.get(RssSparkConfig.RSS_CLIENT_UNREGISTER_REQUEST_TIMEOUT_SEC);
+ long retryIntervalMax = sparkConf.get(RssSparkConfig.RSS_CLIENT_RETRY_INTERVAL_MAX);
+ int heartBeatThreadNum = sparkConf.get(RssSparkConfig.RSS_CLIENT_HEARTBEAT_THREAD_NUM);
+
+ final int retryMax = sparkConf.get(RssSparkConfig.RSS_CLIENT_RETRY_MAX);
+ return RssShuffleClientFactory.getInstance()
+ .createShuffleWriteClient(
+ RssShuffleClientFactory.newWriteBuilder()
+ .blockIdSelfManagedEnabled(blockIdSelfManagedEnabled)
+ .managerClientSupplier(managerClientSupplier)
+ .clientType(clientType)
+ .retryMax(retryMax)
+ .retryIntervalMax(retryIntervalMax)
+ .heartBeatThreadNum(heartBeatThreadNum)
+ .replica(dataReplica)
+ .replicaWrite(dataReplicaWrite)
+ .replicaRead(dataReplicaRead)
+ .replicaSkipEnabled(dataReplicaSkipEnabled)
+ .dataTransferPoolSize(dataTransferPoolSize)
+ .dataCommitPoolSize(dataCommitPoolSize)
+ .unregisterThreadPoolSize(unregisterThreadPoolSize)
+ .unregisterTimeSec(unregisterTimeoutSec)
+ .unregisterRequestTimeSec(unregisterRequestTimeoutSec)
+ .rssConf(rssConf));
+ }
+
+ @VisibleForTesting
+ protected static ShuffleDataDistributionType getDataDistributionType(SparkConf sparkConf) {
+ RssConf rssConf = RssSparkConfig.toRssConf(sparkConf);
+ if ((boolean) sparkConf.get(SQLConf.ADAPTIVE_EXECUTION_ENABLED())
+ && !rssConf.containsKey(RssClientConf.DATA_DISTRIBUTION_TYPE.key())) {
+ return ShuffleDataDistributionType.LOCAL_ORDER;
+ }
+
+ return rssConf.get(RssClientConf.DATA_DISTRIBUTION_TYPE);
+ }
+
+ static class ReadMetrics extends ShuffleReadMetrics {
+ private ShuffleReadMetricsReporter reporter;
+
+ ReadMetrics(ShuffleReadMetricsReporter reporter) {
+ this.reporter = reporter;
+ }
+
+ @Override
+ public void incRemoteBytesRead(long v) {
+ reporter.incRemoteBytesRead(v);
+ }
+
+ @Override
+ public void incFetchWaitTime(long v) {
+ reporter.incFetchWaitTime(v);
+ }
+
+ @Override
+ public void incRecordsRead(long v) {
+ reporter.incRecordsRead(v);
+ }
+ }
+
+ public static class WriteMetrics extends ShuffleWriteMetrics {
+ private ShuffleWriteMetricsReporter reporter;
+
+ public WriteMetrics(ShuffleWriteMetricsReporter reporter) {
+ this.reporter = reporter;
+ }
+
+ @Override
+ public void incBytesWritten(long v) {
+ reporter.incBytesWritten(v);
+ }
+
+ @Override
+ public void incRecordsWritten(long v) {
+ reporter.incRecordsWritten(v);
+ }
+
+ @Override
+ public void incWriteTime(long v) {
+ reporter.incWriteTime(v);
+ }
+ }
+
+ @VisibleForTesting
+ public void setAppId(String appId) {
+ this.id = new AtomicReference<>(appId);
+ }
+
+ private ShuffleResult getShuffleResultForMultiPart(
+ String clientType,
+ Map<ShuffleServerInfo, Set<Integer>> serverToPartitions,
+ String appId,
+ int shuffleId,
+ int stageAttemptId,
+ PartitionDataReplicaRequirementTracking replicaRequirementTracking) {
+ Set<Integer> failedPartitions = Sets.newHashSet();
+ try {
+ return shuffleWriteClient.getShuffleResultForMultiPartV2(
+ clientType,
+ serverToPartitions,
+ appId,
+ shuffleId,
+ failedPartitions,
+ replicaRequirementTracking);
+ } catch (RssFetchFailedException e) {
+ throw RssSparkShuffleUtils.reportRssFetchFailedException(
+ managerClientSupplier, e, sparkConf, appId, shuffleId, stageAttemptId, failedPartitions);
+ }
+ }
+}
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/Spark4VersionUtils.java
similarity index 52%
copy from client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java
copy to client-spark/spark4/src/main/java/org/apache/spark/shuffle/Spark4VersionUtils.java
index a55f0aa7d..d0e003efc 100644
--- a/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/Spark4VersionUtils.java
@@ -20,35 +20,21 @@ package org.apache.spark.shuffle;
import org.apache.spark.package$;
import org.apache.spark.util.VersionUtils;
-public class SparkVersionUtils {
- public static final String SPARK_VERSION = package$.MODULE$.SPARK_VERSION();
- public static final int MAJOR_VERSION;
- public static final int MINOR_VERSION;
+public class Spark4VersionUtils {
+ public static final String SPARK_VERSION_SHORT = package$.MODULE$.SPARK_VERSION_SHORT();
- static {
- int majorVersion;
- int minorVersion;
- try {
- majorVersion = VersionUtils.majorVersion(SPARK_VERSION);
- minorVersion = VersionUtils.minorVersion(SPARK_VERSION);
- } catch (Exception e) {
- // ignore, just in case some wild spark version is passed.
- majorVersion = -1;
- minorVersion = -1;
- }
- MAJOR_VERSION = majorVersion;
- MINOR_VERSION = minorVersion;
- }
-
- public static boolean isSpark2() {
- return MAJOR_VERSION == 2;
- }
+ private Spark4VersionUtils() {}
- public static boolean isSpark3() {
- return MAJOR_VERSION == 3;
- }
+ public static boolean isSparkVersionAtLeast(String target) {
+ int targetMajor = VersionUtils.majorVersion(target);
+ int targetMinor = VersionUtils.minorVersion(target);
- public static boolean isSpark320() {
- return SPARK_VERSION.matches("^3.2.0([^\\d].*)?$");
+ if (SparkVersionUtils.MAJOR_VERSION > targetMajor) {
+ return true;
+ } else if (SparkVersionUtils.MAJOR_VERSION == targetMajor) {
+ return SparkVersionUtils.MINOR_VERSION >= targetMinor;
+ } else {
+ return false;
+ }
}
}
diff --git a/client-spark/spark4/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java
new file mode 100644
index 000000000..4810fec6a
--- /dev/null
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java
@@ -0,0 +1,508 @@
+/*
+ * 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.spark.shuffle.reader;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+import scala.Function0;
+import scala.Function1;
+import scala.Option;
+import scala.Product2;
+import scala.collection.AbstractIterator;
+import scala.collection.Iterator;
+import scala.runtime.AbstractFunction0;
+import scala.runtime.AbstractFunction1;
+import scala.runtime.BoxedUnit;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Lists;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.spark.Aggregator;
+import org.apache.spark.InterruptibleIterator;
+import org.apache.spark.ShuffleDependency;
+import org.apache.spark.TaskContext;
+import org.apache.spark.executor.ShuffleReadMetrics;
+import org.apache.spark.serializer.Serializer;
+import org.apache.spark.shuffle.FunctionUtils;
+import org.apache.spark.shuffle.RssShuffleHandle;
+import org.apache.spark.shuffle.RssShuffleManager;
+import org.apache.spark.shuffle.RssSparkConfig;
+import org.apache.spark.shuffle.ShuffleReader;
+import org.apache.spark.util.CompletionIterator;
+import org.apache.spark.util.CompletionIterator$;
+import org.apache.spark.util.collection.ExternalSorter;
+import org.roaringbitmap.longlong.Roaring64NavigableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleManagerClient;
+import org.apache.uniffle.client.api.ShuffleReadClient;
+import org.apache.uniffle.client.factory.ShuffleClientFactory;
+import org.apache.uniffle.client.request.RssReportShuffleReadMetricRequest;
+import org.apache.uniffle.client.response.RssReportShuffleReadMetricResponse;
+import org.apache.uniffle.client.util.RssClientConfig;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.ShuffleReadTimes;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.compression.Codec;
+import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.common.config.RssConf;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.rpc.StatusCode;
+import org.apache.uniffle.shuffle.ShuffleReadTaskStats;
+import org.apache.uniffle.shuffle.ShuffleWriteTaskStats;
+import org.apache.uniffle.storage.handler.impl.ShuffleServerReadCostTracker;
+
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_READ_OVERLAPPING_DECOMPRESSION_ENABLED;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_READ_OVERLAPPING_DECOMPRESSION_THREADS;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_READ_REORDER_MULTI_SERVERS_ENABLED;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED;
+
+public class RssShuffleReader<K, C> implements ShuffleReader<K, C> {
+ private static final Logger LOG = LoggerFactory.getLogger(RssShuffleReader.class);
+ private final Map<Integer, List<ShuffleServerInfo>> partitionToShuffleServers;
+
+ private String appId;
+ private int shuffleId;
+ private int startPartition;
+ private int endPartition;
+ private TaskContext context;
+ private ShuffleDependency<K, ?, C> shuffleDependency;
+ private int numMaps;
+ private Serializer serializer;
+ private long taskAttemptId;
+ private String taskId;
+ private String basePath;
+ private int partitionNum;
+ private Map<Integer, Roaring64NavigableMap> partitionToExpectBlocks;
+ private Roaring64NavigableMap taskIdBitmap;
+ private Configuration hadoopConf;
+ private int mapStartIndex;
+ private int mapEndIndex;
+ private ShuffleReadMetrics readMetrics;
+ private RssConf rssConf;
+ private ShuffleDataDistributionType dataDistributionType;
+ private Supplier<ShuffleManagerClient> managerClientSupplier;
+ private ShuffleServerReadCostTracker shuffleServerReadCostTracker =
+ new ShuffleServerReadCostTracker();
+
+ private boolean isShuffleReadFailed = false;
+ private Optional<String> shuffleReadReason = Optional.empty();
+
+ private ShuffleReadTimes shuffleReadTimes = new ShuffleReadTimes();
+
+ private long expectedRecordsRead = 0L;
+ private long actualRecordsRead = 0L;
+
+ private Optional<ShuffleReadTaskStats> shuffleReadTaskStats = Optional.empty();
+
+ public RssShuffleReader(
+ int startPartition,
+ int endPartition,
+ int mapStartIndex,
+ int mapEndIndex,
+ TaskContext context,
+ RssShuffleHandle<K, ?, C> rssShuffleHandle,
+ String basePath,
+ Configuration hadoopConf,
+ int partitionNum,
+ Map<Integer, Roaring64NavigableMap> partitionToExpectBlocks,
+ Roaring64NavigableMap taskIdBitmap,
+ ShuffleReadMetrics readMetrics,
+ Supplier<ShuffleManagerClient> managerClientSupplier,
+ RssConf rssConf,
+ ShuffleDataDistributionType dataDistributionType,
+ Map<Integer, List<ShuffleServerInfo>> allPartitionToServers,
+ long expectedRecordsRead) {
+ this(
+ startPartition,
+ endPartition,
+ mapStartIndex,
+ mapEndIndex,
+ context,
+ rssShuffleHandle,
+ basePath,
+ hadoopConf,
+ partitionNum,
+ partitionToExpectBlocks,
+ taskIdBitmap,
+ readMetrics,
+ managerClientSupplier,
+ rssConf,
+ dataDistributionType,
+ allPartitionToServers);
+ this.expectedRecordsRead = expectedRecordsRead;
+ if (RssShuffleManager.isIntegrationValidationFailureAnalysisEnabled(rssConf)) {
+ this.shuffleReadTaskStats = Optional.of(new ShuffleReadTaskStats());
+ }
+ }
+
+ public RssShuffleReader(
+ int startPartition,
+ int endPartition,
+ int mapStartIndex,
+ int mapEndIndex,
+ TaskContext context,
+ RssShuffleHandle<K, ?, C> rssShuffleHandle,
+ String basePath,
+ Configuration hadoopConf,
+ int partitionNum,
+ Map<Integer, Roaring64NavigableMap> partitionToExpectBlocks,
+ Roaring64NavigableMap taskIdBitmap,
+ ShuffleReadMetrics readMetrics,
+ Supplier<ShuffleManagerClient> managerClientSupplier,
+ RssConf rssConf,
+ ShuffleDataDistributionType dataDistributionType,
+ Map<Integer, List<ShuffleServerInfo>> allPartitionToServers) {
+ this.appId = rssShuffleHandle.getAppId();
+ this.startPartition = startPartition;
+ this.endPartition = endPartition;
+ this.mapStartIndex = mapStartIndex;
+ this.mapEndIndex = mapEndIndex;
+ this.context = context;
+ this.numMaps = rssShuffleHandle.getNumMaps();
+ this.shuffleDependency = rssShuffleHandle.getDependency();
+ this.shuffleId = shuffleDependency.shuffleId();
+ this.serializer = rssShuffleHandle.getDependency().serializer();
+ this.taskId = "" + context.taskAttemptId() + "_" + context.attemptNumber();
+ this.taskAttemptId = context.taskAttemptId();
+ this.basePath = basePath;
+ this.partitionNum = partitionNum;
+ this.partitionToExpectBlocks = partitionToExpectBlocks;
+ this.taskIdBitmap = taskIdBitmap;
+ this.hadoopConf = hadoopConf;
+ this.readMetrics = readMetrics;
+ this.partitionToShuffleServers = allPartitionToServers;
+ this.rssConf = rssConf;
+ this.dataDistributionType = dataDistributionType;
+ this.managerClientSupplier = managerClientSupplier;
+ }
+
+ @Override
+ public Iterator<Product2<K, C>> read() {
+ LOG.info("Starting shuffle read: {}", getReadInfo());
+
+ Iterator<Product2<K, C>> resultIter;
+ MultiPartitionIterator rssShuffleDataIterator = new MultiPartitionIterator<K, C>();
+
+ if (shuffleDependency.keyOrdering().isDefined()) {
+ // Create an ExternalSorter to sort the data
+ Option<Aggregator<K, Object, C>> aggregator = Option.empty();
+ if (shuffleDependency.aggregator().isDefined()) {
+ if (shuffleDependency.mapSideCombine()) {
+ aggregator =
+ Option.apply(
+ (Aggregator<K, Object, C>)
+ new Aggregator<K, C, C>(
+ x -> x,
+ shuffleDependency.aggregator().get().mergeCombiners(),
+ shuffleDependency.aggregator().get().mergeCombiners()));
+ } else {
+ aggregator =
+ Option.apply((Aggregator<K, Object, C>) shuffleDependency.aggregator().get());
+ }
+ }
+ ExternalSorter<K, Object, C> sorter =
+ new ExternalSorter<>(
+ context, aggregator, Option.empty(), shuffleDependency.keyOrdering(), serializer);
+ long startTime = System.currentTimeMillis();
+ sorter.insertAll(rssShuffleDataIterator);
+ LOG.info(
+ "Inserted aggregated records to sorter, took {} millis",
+ System.currentTimeMillis() - startTime);
+ context.taskMetrics().incMemoryBytesSpilled(sorter.memoryBytesSpilled());
+ context.taskMetrics().incPeakExecutionMemory(sorter.peakMemoryUsedBytes());
+ context.taskMetrics().incDiskBytesSpilled(sorter.diskBytesSpilled());
+ Function0<BoxedUnit> fn0 =
+ new AbstractFunction0<BoxedUnit>() {
+ @Override
+ public BoxedUnit apply() {
+ sorter.stop();
+ return BoxedUnit.UNIT;
+ }
+ };
+ Function1<TaskContext, Void> fn1 =
+ new AbstractFunction1<TaskContext, Void>() {
+ public Void apply(TaskContext context) {
+ sorter.stop();
+ return null;
+ }
+ };
+ context.addTaskCompletionListener(fn1);
+ resultIter = CompletionIterator$.MODULE$.apply(sorter.iterator(), fn0);
+ } else if (shuffleDependency.aggregator().isDefined()) {
+ Aggregator<K, Object, C> aggregator =
+ (Aggregator<K, Object, C>) shuffleDependency.aggregator().get();
+ if (shuffleDependency.mapSideCombine()) {
+ resultIter = aggregator.combineCombinersByKey(rssShuffleDataIterator, context);
+ } else {
+ resultIter = aggregator.combineValuesByKey(rssShuffleDataIterator, context);
+ }
+ } else {
+ resultIter = rssShuffleDataIterator;
+ }
+
+ if (!(resultIter instanceof InterruptibleIterator)) {
+ resultIter = new InterruptibleIterator<>(context, resultIter);
+ }
+ // resubmit stage and shuffle manager server port are both set
+ if (rssConf.getBoolean(RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED)
+ && rssConf.getInteger(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT, 0) > 0) {
+ resultIter =
+ RssFetchFailedIterator.newBuilder()
+ .appId(appId)
+ .shuffleId(shuffleId)
+ .partitionId(startPartition)
+ .stageAttemptId(context.stageAttemptNumber())
+ .managerClientSupplier(managerClientSupplier)
+ .build(resultIter);
+ }
+ return resultIter;
+ }
+
+ private String getReadInfo() {
+ return "appId="
+ + appId
+ + ", shuffleId="
+ + shuffleId
+ + ", taskId="
+ + taskId
+ + ", partitions: ["
+ + startPartition
+ + ", "
+ + endPartition
+ + ")"
+ + ", maps: ["
+ + mapStartIndex
+ + ", "
+ + mapEndIndex
+ + "]"
+ + ", expected records: "
+ + expectedRecordsRead;
+ }
+
+ @VisibleForTesting
+ public Configuration getHadoopConf() {
+ return hadoopConf;
+ }
+
+ class MultiPartitionIterator<K, C> extends AbstractIterator<Product2<K, C>> {
+ java.util.Iterator<CompletionIterator<Product2<K, C>, RssShuffleDataIterator<K, C>>> iterator;
+ CompletionIterator<Product2<K, C>, RssShuffleDataIterator<K, C>> dataIterator;
+
+ MultiPartitionIterator() {
+ List<CompletionIterator<Product2<K, C>, RssShuffleDataIterator<K, C>>> iterators =
+ Lists.newArrayList();
+ if (numMaps <= 0) {
+ return;
+ }
+ for (int partition = startPartition; partition < endPartition; partition++) {
+ if (partitionToExpectBlocks.get(partition).isEmpty()) {
+ LOG.info("{} partition is empty partition", partition);
+ continue;
+ }
+ List<ShuffleServerInfo> shuffleServerInfoList = partitionToShuffleServers.get(partition);
+ if (shuffleServerInfoList != null
+ && shuffleServerInfoList.size() > 1
+ && rssConf.getBoolean(RSS_READ_REORDER_MULTI_SERVERS_ENABLED)) {
+ Collections.shuffle(shuffleServerInfoList);
+ }
+
+ // This mechanism of taskId filter is to filter out the most of data for AQE skew and multi
+ // replicas cases
+ boolean isReplicaFilterEnabled =
+ rssConf.getInteger(
+ RssClientConfig.RSS_DATA_REPLICA,
+ RssClientConfig.RSS_DATA_REPLICA_DEFAULT_VALUE)
+ > 1
+ && shuffleServerInfoList.size() > 1;
+ boolean expectedTaskIdsBitmapFilterEnable =
+ !(mapStartIndex == 0 && mapEndIndex == Integer.MAX_VALUE) || isReplicaFilterEnabled;
+ int retryMax =
+ rssConf.getInteger(
+ RssClientConfig.RSS_CLIENT_RETRY_MAX,
+ RssClientConfig.RSS_CLIENT_RETRY_MAX_DEFAULT_VALUE);
+ long retryIntervalMax =
+ rssConf.getLong(
+ RssClientConfig.RSS_CLIENT_RETRY_INTERVAL_MAX,
+ RssClientConfig.RSS_CLIENT_RETRY_INTERVAL_MAX_DEFAULT_VALUE);
+ boolean compress =
+ rssConf.getBoolean(
+ RssSparkConfig.SPARK_SHUFFLE_COMPRESS_KEY.substring(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX.length()),
+ RssSparkConfig.SPARK_SHUFFLE_COMPRESS_DEFAULT);
+ Optional<Codec> codec = compress ? Codec.newInstance(rssConf) : Optional.empty();
+ ShuffleClientFactory.ReadClientBuilder builder =
+ ShuffleClientFactory.newReadBuilder()
+ .readCostTracker(shuffleServerReadCostTracker)
+ .appId(appId)
+ .shuffleId(shuffleId)
+ .partitionId(partition)
+ .basePath(basePath)
+ .partitionNumPerRange(1)
+ .partitionNum(partitionNum)
+ .blockIdBitmap(partitionToExpectBlocks.get(partition))
+ .taskIdBitmap(taskIdBitmap)
+ .shuffleServerInfoList(shuffleServerInfoList)
+ .hadoopConf(hadoopConf)
+ .shuffleDataDistributionType(dataDistributionType)
+ .expectedTaskIdsBitmapFilterEnable(expectedTaskIdsBitmapFilterEnable)
+ .retryMax(retryMax)
+ .retryIntervalMax(retryIntervalMax)
+ .rssConf(rssConf)
+ .taskAttemptId(taskAttemptId);
+ if (codec.isPresent() && rssConf.get(RSS_READ_OVERLAPPING_DECOMPRESSION_ENABLED)) {
+ builder
+ .overlappingDecompressionEnabled(true)
+ .codec(codec.get())
+ .overlappingDecompressionThreadNum(
+ rssConf.get(RSS_READ_OVERLAPPING_DECOMPRESSION_THREADS));
+ }
+ ShuffleReadClient shuffleReadClient =
+ ShuffleClientFactory.getInstance().createShuffleReadClient(builder);
+ RssShuffleDataIterator<K, C> iterator =
+ new RssShuffleDataIterator<>(
+ shuffleDependency.serializer(),
+ shuffleReadClient,
+ readMetrics,
+ rssConf,
+ codec,
+ shuffleReadTaskStats,
+ partition);
+ CompletionIterator<Product2<K, C>, RssShuffleDataIterator<K, C>> completionIterator =
+ CompletionIterator$.MODULE$.apply(
+ iterator,
+ FunctionUtils.once(
+ () -> {
+ shuffleReadTimes.merge(iterator.getReadTimes());
+ context.taskMetrics().mergeShuffleReadMetrics();
+ return iterator.cleanup();
+ }));
+ iterators.add(completionIterator);
+ }
+ iterator = iterators.iterator();
+ if (iterator.hasNext()) {
+ dataIterator = iterator.next();
+ iterator.remove();
+ }
+ context.addTaskCompletionListener(
+ (taskContext) -> {
+ if (dataIterator != null) {
+ dataIterator.completion();
+ }
+ iterator.forEachRemaining(CompletionIterator::completion);
+ });
+ }
+
+ @Override
+ public boolean hasNext() {
+ try {
+ if (dataIterator == null) {
+ return false;
+ }
+ while (!dataIterator.hasNext()) {
+ if (!iterator.hasNext()) {
+ validate();
+ postShuffleReadMetricsToDriver();
+ return false;
+ }
+ dataIterator = iterator.next();
+ iterator.remove();
+ }
+ return dataIterator.hasNext();
+ } catch (Exception e) {
+ isShuffleReadFailed = true;
+ shuffleReadReason = Optional.ofNullable(e.getMessage());
+ postShuffleReadMetricsToDriver();
+ throw e;
+ }
+ }
+
+ @Override
+ public Product2<K, C> next() {
+ Product2<K, C> result = dataIterator.next();
+ actualRecordsRead += 1;
+ return result;
+ }
+ }
+
+ private void validate() {
+ if (RssShuffleManager.isIntegrityValidationEnabled(rssConf)
+ && expectedRecordsRead > 0
+ && (expectedRecordsRead != actualRecordsRead)) {
+ // dig to analyze the missing records from the upstream map id
+ if (shuffleReadTaskStats.isPresent()) {
+ ShuffleReadTaskStats readTaskStats = shuffleReadTaskStats.get();
+ Map<Long, ShuffleWriteTaskStats> upstreamWriteTaskStats =
+ RssShuffleManager.getUpstreamWriteTaskStats(
+ rssConf, shuffleId, startPartition, endPartition, mapStartIndex, mapEndIndex);
+ readTaskStats.diff(upstreamWriteTaskStats, startPartition, endPartition);
+ }
+ throw new RssException(
+ "Inconsistent number of records: "
+ + expectedRecordsRead
+ + " written, "
+ + actualRecordsRead
+ + " read");
+ }
+ }
+
+ private void postShuffleReadMetricsToDriver() {
+ if (managerClientSupplier != null) {
+ ShuffleManagerClient client = managerClientSupplier.get();
+ if (client != null) {
+ try {
+ RssReportShuffleReadMetricResponse response =
+ client.reportShuffleReadMetric(
+ new RssReportShuffleReadMetricRequest(
+ context.stageId(),
+ shuffleId,
+ context.taskAttemptId(),
+ shuffleServerReadCostTracker.list().entrySet().stream()
+ .collect(
+ Collectors.toMap(
+ Map.Entry::getKey,
+ x ->
+ new RssReportShuffleReadMetricRequest.TaskShuffleReadMetric(
+ x.getValue().getDurationMillis(),
+ x.getValue().getReadBytes(),
+ x.getValue().getMemoryReadDurationMillis(),
+ x.getValue().getMemoryReadBytes(),
+ x.getValue().getLocalfileReadDurationMillis(),
+ x.getValue().getLocalfileReadBytes(),
+ x.getValue().getHadoopReadLocalFileDurationMillis(),
+ x.getValue().getHadoopReadLocalFileBytes()))),
+ isShuffleReadFailed,
+ shuffleReadReason,
+ shuffleReadTimes,
+ context.attemptNumber()));
+ if (response != null && response.getStatusCode() != StatusCode.SUCCESS) {
+ LOG.error("Errors on reporting shuffle read metrics to driver");
+ }
+ } catch (Exception e) {
+ LOG.error("Errors on post shuffle read metric to driver", e);
+ }
+ }
+ }
+ }
+}
diff --git a/client-spark/spark4/src/main/java/org/apache/spark/shuffle/writer/PartitionLengthStatistic.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/writer/PartitionLengthStatistic.java
new file mode 100644
index 000000000..20dc4074e
--- /dev/null
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/writer/PartitionLengthStatistic.java
@@ -0,0 +1,52 @@
+/*
+ * 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.spark.shuffle.writer;
+
+import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.apache.uniffle.common.ShuffleBlockInfo;
+import org.apache.uniffle.common.exception.RssException;
+
+public class PartitionLengthStatistic {
+ private final AtomicLong[] partitionLens;
+
+ public PartitionLengthStatistic(int numPartitions) {
+ this.partitionLens = new AtomicLong[numPartitions];
+ for (int i = 0; i < numPartitions; i++) {
+ partitionLens[i] = new AtomicLong(0);
+ }
+ }
+
+ public void inc(ShuffleBlockInfo block) {
+ int partitionId = block.getPartitionId();
+ if (partitionId >= partitionLens.length) {
+ throw new RssException(
+ "Partition ID "
+ + partitionId
+ + " is out of bounds (should be less than "
+ + partitionLens.length
+ + ")");
+ }
+ partitionLens[partitionId].addAndGet(block.getLength());
+ }
+
+ public long[] toArray() {
+ return Arrays.stream(this.partitionLens).mapToLong(x -> x.get()).toArray();
+ }
+}
diff --git a/client-spark/spark4/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java
new file mode 100644
index 000000000..5bfb33104
--- /dev/null
+++ b/client-spark/spark4/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java
@@ -0,0 +1,911 @@
+/*
+ * 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.spark.shuffle.writer;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+import scala.Function1;
+import scala.Option;
+import scala.Product2;
+import scala.Tuple2;
+import scala.collection.Iterator;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
+import com.google.common.util.concurrent.Uninterruptibles;
+import org.apache.commons.collections4.CollectionUtils;
+import org.apache.spark.Partitioner;
+import org.apache.spark.ShuffleDependency;
+import org.apache.spark.SparkConf;
+import org.apache.spark.TaskContext;
+import org.apache.spark.executor.ShuffleWriteMetrics;
+import org.apache.spark.scheduler.MapStatus;
+import org.apache.spark.shuffle.FetchFailedException;
+import org.apache.spark.shuffle.RssShuffleHandle;
+import org.apache.spark.shuffle.RssShuffleManager;
+import org.apache.spark.shuffle.RssSparkConfig;
+import org.apache.spark.shuffle.RssSparkShuffleUtils;
+import org.apache.spark.shuffle.ShuffleWriter;
+import org.apache.spark.shuffle.handle.ShuffleHandleInfo;
+import org.apache.spark.storage.BlockManagerId;
+import org.roaringbitmap.longlong.Roaring64NavigableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleManagerClient;
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+import org.apache.uniffle.client.impl.FailedBlockSendTracker;
+import org.apache.uniffle.client.impl.TrackingBlockStatus;
+import org.apache.uniffle.client.request.RssReportShuffleWriteFailureRequest;
+import org.apache.uniffle.client.request.RssReportShuffleWriteMetricRequest;
+import org.apache.uniffle.client.response.RssReportShuffleWriteFailureResponse;
+import org.apache.uniffle.client.response.RssReportShuffleWriteMetricResponse;
+import org.apache.uniffle.common.ShuffleBlockInfo;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.common.config.RssConf;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.exception.RssSendFailedException;
+import org.apache.uniffle.common.exception.RssWaitFailedException;
+import org.apache.uniffle.common.rpc.StatusCode;
+import org.apache.uniffle.common.util.BlockIdLayout;
+import org.apache.uniffle.shuffle.BlockStats;
+import org.apache.uniffle.shuffle.ReassignExecutor;
+import org.apache.uniffle.shuffle.ShuffleWriteTaskStats;
+import org.apache.uniffle.storage.util.StorageType;
+
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_CLIENT_MAP_SIDE_COMBINE_ENABLED;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_PARTITION_REASSIGN_BLOCK_RETRY_MAX_TIMES;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_RESUBMIT_STAGE_WITH_WRITE_FAILURE_ENABLED;
+import static org.apache.spark.shuffle.RssSparkConfig.toRssConf;
+
+public class RssShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(RssShuffleWriter.class);
+ private static final String DUMMY_HOST = "dummy_host";
+ private static final int DUMMY_PORT = 99999;
+ public static final String DEFAULT_ERROR_MESSAGE = "Default Error Message";
+
+ private final String appId;
+ private final int shuffleId;
+ private final ShuffleHandleInfo shuffleHandleInfo;
+ private WriteBufferManager bufferManager;
+ private String taskId;
+ private final int numMaps;
+ private final ShuffleDependency<K, V, C> shuffleDependency;
+ private final Partitioner partitioner;
+ private final RssShuffleManager shuffleManager;
+ private final boolean shouldPartition;
+ private final long sendCheckTimeout;
+ private final long sendCheckInterval;
+ private final int bitmapSplitNum;
+ private final ShuffleWriteClient shuffleWriteClient;
+ private final Set<ShuffleServerInfo> shuffleServersForData;
+ // Gluten needs this variable
+ protected final boolean isMemoryShuffleEnabled;
+ private final Function<String, Boolean> taskFailureCallback;
+ private final Set<Long> blockIds = Sets.newConcurrentHashSet();
+ private TaskContext taskContext;
+ private SparkConf sparkConf;
+ private RssConf rssConf;
+ private boolean blockFailSentRetryEnabled;
+
+ /** used by columnar rss shuffle writer implementation */
+ protected final long taskAttemptId;
+
+ protected final ShuffleWriteMetrics shuffleWriteMetrics;
+
+ private final BlockingQueue<Object> finishEventQueue = new LinkedBlockingQueue<>();
+
+ // Will be updated when the reassignment is triggered.
+ private TaskAttemptAssignment taskAttemptAssignment;
+
+ private final Supplier<ShuffleManagerClient> managerClientSupplier;
+ private boolean enableWriteFailureRetry;
+ private Set<ShuffleServerInfo> recordReportFailedShuffleservers;
+
+ private long totalShuffleWriteMills = 0L;
+ private long checkSendResultMills = 0L;
+
+ private boolean isShuffleWriteFailed = false;
+ private Optional<String> shuffleWriteFailureReason = Optional.empty();
+
+ // Visible for the Gluten
+ protected ShuffleWriteTaskStats shuffleTaskStats;
+
+ private boolean isIntegrityValidationClientManagementEnabled = false;
+
+ private ReassignExecutor reassignExecutor;
+
+ // Only for tests
+ @VisibleForTesting
+ public RssShuffleWriter(
+ String appId,
+ int shuffleId,
+ String taskId,
+ long taskAttemptId,
+ WriteBufferManager bufferManager,
+ ShuffleWriteMetrics shuffleWriteMetrics,
+ RssShuffleManager shuffleManager,
+ SparkConf sparkConf,
+ ShuffleWriteClient shuffleWriteClient,
+ Supplier<ShuffleManagerClient> managerClientSupplier,
+ RssShuffleHandle<K, V, C> rssHandle,
+ ShuffleHandleInfo shuffleHandleInfo,
+ TaskContext context) {
+ this(
+ appId,
+ shuffleId,
+ taskId,
+ taskAttemptId,
+ shuffleWriteMetrics,
+ shuffleManager,
+ sparkConf,
+ shuffleWriteClient,
+ managerClientSupplier,
+ rssHandle,
+ (tid) -> true,
+ shuffleHandleInfo,
+ context);
+ this.bufferManager = bufferManager;
+ this.taskAttemptAssignment = new TaskAttemptAssignment(taskAttemptId, shuffleHandleInfo);
+ this.reassignExecutor =
+ new ReassignExecutor(
+ shuffleManager.getTaskToFailedBlockSendTracker(),
+ taskId,
+ taskAttemptAssignment,
+ block -> clearFailedBlockState(block),
+ blocks -> processShuffleBlockInfos(blocks),
+ managerClientSupplier,
+ taskContext,
+ shuffleId,
+ rssConf.get(RSS_PARTITION_REASSIGN_BLOCK_RETRY_MAX_TIMES));
+ }
+
+ private RssShuffleWriter(
+ String appId,
+ int shuffleId,
+ String taskId,
+ long taskAttemptId,
+ ShuffleWriteMetrics shuffleWriteMetrics,
+ RssShuffleManager shuffleManager,
+ SparkConf sparkConf,
+ ShuffleWriteClient shuffleWriteClient,
+ Supplier<ShuffleManagerClient> managerClientSupplier,
+ RssShuffleHandle<K, V, C> rssHandle,
+ Function<String, Boolean> taskFailureCallback,
+ ShuffleHandleInfo shuffleHandleInfo,
+ TaskContext context) {
+ LOG.info(
+ "Starting shuffle write: appId={}, taskId={}, taskAttemptId={}, shuffleId={}",
+ rssHandle.getAppId(),
+ taskId,
+ taskAttemptId,
+ shuffleId);
+ this.shuffleManager = shuffleManager;
+ this.appId = appId;
+ this.shuffleId = shuffleId;
+ this.taskId = taskId;
+ this.taskAttemptId = taskAttemptId;
+ this.numMaps = rssHandle.getNumMaps();
+ this.shuffleWriteMetrics = shuffleWriteMetrics;
+ this.shuffleDependency = rssHandle.getDependency();
+ this.partitioner = shuffleDependency.partitioner();
+ this.shouldPartition = partitioner.numPartitions() > 1;
+ this.sendCheckTimeout = sparkConf.get(RssSparkConfig.RSS_CLIENT_SEND_CHECK_TIMEOUT_MS);
+ this.sendCheckInterval = sparkConf.get(RssSparkConfig.RSS_CLIENT_SEND_CHECK_INTERVAL_MS);
+ this.bitmapSplitNum = sparkConf.get(RssSparkConfig.RSS_CLIENT_BITMAP_SPLIT_NUM);
+ this.shuffleWriteClient = shuffleWriteClient;
+ this.shuffleServersForData = shuffleHandleInfo.getServers();
+ this.isMemoryShuffleEnabled =
+ isMemoryShuffleEnabled(sparkConf.get(RssSparkConfig.RSS_STORAGE_TYPE.key()));
+ this.taskFailureCallback = taskFailureCallback;
+ this.shuffleHandleInfo = shuffleHandleInfo;
+ this.taskContext = context;
+ this.sparkConf = sparkConf;
+ this.rssConf = toRssConf(sparkConf);
+ this.managerClientSupplier = managerClientSupplier;
+ this.blockFailSentRetryEnabled =
+ sparkConf.getBoolean(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX
+ + RssClientConf.RSS_CLIENT_REASSIGN_ENABLED.key(),
+ RssClientConf.RSS_CLIENT_REASSIGN_ENABLED.defaultValue());
+ this.enableWriteFailureRetry = rssConf.get(RSS_RESUBMIT_STAGE_WITH_WRITE_FAILURE_ENABLED);
+ this.recordReportFailedShuffleservers = Sets.newConcurrentHashSet();
+ this.isIntegrityValidationClientManagementEnabled =
+ RssShuffleManager.isIntegrityValidationClientManagementEnabled(rssConf);
+ this.shuffleTaskStats =
+ new ShuffleWriteTaskStats(
+ rssConf, partitioner.numPartitions(), taskAttemptId, taskContext.taskAttemptId());
+ }
+
+ // Gluten needs this method
+ public RssShuffleWriter(
+ String appId,
+ int shuffleId,
+ String taskId,
+ long taskAttemptId,
+ ShuffleWriteMetrics shuffleWriteMetrics,
+ RssShuffleManager shuffleManager,
+ SparkConf sparkConf,
+ ShuffleWriteClient shuffleWriteClient,
+ RssShuffleHandle<K, V, C> rssHandle,
+ Function<String, Boolean> taskFailureCallback,
+ TaskContext context) {
+ this(
+ appId,
+ shuffleId,
+ taskId,
+ taskAttemptId,
+ shuffleWriteMetrics,
+ shuffleManager,
+ sparkConf,
+ shuffleWriteClient,
+ shuffleManager.getShuffleManagerClientSupplier(),
+ rssHandle,
+ taskFailureCallback,
+ context);
+ }
+
+ public RssShuffleWriter(
+ String appId,
+ int shuffleId,
+ String taskId,
+ long taskAttemptId,
+ ShuffleWriteMetrics shuffleWriteMetrics,
+ RssShuffleManager shuffleManager,
+ SparkConf sparkConf,
+ ShuffleWriteClient shuffleWriteClient,
+ Supplier<ShuffleManagerClient> managerClientSupplier,
+ RssShuffleHandle<K, V, C> rssHandle,
+ Function<String, Boolean> taskFailureCallback,
+ TaskContext context) {
+ this(
+ appId,
+ shuffleId,
+ taskId,
+ taskAttemptId,
+ shuffleWriteMetrics,
+ shuffleManager,
+ sparkConf,
+ shuffleWriteClient,
+ managerClientSupplier,
+ rssHandle,
+ taskFailureCallback,
+ shuffleManager.getShuffleHandleInfo(
+ context.stageId(), context.stageAttemptNumber(), rssHandle, true),
+ context);
+ this.taskAttemptAssignment = new TaskAttemptAssignment(taskAttemptId, shuffleHandleInfo);
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(sparkConf);
+ final WriteBufferManager bufferManager =
+ new WriteBufferManager(
+ shuffleId,
+ taskId,
+ taskAttemptId,
+ bufferOptions,
+ rssHandle.getDependency().serializer(),
+ context.taskMemoryManager(),
+ shuffleWriteMetrics,
+ RssSparkConfig.toRssConf(sparkConf),
+ this::processShuffleBlockInfos,
+ this::getPartitionAssignedServers,
+ context.stageAttemptNumber());
+ this.bufferManager = bufferManager;
+ this.reassignExecutor =
+ new ReassignExecutor(
+ shuffleManager.getTaskToFailedBlockSendTracker(),
+ taskId,
+ taskAttemptAssignment,
+ block -> clearFailedBlockState(block),
+ blocks -> processShuffleBlockInfos(blocks),
+ managerClientSupplier,
+ taskContext,
+ shuffleId,
+ rssConf.get(RSS_PARTITION_REASSIGN_BLOCK_RETRY_MAX_TIMES));
+ }
+
+ @VisibleForTesting
+ protected List<ShuffleServerInfo> getPartitionAssignedServers(int partitionId) {
+ return this.taskAttemptAssignment.retrieve(partitionId);
+ }
+
+ private boolean isMemoryShuffleEnabled(String storageType) {
+ return StorageType.withMemory(StorageType.valueOf(storageType));
+ }
+
+ @Override
+ public void write(Iterator<Product2<K, V>> records) {
+ try {
+ writeImpl(records);
+ } catch (Exception e) {
+ if (e instanceof RssException) {
+ isShuffleWriteFailed = true;
+ shuffleWriteFailureReason = Optional.ofNullable(e.getMessage());
+ }
+ taskFailureCallback.apply(taskId);
+ if (enableWriteFailureRetry) {
+ throwFetchFailedIfNecessary(e, Sets.newConcurrentHashSet());
+ } else {
+ throw e;
+ }
+ }
+ }
+
+ // Gluten needs this method.
+ protected void writeImpl(Iterator<Product2<K, V>> records) {
+ List<ShuffleBlockInfo> shuffleBlockInfos;
+ boolean isCombine = shuffleDependency.mapSideCombine();
+
+ Iterator<? extends Product2<K, ?>> iterator = records;
+ if (isCombine) {
+ if (RssSparkConfig.toRssConf(sparkConf).get(RSS_CLIENT_MAP_SIDE_COMBINE_ENABLED)) {
+ iterator = shuffleDependency.aggregator().get().combineValuesByKey(records, taskContext);
+ } else {
+ Function1<V, C> combiner = shuffleDependency.aggregator().get().createCombiner();
+ iterator =
+ records.map(
+ (Function1<Product2<K, V>, Product2<K, C>>)
+ x -> new Tuple2<>(x._1(), combiner.apply(x._2())));
+ }
+ }
+ long recordCount = 0;
+ while (iterator.hasNext()) {
+ recordCount++;
+ checkDataIfAnyFailure();
+ Product2<K, ?> record = iterator.next();
+ K key = record._1();
+ int partition = getPartition(key);
+ shuffleBlockInfos = bufferManager.addRecord(partition, record._1(), record._2());
+ if (shuffleBlockInfos != null && !shuffleBlockInfos.isEmpty()) {
+ processShuffleBlockInfos(shuffleBlockInfos);
+ }
+ shuffleTaskStats.incPartitionRecord(partition);
+ }
+ final long start = System.currentTimeMillis();
+ shuffleBlockInfos = bufferManager.clear(1.0);
+ if (shuffleBlockInfos != null && !shuffleBlockInfos.isEmpty()) {
+ processShuffleBlockInfos(shuffleBlockInfos);
+ }
+ @SuppressWarnings("checkstyle:VariableDeclarationUsageDistance")
+ long checkStartTs = System.currentTimeMillis();
+ checkAllBufferSpilled();
+ checkSentRecordCount(recordCount);
+ checkBlockSendResult(blockIds);
+ checkSentBlockCount();
+ bufferManager.getShuffleServerPushCostTracker().statistics();
+ long commitStartTs = System.currentTimeMillis();
+ long checkDuration = commitStartTs - checkStartTs;
+ if (!isMemoryShuffleEnabled) {
+ sendCommit();
+ }
+ long writeDurationMs = bufferManager.getWriteTime() + (System.currentTimeMillis() - start);
+ this.totalShuffleWriteMills = writeDurationMs;
+ this.checkSendResultMills = checkDuration;
+ shuffleWriteMetrics.incWriteTime(TimeUnit.MILLISECONDS.toNanos(writeDurationMs));
+ LOG.info(
+ "Finished shuffle writing for shuffleId[{}], taskId[{}], taskAttemptId[{}] in {} ms, including blockWait[{}], commit[{}], {}",
+ shuffleId,
+ taskId,
+ taskAttemptId,
+ writeDurationMs,
+ checkDuration,
+ System.currentTimeMillis() - commitStartTs,
+ bufferManager.getManagerCostInfo());
+ }
+
+ private void checkAllBufferSpilled() {
+ if (bufferManager.getBuffers().size() > 0) {
+ throw new RssSendFailedException(
+ "Potential data loss due to existing remaining data buffers that are not flushed. This should not happen.");
+ }
+ }
+
+ private void checkSentRecordCount(long recordCount) {
+ if (recordCount != bufferManager.getRecordCount()) {
+ String message =
+ String.format(
+ "Inconsistent records number for taskId[%s]. expected: %d, actual: %d.",
+ taskId, recordCount, bufferManager.getRecordCount());
+ throw new RssSendFailedException(message);
+ }
+ }
+
+ private void checkSentBlockCount() {
+ long expected = blockIds.size();
+ long bufferManagerTracked = bufferManager.getBlockCount();
+
+ // to filter the multiple replica's duplicate blockIds
+ Roaring64NavigableMap blockIdBitmap = Roaring64NavigableMap.bitmapOf();
+ for (Map<Integer, BlockStats> partitionBlockStats :
+ shuffleTaskStats.getAllBlockStats().values()) {
+ partitionBlockStats.values().forEach(x -> x.getBlockIds().forEach(blockIdBitmap::addLong));
+ }
+ long serverTracked = blockIdBitmap.getLongCardinality();
+ if (expected != serverTracked || expected != bufferManagerTracked) {
+ throw new RssSendFailedException(
+ "Potential block loss may occur for task["
+ + taskId
+ + "]. BlockId number expected: "
+ + expected
+ + ", serverTracked: "
+ + serverTracked
+ + ", bufferManagerTracked: "
+ + bufferManagerTracked);
+ }
+ }
+
+ // Required by ShuffleWriter interface; returns empty since RSS does not support push-based shuffle.
+ public long[] getPartitionLengths() {
+ return new long[0];
+ }
+
+ @VisibleForTesting
+ protected List<CompletableFuture<Long>> processShuffleBlockInfos(
+ List<ShuffleBlockInfo> shuffleBlockInfoList) {
+ if (shuffleBlockInfoList != null && !shuffleBlockInfoList.isEmpty()) {
+ shuffleBlockInfoList.forEach(
+ sbi -> {
+ long blockId = sbi.getBlockId();
+ long recordNumber = sbi.getRecordNumber();
+ blockIds.add(blockId);
+ int partitionId = sbi.getPartitionId();
+ shuffleTaskStats.incPartitionBlock(partitionId);
+ sbi.getShuffleServerInfos()
+ .forEach(
+ s ->
+ shuffleTaskStats.mergeBlockStats(
+ s, partitionId, new BlockStats(recordNumber, blockId)));
+ });
+ return postBlockEvent(shuffleBlockInfoList);
+ }
+ return Collections.emptyList();
+ }
+
+ protected List<CompletableFuture<Long>> postBlockEvent(
+ List<ShuffleBlockInfo> shuffleBlockInfoList) {
+ List<CompletableFuture<Long>> futures = new ArrayList<>();
+ for (AddBlockEvent event : bufferManager.buildBlockEvents(shuffleBlockInfoList)) {
+ for (ShuffleBlockInfo block : event.getShuffleDataInfoList()) {
+ block.withCompletionCallback(
+ (b, isSuccessful) -> {
+ // If partition reassignment is enabled, the block is only released upon successful
+ // completion.
+ // Otherwise, the block is released immediately once completed.
+ if (!blockFailSentRetryEnabled || isSuccessful) {
+ bufferManager.releaseBlockResource(b);
+ shuffleTaskStats.incPartitionLength(b.getPartitionId(), b.getLength());
+ }
+ });
+ }
+ event.addCallback(
+ () -> {
+ if (!finishEventQueue.add(new Object())) {
+ LOG.error("Add event {} to finishEventQueue fail", event);
+ }
+ });
+ futures.add(shuffleManager.sendData(event));
+ }
+ return futures;
+ }
+
+ // Gluten needs this method
+ protected void internalCheckBlockSendResult() {
+ this.checkBlockSendResult(this.blockIds);
+ }
+
+ @VisibleForTesting
+ protected void checkBlockSendResult(Set<Long> blockIds) {
+ boolean interrupted = false;
+
+ try {
+ long remainingMs = sendCheckTimeout;
+ long end = System.currentTimeMillis() + remainingMs;
+ long currentAckValue = 0;
+ for (Long blockId : blockIds) {
+ currentAckValue ^= blockId;
+ }
+ while (true) {
+ try {
+ finishEventQueue.clear();
+ checkDataIfAnyFailure();
+ Set<Long> successBlockIds = shuffleManager.getSuccessBlockIds(taskId);
+ if (blockIds.size() == successBlockIds.size()) {
+ for (Long successBlockId : successBlockIds) {
+ currentAckValue ^= successBlockId;
+ }
+ if (currentAckValue != 0) {
+ String errorMsg = "Ack value is not equal to 0, it should not happen!";
+ throw new RssSendFailedException(errorMsg);
+ }
+ break;
+ }
+ if (finishEventQueue.isEmpty()) {
+ remainingMs = Math.max(end - System.currentTimeMillis(), 0);
+ Object event = finishEventQueue.poll(remainingMs, TimeUnit.MILLISECONDS);
+ if (event == null) {
+ break;
+ }
+ }
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ Set<Long> successBlockIds = shuffleManager.getSuccessBlockIds(taskId);
+ if (currentAckValue != 0 || blockIds.size() != successBlockIds.size()) {
+ int missing = blockIds.size() - successBlockIds.size();
+ int failed =
+ Optional.ofNullable(shuffleManager.getFailedBlockIds(taskId)).map(Set::size).orElse(0);
+ String message =
+ String.format(
+ "TaskId[%s] failed because %d blocks (failed: %d) can't be sent to shuffle server in %d ms",
+ taskId, missing, failed, sendCheckTimeout);
+
+ // detailed error message
+ Set<Long> missingBlockIds = new HashSet<>(blockIds);
+ missingBlockIds.removeAll(successBlockIds);
+ BlockIdLayout layout = BlockIdLayout.from(rssConf);
+ LOG.error(
+ "{}, includes partitions: {}",
+ message,
+ missingBlockIds.stream()
+ .map(x -> layout.getPartitionId(x))
+ .collect(Collectors.toSet()));
+
+ throw new RssWaitFailedException(message);
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // This method should remain protected so that Gluten can invoke it
+ protected void checkDataIfAnyFailure() {
+ if (blockFailSentRetryEnabled) {
+ reassignExecutor.reassign();
+ } else {
+ String errorMsg = getFirstBlockFailure();
+ if (errorMsg != null) {
+ throw new RssSendFailedException("Fail to send the block. Error: " + errorMsg);
+ }
+ }
+ }
+
+ private String getFirstBlockFailure() {
+ Set<Long> failedBlockIds = shuffleManager.getFailedBlockIds(taskId);
+ if (!failedBlockIds.isEmpty()) {
+ List<TrackingBlockStatus> trackingBlockStatues =
+ shuffleManager
+ .getBlockIdsFailedSendTracker(taskId)
+ .getFailedBlockStatus(failedBlockIds.iterator().next());
+ String errorMsg = DEFAULT_ERROR_MESSAGE;
+ if (CollectionUtils.isNotEmpty(trackingBlockStatues)) {
+ errorMsg = trackingBlockStatues.get(0).getStatusCode().name();
+ }
+ LOG.error(
+ "Errors on sending blocks for task[{}]. {} blocks can't be sent to remote servers: {}",
+ taskId,
+ failedBlockIds.size(),
+ shuffleManager.getBlockIdsFailedSendTracker(taskId).getFaultyShuffleServers());
+ return errorMsg;
+ }
+ return null;
+ }
+
+ private void clearFailedBlockState(ShuffleBlockInfo block) {
+ shuffleManager.getBlockIdsFailedSendTracker(taskId).remove(block.getBlockId());
+ shuffleTaskStats.decPartitionBlock(block.getPartitionId());
+ block.getShuffleServerInfos().stream()
+ .forEach(
+ s ->
+ shuffleTaskStats.removeBlockStats(
+ s,
+ block.getPartitionId(),
+ new BlockStats(block.getRecordNumber(), block.getBlockId())));
+ blockIds.remove(block.getBlockId());
+ }
+
+ @VisibleForTesting
+ protected void sendCommit() {
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ Future<Boolean> future =
+ executor.submit(
+ () -> shuffleWriteClient.sendCommit(shuffleServersForData, appId, shuffleId, numMaps));
+ int maxWait = 5000;
+ int currentWait = 200;
+ long start = System.currentTimeMillis();
+ while (!future.isDone()) {
+ LOG.info(
+ "Wait commit to shuffle server for task[{}] cost {} ms",
+ taskAttemptId,
+ System.currentTimeMillis() - start);
+ Uninterruptibles.sleepUninterruptibly(currentWait, TimeUnit.MILLISECONDS);
+ currentWait = Math.min(currentWait * 2, maxWait);
+ }
+ try {
+ if (!future.get()) {
+ throw new RssException("Failed to commit task to shuffle server");
+ }
+ } catch (InterruptedException ie) {
+ LOG.warn("Ignore the InterruptedException which should be caused by internal killed");
+ } catch (Exception e) {
+ throw new RssException("Exception happened when get commit status", e);
+ } finally {
+ executor.shutdown();
+ }
+ }
+
+ @VisibleForTesting
+ protected <T> int getPartition(T key) {
+ int result = 0;
+ if (shouldPartition) {
+ result = partitioner.getPartition(key);
+ }
+ return result;
+ }
+
+ @Override
+ public Option<MapStatus> stop(boolean success) {
+ try {
+ if (success) {
+ long start = System.currentTimeMillis();
+ shuffleWriteClient.reportShuffleResult(
+ getServerToPartitionToBlockIds(),
+ appId,
+ shuffleId,
+ taskAttemptId,
+ bitmapSplitNum,
+ recordReportFailedShuffleservers,
+ enableWriteFailureRetry,
+ isIntegrityValidationClientManagementEnabled
+ ? null
+ : getServerToPartitionToRecordNumbers());
+ long reportDuration = System.currentTimeMillis() - start;
+ LOG.info(
+ "Reported all shuffle result for shuffleId[{}] task[{}] with bitmapNum[{}] cost {} ms",
+ shuffleId,
+ taskAttemptId,
+ bitmapSplitNum,
+ reportDuration);
+ shuffleWriteMetrics.incWriteTime(TimeUnit.MILLISECONDS.toNanos(reportDuration));
+
+ if (RssShuffleManager.isIntegrationValidationFailureAnalysisEnabled(rssConf)) {
+ shuffleTaskStats.log();
+ }
+
+ shuffleTaskStats.check();
+
+ // TODO: Use preferred shuffle server host/port instead of dummy values to optimize read locality
+ final BlockManagerId blockManagerId =
+ BlockManagerId.apply(
+ appId + "_" + taskId,
+ DUMMY_HOST,
+ DUMMY_PORT,
+ Option.apply(
+ isIntegrityValidationClientManagementEnabled
+ ? shuffleTaskStats.encode()
+ : Long.toString(taskAttemptId)));
+ MapStatus mapStatus =
+ MapStatus.apply(blockManagerId, shuffleTaskStats.getPartitionLengths(), taskAttemptId);
+ return Option.apply(mapStatus);
+ } else {
+ return Option.empty();
+ }
+ } catch (Exception e) {
+ // If an exception is thrown during the reporting process, it should be judged as a failure
+ // and Stage retry should be triggered.
+ if (enableWriteFailureRetry) {
+ throw throwFetchFailedIfNecessary(e, recordReportFailedShuffleservers);
+ } else {
+ throw e;
+ }
+ } finally {
+ // report shuffle write metrics to driver
+ if (managerClientSupplier != null) {
+ ShuffleManagerClient shuffleManagerClient = managerClientSupplier.get();
+ if (shuffleManagerClient != null) {
+ RssReportShuffleWriteMetricRequest.TaskShuffleWriteTimes writeTimes =
+ new RssReportShuffleWriteMetricRequest.TaskShuffleWriteTimes(
+ totalShuffleWriteMills,
+ bufferManager.getCopyTime(),
+ bufferManager.getSerializeTime(),
+ bufferManager.getCompressTime(),
+ bufferManager.getSortTime(),
+ bufferManager.getRequireMemoryTime(),
+ checkSendResultMills);
+ try {
+ RssReportShuffleWriteMetricResponse response =
+ shuffleManagerClient.reportShuffleWriteMetric(
+ new RssReportShuffleWriteMetricRequest(
+ taskContext.stageId(),
+ shuffleId,
+ taskContext.taskAttemptId(),
+ bufferManager.getShuffleServerPushCostTracker().toMetric(),
+ writeTimes,
+ isShuffleWriteFailed,
+ shuffleWriteFailureReason,
+ bufferManager.getUncompressedDataLen(),
+ taskContext.attemptNumber()));
+ if (response.getStatusCode() != StatusCode.SUCCESS) {
+ LOG.error(
+ "Errors on reporting shuffle write metrics to driver. status_code: {}",
+ response.getStatusCode());
+ }
+ } catch (Exception e) {
+ LOG.error("Errors on reporting shuffle write metrics to driver", e);
+ }
+ }
+ }
+
+ if (blockFailSentRetryEnabled) {
+ if (success) {
+ if (CollectionUtils.isNotEmpty(shuffleManager.getFailedBlockIds(taskId))) {
+ LOG.error(
+ "Errors on stopping writer due to the remaining failed blockIds. This should not happen.");
+ return Option.empty();
+ }
+ } else {
+ shuffleManager.getBlockIdsFailedSendTracker(taskId).clearAndReleaseBlockResources();
+ }
+ }
+ // Release memory and metadata to prevent executor memory leaks
+ if (bufferManager != null) {
+ bufferManager.freeAllMemory();
+ bufferManager.close();
+ }
+ if (shuffleManager != null) {
+ shuffleManager.clearTaskMeta(taskId);
+ }
+ }
+ }
+
+ @VisibleForTesting
+ Map<Integer, Set<Long>> getPartitionToBlockIds() {
+ return getServerToPartitionToBlockIds().values().stream()
+ .flatMap(s -> s.entrySet().stream())
+ .collect(
+ Collectors.toMap(
+ Map.Entry::getKey,
+ Map.Entry::getValue,
+ (existingSet, newSet) -> {
+ Set<Long> mergedSet = new HashSet<>(existingSet);
+ mergedSet.addAll(newSet);
+ return mergedSet;
+ }));
+ }
+
+ @VisibleForTesting
+ public WriteBufferManager getBufferManager() {
+ return bufferManager;
+ }
+
+ private RssException throwFetchFailedIfNecessary(
+ Exception e, Set<ShuffleServerInfo> reportFailuredServers) {
+ // The shuffleServer is registered only when a Block fails to be sent
+ if (e instanceof RssSendFailedException) {
+ FailedBlockSendTracker blockIdsFailedSendTracker =
+ shuffleManager.getBlockIdsFailedSendTracker(taskId);
+ List<ShuffleServerInfo> shuffleServerInfos =
+ Lists.newArrayList(blockIdsFailedSendTracker.getFaultyShuffleServers());
+ shuffleServerInfos.addAll(reportFailuredServers);
+ RssReportShuffleWriteFailureRequest req =
+ new RssReportShuffleWriteFailureRequest(
+ appId,
+ shuffleId,
+ taskContext.stageId(),
+ taskContext.stageAttemptNumber(),
+ shuffleServerInfos,
+ e.getMessage());
+ RssReportShuffleWriteFailureResponse response =
+ managerClientSupplier.get().reportShuffleWriteFailure(req);
+ if (response.getReSubmitWholeStage()) {
+ LOG.warn(response.getMessage());
+ FetchFailedException ffe =
+ RssSparkShuffleUtils.createFetchFailedException(
+ shuffleId, -1, taskContext.stageAttemptNumber(), e);
+ throw new RssException(ffe);
+ }
+ }
+ throw new RssException(e);
+ }
+
+ @VisibleForTesting
+ protected void enableBlockFailSentRetry() {
+ this.blockFailSentRetryEnabled = true;
+ }
+
+ @VisibleForTesting
+ protected void resetBlockFailSentRetryMaxTimes(int times) {
+ reassignExecutor.resetBlockRetryMaxTimes(times);
+ }
+
+ @VisibleForTesting
+ protected void setTaskId(String taskId) {
+ this.taskId = taskId;
+ }
+
+ @VisibleForTesting
+ protected Map<ShuffleServerInfo, Map<Integer, Set<Long>>> getServerToPartitionToBlockIds() {
+ Map<ShuffleServerInfo, Map<Integer, Set<Long>>> serverToPartitionToBlockIds = new HashMap<>();
+ Map<ShuffleServerInfo, Map<Integer, BlockStats>> allBlockStats =
+ shuffleTaskStats.getAllBlockStats();
+ for (Map.Entry<ShuffleServerInfo, Map<Integer, BlockStats>> entry : allBlockStats.entrySet()) {
+ ShuffleServerInfo server = entry.getKey();
+ for (Map.Entry<Integer, BlockStats> childEntry : entry.getValue().entrySet()) {
+ int partitionId = childEntry.getKey();
+ BlockStats stats = childEntry.getValue();
+ serverToPartitionToBlockIds
+ .computeIfAbsent(server, k -> new HashMap<>())
+ .computeIfAbsent(partitionId, z -> new HashSet<>())
+ .addAll(stats.getBlockIds());
+ }
+ }
+ return serverToPartitionToBlockIds;
+ }
+
+ @VisibleForTesting
+ protected Map<ShuffleServerInfo, Map<Integer, Long>> getServerToPartitionToRecordNumbers() {
+ Map<ShuffleServerInfo, Map<Integer, Long>> serverToPartitionToRecordNumbers = new HashMap<>();
+ Map<ShuffleServerInfo, Map<Integer, BlockStats>> allBlockStats =
+ shuffleTaskStats.getAllBlockStats();
+ for (Map.Entry<ShuffleServerInfo, Map<Integer, BlockStats>> entry : allBlockStats.entrySet()) {
+ ShuffleServerInfo server = entry.getKey();
+ for (Map.Entry<Integer, BlockStats> childEntry : entry.getValue().entrySet()) {
+ int partitionId = childEntry.getKey();
+ BlockStats stats = childEntry.getValue();
+ serverToPartitionToRecordNumbers
+ .computeIfAbsent(server, k -> new HashMap<>())
+ .merge(partitionId, stats.getRecordNumber(), Long::sum);
+ }
+ }
+ return serverToPartitionToRecordNumbers;
+ }
+
+ @VisibleForTesting
+ protected RssShuffleManager getShuffleManager() {
+ return shuffleManager;
+ }
+
+ @VisibleForTesting
+ public ReassignExecutor getReassignExecutor() {
+ return reassignExecutor;
+ }
+
+ public TaskAttemptAssignment getTaskAttemptAssignment() {
+ return taskAttemptAssignment;
+ }
+}
diff --git a/client-spark/spark4/src/test/java/org/apache/spark/shuffle/DelegationRssShuffleManagerTest.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/DelegationRssShuffleManagerTest.java
new file mode 100644
index 000000000..a38bfad7b
--- /dev/null
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/DelegationRssShuffleManagerTest.java
@@ -0,0 +1,237 @@
+/*
+ * 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.spark.shuffle;
+
+import java.util.NoSuchElementException;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.shuffle.sort.SortShuffleManager;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+import org.apache.uniffle.client.api.CoordinatorClient;
+import org.apache.uniffle.client.request.RssAccessClusterRequest;
+import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.storage.util.StorageType;
+
+import static org.apache.uniffle.common.rpc.StatusCode.ACCESS_DENIED;
+import static org.apache.uniffle.common.rpc.StatusCode.SUCCESS;
+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.mockito.Mockito.verify;
+
+public class DelegationRssShuffleManagerTest extends RssShuffleManagerTestBase {
+
+ @Test
+ public void testCreateInDriverDenied() throws Exception {
+ setupMockedRssShuffleUtils(ACCESS_DENIED);
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ assertCreateSortShuffleManager(conf);
+ }
+
+ @Test
+ public void testCreateInDriver() throws Exception {
+ setupMockedRssShuffleUtils(SUCCESS);
+
+ SparkConf conf = new SparkConf();
+ assertCreateSortShuffleManager(conf);
+ conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set("spark.foo.bar.key", "mockId");
+ conf.set(RssSparkConfig.RSS_ACCESS_ID_PROVIDER_KEY.key(), "spark.foo.bar.key");
+ assertCreateSortShuffleManager(conf);
+
+ conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_ACCESS_ID.key(), "mockId");
+ assertCreateSortShuffleManager(conf);
+
+ conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_ACCESS_ID.key(), "mockId");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ assertCreateRssShuffleManager(conf);
+
+ conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ assertCreateSortShuffleManager(conf);
+ }
+
+ @Test
+ public void testCreateInExecutor() throws Exception {
+ DelegationRssShuffleManager delegationRssShuffleManager;
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ delegationRssShuffleManager = new DelegationRssShuffleManager(conf, false);
+ assertFalse(delegationRssShuffleManager.getDelegate() instanceof RssShuffleManager);
+ assertTrue(delegationRssShuffleManager.getDelegate() instanceof SortShuffleManager);
+ }
+
+ @Test
+ public void testCreateFallback() throws Exception {
+ setupMockedRssShuffleUtils(SUCCESS);
+
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_ACCESS_ID.key(), "mockId");
+ conf.set(RssSparkConfig.RSS_ENABLED.key(), "true");
+ conf.set(RssSparkConfig.RSS_STORAGE_TYPE.key(), "MEMORY_LOCALFILE");
+
+ // fall back to SortShuffleManager in driver
+ assertCreateSortShuffleManager(conf);
+
+ // No fall back in executor
+ conf.set(RssSparkConfig.RSS_ENABLED.key(), "true");
+ boolean hasException = false;
+ try {
+ new DelegationRssShuffleManager(conf, false);
+ } catch (NoSuchElementException e) {
+ assertTrue(e.getMessage().startsWith("spark.rss.coordinator.quorum"));
+ hasException = true;
+ }
+ assertTrue(hasException);
+ }
+
+ @Test
+ public void testTryAccessCluster() throws Exception {
+ setupMockedRssShuffleUtils(SUCCESS);
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_INTERVAL_MS, 3000L);
+ conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_TIMES, 3);
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_ACCESS_ID.key(), "mockId");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ assertCreateRssShuffleManager(conf);
+
+ setupMockedRssShuffleUtils(ACCESS_DENIED);
+ SparkConf secondConf = new SparkConf();
+ secondConf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_INTERVAL_MS, 3000L);
+ secondConf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_TIMES, 3);
+ secondConf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ secondConf.set(RssSparkConfig.RSS_ACCESS_ID.key(), "mockId");
+ secondConf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ secondConf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ assertCreateSortShuffleManager(secondConf);
+ }
+
+ @Test
+ public void testDefaultIncludeExcludeProperties() throws Exception {
+ final CoordinatorClient mockClient = setupMockedRssShuffleUtils(SUCCESS);
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_INTERVAL_MS, 3000L);
+ conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_TIMES, 3);
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_ACCESS_ID.key(), "mockId");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ final int confInitKeyCount = conf.getAll().length;
+ assertCreateRssShuffleManager(conf);
+
+ // default case: access cluster should include all properties in conf and an extra one.
+ ArgumentCaptor<RssAccessClusterRequest> argumentCaptor =
+ ArgumentCaptor.forClass(RssAccessClusterRequest.class);
+ verify(mockClient).accessCluster(argumentCaptor.capture());
+ RssAccessClusterRequest request = argumentCaptor.getValue();
+ assertEquals(confInitKeyCount + 1, request.getExtraProperties().size());
+ }
+
+ @Test
+ public void testIncludeProperties() throws Exception {
+ final CoordinatorClient mockClient = setupMockedRssShuffleUtils(SUCCESS);
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_INTERVAL_MS, 3000L);
+ conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_TIMES, 3);
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_ACCESS_ID.key(), "mockId");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ // test include properties
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX
+ + RssClientConf.RSS_CLIENT_REPORT_INCLUDE_PROPERTIES.key(),
+ RssSparkConfig.RSS_ACCESS_ID
+ .key()
+ .substring(RssSparkConfig.SPARK_RSS_CONFIG_PREFIX.length()));
+ assertCreateRssShuffleManager(conf);
+
+ ArgumentCaptor<RssAccessClusterRequest> argumentCaptor =
+ ArgumentCaptor.forClass(RssAccessClusterRequest.class);
+
+ verify(mockClient).accessCluster(argumentCaptor.capture());
+ RssAccessClusterRequest request = argumentCaptor.getValue();
+ // only accessId and extra one
+ assertEquals(1 + 1, request.getExtraProperties().size());
+ }
+
+ @Test
+ public void testExcludeProperties() throws Exception {
+ final CoordinatorClient mockClient = setupMockedRssShuffleUtils(SUCCESS);
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_INTERVAL_MS, 3000L);
+ conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_TIMES, 3);
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_ACCESS_ID.key(), "mockId");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ // test exclude properties
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX
+ + RssClientConf.RSS_CLIENT_REPORT_EXCLUDE_PROPERTIES.key(),
+ RssSparkConfig.RSS_ACCESS_ID
+ .key()
+ .substring(RssSparkConfig.SPARK_RSS_CONFIG_PREFIX.length()));
+ final int confInitKeyCount = conf.getAll().length;
+ assertCreateRssShuffleManager(conf);
+
+ ArgumentCaptor<RssAccessClusterRequest> argumentCaptor =
+ ArgumentCaptor.forClass(RssAccessClusterRequest.class);
+
+ verify(mockClient).accessCluster(argumentCaptor.capture());
+ RssAccessClusterRequest request = argumentCaptor.getValue();
+ // all accessId and extra one except the excluded one
+ assertEquals(confInitKeyCount + 1 - 1, request.getExtraProperties().size());
+ }
+
+ private void assertCreateSortShuffleManager(SparkConf conf) throws Exception {
+ DelegationRssShuffleManager delegationRssShuffleManager =
+ new DelegationRssShuffleManager(conf, true);
+ assertTrue(delegationRssShuffleManager.getDelegate() instanceof SortShuffleManager);
+ assertFalse(delegationRssShuffleManager.getDelegate() instanceof RssShuffleManager);
+ assertFalse(conf.getBoolean(RssSparkConfig.RSS_ENABLED.key(), false));
+ assertEquals("sort", conf.get("spark.shuffle.manager"));
+ }
+
+ private void assertCreateRssShuffleManager(SparkConf conf) throws Exception {
+ DelegationRssShuffleManager delegationRssShuffleManager =
+ new DelegationRssShuffleManager(conf, true);
+ assertFalse(delegationRssShuffleManager.getDelegate() instanceof SortShuffleManager);
+ assertTrue(delegationRssShuffleManager.getDelegate() instanceof RssShuffleManager);
+ assertTrue(Boolean.parseBoolean(conf.get(RssSparkConfig.RSS_ENABLED.key())));
+ assertEquals(RssShuffleManager.class.getCanonicalName(), conf.get("spark.shuffle.manager"));
+ }
+}
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/FunctionUtilsTests.java
similarity index 54%
copy from storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
copy to client-spark/spark4/src/test/java/org/apache/spark/shuffle/FunctionUtilsTests.java
index d84bdfead..7640f504e 100644
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/FunctionUtilsTests.java
@@ -15,30 +15,25 @@
* limitations under the License.
*/
-package org.apache.uniffle.storage.handler.impl;
+package org.apache.spark.shuffle;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
+import java.util.concurrent.atomic.AtomicInteger;
-import org.apache.uniffle.common.KerberizedHadoopBase;
+import scala.Function0;
-public class KerberizedHadoopClientReadHandlerTest extends KerberizedHadoopBase {
+import org.junit.jupiter.api.Test;
- @BeforeAll
- public static void beforeAll() throws Exception {
- testRunner = KerberizedHadoopClientReadHandlerTest.class;
- KerberizedHadoopBase.init();
- }
+import static org.junit.jupiter.api.Assertions.assertEquals;
- @BeforeEach
- public void beforeEach() throws Exception {
- initHadoopSecurityContext();
- }
+public class FunctionUtilsTests {
@Test
- public void test() throws Exception {
- HadoopClientReadHandlerTest.createAndRunCases(
- kerberizedHadoop.getSchemeAndAuthorityPrefix(), kerberizedHadoop.getConf(), "alex");
+ public void testOnceFunction0() {
+ AtomicInteger count = new AtomicInteger(0);
+ Function0<Integer> once = FunctionUtils.once(count::incrementAndGet);
+ once.apply();
+ once.apply();
+ once.apply();
+ assertEquals(1, count.get());
}
}
diff --git a/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTest.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTest.java
new file mode 100644
index 000000000..81b4676e8
--- /dev/null
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTest.java
@@ -0,0 +1,300 @@
+/*
+ * 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.spark.shuffle;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Supplier;
+
+import org.apache.spark.Partitioner;
+import org.apache.spark.ShuffleDependency;
+import org.apache.spark.SparkConf;
+import org.apache.spark.shuffle.handle.ShuffleHandleInfo;
+import org.apache.spark.shuffle.handle.SimpleShuffleHandleInfo;
+import org.apache.spark.sql.internal.SQLConf;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import org.apache.uniffle.client.impl.FailedBlockSendTracker;
+import org.apache.uniffle.client.util.RssClientConfig;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.rpc.StatusCode;
+import org.apache.uniffle.common.util.BlockIdLayout;
+import org.apache.uniffle.common.util.JavaUtils;
+import org.apache.uniffle.storage.util.StorageType;
+
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_SHUFFLE_MANAGER_GRPC_PORT;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrowsExactly;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class RssShuffleManagerTest extends RssShuffleManagerTestBase {
+ private static final String SPARK_ADAPTIVE_EXECUTION_ENABLED_KEY = "spark.sql.adaptive.enabled";
+
+ @Test
+ public void testGetDataDistributionType() {
+ // case1
+ SparkConf sparkConf = new SparkConf();
+ sparkConf.set(SPARK_ADAPTIVE_EXECUTION_ENABLED_KEY, "true");
+ assertEquals(
+ ShuffleDataDistributionType.LOCAL_ORDER,
+ RssShuffleManager.getDataDistributionType(sparkConf));
+
+ // case2
+ sparkConf = new SparkConf();
+ sparkConf.set(SPARK_ADAPTIVE_EXECUTION_ENABLED_KEY, "false");
+ assertEquals(
+ RssClientConf.DATA_DISTRIBUTION_TYPE.defaultValue(),
+ RssShuffleManager.getDataDistributionType(sparkConf));
+
+ // case3
+ sparkConf = new SparkConf();
+ sparkConf.set(SPARK_ADAPTIVE_EXECUTION_ENABLED_KEY, "true");
+ sparkConf.set(
+ "spark." + RssClientConf.DATA_DISTRIBUTION_TYPE.key(),
+ ShuffleDataDistributionType.NORMAL.name());
+ assertEquals(
+ ShuffleDataDistributionType.NORMAL, RssShuffleManager.getDataDistributionType(sparkConf));
+
+ // case4
+ sparkConf = new SparkConf();
+ sparkConf.set(SPARK_ADAPTIVE_EXECUTION_ENABLED_KEY, "true");
+ sparkConf.set(
+ "spark." + RssClientConf.DATA_DISTRIBUTION_TYPE.key(),
+ ShuffleDataDistributionType.LOCAL_ORDER.name());
+ assertEquals(
+ ShuffleDataDistributionType.LOCAL_ORDER,
+ RssShuffleManager.getDataDistributionType(sparkConf));
+
+ // case5
+ sparkConf = new SparkConf();
+ boolean aqeEnable = (boolean) sparkConf.get(SQLConf.ADAPTIVE_EXECUTION_ENABLED());
+ if (aqeEnable) {
+ assertEquals(
+ ShuffleDataDistributionType.LOCAL_ORDER,
+ RssShuffleManager.getDataDistributionType(sparkConf));
+ } else {
+ assertEquals(
+ RssClientConf.DATA_DISTRIBUTION_TYPE.defaultValue(),
+ RssShuffleManager.getDataDistributionType(sparkConf));
+ }
+ }
+
+ @Test
+ public void testGetRemoteStorageInfo() {
+ setupMockedRssShuffleUtils(StatusCode.SUCCESS);
+
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.driver.host", "localhost");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+
+ // inject some hadoop configs
+ conf.set("spark.rss.hadoop.k1", "v1");
+ conf.set("spark.rss.hadoop.k2", "v2");
+
+ RssShuffleManager shuffleManager = new RssShuffleManager(conf, true);
+ RemoteStorageInfo remoteStorageInfo = shuffleManager.getRemoteStorageInfo();
+ assertEquals(2, remoteStorageInfo.getConfItems().size());
+ }
+
+ @Test
+ public void testCreateShuffleManagerServer() {
+ setupMockedRssShuffleUtils(StatusCode.SUCCESS);
+
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.driver.host", "localhost");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ // enable stage recompute
+ conf.set("spark." + RssClientConfig.RSS_RESUBMIT_STAGE, "true");
+
+ RssShuffleManager shuffleManager = new RssShuffleManager(conf, true);
+
+ assertTrue(conf.get(RSS_SHUFFLE_MANAGER_GRPC_PORT) > 0);
+ }
+
+ @Test
+ public void testRssShuffleManagerInterface() throws Exception {
+ setupMockedRssShuffleUtils(StatusCode.SUCCESS);
+
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+
+ conf.set("spark.task.maxFailures", "3");
+ RssShuffleManager shuffleManager = new RssShuffleManager(conf, true);
+ assertEquals(shuffleManager.getMaxFetchFailures(), 2);
+ // by default, the appId is null
+ assertNull(shuffleManager.getAppId());
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {16, 20, 24})
+ public void testRssShuffleManagerRegisterShuffle(int partitionIdBits) {
+ BlockIdLayout layout =
+ BlockIdLayout.from(
+ 63 - partitionIdBits - partitionIdBits - 2, partitionIdBits, partitionIdBits + 2);
+
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ conf.set("spark.task.maxFailures", "4");
+
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX + RssClientConf.BLOCKID_SEQUENCE_NO_BITS.key(),
+ String.valueOf(layout.sequenceNoBits));
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX + RssClientConf.BLOCKID_PARTITION_ID_BITS.key(),
+ String.valueOf(layout.partitionIdBits));
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX + RssClientConf.BLOCKID_TASK_ATTEMPT_ID_BITS.key(),
+ String.valueOf(layout.taskAttemptIdBits));
+
+ // register a shuffle with too many partitions should fail
+ Partitioner mockPartitioner = mock(Partitioner.class);
+ when(mockPartitioner.numPartitions()).thenReturn(layout.maxNumPartitions + 1);
+ ShuffleDependency<String, String, String> mockDependency = mock(ShuffleDependency.class);
+ when(mockDependency.partitioner()).thenReturn(mockPartitioner);
+
+ RssShuffleManager shuffleManager = new RssShuffleManager(conf, true);
+
+ RssException e =
+ assertThrowsExactly(
+ RssException.class, () -> shuffleManager.registerShuffle(0, mockDependency));
+ assertEquals(
+ "Cannot register shuffle with "
+ + (layout.maxNumPartitions + 1)
+ + " partitions because the configured block id layout supports at most "
+ + layout.maxNumPartitions
+ + " partitions.",
+ e.getMessage());
+ }
+
+ @Test
+ public void testWithStageRetry() {
+ // case1: disable the stage retry
+ SparkConf conf = createSparkConf();
+ RssShuffleManager shuffleManager = new RssShuffleManager(conf, true);
+ assertFalse(shuffleManager.isRssStageRetryEnabled());
+ assertFalse(shuffleManager.isRssStageRetryForFetchFailureEnabled());
+ assertFalse(shuffleManager.isRssStageRetryForWriteFailureEnabled());
+ shuffleManager.stop();
+
+ // case2: enable the stage retry
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX + RssSparkConfig.RSS_RESUBMIT_STAGE_ENABLED.key(),
+ "true");
+ shuffleManager = new RssShuffleManager(conf, true);
+ assertTrue(shuffleManager.isRssStageRetryEnabled());
+ assertTrue(shuffleManager.isRssStageRetryForFetchFailureEnabled());
+ assertTrue(shuffleManager.isRssStageRetryForWriteFailureEnabled());
+ shuffleManager.stop();
+
+ // case3: overwrite the stage retry
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX
+ + RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED.key(),
+ "false");
+ shuffleManager = new RssShuffleManager(conf, true);
+ assertTrue(shuffleManager.isRssStageRetryEnabled());
+ assertFalse(shuffleManager.isRssStageRetryForFetchFailureEnabled());
+ assertTrue(shuffleManager.isRssStageRetryForWriteFailureEnabled());
+ shuffleManager.stop();
+
+ // case4: enable the partial stage retry of fetch failure
+ conf = createSparkConf();
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX
+ + RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED.key(),
+ "true");
+ shuffleManager = new RssShuffleManager(conf, true);
+ assertTrue(shuffleManager.isRssStageRetryEnabled());
+ assertTrue(shuffleManager.isRssStageRetryForFetchFailureEnabled());
+ assertFalse(shuffleManager.isRssStageRetryForWriteFailureEnabled());
+ shuffleManager.stop();
+ }
+
+ private SparkConf createSparkConf() {
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ conf.set("spark.task.maxFailures", "4");
+ conf.set("spark.driver.host", "localhost");
+ return conf;
+ }
+
+ @Test
+ public void testReadCacheShuffleInfo() {
+ SparkConf conf = new SparkConf();
+ conf.setAppName("testApp")
+ .setMaster("local[2]")
+ .set(RssSparkConfig.RSS_TEST_FLAG.key(), "true")
+ .set(RssSparkConfig.RSS_TEST_MODE_ENABLE.key(), "true")
+ .set(RssSparkConfig.RSS_CLIENT_SEND_CHECK_TIMEOUT_MS.key(), "10000")
+ .set(RssSparkConfig.RSS_CLIENT_RETRY_MAX.key(), "10")
+ .set(RssSparkConfig.RSS_CLIENT_SEND_CHECK_INTERVAL_MS.key(), "1000")
+ .set(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name())
+ .set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "127.0.0.1:12345,127.0.0.1:12346");
+ Map<String, Set<Long>> successBlocks = JavaUtils.newConcurrentMap();
+ Map<String, FailedBlockSendTracker> taskToFailedBlockSendTracker = JavaUtils.newConcurrentMap();
+ RssShuffleManager manager =
+ TestUtils.createShuffleManager(
+ conf, false, null, successBlocks, taskToFailedBlockSendTracker);
+
+ // case1: legal fetch and cache
+ Supplier<ShuffleHandleInfo> func1 =
+ () ->
+ new SimpleShuffleHandleInfo(
+ 1, Collections.emptyMap(), RemoteStorageInfo.EMPTY_REMOTE_STORAGE);
+ ShuffleHandleInfo handle1 = manager.getOrFetchShuffleHandle(1, func1);
+ ShuffleHandleInfo handle2 = manager.getOrFetchShuffleHandle(1, func1);
+ assertEquals(handle1, handle2);
+
+ // case2: illegal fetch
+ manager.clearShuffleHandleCache();
+ Supplier<ShuffleHandleInfo> func2 = () -> null;
+ try {
+ ShuffleHandleInfo handle3 = manager.getOrFetchShuffleHandle(1, func2);
+ fail();
+ } catch (Exception e) {
+ // ignore
+ }
+ }
+}
diff --git a/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTestBase.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTestBase.java
new file mode 100644
index 000000000..0cfcffbed
--- /dev/null
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTestBase.java
@@ -0,0 +1,70 @@
+/*
+ * 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.spark.shuffle;
+
+import java.util.List;
+
+import com.google.common.collect.Lists;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import org.apache.uniffle.client.api.CoordinatorClient;
+import org.apache.uniffle.client.impl.grpc.CoordinatorGrpcRetryableClient;
+import org.apache.uniffle.client.response.RssAccessClusterResponse;
+import org.apache.uniffle.common.rpc.StatusCode;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.when;
+
+public class RssShuffleManagerTestBase {
+ protected static MockedStatic<RssSparkShuffleUtils> mockedStaticRssShuffleUtils;
+
+ @BeforeAll
+ public static void setUp() {
+ mockedStaticRssShuffleUtils =
+ mockStatic(RssSparkShuffleUtils.class, Mockito.CALLS_REAL_METHODS);
+ }
+
+ @AfterAll
+ public static void tearDown() {
+ mockedStaticRssShuffleUtils.close();
+ }
+
+ protected CoordinatorClient createCoordinatorClient(StatusCode status) {
+ CoordinatorClient mockedCoordinatorClient = mock(CoordinatorClient.class);
+ when(mockedCoordinatorClient.accessCluster(any()))
+ .thenReturn(new RssAccessClusterResponse(status, ""));
+ return mockedCoordinatorClient;
+ }
+
+ CoordinatorClient setupMockedRssShuffleUtils(StatusCode status) {
+ CoordinatorClient mockCoordinatorClient = createCoordinatorClient(status);
+ List<CoordinatorClient> coordinatorClients = Lists.newArrayList();
+ coordinatorClients.add(mockCoordinatorClient);
+ CoordinatorGrpcRetryableClient client =
+ new CoordinatorGrpcRetryableClient(coordinatorClients, 0, 1, 1);
+ mockedStaticRssShuffleUtils
+ .when(() -> RssSparkShuffleUtils.createCoordinatorClientsForAccessCluster(any()))
+ .thenReturn(client);
+ return mockCoordinatorClient;
+ }
+}
diff --git a/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssSpark4ShuffleUtilsTest.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssSpark4ShuffleUtilsTest.java
new file mode 100644
index 000000000..73af4a5da
--- /dev/null
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/RssSpark4ShuffleUtilsTest.java
@@ -0,0 +1,62 @@
+/*
+ * 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.spark.shuffle;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.SparkContext;
+import org.junit.jupiter.api.Test;
+
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.exception.RssFetchFailedException;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class RssSpark4ShuffleUtilsTest {
+
+ @Test
+ public void testGetActiveSparkContext() {
+ assertThrows(RssException.class, RssSparkShuffleUtils::getActiveSparkContext);
+ SparkConf conf = new SparkConf();
+ conf.setMaster("local[1]");
+ conf.setAppName("test");
+ SparkContext sc = null;
+ try {
+ sc = SparkContext.getOrCreate(conf);
+ assertEquals(sc, RssSparkShuffleUtils.getActiveSparkContext());
+ } finally {
+ if (sc != null) {
+ sc.stop();
+ }
+ }
+ }
+
+ @Test
+ public void testCreateFetchFailedException() {
+ FetchFailedException ffe = RssSparkShuffleUtils.createFetchFailedException(0, -1, 10, null);
+ ffe =
+ RssSparkShuffleUtils.createFetchFailedException(
+ 0, -1, 100, new RssFetchFailedException("xx"));
+ }
+
+ @Test
+ public void testIsStageResubmitSupported() {
+ assertTrue(RssSparkShuffleUtils.isStageResubmitSupported());
+ }
+}
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java
similarity index 54%
copy from storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
copy to client-spark/spark4/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java
index d84bdfead..d563241bb 100644
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java
@@ -15,30 +15,24 @@
* limitations under the License.
*/
-package org.apache.uniffle.storage.handler.impl;
+package org.apache.spark.shuffle;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
-import org.apache.uniffle.common.KerberizedHadoopBase;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
-public class KerberizedHadoopClientReadHandlerTest extends KerberizedHadoopBase {
-
- @BeforeAll
- public static void beforeAll() throws Exception {
- testRunner = KerberizedHadoopClientReadHandlerTest.class;
- KerberizedHadoopBase.init();
- }
-
- @BeforeEach
- public void beforeEach() throws Exception {
- initHadoopSecurityContext();
+public class SparkVersionUtilsTest {
+ @Test
+ public void testSparkVersion() {
+ assertFalse(SparkVersionUtils.isSpark2());
+ assertFalse(SparkVersionUtils.isSpark3());
+ assertTrue(SparkVersionUtils.isSpark4());
}
@Test
- public void test() throws Exception {
- HadoopClientReadHandlerTest.createAndRunCases(
- kerberizedHadoop.getSchemeAndAuthorityPrefix(), kerberizedHadoop.getConf(), "alex");
+ public void testSpark4Version() {
+ assertTrue(Spark4VersionUtils.isSparkVersionAtLeast("3.5.0"));
+ assertTrue(Spark4VersionUtils.isSparkVersionAtLeast("4.0.0"));
}
}
diff --git a/client-spark/spark4/src/test/java/org/apache/spark/shuffle/TestUtils.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/TestUtils.java
new file mode 100644
index 000000000..cb2a7f909
--- /dev/null
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/TestUtils.java
@@ -0,0 +1,51 @@
+/*
+ * 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.spark.shuffle;
+
+import java.util.Map;
+import java.util.Set;
+
+import org.apache.commons.lang3.SystemUtils;
+import org.apache.spark.SparkConf;
+import org.apache.spark.shuffle.writer.DataPusher;
+
+import org.apache.uniffle.client.impl.FailedBlockSendTracker;
+import org.apache.uniffle.common.ShuffleBlockInfo;
+
+public class TestUtils {
+
+ private TestUtils() {}
+
+ public static RssShuffleManager createShuffleManager(
+ SparkConf conf,
+ Boolean isDriver,
+ DataPusher dataPusher,
+ Map<String, Set<Long>> successBlockIds,
+ Map<String, FailedBlockSendTracker> taskToFailedBlockSendTracker) {
+ return new RssShuffleManager(
+ conf, isDriver, dataPusher, successBlockIds, taskToFailedBlockSendTracker);
+ }
+
+ public static boolean isMacOnAppleSilicon() {
+ return SystemUtils.IS_OS_MAC_OSX && SystemUtils.OS_ARCH.equals("aarch64");
+ }
+
+ public static ShuffleBlockInfo createMockBlockOnlyBlockId(long blockId) {
+ return new ShuffleBlockInfo(1, 1, blockId, 1, 1, new byte[1], null, 1, 100, 1);
+ }
+}
diff --git a/client-spark/spark4/src/test/java/org/apache/spark/shuffle/reader/RssShuffleReaderTest.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/reader/RssShuffleReaderTest.java
new file mode 100644
index 000000000..bc77f7192
--- /dev/null
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/reader/RssShuffleReaderTest.java
@@ -0,0 +1,165 @@
+/*
+ * 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.spark.shuffle.reader;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import scala.Option;
+
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import org.apache.spark.ShuffleDependency;
+import org.apache.spark.SparkConf;
+import org.apache.spark.TaskContext;
+import org.apache.spark.executor.ShuffleReadMetrics;
+import org.apache.spark.executor.TaskMetrics;
+import org.apache.spark.serializer.KryoSerializer;
+import org.apache.spark.serializer.Serializer;
+import org.apache.spark.shuffle.RssShuffleHandle;
+import org.junit.jupiter.api.Test;
+import org.roaringbitmap.longlong.Roaring64NavigableMap;
+
+import org.apache.uniffle.client.api.ShuffleManagerClient;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.common.config.RssConf;
+import org.apache.uniffle.common.util.ExpiringCloseableSupplier;
+import org.apache.uniffle.storage.handler.impl.HadoopShuffleWriteHandler;
+import org.apache.uniffle.storage.util.StorageType;
+
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+public class RssShuffleReaderTest extends AbstractRssReaderTest {
+
+ private static final Serializer KRYO_SERIALIZER = new KryoSerializer(new SparkConf(false));
+
+ @Test
+ public void readTest() throws Exception {
+ ShuffleServerInfo ssi = new ShuffleServerInfo("127.0.0.1", 0);
+ String basePath = HDFS_URI + "readTest1";
+ HadoopShuffleWriteHandler writeHandler =
+ new HadoopShuffleWriteHandler("appId", 0, 0, 0, basePath, ssi.getId(), conf);
+ final HadoopShuffleWriteHandler writeHandler1 =
+ new HadoopShuffleWriteHandler("appId", 0, 1, 1, basePath, ssi.getId(), conf);
+
+ Roaring64NavigableMap blockIdBitmap = Roaring64NavigableMap.bitmapOf();
+ final Roaring64NavigableMap taskIdBitmap = Roaring64NavigableMap.bitmapOf(0);
+ Map<String, String> expectedData = Maps.newHashMap();
+ final Roaring64NavigableMap blockIdBitmap1 = Roaring64NavigableMap.bitmapOf();
+ writeTestData(writeHandler, 2, 5, expectedData, blockIdBitmap, "key", KRYO_SERIALIZER, 0);
+
+ RssShuffleHandle<String, String, String> handleMock = mock(RssShuffleHandle.class);
+ ShuffleDependency<String, String, String> dependencyMock = mock(ShuffleDependency.class);
+ when(handleMock.getAppId()).thenReturn("appId");
+ when(handleMock.getDependency()).thenReturn(dependencyMock);
+ when(handleMock.getShuffleId()).thenReturn(1);
+ when(handleMock.getNumMaps()).thenReturn(1);
+ Map<Integer, List<ShuffleServerInfo>> partitionToServers = new HashMap<>();
+ partitionToServers.put(0, Lists.newArrayList(ssi));
+ partitionToServers.put(1, Lists.newArrayList(ssi));
+ when(handleMock.getPartitionToServers()).thenReturn(partitionToServers);
+ when(dependencyMock.serializer()).thenReturn(KRYO_SERIALIZER);
+ TaskContext contextMock = mock(TaskContext.class);
+ when(contextMock.attemptNumber()).thenReturn(1);
+ when(contextMock.taskAttemptId()).thenReturn(1L);
+ when(contextMock.taskMetrics()).thenReturn(new TaskMetrics());
+ doNothing().when(contextMock).killTaskIfInterrupted();
+ when(dependencyMock.aggregator()).thenReturn(Option.empty());
+ when(dependencyMock.keyOrdering()).thenReturn(Option.empty());
+ when(dependencyMock.mapSideCombine()).thenReturn(false);
+
+ Map<Integer, Roaring64NavigableMap> partitionToExpectBlocks = Maps.newHashMap();
+ partitionToExpectBlocks.put(0, blockIdBitmap);
+ RssConf rssConf = new RssConf();
+ rssConf.set(RssClientConf.RSS_STORAGE_TYPE, StorageType.HDFS.name());
+ rssConf.set(RssClientConf.RSS_INDEX_READ_LIMIT, 1000);
+ rssConf.set(RssClientConf.RSS_CLIENT_READ_BUFFER_SIZE, "1000");
+ ShuffleManagerClient mockShuffleManagerClient = mock(ShuffleManagerClient.class);
+ RssShuffleReader<String, String> rssShuffleReaderSpy =
+ spy(
+ new RssShuffleReader<>(
+ 0,
+ 1,
+ 0,
+ Integer.MAX_VALUE,
+ contextMock,
+ handleMock,
+ basePath,
+ conf,
+ 1,
+ partitionToExpectBlocks,
+ taskIdBitmap,
+ new ShuffleReadMetrics(),
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ rssConf,
+ ShuffleDataDistributionType.NORMAL,
+ partitionToServers));
+ validateResult(rssShuffleReaderSpy.read(), expectedData, 10);
+
+ writeTestData(
+ writeHandler1, 2, 4, expectedData, blockIdBitmap1, "another_key", KRYO_SERIALIZER, 1);
+ partitionToExpectBlocks.put(1, blockIdBitmap1);
+ RssShuffleReader<String, String> rssShuffleReaderSpy1 =
+ spy(
+ new RssShuffleReader<>(
+ 0,
+ 2,
+ 0,
+ Integer.MAX_VALUE,
+ contextMock,
+ handleMock,
+ basePath,
+ conf,
+ 2,
+ partitionToExpectBlocks,
+ taskIdBitmap,
+ new ShuffleReadMetrics(),
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ rssConf,
+ ShuffleDataDistributionType.NORMAL,
+ partitionToServers));
+ validateResult(rssShuffleReaderSpy1.read(), expectedData, 18);
+
+ RssShuffleReader<String, String> rssShuffleReaderSpy2 =
+ spy(
+ new RssShuffleReader<>(
+ 0,
+ 2,
+ 0,
+ Integer.MAX_VALUE,
+ contextMock,
+ handleMock,
+ basePath,
+ conf,
+ 2,
+ partitionToExpectBlocks,
+ Roaring64NavigableMap.bitmapOf(),
+ new ShuffleReadMetrics(),
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ rssConf,
+ ShuffleDataDistributionType.NORMAL,
+ partitionToServers));
+ validateResult(rssShuffleReaderSpy2.read(), Maps.newHashMap(), 0);
+ }
+}
diff --git a/client-spark/spark4/src/test/java/org/apache/spark/shuffle/writer/PartitionLengthStatisticTest.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/writer/PartitionLengthStatisticTest.java
new file mode 100644
index 000000000..66f3d58aa
--- /dev/null
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/writer/PartitionLengthStatisticTest.java
@@ -0,0 +1,67 @@
+/*
+ * 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.spark.shuffle.writer;
+
+import io.netty.buffer.Unpooled;
+import org.junit.jupiter.api.Test;
+
+import org.apache.uniffle.common.ShuffleBlockInfo;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.fail;
+
+public class PartitionLengthStatisticTest {
+
+ @Test
+ public void test() {
+ int numPartitions = 10;
+ PartitionLengthStatistic statistic = new PartitionLengthStatistic(numPartitions);
+ // case1
+ assertEquals(numPartitions, statistic.toArray().length);
+ assertEquals(0, statistic.toArray()[0]);
+
+ // case2
+ ShuffleBlockInfo blockInfo =
+ new ShuffleBlockInfo(
+ 1, 1, 1, 100, 123, Unpooled.wrappedBuffer(new byte[100]).retain(), null, 5, 0, 1);
+ statistic.inc(blockInfo);
+ assertEquals(numPartitions, statistic.toArray().length);
+ assertEquals(0, statistic.toArray()[0]);
+ assertEquals(100, statistic.toArray()[1]);
+
+ // case3
+ blockInfo =
+ new ShuffleBlockInfo(
+ 1,
+ numPartitions,
+ 1,
+ 100,
+ 123,
+ Unpooled.wrappedBuffer(new byte[100]).retain(),
+ null,
+ 5,
+ 0,
+ 1);
+ try {
+ statistic.inc(blockInfo);
+ fail();
+ } catch (Exception e) {
+ // ignore
+ }
+ }
+}
diff --git a/client-spark/spark4/src/test/java/org/apache/spark/shuffle/writer/RssShuffleWriterTest.java b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/writer/RssShuffleWriterTest.java
new file mode 100644
index 000000000..dcabaa0a0
--- /dev/null
+++ b/client-spark/spark4/src/test/java/org/apache/spark/shuffle/writer/RssShuffleWriterTest.java
@@ -0,0 +1,1134 @@
+/*
+ * 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.spark.shuffle.writer;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import scala.Product2;
+import scala.Tuple2;
+import scala.collection.immutable.Nil;
+
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
+import org.apache.spark.Partitioner;
+import org.apache.spark.ShuffleDependency;
+import org.apache.spark.SparkConf;
+import org.apache.spark.TaskContext;
+import org.apache.spark.executor.ShuffleWriteMetrics;
+import org.apache.spark.executor.TaskMetrics;
+import org.apache.spark.memory.TaskMemoryManager;
+import org.apache.spark.serializer.KryoSerializer;
+import org.apache.spark.serializer.Serializer;
+import org.apache.spark.shuffle.RssShuffleHandle;
+import org.apache.spark.shuffle.RssShuffleManager;
+import org.apache.spark.shuffle.RssSparkConfig;
+import org.apache.spark.shuffle.TestUtils;
+import org.apache.spark.shuffle.handle.MutableShuffleHandleInfo;
+import org.apache.spark.shuffle.handle.SimpleShuffleHandleInfo;
+import org.awaitility.Awaitility;
+import org.junit.jupiter.api.Test;
+import org.roaringbitmap.longlong.Roaring64NavigableMap;
+
+import org.apache.uniffle.client.api.ShuffleManagerClient;
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+import org.apache.uniffle.client.impl.FailedBlockSendTracker;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleBlockInfo;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.rpc.StatusCode;
+import org.apache.uniffle.common.util.ExpiringCloseableSupplier;
+import org.apache.uniffle.common.util.JavaUtils;
+import org.apache.uniffle.storage.util.StorageType;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+public class RssShuffleWriterTest {
+
+ private scala.collection.immutable.List<Product2<String, String>> createMockRecords() {
+ return Nil.$colon$colon(new Tuple2<>("testKey2", "testValue2"))
+ .$colon$colon(new Tuple2<>("testKey3", "testValue3"))
+ .$colon$colon(new Tuple2<>("testKey4", "testValue4"))
+ .$colon$colon(new Tuple2<>("testKey6", "testValue6"))
+ .$colon$colon(new Tuple2<>("testKey1", "testValue1"))
+ .$colon$colon(new Tuple2<>("testKey5", "testValue5"));
+ }
+
+ private MutableShuffleHandleInfo createMutableShuffleHandle() {
+ Map<Integer, List<ShuffleServerInfo>> partitionToServers = Maps.newHashMap();
+ List<ShuffleServerInfo> ssi12 =
+ Arrays.asList(
+ new ShuffleServerInfo("id1", "0.0.0.1", 100),
+ new ShuffleServerInfo("id2", "0.0.0.2", 100));
+ partitionToServers.put(0, ssi12);
+ List<ShuffleServerInfo> ssi34 =
+ Arrays.asList(
+ new ShuffleServerInfo("id3", "0.0.0.3", 100),
+ new ShuffleServerInfo("id4", "0.0.0.4", 100));
+ partitionToServers.put(1, ssi34);
+ List<ShuffleServerInfo> ssi56 =
+ Arrays.asList(
+ new ShuffleServerInfo("id5", "0.0.0.5", 100),
+ new ShuffleServerInfo("id6", "0.0.0.6", 100));
+ partitionToServers.put(2, ssi56);
+
+ MutableShuffleHandleInfo shuffleHandleInfo =
+ new MutableShuffleHandleInfo(0, partitionToServers, RemoteStorageInfo.EMPTY_REMOTE_STORAGE);
+ return shuffleHandleInfo;
+ }
+
+ private RssShuffleWriter createMockWriter(MutableShuffleHandleInfo shuffleHandle, String taskId) {
+ SparkConf conf = new SparkConf();
+ conf.setAppName("testApp")
+ .setMaster("local[2]")
+ .set(RssSparkConfig.RSS_WRITER_SERIALIZER_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_TEST_FLAG.key(), "true")
+ .set(RssSparkConfig.RSS_TEST_MODE_ENABLE.key(), "true")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SEGMENT_SIZE.key(), "64")
+ .set(RssSparkConfig.RSS_CLIENT_SEND_CHECK_INTERVAL_MS.key(), "1000")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SPILL_SIZE.key(), "128")
+ .set(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name());
+
+ Map<String, Set<Long>> successBlockIds = JavaUtils.newConcurrentMap();
+ Map<String, FailedBlockSendTracker> taskToFailedBlockSendTracker = JavaUtils.newConcurrentMap();
+ taskToFailedBlockSendTracker.put(taskId, new FailedBlockSendTracker());
+
+ FakedDataPusher dataPusher = null;
+ final RssShuffleManager manager =
+ TestUtils.createShuffleManager(
+ conf, false, dataPusher, successBlockIds, taskToFailedBlockSendTracker);
+ Serializer kryoSerializer = new KryoSerializer(conf);
+ Partitioner mockPartitioner = mock(Partitioner.class);
+ final ShuffleWriteClient mockShuffleWriteClient = mock(ShuffleWriteClient.class);
+ final ShuffleManagerClient mockShuffleManagerClient = mock(ShuffleManagerClient.class);
+ ShuffleDependency<String, String, String> mockDependency = mock(ShuffleDependency.class);
+ RssShuffleHandle<String, String, String> mockHandle = mock(RssShuffleHandle.class);
+ when(mockHandle.getDependency()).thenReturn(mockDependency);
+ when(mockDependency.serializer()).thenReturn(kryoSerializer);
+ when(mockDependency.partitioner()).thenReturn(mockPartitioner);
+ when(mockPartitioner.numPartitions()).thenReturn(3);
+
+ when(mockPartitioner.getPartition("testKey1")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey2")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey4")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey5")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey3")).thenReturn(2);
+ when(mockPartitioner.getPartition("testKey7")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey8")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey9")).thenReturn(2);
+ when(mockPartitioner.getPartition("testKey6")).thenReturn(2);
+
+ TaskMemoryManager mockTaskMemoryManager = mock(TaskMemoryManager.class);
+
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
+ ShuffleWriteMetrics shuffleWriteMetrics = new ShuffleWriteMetrics();
+ WriteBufferManager bufferManager =
+ new WriteBufferManager(
+ 0,
+ 0,
+ bufferOptions,
+ kryoSerializer,
+ shuffleHandle.getAvailablePartitionServersForWriter(),
+ mockTaskMemoryManager,
+ shuffleWriteMetrics,
+ RssSparkConfig.toRssConf(conf));
+ bufferManager.setTaskId(taskId);
+
+ WriteBufferManager bufferManagerSpy = spy(bufferManager);
+ TaskContext contextMock = mock(TaskContext.class);
+ RssShuffleWriter<String, String, String> rssShuffleWriter =
+ new RssShuffleWriter<>(
+ "appId",
+ 0,
+ taskId,
+ 1L,
+ bufferManagerSpy,
+ shuffleWriteMetrics,
+ manager,
+ conf,
+ mockShuffleWriteClient,
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ mockHandle,
+ shuffleHandle,
+ contextMock);
+ rssShuffleWriter.enableBlockFailSentRetry();
+ doReturn(100000L).when(bufferManagerSpy).acquireMemory(anyLong());
+
+ RssShuffleWriter<String, String, String> rssShuffleWriterSpy = spy(rssShuffleWriter);
+ doNothing().when(rssShuffleWriterSpy).sendCommit();
+
+ return rssShuffleWriterSpy;
+ }
+
+ private void updateShuffleHandleAssignment(
+ MutableShuffleHandleInfo handle,
+ Set<Integer> partitionIds,
+ String receivingFailureServerId,
+ Set<ShuffleServerInfo> replacements) {
+ for (int partitionId : partitionIds) {
+ handle.updateAssignment(partitionId, receivingFailureServerId, replacements);
+ }
+ }
+
+ /** Test the reassign multi times for one partitionId. */
+ @Test
+ public void reassignMultiTimesForOnePartitionIdTest() {
+ String taskId = "taskId";
+ MutableShuffleHandleInfo shuffleHandle = createMutableShuffleHandle();
+ RssShuffleWriter writer = createMockWriter(shuffleHandle, taskId);
+ writer.resetBlockFailSentRetryMaxTimes(10);
+
+ // Make the id1 + id10 + id11 broken, and then finally, it will use the id12 successfully
+ AtomicInteger failureCnt = new AtomicInteger();
+ RssShuffleManager shuffleManager = writer.getShuffleManager();
+ Map<String, Set<Long>> taskToSuccessBlockIds = shuffleManager.getTaskToSuccessBlockIds();
+ Map<String, FailedBlockSendTracker> taskToFailedBlockSendTracker =
+ shuffleManager.getTaskToFailedBlockSendTracker();
+ TaskAttemptAssignment taskAssignment = writer.getTaskAttemptAssignment();
+ FakedDataPusher pusher =
+ new FakedDataPusher(
+ addBlockEvent -> {
+ List<ShuffleBlockInfo> blocks = addBlockEvent.getShuffleDataInfoList();
+ for (ShuffleBlockInfo block : blocks) {
+ ShuffleServerInfo server = block.getShuffleServerInfos().get(0);
+ String serverId = server.getId();
+ if (Arrays.asList("id1", "id10", "id11").contains(serverId)) {
+ taskToFailedBlockSendTracker
+ .computeIfAbsent(taskId, x -> new FailedBlockSendTracker())
+ .add(block, server, StatusCode.NO_BUFFER);
+ failureCnt.incrementAndGet();
+
+ // refresh the assignment to simulate the reassign rpc.
+ if (serverId.equals("id1")) {
+ ShuffleServerInfo replacement1 = new ShuffleServerInfo("id10", "0.0.0.10", 100);
+ updateShuffleHandleAssignment(
+ shuffleHandle,
+ Sets.newHashSet(0, 1, 2),
+ "id1",
+ Sets.newHashSet(replacement1));
+ taskAssignment.update(shuffleHandle);
+ } else if (serverId.equals("id10")) {
+ ShuffleServerInfo replacement2 = new ShuffleServerInfo("id11", "0.0.0.10", 100);
+ updateShuffleHandleAssignment(
+ shuffleHandle,
+ Sets.newHashSet(0, 1, 2),
+ "id10",
+ Sets.newHashSet(replacement2));
+ taskAssignment.update(shuffleHandle);
+ } else if (serverId.equals("id11")) {
+ ShuffleServerInfo replacement3 = new ShuffleServerInfo("id12", "0.0.0.10", 100);
+ updateShuffleHandleAssignment(
+ shuffleHandle,
+ Sets.newHashSet(0, 1, 2),
+ "id11",
+ Sets.newHashSet(replacement3));
+ taskAssignment.update(shuffleHandle);
+ }
+
+ } else {
+ taskToSuccessBlockIds
+ .computeIfAbsent(taskId, x -> new HashSet<>())
+ .add(block.getBlockId());
+ }
+ }
+ return new CompletableFuture<>();
+ });
+ shuffleManager.setDataPusher(pusher);
+
+ writer
+ .getBufferManager()
+ .setPartitionAssignmentRetrieveFunc(
+ partitionId -> writer.getPartitionAssignedServers(partitionId));
+
+ // case1: the reassignment will refresh the following plan. So the failure will only occur one
+ // time.
+ scala.collection.immutable.List<Product2<String, String>> mockedData = createMockRecords();
+ writer.write(mockedData.iterator());
+
+ Awaitility.await()
+ .timeout(Duration.ofSeconds(5))
+ .until(() -> taskToSuccessBlockIds.get(taskId).size() == mockedData.size());
+ assertEquals(3, failureCnt.get());
+ }
+
+ /** Once the reassignment occurs, the following AddBlockEvents will use the latest assignment. */
+ @Test
+ public void refreshAssignmentTest() {
+ String taskId = "taskId";
+ MutableShuffleHandleInfo shuffleHandle = createMutableShuffleHandle();
+ RssShuffleWriter writer = createMockWriter(shuffleHandle, taskId);
+
+ ShuffleServerInfo replacement = new ShuffleServerInfo("id10", "0.0.0.10", 100);
+ updateShuffleHandleAssignment(
+ shuffleHandle, Sets.newHashSet(0, 1, 2), "id1", Sets.newHashSet(replacement));
+
+ AtomicInteger failureCnt = new AtomicInteger();
+ RssShuffleManager shuffleManager = writer.getShuffleManager();
+ Map<String, Set<Long>> taskToSuccessBlockIds = shuffleManager.getTaskToSuccessBlockIds();
+ Map<String, FailedBlockSendTracker> taskToFailedBlockSendTracker =
+ shuffleManager.getTaskToFailedBlockSendTracker();
+ FakedDataPusher pusher =
+ new FakedDataPusher(
+ addBlockEvent -> {
+ List<ShuffleBlockInfo> blocks = addBlockEvent.getShuffleDataInfoList();
+ for (ShuffleBlockInfo block : blocks) {
+ ShuffleServerInfo server = block.getShuffleServerInfos().get(0);
+ if (server.getId().equals("id1")) {
+ taskToFailedBlockSendTracker
+ .computeIfAbsent(taskId, x -> new FailedBlockSendTracker())
+ .add(block, server, StatusCode.NO_BUFFER);
+ failureCnt.incrementAndGet();
+ // refresh the assignment to simulate the reassign rpc.
+ writer.getTaskAttemptAssignment().update(shuffleHandle);
+ } else {
+ taskToSuccessBlockIds
+ .computeIfAbsent(taskId, x -> new HashSet<>())
+ .add(block.getBlockId());
+ }
+ }
+ return new CompletableFuture<>();
+ });
+ shuffleManager.setDataPusher(pusher);
+
+ writer
+ .getBufferManager()
+ .setPartitionAssignmentRetrieveFunc(
+ partitionId -> writer.getPartitionAssignedServers(partitionId));
+
+ // case1: the reassignment will refresh the following plan. So the failure will only occur one
+ // time.
+ scala.collection.immutable.List<Product2<String, String>> mockedData = createMockRecords();
+ writer.write(mockedData.iterator());
+
+ Awaitility.await()
+ .timeout(Duration.ofSeconds(5))
+ .until(() -> taskToSuccessBlockIds.get(taskId).size() == mockedData.size());
+ assertEquals(1, failureCnt.get());
+ }
+
+ @Test
+ public void blockFailureResendTest() {
+ SparkConf conf = new SparkConf();
+ conf.setAppName("testApp")
+ .setMaster("local[2]")
+ .set(RssSparkConfig.RSS_WRITER_SERIALIZER_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_TEST_FLAG.key(), "true")
+ .set(RssSparkConfig.RSS_TEST_MODE_ENABLE.key(), "true")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SEGMENT_SIZE.key(), "64")
+ .set(RssSparkConfig.RSS_CLIENT_SEND_CHECK_INTERVAL_MS.key(), "1000")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SPILL_SIZE.key(), "128")
+ .set(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name());
+
+ List<ShuffleBlockInfo> shuffleBlockInfos = Lists.newArrayList();
+ Map<String, Set<Long>> successBlockIds = JavaUtils.newConcurrentMap();
+ Map<String, FailedBlockSendTracker> taskToFailedBlockSendTracker = JavaUtils.newConcurrentMap();
+ taskToFailedBlockSendTracker.put("taskId", new FailedBlockSendTracker());
+
+ AtomicInteger sentFailureCnt = new AtomicInteger();
+ FakedDataPusher dataPusher =
+ new FakedDataPusher(
+ event -> {
+ assertEquals("taskId", event.getTaskId());
+ FailedBlockSendTracker tracker = taskToFailedBlockSendTracker.get(event.getTaskId());
+ for (ShuffleBlockInfo block : event.getShuffleDataInfoList()) {
+ boolean isSuccessful = true;
+ ShuffleServerInfo shuffleServer = block.getShuffleServerInfos().get(0);
+ if (shuffleServer.getId().equals("id1") && block.getRetryCnt() == 0) {
+ tracker.add(block, shuffleServer, StatusCode.NO_BUFFER);
+ sentFailureCnt.addAndGet(1);
+ isSuccessful = false;
+ } else {
+ successBlockIds.putIfAbsent(event.getTaskId(), Sets.newConcurrentHashSet());
+ successBlockIds.get(event.getTaskId()).add(block.getBlockId());
+ shuffleBlockInfos.add(block);
+ }
+ block.executeCompletionCallback(isSuccessful);
+ }
+ return new CompletableFuture<>();
+ });
+
+ final RssShuffleManager manager =
+ TestUtils.createShuffleManager(
+ conf, false, dataPusher, successBlockIds, taskToFailedBlockSendTracker);
+ Serializer kryoSerializer = new KryoSerializer(conf);
+ Partitioner mockPartitioner = mock(Partitioner.class);
+ final ShuffleWriteClient mockShuffleWriteClient = mock(ShuffleWriteClient.class);
+ final ShuffleManagerClient mockShuffleManagerClient = mock(ShuffleManagerClient.class);
+ ShuffleDependency<String, String, String> mockDependency = mock(ShuffleDependency.class);
+ RssShuffleHandle<String, String, String> mockHandle = mock(RssShuffleHandle.class);
+ when(mockHandle.getDependency()).thenReturn(mockDependency);
+ when(mockDependency.serializer()).thenReturn(kryoSerializer);
+ when(mockDependency.partitioner()).thenReturn(mockPartitioner);
+ when(mockPartitioner.numPartitions()).thenReturn(3);
+
+ Map<Integer, List<ShuffleServerInfo>> partitionToServers = Maps.newHashMap();
+ List<ShuffleServerInfo> ssi12 =
+ Arrays.asList(
+ new ShuffleServerInfo("id1", "0.0.0.1", 100),
+ new ShuffleServerInfo("id2", "0.0.0.2", 100));
+ partitionToServers.put(0, ssi12);
+ List<ShuffleServerInfo> ssi34 =
+ Arrays.asList(
+ new ShuffleServerInfo("id3", "0.0.0.3", 100),
+ new ShuffleServerInfo("id4", "0.0.0.4", 100));
+ partitionToServers.put(1, ssi34);
+ List<ShuffleServerInfo> ssi56 =
+ Arrays.asList(
+ new ShuffleServerInfo("id5", "0.0.0.5", 100),
+ new ShuffleServerInfo("id6", "0.0.0.6", 100));
+ partitionToServers.put(2, ssi56);
+
+ when(mockPartitioner.getPartition("testKey1")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey2")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey4")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey5")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey3")).thenReturn(2);
+ when(mockPartitioner.getPartition("testKey7")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey8")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey9")).thenReturn(2);
+ when(mockPartitioner.getPartition("testKey6")).thenReturn(2);
+
+ TaskMemoryManager mockTaskMemoryManager = mock(TaskMemoryManager.class);
+
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
+ ShuffleWriteMetrics shuffleWriteMetrics = new ShuffleWriteMetrics();
+ WriteBufferManager bufferManager =
+ new WriteBufferManager(
+ 0,
+ 0,
+ bufferOptions,
+ kryoSerializer,
+ partitionToServers,
+ mockTaskMemoryManager,
+ shuffleWriteMetrics,
+ RssSparkConfig.toRssConf(conf));
+ bufferManager.setTaskId("taskId");
+
+ WriteBufferManager bufferManagerSpy = spy(bufferManager);
+ TaskContext contextMock = mock(TaskContext.class);
+ MutableShuffleHandleInfo shuffleHandleInfo =
+ new MutableShuffleHandleInfo(0, partitionToServers, RemoteStorageInfo.EMPTY_REMOTE_STORAGE);
+ RssShuffleWriter<String, String, String> rssShuffleWriter =
+ new RssShuffleWriter<>(
+ "appId",
+ 0,
+ "taskId",
+ 1L,
+ bufferManagerSpy,
+ shuffleWriteMetrics,
+ manager,
+ conf,
+ mockShuffleWriteClient,
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ mockHandle,
+ shuffleHandleInfo,
+ contextMock);
+ rssShuffleWriter.enableBlockFailSentRetry();
+ doReturn(100000L).when(bufferManagerSpy).acquireMemory(anyLong());
+
+ ShuffleServerInfo replacement = new ShuffleServerInfo("id10", "0.0.0.10", 100);
+ shuffleHandleInfo.updateAssignment(0, "id1", Sets.newHashSet(replacement));
+ shuffleHandleInfo.updateAssignment(1, "id1", Sets.newHashSet(replacement));
+ shuffleHandleInfo.updateAssignment(2, "id1", Sets.newHashSet(replacement));
+
+ rssShuffleWriter.getTaskAttemptAssignment().update(shuffleHandleInfo);
+
+ RssShuffleWriter<String, String, String> rssShuffleWriterSpy = spy(rssShuffleWriter);
+ doNothing().when(rssShuffleWriterSpy).sendCommit();
+
+ // case 1. failed blocks will be resent
+ scala.collection.immutable.List<Product2<String, String>> data = createMockRecords();
+ rssShuffleWriterSpy.write(data.iterator());
+
+ Awaitility.await()
+ .timeout(Duration.ofSeconds(5))
+ .until(() -> successBlockIds.get("taskId").size() == data.size());
+ assertEquals(2, sentFailureCnt.get());
+ assertEquals(0, taskToFailedBlockSendTracker.get("taskId").getFailedBlockIds().size());
+ assertEquals(6, shuffleWriteMetrics.recordsWritten());
+ assertEquals(
+ shuffleBlockInfos.stream().mapToInt(ShuffleBlockInfo::getLength).sum(),
+ shuffleWriteMetrics.bytesWritten());
+ assertEquals(6, shuffleBlockInfos.size());
+
+ assertEquals(0, bufferManagerSpy.getUsedBytes());
+ assertEquals(0, bufferManagerSpy.getInSendListBytes());
+
+ // check the blockId -> servers mapping.
+ // server -> partitionId -> blockIds
+ Map<ShuffleServerInfo, Map<Integer, Set<Long>>> serverToPartitionToBlockIds =
+ rssShuffleWriterSpy.getServerToPartitionToBlockIds();
+ assertEquals(2, serverToPartitionToBlockIds.get(replacement).get(0).size());
+
+ // case2. If exceeding the max retry times, it will fast fail.
+ String taskId = "t2";
+ rssShuffleWriter.getReassignExecutor().resetTaskId(taskId);
+ bufferManagerSpy.resetRecordCount();
+ rssShuffleWriter.resetBlockFailSentRetryMaxTimes(1);
+ rssShuffleWriter.setTaskId(taskId);
+ rssShuffleWriter.getBufferManager().setTaskId(taskId);
+ FailedBlockSendTracker tracker = new FailedBlockSendTracker();
+ taskToFailedBlockSendTracker.put(taskId, tracker);
+ FakedDataPusher alwaysFailedDataPusher =
+ new FakedDataPusher(
+ event -> {
+ assertEquals(taskId, event.getTaskId());
+ for (ShuffleBlockInfo block : event.getShuffleDataInfoList()) {
+ tracker.add(block, block.getShuffleServerInfos().get(0), StatusCode.NO_BUFFER);
+ block.executeCompletionCallback(false);
+ }
+ List<Runnable> callbackChain =
+ Optional.of(event.getProcessedCallbackChain()).orElse(Collections.EMPTY_LIST);
+ for (Runnable runnable : callbackChain) {
+ runnable.run();
+ }
+ return new CompletableFuture<>();
+ });
+ manager.setDataPusher(alwaysFailedDataPusher);
+
+ scala.collection.immutable.List<Product2<String, String>> mockedData = createMockRecords();
+
+ try {
+ rssShuffleWriter.write(mockedData.iterator());
+ fail();
+ } catch (Exception e) {
+ // ignore
+ }
+ }
+
+ @Test
+ public void checkBlockSendResultTest() {
+ SparkConf conf = new SparkConf();
+ conf.setAppName("testApp")
+ .setMaster("local[2]")
+ .set(RssSparkConfig.RSS_TEST_FLAG.key(), "true")
+ .set(RssSparkConfig.RSS_TEST_MODE_ENABLE.key(), "true")
+ .set(RssSparkConfig.RSS_CLIENT_SEND_CHECK_TIMEOUT_MS.key(), "10000")
+ .set(RssSparkConfig.RSS_CLIENT_RETRY_MAX.key(), "10")
+ .set(RssSparkConfig.RSS_CLIENT_SEND_CHECK_INTERVAL_MS.key(), "1000")
+ .set(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name())
+ .set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "127.0.0.1:12345,127.0.0.1:12346");
+ Map<String, Set<Long>> successBlocks = JavaUtils.newConcurrentMap();
+ Map<String, FailedBlockSendTracker> taskToFailedBlockSendTracker = JavaUtils.newConcurrentMap();
+ Serializer kryoSerializer = new KryoSerializer(conf);
+ RssShuffleManager manager =
+ TestUtils.createShuffleManager(
+ conf, false, null, successBlocks, taskToFailedBlockSendTracker);
+
+ ShuffleWriteClient mockShuffleWriteClient = mock(ShuffleWriteClient.class);
+ ShuffleManagerClient mockShuffleManagerClient = mock(ShuffleManagerClient.class);
+ Partitioner mockPartitioner = mock(Partitioner.class);
+ RssShuffleHandle<String, String, String> mockHandle = mock(RssShuffleHandle.class);
+ ShuffleDependency<String, String, String> mockDependency = mock(ShuffleDependency.class);
+ when(mockHandle.getDependency()).thenReturn(mockDependency);
+ when(mockPartitioner.numPartitions()).thenReturn(2);
+ TaskMemoryManager mockTaskMemoryManager = mock(TaskMemoryManager.class);
+ when(mockHandle.getPartitionToServers()).thenReturn(Maps.newHashMap());
+ when(mockDependency.partitioner()).thenReturn(mockPartitioner);
+
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
+ WriteBufferManager bufferManager =
+ new WriteBufferManager(
+ 0,
+ 0,
+ bufferOptions,
+ kryoSerializer,
+ Maps.newHashMap(),
+ mockTaskMemoryManager,
+ new ShuffleWriteMetrics(),
+ RssSparkConfig.toRssConf(conf));
+ WriteBufferManager bufferManagerSpy = spy(bufferManager);
+
+ TaskContext contextMock = mock(TaskContext.class);
+ SimpleShuffleHandleInfo mockShuffleHandleInfo = mock(SimpleShuffleHandleInfo.class);
+ RssShuffleWriter<String, String, String> rssShuffleWriter =
+ new RssShuffleWriter<>(
+ "appId",
+ 0,
+ "taskId",
+ 1L,
+ bufferManagerSpy,
+ (new TaskMetrics()).shuffleWriteMetrics(),
+ manager,
+ conf,
+ mockShuffleWriteClient,
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ mockHandle,
+ mockShuffleHandleInfo,
+ contextMock);
+ doReturn(1000000L).when(bufferManagerSpy).acquireMemory(anyLong());
+
+ // case 1: all blocks are sent successfully
+ successBlocks.put("taskId", Sets.newHashSet(1L, 2L, 3L));
+ rssShuffleWriter.checkBlockSendResult(Sets.newHashSet(1L, 2L, 3L));
+ successBlocks.clear();
+
+ // case 2: partial blocks aren't sent before spark.rss.client.send.check.timeout.ms,
+ // Runtime exception will be thrown
+ successBlocks.put("taskId", Sets.newHashSet(1L, 2L));
+ Throwable e2 =
+ assertThrows(
+ RuntimeException.class,
+ () -> rssShuffleWriter.checkBlockSendResult(Sets.newHashSet(1L, 2L, 3L)));
+ assertTrue(e2.getMessage().contains("failed because"));
+ successBlocks.clear();
+
+ // case 3: partial blocks are sent failed, Runtime exception will be thrown
+ successBlocks.put("taskId", Sets.newHashSet(1L, 2L));
+ FailedBlockSendTracker failedBlockSendTracker = new FailedBlockSendTracker();
+ taskToFailedBlockSendTracker.put("taskId", failedBlockSendTracker);
+ ShuffleServerInfo shuffleServerInfo = new ShuffleServerInfo("127.0.0.1", 20001);
+ failedBlockSendTracker.add(
+ TestUtils.createMockBlockOnlyBlockId(3L), shuffleServerInfo, StatusCode.INTERNAL_ERROR);
+ Throwable e3 =
+ assertThrows(
+ RuntimeException.class,
+ () -> rssShuffleWriter.checkBlockSendResult(Sets.newHashSet(1L, 2L, 3L)));
+ assertTrue(e3.getMessage().startsWith("Fail to send the block"));
+ successBlocks.clear();
+ taskToFailedBlockSendTracker.clear();
+ }
+
+ static class FakedDataPusher extends DataPusher {
+ private final Function<AddBlockEvent, CompletableFuture<Long>> sendFunc;
+
+ FakedDataPusher(Function<AddBlockEvent, CompletableFuture<Long>> sendFunc) {
+ this(null, null, null, null, null, 1, 1, sendFunc);
+ }
+
+ private FakedDataPusher(
+ ShuffleWriteClient shuffleWriteClient,
+ Map<String, Set<Long>> taskToSuccessBlockIds,
+ Map<String, Set<Long>> taskToFailedBlockIds,
+ Map<String, FailedBlockSendTracker> failedBlockSendTracker,
+ Set<String> failedTaskIds,
+ int threadPoolSize,
+ int threadKeepAliveTime,
+ Function<AddBlockEvent, CompletableFuture<Long>> sendFunc) {
+ super(
+ shuffleWriteClient,
+ taskToSuccessBlockIds,
+ failedBlockSendTracker,
+ failedTaskIds,
+ threadPoolSize,
+ threadKeepAliveTime);
+ this.sendFunc = sendFunc;
+ }
+
+ @Override
+ public CompletableFuture<Long> send(AddBlockEvent event) {
+ return sendFunc.apply(event);
+ }
+ }
+
+ @Test
+ public void dataConsistencyWhenSpillTriggeredTest() throws Exception {
+ SparkConf conf = new SparkConf();
+ conf.set("spark.rss.client.memory.spill.enabled", "true");
+ conf.setAppName("dataConsistencyWhenSpillTriggeredTest_app")
+ .setMaster("local[2]")
+ .set(RssSparkConfig.RSS_WRITER_SERIALIZER_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_WRITER_PRE_ALLOCATED_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_TEST_FLAG.key(), "true")
+ .set(RssSparkConfig.RSS_TEST_MODE_ENABLE.key(), "true")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SEGMENT_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_CLIENT_SEND_CHECK_INTERVAL_MS.key(), "1000")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SPILL_SIZE.key(), "100000")
+ .set(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.MEMORY.name())
+ .set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "127.0.0.1:12345,127.0.0.1:12346");
+
+ Map<String, Set<Long>> successBlockIds = Maps.newConcurrentMap();
+
+ List<Long> freeMemoryList = new ArrayList<>();
+ FakedDataPusher dataPusher =
+ new FakedDataPusher(
+ event -> {
+ event.getProcessedCallbackChain().stream().forEach(x -> x.run());
+ long sum =
+ event.getShuffleDataInfoList().stream().mapToLong(x -> x.getFreeMemory()).sum();
+ freeMemoryList.add(sum);
+ successBlockIds.putIfAbsent(event.getTaskId(), new HashSet<>());
+ successBlockIds
+ .get(event.getTaskId())
+ .add(event.getShuffleDataInfoList().get(0).getBlockId());
+ return CompletableFuture.completedFuture(sum);
+ });
+
+ final RssShuffleManager manager =
+ TestUtils.createShuffleManager(
+ conf, false, dataPusher, successBlockIds, JavaUtils.newConcurrentMap());
+
+ WriteBufferManagerTest.FakedTaskMemoryManager fakedTaskMemoryManager =
+ new WriteBufferManagerTest.FakedTaskMemoryManager();
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
+ Map<Integer, List<ShuffleServerInfo>> partitionToServers = Maps.newConcurrentMap();
+ partitionToServers.put(0, Lists.newArrayList(new ShuffleServerInfo("127.0.0.1", 1111)));
+ WriteBufferManager bufferManager =
+ new WriteBufferManager(
+ 0,
+ "taskId",
+ 0,
+ bufferOptions,
+ new KryoSerializer(conf),
+ partitionToServers,
+ fakedTaskMemoryManager,
+ new ShuffleWriteMetrics(),
+ RssSparkConfig.toRssConf(conf),
+ null,
+ 0);
+
+ Serializer kryoSerializer = new KryoSerializer(conf);
+ Partitioner mockPartitioner = mock(Partitioner.class);
+ final ShuffleWriteClient mockShuffleWriteClient = mock(ShuffleWriteClient.class);
+ final ShuffleManagerClient mockShuffleManagerClient = mock(ShuffleManagerClient.class);
+ ShuffleDependency<String, String, String> mockDependency = mock(ShuffleDependency.class);
+ RssShuffleHandle<String, String, String> mockHandle = mock(RssShuffleHandle.class);
+ when(mockHandle.getDependency()).thenReturn(mockDependency);
+ when(mockDependency.serializer()).thenReturn(kryoSerializer);
+ when(mockDependency.partitioner()).thenReturn(mockPartitioner);
+ when(mockPartitioner.numPartitions()).thenReturn(1);
+ TaskContext contextMock = mock(TaskContext.class);
+ SimpleShuffleHandleInfo mockShuffleHandleInfo = mock(SimpleShuffleHandleInfo.class);
+
+ RssShuffleWriter<String, String, String> rssShuffleWriter =
+ new RssShuffleWriter<>(
+ "appId",
+ 0,
+ "taskId",
+ 1L,
+ bufferManager,
+ new ShuffleWriteMetrics(),
+ manager,
+ conf,
+ mockShuffleWriteClient,
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ mockHandle,
+ mockShuffleHandleInfo,
+ contextMock);
+ rssShuffleWriter.getBufferManager().setSpillFunc(rssShuffleWriter::processShuffleBlockInfos);
+
+ scala.collection.immutable.List<Product2<String, String>> data =
+ Nil
+ // One record is 26 bytes
+ .$colon$colon(new Tuple2<>("Key", "Value11111111111111"))
+ .$colon$colon(new Tuple2<>("Key", "Value11111111111111"))
+ .$colon$colon(new Tuple2<>("Key", "Value11111111111111"))
+ .$colon$colon(new Tuple2<>("Key", "Value11111111111111"));
+
+ // case1: all blocks are sent and pass the blocks check when spill is triggered
+ rssShuffleWriter.write(data.iterator());
+ assertEquals(4, successBlockIds.get("taskId").size());
+ for (int i = 0; i < 4; i++) {
+ assertEquals(32, freeMemoryList.get(i));
+ }
+ }
+
+ @Test
+ public void writeTest() throws Exception {
+ SparkConf conf = new SparkConf();
+ conf.setAppName("testApp")
+ .setMaster("local[2]")
+ .set(RssSparkConfig.RSS_WRITER_SERIALIZER_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SIZE.key(), "32")
+ .set(RssSparkConfig.RSS_TEST_FLAG.key(), "true")
+ .set(RssSparkConfig.RSS_TEST_MODE_ENABLE.key(), "true")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SEGMENT_SIZE.key(), "64")
+ .set(RssSparkConfig.RSS_CLIENT_SEND_CHECK_INTERVAL_MS.key(), "1000")
+ .set(RssSparkConfig.RSS_WRITER_BUFFER_SPILL_SIZE.key(), "128")
+ .set(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.LOCALFILE.name())
+ .set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "127.0.0.1:12345,127.0.0.1:12346");
+ List<ShuffleBlockInfo> shuffleBlockInfos = Lists.newArrayList();
+ Map<String, Set<Long>> successBlockIds = Maps.newConcurrentMap();
+
+ FakedDataPusher dataPusher =
+ new FakedDataPusher(
+ event -> {
+ assertEquals("taskId", event.getTaskId());
+ shuffleBlockInfos.addAll(event.getShuffleDataInfoList());
+ Set<Long> blockIds =
+ event
+ .getShuffleDataInfoList()
+ .parallelStream()
+ .map(sdi -> sdi.getBlockId())
+ .collect(Collectors.toSet());
+ successBlockIds.putIfAbsent(event.getTaskId(), Sets.newConcurrentHashSet());
+ successBlockIds.get(event.getTaskId()).addAll(blockIds);
+ return new CompletableFuture<>();
+ });
+
+ final RssShuffleManager manager =
+ TestUtils.createShuffleManager(
+ conf, false, dataPusher, successBlockIds, JavaUtils.newConcurrentMap());
+ Serializer kryoSerializer = new KryoSerializer(conf);
+ Partitioner mockPartitioner = mock(Partitioner.class);
+ final ShuffleWriteClient mockShuffleWriteClient = mock(ShuffleWriteClient.class);
+ final ShuffleManagerClient mockShuffleManagerClient = mock(ShuffleManagerClient.class);
+ ShuffleDependency<String, String, String> mockDependency = mock(ShuffleDependency.class);
+ RssShuffleHandle<String, String, String> mockHandle = mock(RssShuffleHandle.class);
+ when(mockHandle.getDependency()).thenReturn(mockDependency);
+ when(mockDependency.serializer()).thenReturn(kryoSerializer);
+ when(mockDependency.partitioner()).thenReturn(mockPartitioner);
+ when(mockPartitioner.numPartitions()).thenReturn(3);
+
+ Map<Integer, List<ShuffleServerInfo>> partitionToServers = Maps.newHashMap();
+ List<ShuffleServerInfo> ssi34 =
+ Arrays.asList(
+ new ShuffleServerInfo("id3", "0.0.0.3", 100),
+ new ShuffleServerInfo("id4", "0.0.0.4", 100));
+ partitionToServers.put(1, ssi34);
+ List<ShuffleServerInfo> ssi56 =
+ Arrays.asList(
+ new ShuffleServerInfo("id5", "0.0.0.5", 100),
+ new ShuffleServerInfo("id6", "0.0.0.6", 100));
+ partitionToServers.put(2, ssi56);
+ List<ShuffleServerInfo> ssi12 =
+ Arrays.asList(
+ new ShuffleServerInfo("id1", "0.0.0.1", 100),
+ new ShuffleServerInfo("id2", "0.0.0.2", 100));
+ partitionToServers.put(0, ssi12);
+ when(mockPartitioner.getPartition("testKey1")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey2")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey4")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey5")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey3")).thenReturn(2);
+ when(mockPartitioner.getPartition("testKey7")).thenReturn(0);
+ when(mockPartitioner.getPartition("testKey8")).thenReturn(1);
+ when(mockPartitioner.getPartition("testKey9")).thenReturn(2);
+ when(mockPartitioner.getPartition("testKey6")).thenReturn(2);
+
+ TaskMemoryManager mockTaskMemoryManager = mock(TaskMemoryManager.class);
+
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
+ ShuffleWriteMetrics shuffleWriteMetrics = new ShuffleWriteMetrics();
+ WriteBufferManager bufferManager =
+ new WriteBufferManager(
+ 0,
+ 0,
+ bufferOptions,
+ kryoSerializer,
+ partitionToServers,
+ mockTaskMemoryManager,
+ shuffleWriteMetrics,
+ RssSparkConfig.toRssConf(conf));
+ bufferManager.setTaskId("taskId");
+
+ WriteBufferManager bufferManagerSpy = spy(bufferManager);
+ TaskContext contextMock = mock(TaskContext.class);
+ SimpleShuffleHandleInfo mockShuffleHandleInfo = mock(SimpleShuffleHandleInfo.class);
+ RssShuffleWriter<String, String, String> rssShuffleWriter =
+ new RssShuffleWriter<>(
+ "appId",
+ 0,
+ "taskId",
+ 1L,
+ bufferManagerSpy,
+ shuffleWriteMetrics,
+ manager,
+ conf,
+ mockShuffleWriteClient,
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ mockHandle,
+ mockShuffleHandleInfo,
+ contextMock);
+ doReturn(1000000L).when(bufferManagerSpy).acquireMemory(anyLong());
+
+ RssShuffleWriter<String, String, String> rssShuffleWriterSpy = spy(rssShuffleWriter);
+ doNothing().when(rssShuffleWriterSpy).sendCommit();
+
+ // case 1
+ scala.collection.immutable.List<Product2<String, String>> data =
+ Nil.$colon$colon(new Tuple2<>("testKey2", "testValue2"))
+ .$colon$colon(new Tuple2<>("testKey3", "testValue3"))
+ .$colon$colon(new Tuple2<>("testKey4", "testValue4"))
+ .$colon$colon(new Tuple2<>("testKey6", "testValue6"))
+ .$colon$colon(new Tuple2<>("testKey1", "testValue1"))
+ .$colon$colon(new Tuple2<>("testKey5", "testValue5"));
+ rssShuffleWriterSpy.write(data.iterator());
+
+ assertTrue(shuffleWriteMetrics.writeTime() > 0);
+ assertEquals(6, shuffleWriteMetrics.recordsWritten());
+
+ assertEquals(
+ shuffleBlockInfos.stream().mapToInt(ShuffleBlockInfo::getLength).sum(),
+ shuffleWriteMetrics.bytesWritten());
+
+ assertEquals(6, shuffleBlockInfos.size());
+ for (ShuffleBlockInfo shuffleBlockInfo : shuffleBlockInfos) {
+ assertEquals(22, shuffleBlockInfo.getUncompressLength());
+ assertEquals(0, shuffleBlockInfo.getShuffleId());
+ if (shuffleBlockInfo.getPartitionId() == 0) {
+ assertEquals(shuffleBlockInfo.getShuffleServerInfos(), ssi12);
+ }
+ if (shuffleBlockInfo.getPartitionId() == 1) {
+ assertEquals(shuffleBlockInfo.getShuffleServerInfos(), ssi34);
+ }
+ if (shuffleBlockInfo.getPartitionId() == 2) {
+ assertEquals(shuffleBlockInfo.getShuffleServerInfos(), ssi56);
+ }
+ if (shuffleBlockInfo.getPartitionId() < 0 || shuffleBlockInfo.getPartitionId() > 2) {
+ throw new Exception("Shouldn't be here");
+ }
+ }
+ Map<Integer, Set<Long>> partitionToBlockIds = rssShuffleWriterSpy.getPartitionToBlockIds();
+ assertEquals(2, partitionToBlockIds.get(1).size());
+ assertEquals(2, partitionToBlockIds.get(0).size());
+ assertEquals(2, partitionToBlockIds.get(2).size());
+ partitionToBlockIds.clear();
+ }
+
+ @Test
+ public void postBlockEventTest() throws Exception {
+ SparkConf conf = new SparkConf();
+ conf.set(
+ RssSparkConfig.SPARK_RSS_CONFIG_PREFIX
+ + RssSparkConfig.RSS_CLIENT_SEND_SIZE_LIMITATION.key(),
+ "64")
+ .set(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.MEMORY_LOCALFILE.name());
+
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
+ WriteBufferManager bufferManager =
+ new WriteBufferManager(
+ 0,
+ 0,
+ bufferOptions,
+ new KryoSerializer(conf),
+ Maps.newHashMap(),
+ mock(TaskMemoryManager.class),
+ new ShuffleWriteMetrics(),
+ RssSparkConfig.toRssConf(conf));
+ WriteBufferManager bufferManagerSpy = spy(bufferManager);
+
+ ShuffleDependency<String, String, String> mockDependency = mock(ShuffleDependency.class);
+ ShuffleWriteMetrics mockMetrics = mock(ShuffleWriteMetrics.class);
+ Partitioner mockPartitioner = mock(Partitioner.class);
+ when(mockDependency.partitioner()).thenReturn(mockPartitioner);
+ SparkConf sparkConf = new SparkConf();
+ when(mockPartitioner.numPartitions()).thenReturn(2);
+ List<AddBlockEvent> events = Lists.newArrayList();
+
+ FakedDataPusher dataPusher =
+ new FakedDataPusher(
+ event -> {
+ events.add(event);
+ return new CompletableFuture<>();
+ });
+
+ RssShuffleManager mockShuffleManager =
+ spy(
+ TestUtils.createShuffleManager(
+ sparkConf,
+ false,
+ dataPusher,
+ Maps.newConcurrentMap(),
+ JavaUtils.newConcurrentMap()));
+
+ RssShuffleHandle<String, String, String> mockHandle = mock(RssShuffleHandle.class);
+ when(mockHandle.getDependency()).thenReturn(mockDependency);
+ TaskContext contextMock = mock(TaskContext.class);
+ SimpleShuffleHandleInfo mockShuffleHandleInfo = mock(SimpleShuffleHandleInfo.class);
+ ShuffleWriteClient mockWriteClient = mock(ShuffleWriteClient.class);
+ ShuffleManagerClient mockShuffleManagerClient = mock(ShuffleManagerClient.class);
+
+ List<ShuffleBlockInfo> shuffleBlockInfoList = createShuffleBlockList(1, 31);
+ RssShuffleWriter<String, String, String> writer =
+ new RssShuffleWriter<>(
+ "appId",
+ 0,
+ "taskId",
+ 1L,
+ bufferManagerSpy,
+ mockMetrics,
+ mockShuffleManager,
+ conf,
+ mockWriteClient,
+ ExpiringCloseableSupplier.of(() -> mockShuffleManagerClient),
+ mockHandle,
+ mockShuffleHandleInfo,
+ contextMock);
+ writer.postBlockEvent(shuffleBlockInfoList);
+ Awaitility.await().timeout(Duration.ofSeconds(1)).until(() -> events.size() == 1);
+ assertEquals(1, events.get(0).getShuffleDataInfoList().size());
+ events.clear();
+
+ testSingleEvent(events, writer, 2, 15);
+
+ testSingleEvent(events, writer, 1, 33);
+
+ testSingleEvent(events, writer, 2, 16);
+
+ testSingleEvent(events, writer, 2, 15);
+
+ testSingleEvent(events, writer, 2, 17);
+
+ testSingleEvent(events, writer, 2, 32);
+
+ testTwoEvents(events, writer, 2, 33, 1, 1);
+
+ testTwoEvents(events, writer, 3, 17, 2, 1);
+ }
+
+ private void testTwoEvents(
+ List<AddBlockEvent> events,
+ RssShuffleWriter<String, String, String> writer,
+ int blockNum,
+ int blockLength,
+ int firstEventSize,
+ int secondEventSize)
+ throws InterruptedException {
+ List<ShuffleBlockInfo> shuffleBlockInfoList;
+ shuffleBlockInfoList = createShuffleBlockList(blockNum, blockLength);
+ writer.postBlockEvent(shuffleBlockInfoList);
+ Thread.sleep(500);
+ assertEquals(2, events.size());
+ assertEquals(firstEventSize, events.get(0).getShuffleDataInfoList().size());
+ assertEquals(secondEventSize, events.get(1).getShuffleDataInfoList().size());
+ events.clear();
+ }
+
+ private void testSingleEvent(
+ List<AddBlockEvent> events,
+ RssShuffleWriter<String, String, String> writer,
+ int blockNum,
+ int blockLength)
+ throws InterruptedException {
+ List<ShuffleBlockInfo> shuffleBlockInfoList;
+ shuffleBlockInfoList = createShuffleBlockList(blockNum, blockLength);
+ writer.postBlockEvent(shuffleBlockInfoList);
+ Thread.sleep(500);
+ assertEquals(1, events.size());
+ assertEquals(blockNum, events.get(0).getShuffleDataInfoList().size());
+ events.clear();
+ }
+
+ private List<ShuffleBlockInfo> createShuffleBlockList(int blockNum, int blockLength) {
+ List<ShuffleServerInfo> shuffleServerInfoList =
+ Lists.newArrayList(new ShuffleServerInfo("id", "host", 0));
+ List<ShuffleBlockInfo> shuffleBlockInfoList = Lists.newArrayList();
+ for (int i = 0; i < blockNum; i++) {
+ shuffleBlockInfoList.add(
+ new ShuffleBlockInfo(
+ 0,
+ 0,
+ 10,
+ blockLength,
+ 10,
+ new byte[] {1},
+ shuffleServerInfoList,
+ blockLength,
+ 10,
+ 0));
+ }
+ return shuffleBlockInfoList;
+ }
+
+ /**
+ * Test that Roaring64NavigableMap correctly filters duplicate blockIds from multiple replicas.
+ * This tests the optimization introduced in issue #2675.
+ */
+ @Test
+ public void testRoaring64NavigableMapDeduplication() {
+ // Simulate serverToPartitionToBlockIds with duplicate blockIds across multiple servers
+ // (representing replicas)
+ Map<ShuffleServerInfo, Map<Integer, Set<Long>>> serverToPartitionToBlockIds = Maps.newHashMap();
+
+ final ShuffleServerInfo server1 = new ShuffleServerInfo("id1", "host1", 100);
+ final ShuffleServerInfo server2 = new ShuffleServerInfo("id2", "host2", 100);
+ final ShuffleServerInfo server3 = new ShuffleServerInfo("id3", "host3", 100);
+
+ // Server1 has blockIds: 1, 2, 3 for partition 0
+ Map<Integer, Set<Long>> server1Partitions = Maps.newHashMap();
+ server1Partitions.put(0, Sets.newHashSet(1L, 2L, 3L));
+ serverToPartitionToBlockIds.put(server1, server1Partitions);
+
+ // Server2 has same blockIds (replica): 1, 2, 3 for partition 0
+ Map<Integer, Set<Long>> server2Partitions = Maps.newHashMap();
+ server2Partitions.put(0, Sets.newHashSet(1L, 2L, 3L));
+ serverToPartitionToBlockIds.put(server2, server2Partitions);
+
+ // Server3 has blockIds: 4, 5 for partition 1
+ Map<Integer, Set<Long>> server3Partitions = Maps.newHashMap();
+ server3Partitions.put(1, Sets.newHashSet(4L, 5L));
+ serverToPartitionToBlockIds.put(server3, server3Partitions);
+
+ // Using Roaring64NavigableMap to filter duplicates (as implemented in checkSentBlockCount)
+ Roaring64NavigableMap blockIdBitmap = Roaring64NavigableMap.bitmapOf();
+ for (Map<Integer, Set<Long>> partitionBlockIds : serverToPartitionToBlockIds.values()) {
+ partitionBlockIds.values().forEach(x -> x.forEach(blockIdBitmap::addLong));
+ }
+
+ // Expected unique blockIds: 1, 2, 3, 4, 5 = 5 unique blocks
+ // Without deduplication: 1,2,3 + 1,2,3 + 4,5 = 8 blocks
+ assertEquals(5, blockIdBitmap.getLongCardinality());
+
+ // Verify the bitmap contains all expected blockIds
+ assertTrue(blockIdBitmap.contains(1L));
+ assertTrue(blockIdBitmap.contains(2L));
+ assertTrue(blockIdBitmap.contains(3L));
+ assertTrue(blockIdBitmap.contains(4L));
+ assertTrue(blockIdBitmap.contains(5L));
+ }
+
+ /**
+ * Test Roaring64NavigableMap with large number of consecutive blockIds. This verifies the memory
+ * efficiency benefit of using Roaring64NavigableMap.
+ */
+ @Test
+ public void testRoaring64NavigableMapWithLargeBlockIds() {
+ Roaring64NavigableMap blockIdBitmap = Roaring64NavigableMap.bitmapOf();
+
+ // Add a large range of consecutive blockIds (simulating real shuffle scenario)
+ int numBlocks = 100000;
+ for (long i = 0; i < numBlocks; i++) {
+ blockIdBitmap.addLong(i);
+ }
+
+ assertEquals(numBlocks, blockIdBitmap.getLongCardinality());
+
+ // Add same blockIds again (simulating replicas)
+ for (long i = 0; i < numBlocks; i++) {
+ blockIdBitmap.addLong(i);
+ }
+
+ // Count should remain the same after adding duplicates
+ assertEquals(numBlocks, blockIdBitmap.getLongCardinality());
+ }
+}
diff --git a/common/src/main/java/org/apache/uniffle/common/web/JettyServer.java b/common/src/main/java/org/apache/uniffle/common/web/JettyServer.java
index ac357a6cb..45ca74896 100644
--- a/common/src/main/java/org/apache/uniffle/common/web/JettyServer.java
+++ b/common/src/main/java/org/apache/uniffle/common/web/JettyServer.java
@@ -175,13 +175,27 @@ public class JettyServer {
try {
server.start();
httpPort = ((ServerConnector) server.getConnectors()[0]).getLocalPort();
- } catch (BindException e) {
- ExitUtils.terminate(1, "Fail to start jetty http server, port is " + httpPort, e, LOG);
+ } catch (Exception e) {
+ if (isBindException(e)) {
+ ExitUtils.terminate(1, "Fail to start jetty http server, port is " + httpPort, e, LOG);
+ } else {
+ throw e;
+ }
}
LOG.info("Jetty http server started, listening on port {}", httpPort);
return httpPort;
}
+ private static boolean isBindException(Throwable t) {
+ while (t != null) {
+ if (t instanceof BindException) {
+ return true;
+ }
+ t = t.getCause();
+ }
+ return false;
+ }
+
public void stop() throws Exception {
server.stop();
}
diff --git a/common/src/test/java/org/apache/uniffle/common/KerberizedHadoop.java b/common/src/test/java/org/apache/uniffle/common/KerberizedHadoop.java
index 4b57248c1..27836eaf5 100644
--- a/common/src/test/java/org/apache/uniffle/common/KerberizedHadoop.java
+++ b/common/src/test/java/org/apache/uniffle/common/KerberizedHadoop.java
@@ -24,6 +24,7 @@ import java.io.OutputStreamWriter;
import java.io.Serializable;
import java.lang.reflect.Method;
import java.net.BindException;
+import java.net.InetAddress;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
@@ -347,6 +348,16 @@ public class KerberizedHadoop implements Serializable {
LOGGER.info("Proxy: {}", superUser.getShortUserName());
}
+ /**
+ * Added in Hadoop 3.4 as a new abstract method in ImpersonationProvider. This override is
+ * needed for compilation with Hadoop 3.4+ while remaining compatible with older versions.
+ */
+ public void authorize(UserGroupInformation userGroupInformation, InetAddress address)
+ throws AuthorizationException {
+ UserGroupInformation superUser = userGroupInformation.getRealUser();
+ LOGGER.info("Proxy: {}", superUser.getShortUserName());
+ }
+
@Override
public void setConf(Configuration conf) {
// ignore
diff --git a/common/src/test/java/org/apache/uniffle/common/filesystem/HadoopFilesystemProviderTest.java b/common/src/test/java/org/apache/uniffle/common/filesystem/HadoopFilesystemProviderTest.java
index 9689629a6..bbed61b6a 100644
--- a/common/src/test/java/org/apache/uniffle/common/filesystem/HadoopFilesystemProviderTest.java
+++ b/common/src/test/java/org/apache/uniffle/common/filesystem/HadoopFilesystemProviderTest.java
@@ -30,12 +30,17 @@ import org.apache.hadoop.security.AccessControlException;
import org.apache.hadoop.security.UserGroupInformation;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.apache.uniffle.common.KerberizedHadoopBase;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class HadoopFilesystemProviderTest extends KerberizedHadoopBase {
@BeforeAll
diff --git a/common/src/test/java/org/apache/uniffle/common/security/HadoopSecurityContextTest.java b/common/src/test/java/org/apache/uniffle/common/security/HadoopSecurityContextTest.java
index a0aae828a..f6e06926a 100644
--- a/common/src/test/java/org/apache/uniffle/common/security/HadoopSecurityContextTest.java
+++ b/common/src/test/java/org/apache/uniffle/common/security/HadoopSecurityContextTest.java
@@ -27,6 +27,7 @@ import org.apache.hadoop.fs.Path;
import org.apache.hadoop.security.UserGroupInformation;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.apache.uniffle.common.KerberizedHadoopBase;
@@ -34,6 +35,10 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class HadoopSecurityContextTest extends KerberizedHadoopBase {
@BeforeAll
diff --git a/common/src/test/java/org/apache/uniffle/common/security/SecurityContextFactoryTest.java b/common/src/test/java/org/apache/uniffle/common/security/SecurityContextFactoryTest.java
index 7b6116f9f..61c1c4d2c 100644
--- a/common/src/test/java/org/apache/uniffle/common/security/SecurityContextFactoryTest.java
+++ b/common/src/test/java/org/apache/uniffle/common/security/SecurityContextFactoryTest.java
@@ -20,12 +20,17 @@ package org.apache.uniffle.common.security;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.apache.uniffle.common.KerberizedHadoopBase;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.fail;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class SecurityContextFactoryTest extends KerberizedHadoopBase {
@BeforeAll
diff --git a/integration-test/common/pom.xml b/integration-test/common/pom.xml
index eef340e48..9af6514b2 100644
--- a/integration-test/common/pom.xml
+++ b/integration-test/common/pom.xml
@@ -177,7 +177,7 @@
</systemPropertyVariables>
<redirectTestOutputToFile>${test.redirectToFile}</redirectTestOutputToFile>
<useFile>${test.redirectToFile}</useFile>
- <argLine>-ea -Xmx5g -Xss4m -XX:MaxMetaspaceSize=2g -XX:ReservedCodeCacheSize=1g</argLine>
+ <argLine>-ea -Xmx5g -Xss4m -XX:MaxMetaspaceSize=2g -XX:ReservedCodeCacheSize=1g ${extraJavaTestArgs}</argLine>
<failIfNoTests>false</failIfNoTests>
</configuration>
</plugin>
diff --git a/integration-test/common/src/test/java/org/apache/uniffle/test/AccessCandidatesCheckerKerberizedHadoopTest.java b/integration-test/common/src/test/java/org/apache/uniffle/test/AccessCandidatesCheckerKerberizedHadoopTest.java
index baa6b1d24..cacc36237 100644
--- a/integration-test/common/src/test/java/org/apache/uniffle/test/AccessCandidatesCheckerKerberizedHadoopTest.java
+++ b/integration-test/common/src/test/java/org/apache/uniffle/test/AccessCandidatesCheckerKerberizedHadoopTest.java
@@ -21,10 +21,15 @@ import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.apache.uniffle.common.KerberizedHadoopBase;
import org.apache.uniffle.coordinator.metric.CoordinatorMetrics;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class AccessCandidatesCheckerKerberizedHadoopTest extends KerberizedHadoopBase {
@BeforeAll
diff --git a/integration-test/common/src/test/java/org/apache/uniffle/test/DynamicClientConfServiceKerberlizedHadoopTest.java b/integration-test/common/src/test/java/org/apache/uniffle/test/DynamicClientConfServiceKerberlizedHadoopTest.java
index c4c2b575f..00de1e3b8 100644
--- a/integration-test/common/src/test/java/org/apache/uniffle/test/DynamicClientConfServiceKerberlizedHadoopTest.java
+++ b/integration-test/common/src/test/java/org/apache/uniffle/test/DynamicClientConfServiceKerberlizedHadoopTest.java
@@ -19,9 +19,14 @@ package org.apache.uniffle.test;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.apache.uniffle.common.KerberizedHadoopBase;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class DynamicClientConfServiceKerberlizedHadoopTest extends KerberizedHadoopBase {
@BeforeAll
diff --git a/integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleServerWithKerberizedHadoopTest.java b/integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleServerWithKerberizedHadoopTest.java
index 3a305d2bd..4f272de76 100644
--- a/integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleServerWithKerberizedHadoopTest.java
+++ b/integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleServerWithKerberizedHadoopTest.java
@@ -30,6 +30,7 @@ import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.junit.jupiter.api.io.TempDir;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
@@ -69,6 +70,10 @@ import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class ShuffleServerWithKerberizedHadoopTest extends KerberizedHadoopBase {
protected static final String LOCALHOST;
diff --git a/integration-test/spark-common/pom.xml b/integration-test/spark-common/pom.xml
index e5a81a8ad..028e59b61 100644
--- a/integration-test/spark-common/pom.xml
+++ b/integration-test/spark-common/pom.xml
@@ -194,7 +194,7 @@
</systemPropertyVariables>
<redirectTestOutputToFile>${test.redirectToFile}</redirectTestOutputToFile>
<useFile>${test.redirectToFile}</useFile>
- <argLine>-ea -Xmx5g -Xss4m -XX:MaxMetaspaceSize=2g -XX:ReservedCodeCacheSize=1g</argLine>
+ <argLine>-ea -Xmx5g -Xss4m -XX:MaxMetaspaceSize=2g -XX:ReservedCodeCacheSize=1g ${extraJavaTestArgs}</argLine>
<failIfNoTests>false</failIfNoTests>
</configuration>
</plugin>
diff --git a/integration-test/common/pom.xml b/integration-test/spark4/pom.xml
similarity index 54%
copy from integration-test/common/pom.xml
copy to integration-test/spark4/pom.xml
index eef340e48..801c1a1c4 100644
--- a/integration-test/common/pom.xml
+++ b/integration-test/spark4/pom.xml
@@ -19,40 +19,59 @@
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
- <modelVersion>4.0.0</modelVersion>
<parent>
- <groupId>org.apache.uniffle</groupId>
<artifactId>uniffle-parent</artifactId>
+ <groupId>org.apache.uniffle</groupId>
<version>0.11.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
+ <modelVersion>4.0.0</modelVersion>
- <groupId>org.apache.uniffle</groupId>
- <artifactId>rss-integration-common-test</artifactId>
- <version>0.11.0-SNAPSHOT</version>
+ <artifactId>rss-integration-spark4-test</artifactId>
<packaging>jar</packaging>
- <name>Apache Uniffle Integration Test (Common)</name>
+ <name>Apache Uniffle Integration Test (Spark 4)</name>
<dependencies>
+
<dependency>
<groupId>org.apache.uniffle</groupId>
- <artifactId>rss-internal-client</artifactId>
+ <artifactId>rss-integration-spark-common-test</artifactId>
+ <version>${project.version}</version>
<scope>test</scope>
+ <type>test-jar</type>
</dependency>
<dependency>
<groupId>org.apache.uniffle</groupId>
- <artifactId>rss-common</artifactId>
- <type>test-jar</type>
+ <artifactId>rss-client-spark4</artifactId>
+ <version>${project.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.uniffle</groupId>
- <artifactId>shuffle-server</artifactId>
+ <artifactId>rss-client-spark-common</artifactId>
+ <version>${project.version}</version>
+ <scope>test</scope>
+ <exclusions>
+ <exclusion>
+ <groupId>org.eclipse.jetty.orbit</groupId>
+ <artifactId>javax.servlet</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-internal-client</artifactId>
<scope>test</scope>
+ <exclusions>
+ <exclusion>
+ <groupId>org.eclipse.jetty.orbit</groupId>
+ <artifactId>javax.servlet</artifactId>
+ </exclusion>
+ </exclusions>
</dependency>
<dependency>
<groupId>org.apache.uniffle</groupId>
- <artifactId>cli</artifactId>
+ <artifactId>shuffle-server</artifactId>
<scope>test</scope>
</dependency>
<dependency>
@@ -71,26 +90,21 @@
<artifactId>rss-client</artifactId>
<scope>test</scope>
</dependency>
+
<dependency>
<groupId>org.apache.uniffle</groupId>
- <artifactId>shuffle-storage</artifactId>
+ <artifactId>rss-integration-common-test</artifactId>
+ <version>${project.version}</version>
<type>test-jar</type>
<scope>test</scope>
</dependency>
+
<dependency>
- <groupId>org.apache.uniffle</groupId>
- <artifactId>rss-client</artifactId>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scala.binary.version}</artifactId>
+ <version>${spark.version}</version>
<scope>test</scope>
- <type>test-jar</type>
- </dependency>
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-minicluster</artifactId>
<exclusions>
- <exclusion>
- <groupId>com.sun.jersey</groupId>
- <artifactId>jersey-json</artifactId>
- </exclusion>
<exclusion>
<groupId>org.slf4j</groupId>
<artifactId>slf4j-log4j12</artifactId>
@@ -98,8 +112,22 @@
</exclusions>
</dependency>
<dependency>
- <groupId>com.google.guava</groupId>
- <artifactId>guava</artifactId>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>shuffle-storage</artifactId>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-databind</artifactId>
+ <version>${jackson.version}</version>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-core</artifactId>
+ <version>${jackson.version}</version>
<scope>test</scope>
</dependency>
<dependency>
@@ -109,12 +137,8 @@
</dependency>
<dependency>
<groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-minikdc</artifactId>
+ <artifactId>hadoop-minicluster</artifactId>
<exclusions>
- <exclusion>
- <groupId>junit</groupId>
- <artifactId>junit</artifactId>
- </exclusion>
<exclusion>
<groupId>org.slf4j</groupId>
<artifactId>slf4j-log4j12</artifactId>
@@ -122,66 +146,5 @@
</exclusions>
</dependency>
</dependencies>
- <build>
- <plugins>
- <plugin>
- <groupId>org.apache.maven.plugins</groupId>
- <artifactId>maven-enforcer-plugin</artifactId>
- <version>1.4.1</version>
- <executions>
- <execution>
- <id>enforce</id>
- <goals>
- <goal>enforce</goal>
- </goals>
- <configuration>
- <rules>
- <requireUpperBoundDeps/>
- </rules>
- </configuration>
- </execution>
- </executions>
- </plugin>
- <plugin>
- <groupId>org.apache.maven.plugins</groupId>
- <artifactId>maven-jar-plugin</artifactId>
- <version>2.4</version>
- <executions>
- <execution>
- <goals>
- <goal>test-jar</goal>
- </goals>
- </execution>
- </executions>
- </plugin>
- <plugin>
- <groupId>org.jacoco</groupId>
- <artifactId>jacoco-maven-plugin</artifactId>
- <executions>
- </executions>
- </plugin>
- </plugins>
- <pluginManagement>
- <plugins>
- <plugin>
- <groupId>org.apache.maven.plugins</groupId>
- <artifactId>maven-surefire-plugin</artifactId>
- <version>2.22.2</version>
- <configuration>
- <systemPropertyVariables>
- <java.awt.headless>true</java.awt.headless>
- <java.io.tmpdir>${project.build.directory}/tmp</java.io.tmpdir>
- <project.version>${project.version}</project.version>
- <test.mode>true</test.mode>
- </systemPropertyVariables>
- <redirectTestOutputToFile>${test.redirectToFile}</redirectTestOutputToFile>
- <useFile>${test.redirectToFile}</useFile>
- <argLine>-ea -Xmx5g -Xss4m -XX:MaxMetaspaceSize=2g -XX:ReservedCodeCacheSize=1g</argLine>
- <failIfNoTests>false</failIfNoTests>
- </configuration>
- </plugin>
- </plugins>
- </pluginManagement>
- </build>
</project>
diff --git a/integration-test/spark4/src/test/java/org/apache/uniffle/test/AQERepartitionTest.java b/integration-test/spark4/src/test/java/org/apache/uniffle/test/AQERepartitionTest.java
new file mode 100644
index 000000000..4066201ba
--- /dev/null
+++ b/integration-test/spark4/src/test/java/org/apache/uniffle/test/AQERepartitionTest.java
@@ -0,0 +1,104 @@
+/*
+ * 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.util.Comparator;
+import java.util.List;
+import java.util.Map;
+
+import scala.jdk.javaapi.CollectionConverters;
+
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import org.apache.spark.SparkConf;
+import org.apache.spark.shuffle.RssSparkConfig;
+import org.apache.spark.sql.Column;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.internal.SQLConf;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import org.apache.uniffle.common.rpc.ServerType;
+import org.apache.uniffle.coordinator.CoordinatorConf;
+import org.apache.uniffle.storage.util.StorageType;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class AQERepartitionTest extends SparkIntegrationTestBase {
+
+ @BeforeAll
+ public static void setupServers() throws Exception {
+ CoordinatorConf coordinatorConf = coordinatorConfWithoutPort();
+ Map<String, String> dynamicConf = Maps.newHashMap();
+ dynamicConf.put(CoordinatorConf.COORDINATOR_REMOTE_STORAGE_PATH.key(), HDFS_URI + "rss/test");
+ dynamicConf.put(
+ RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.MEMORY_LOCALFILE_HDFS.name());
+ addDynamicConf(coordinatorConf, dynamicConf);
+ storeCoordinatorConf(coordinatorConf);
+
+ storeShuffleServerConf(shuffleServerConfWithoutPort(0, null, ServerType.GRPC));
+ storeShuffleServerConf(shuffleServerConfWithoutPort(1, null, ServerType.GRPC_NETTY));
+ startServersWithRandomPorts();
+ }
+
+ @Override
+ public void updateCommonSparkConf(SparkConf sparkConf) {
+ sparkConf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED().key(), "true");
+ sparkConf.set(SQLConf.COALESCE_PARTITIONS_ENABLED(), "true");
+ sparkConf.set(SQLConf.COALESCE_PARTITIONS_INITIAL_PARTITION_NUM().key(), "6");
+ sparkConf.set(SQLConf.SHUFFLE_PARTITIONS().key(), "10");
+ }
+
+ @Override
+ public void updateSparkConfCustomer(SparkConf sparkConf) {}
+
+ @Test
+ public void resultCompareTest() throws Exception {
+ run();
+ }
+
+ @Override
+ Map runTest(SparkSession spark, String fileName) throws Exception {
+ Thread.sleep(4000);
+ List<Column> repartitionCols = Lists.newArrayList();
+ repartitionCols.add(new Column("id"));
+ Dataset<Long> df =
+ spark.range(10).repartition(CollectionConverters.asScala(repartitionCols).toList());
+ Long[][] result = (Long[][]) df.rdd().collectPartitions();
+ Map<Integer, List<Long>> map = Maps.newHashMap();
+ for (int i = 0; i < result.length; i++) {
+ map.putIfAbsent(i, Lists.newArrayList());
+ for (int j = 0; j < result[i].length; j++) {
+ map.get(i).add(result[i][j]);
+ }
+ }
+ for (int i = 0; i < result.length; i++) {
+ map.get(i)
+ .sort(
+ new Comparator<Long>() {
+ @Override
+ public int compare(Long o1, Long o2) {
+ return Long.compare(o1, o2);
+ }
+ });
+ }
+ assertTrue(result.length < 10);
+ return map;
+ }
+}
diff --git a/integration-test/spark4/src/test/java/org/apache/uniffle/test/MapSideCombineTest.java b/integration-test/spark4/src/test/java/org/apache/uniffle/test/MapSideCombineTest.java
new file mode 100644
index 000000000..46f96eccf
--- /dev/null
+++ b/integration-test/spark4/src/test/java/org/apache/uniffle/test/MapSideCombineTest.java
@@ -0,0 +1,148 @@
+/*
+ * 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.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import scala.Tuple2;
+
+import com.google.common.collect.Maps;
+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 org.apache.uniffle.common.StorageType;
+import org.apache.uniffle.common.rpc.ServerType;
+import org.apache.uniffle.test.listener.WriteAndReadMetricsSparkListener;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class MapSideCombineTest extends SparkIntegrationTestBase {
+
+ @BeforeAll
+ public static void setupServers() throws Exception {
+ storeCoordinatorConf(coordinatorConfWithoutPort());
+
+ storeShuffleServerConf(shuffleServerConfWithoutPort(0, null, ServerType.GRPC));
+ storeShuffleServerConf(shuffleServerConfWithoutPort(1, null, ServerType.GRPC_NETTY));
+
+ startServersWithRandomPorts();
+ }
+
+ @Override
+ public void updateSparkConfCustomer(SparkConf sparkConf) {
+ sparkConf.set(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.MEMORY_LOCALFILE_HDFS.name());
+ sparkConf.set(RssSparkConfig.RSS_REMOTE_STORAGE_PATH.key(), HDFS_URI + "rss/test");
+ sparkConf.set("spark." + RssSparkConfig.RSS_CLIENT_MAP_SIDE_COMBINE_ENABLED.key(), "true");
+ sparkConf.set("spark." + RssSparkConfig.RSS_CLIENT_INTEGRITY_VALIDATION_ENABLED.key(), "true");
+ }
+
+ @Test
+ public void resultCompareTest() throws Exception {
+ run();
+ }
+
+ @Override
+ Map runTest(SparkSession spark, String fileName) throws Exception {
+ Thread.sleep(4000);
+
+ WriteAndReadMetricsSparkListener listener = new WriteAndReadMetricsSparkListener();
+ spark.sparkContext().addSparkListener(listener);
+
+ Map<String, Object> result = Maps.newHashMap();
+ JavaSparkContext sc = new JavaSparkContext(spark.sparkContext());
+ List<Integer> data = Stream.iterate(1, n -> n + 1).limit(10000).collect(Collectors.toList());
+ JavaPairRDD<Integer, Integer> dataSourceRdd =
+ sc.parallelize(data, 10).mapToPair(x -> new Tuple2<>(x % 10, 1));
+ int jobId = -1;
+
+ // reduceByKey
+ checkMapSideCombine(
+ spark,
+ listener,
+ "reduceByKey",
+ dataSourceRdd.reduceByKey(Integer::sum).collectAsMap(),
+ result,
+ ++jobId);
+
+ // combineByKey
+ checkMapSideCombine(
+ spark,
+ listener,
+ "combineByKey",
+ dataSourceRdd.combineByKey(x -> 1, Integer::sum, Integer::sum).collectAsMap(),
+ result,
+ ++jobId);
+
+ // aggregateByKey
+ checkMapSideCombine(
+ spark,
+ listener,
+ "aggregateByKey",
+ dataSourceRdd.aggregateByKey(10, Integer::sum, Integer::sum).collectAsMap(),
+ result,
+ ++jobId);
+
+ // foldByKey
+ checkMapSideCombine(
+ spark,
+ listener,
+ "foldByKey",
+ dataSourceRdd.foldByKey(10, Integer::sum).collectAsMap(),
+ result,
+ ++jobId);
+
+ // countByKey
+ checkMapSideCombine(spark, listener, "countByKey", dataSourceRdd.countByKey(), result, ++jobId);
+
+ return result;
+ }
+
+ private <K, V> void checkMapSideCombine(
+ SparkSession spark,
+ WriteAndReadMetricsSparkListener listener,
+ String method,
+ Map<K, V> rddResult,
+ Map<String, Object> result,
+ int jobId)
+ throws TimeoutException {
+ rddResult.forEach((key, value) -> result.put(method + "-result-value-" + key, value));
+
+ spark.sparkContext().listenerBus().waitUntilEmpty();
+
+ for (int stageId : spark.sparkContext().statusTracker().getJobInfo(jobId).get().stageIds()) {
+ long writeRecords = listener.getWriteRecords(stageId);
+ long readRecords = listener.getReadRecords(stageId);
+ result.put(stageId + "-write-records", writeRecords);
+ result.put(stageId + "-read-records", readRecords);
+ }
+
+ // each job has two stages, so each job start stageId = jobId * 2
+ int shuffleStageId = jobId * 2;
+ // check map side combine
+ assertEquals(100L, result.get(shuffleStageId + "-write-records"));
+ }
+}
diff --git a/integration-test/spark4/src/test/resources/log4j2.xml b/integration-test/spark4/src/test/resources/log4j2.xml
new file mode 100644
index 000000000..7bf322857
--- /dev/null
+++ b/integration-test/spark4/src/test/resources/log4j2.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ ~ 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.
+-->
+<Configuration status="WARN" monitorInterval="30">
+ <Appenders>
+ <Console name="Console" target="SYSTEM_OUT">
+ <PatternLayout pattern="[%d{yyyy-MM-dd HH:mm:ss.SSS}] [%t] [%p] %c{1}.%M - %m%n%ex"/>
+ </Console>
+ <File name="file" fileName="target/unit-tests.log">
+ <PatternLayout pattern="[%d{yyyy-MM-dd HH:mm:ss.SSS}] [%t] [%p] %c{1}.%M - %m%n%ex"/>
+ </File>
+ </Appenders>
+ <Loggers>
+ <Root level="info">
+ <AppenderRef ref="Console"/>
+ <AppenderRef ref="file"/>
+ </Root>
+ </Loggers>
+</Configuration>
diff --git a/pom.xml b/pom.xml
index fdc0dc20c..ee2faa04f 100644
--- a/pom.xml
+++ b/pom.xml
@@ -109,6 +109,7 @@
<kryo.version>4.0.2</kryo.version>
<scala.version>2.12.18</scala.version>
<scala.maven.plugin.version>3.2.2</scala.maven.plugin.version>
+ <extraJavaTestArgs/>
</properties>
<repositories>
@@ -982,7 +983,7 @@
</systemPropertyVariables>
<redirectTestOutputToFile>${test.redirectToFile}</redirectTestOutputToFile>
<useFile>${test.redirectToFile}</useFile>
- <argLine>${argLine} -ea -Xmx5g -Xss4m -XX:MaxMetaspaceSize=2g -XX:ReservedCodeCacheSize=1g</argLine>
+ <argLine>${argLine} -ea -Xmx5g -Xss4m -XX:MaxMetaspaceSize=2g -XX:ReservedCodeCacheSize=1g ${extraJavaTestArgs}</argLine>
<failIfNoTests>false</failIfNoTests>
<trimStackTrace>${trimStackTrace}</trimStackTrace>
<skipTests>${skipUTs}</skipTests>
@@ -1964,6 +1965,143 @@
</dependencies>
</profile>
+ <profile>
+ <id>spark4</id>
+ <properties>
+ <scala.binary.version>2.13</scala.binary.version>
+ <scala.version>2.13.16</scala.version>
+ <spark.version>4.0.2</spark.version>
+ <client.type>4</client.type>
+ <hadoop.version>3.4.1</hadoop.version>
+ <fasterxml.jackson.version>2.12.7</fasterxml.jackson.version>
+ <jetty.version>9.4.53.v20231009</jetty.version>
+ <enforcer.skip>true</enforcer.skip>
+ <jackson.version>2.18.2</jackson.version>
+ <log4j.version>2.24.3</log4j.version>
+ <slf4j.version>2.0.16</slf4j.version>
+ <java.version>17</java.version>
+ <maven.compiler.source>17</maven.compiler.source>
+ <maven.compiler.target>17</maven.compiler.target>
+ <commons-lang3.version>3.17.0</commons-lang3.version>
+ <extraJavaTestArgs>
+ -XX:+IgnoreUnrecognizedVMOptions
+ --add-opens=java.base/java.lang=ALL-UNNAMED
+ --add-opens=java.base/java.lang.invoke=ALL-UNNAMED
+ --add-opens=java.base/java.lang.reflect=ALL-UNNAMED
+ --add-opens=java.base/java.io=ALL-UNNAMED
+ --add-opens=java.base/java.net=ALL-UNNAMED
+ --add-opens=java.base/java.nio=ALL-UNNAMED
+ --add-opens=java.base/java.util=ALL-UNNAMED
+ --add-opens=java.base/java.util.concurrent=ALL-UNNAMED
+ --add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED
+ --add-opens=java.base/sun.nio.ch=ALL-UNNAMED
+ --add-opens=java.base/sun.security.action=ALL-UNNAMED
+ --add-opens=java.base/sun.util.calendar=ALL-UNNAMED
+ --add-opens=java.security.jgss/sun.security.krb5=ALL-UNNAMED
+ -Djdk.reflect.useDirectMethodHandle=false
+ -Duniffle.test.skip.kerberos=true
+ </extraJavaTestArgs>
+ </properties>
+ <modules>
+ <module>client-spark/common</module>
+ <module>client-spark/spark4</module>
+ <module>client-spark/spark4-shaded</module>
+ <module>integration-test/spark-common</module>
+ <module>integration-test/spark4</module>
+ </modules>
+ <dependencyManagement>
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-client-spark4</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-core_${scala.binary.version}</artifactId>
+ <version>${spark.version}</version>
+ <exclusions>
+ <exclusion>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>*</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-databind</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-core</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scala.binary.version}</artifactId>
+ <version>${spark.version}</version>
+ <exclusions>
+ <exclusion>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>*</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-databind</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-core</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-client-spark-common</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-client-spark-common</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-integration-common-test</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>rss-integration-spark-common-test</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+ </dependencies>
+ </dependencyManagement>
+ <dependencies>
+ <!-- Remove log4j-slf4j-impl dependency -->
+ <dependency>
+ <groupId>org.apache.logging.log4j</groupId>
+ <artifactId>log4j-slf4j-impl</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Add log4j-slf4j2-impl dependency -->
+ <dependency>
+ <groupId>org.apache.logging.log4j</groupId>
+ <artifactId>log4j-slf4j2-impl</artifactId>
+ </dependency>
+ <!-- Hadoop 3.4.1 needs commons-logging explicitly since it's excluded from hadoop-client -->
+ <dependency>
+ <groupId>commons-logging</groupId>
+ <artifactId>commons-logging</artifactId>
+ <scope>test</scope>
+ </dependency>
+ </dependencies>
+ </profile>
+
<profile>
<id>spark3.2.0</id>
<properties>
diff --git a/server/src/main/java/org/apache/uniffle/server/HealthScriptChecker.java b/server/src/main/java/org/apache/uniffle/server/HealthScriptChecker.java
index c5c58a890..530f69501 100644
--- a/server/src/main/java/org/apache/uniffle/server/HealthScriptChecker.java
+++ b/server/src/main/java/org/apache/uniffle/server/HealthScriptChecker.java
@@ -17,7 +17,8 @@
package org.apache.uniffle.server;
-import org.apache.hadoop.util.NodeHealthScriptRunner;
+import java.io.File;
+
import org.apache.hadoop.util.Shell;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -33,7 +34,7 @@ public class HealthScriptChecker extends Checker {
HealthScriptChecker(ShuffleServerConf conf) {
super(conf);
this.healthScriptPath = conf.getString(ShuffleServerConf.HEALTH_CHECKER_SCRIPT_PATH);
- if (!NodeHealthScriptRunner.shouldRun(healthScriptPath)) {
+ if (!shouldRun(healthScriptPath)) {
LOG.error(
"Rss health check script:"
+ healthScriptPath
@@ -44,10 +45,19 @@ public class HealthScriptChecker extends Checker {
this.scriptTimeout = conf.getLong(ShuffleServerConf.HEALTH_CHECKER_SCRIPT_EXECUTE_TIMEOUT);
}
+ /** Returns true if the script path is non-empty, the file exists, and is executable. */
+ private static boolean shouldRun(String scriptPath) {
+ if (scriptPath == null || scriptPath.isEmpty()) {
+ return false;
+ }
+ File scriptFile = new File(scriptPath);
+ return scriptFile.exists() && scriptFile.canExecute();
+ }
+
@Override
public boolean checkIsHealthy() {
HealthCheckerExitStatus status = HealthCheckerExitStatus.SUCCESS;
- // always build executor is to load the latest script, the script file can update in need.
+ // Rebuild the executor each time to pick up any script file changes
Shell.ShellCommandExecutor commandExecutor =
new Shell.ShellCommandExecutor(new String[] {healthScriptPath}, null, null, scriptTimeout);
try {
@@ -55,8 +65,7 @@ public class HealthScriptChecker extends Checker {
} catch (Shell.ExitCodeException e) {
// ignore the exit code of the script
status = HealthCheckerExitStatus.FAILED_WITH_EXIT_CODE;
- // On Windows, we will not hit the Stream closed IOException
- // thrown by stdout buffered reader for timeout event.
+ // Windows does not raise IOException on timeout, so check explicitly
if (Shell.WINDOWS && commandExecutor.isTimedOut()) {
status = HealthCheckerExitStatus.TIMED_OUT;
}
diff --git a/server/src/test/java/org/apache/uniffle/server/KerberizedShuffleTaskManagerTest.java b/server/src/test/java/org/apache/uniffle/server/KerberizedShuffleTaskManagerTest.java
index 0a5508d8f..6374418d8 100644
--- a/server/src/test/java/org/apache/uniffle/server/KerberizedShuffleTaskManagerTest.java
+++ b/server/src/test/java/org/apache/uniffle/server/KerberizedShuffleTaskManagerTest.java
@@ -30,6 +30,7 @@ import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.apache.uniffle.common.KerberizedHadoopBase;
import org.apache.uniffle.common.PartitionRange;
@@ -47,6 +48,10 @@ import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class KerberizedShuffleTaskManagerTest extends KerberizedHadoopBase {
private static final AtomicInteger ATOMIC_INT = new AtomicInteger(0);
diff --git a/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerOnKerberizedHadoopTest.java b/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerOnKerberizedHadoopTest.java
index b4686b0d6..b165a6066 100644
--- a/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerOnKerberizedHadoopTest.java
+++ b/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerOnKerberizedHadoopTest.java
@@ -32,6 +32,7 @@ import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -54,6 +55,10 @@ import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class ShuffleFlushManagerOnKerberizedHadoopTest extends KerberizedHadoopBase {
private static final Logger LOGGER =
LoggerFactory.getLogger(ShuffleFlushManagerOnKerberizedHadoopTest.class);
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java b/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
index d84bdfead..3a7d89e79 100644
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
+++ b/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopClientReadHandlerTest.java
@@ -20,9 +20,14 @@ package org.apache.uniffle.storage.handler.impl;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.apache.uniffle.common.KerberizedHadoopBase;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class KerberizedHadoopClientReadHandlerTest extends KerberizedHadoopBase {
@BeforeAll
diff --git a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopShuffleReadHandlerTest.java b/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopShuffleReadHandlerTest.java
index f314951b9..0b010ec29 100644
--- a/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopShuffleReadHandlerTest.java
+++ b/storage/src/test/java/org/apache/uniffle/storage/handler/impl/KerberizedHadoopShuffleReadHandlerTest.java
@@ -20,9 +20,14 @@ package org.apache.uniffle.storage.handler.impl;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.apache.uniffle.common.KerberizedHadoopBase;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class KerberizedHadoopShuffleReadHandlerTest extends KerberizedHadoopBase {
@BeforeAll
diff --git a/storage/src/test/java/org/apache/uniffle/storage/util/ShuffleKerberizedHadoopStorageUtilsTest.java b/storage/src/test/java/org/apache/uniffle/storage/util/ShuffleKerberizedHadoopStorageUtilsTest.java
index a8459f71a..47d7bb0c9 100644
--- a/storage/src/test/java/org/apache/uniffle/storage/util/ShuffleKerberizedHadoopStorageUtilsTest.java
+++ b/storage/src/test/java/org/apache/uniffle/storage/util/ShuffleKerberizedHadoopStorageUtilsTest.java
@@ -21,10 +21,15 @@ import java.io.File;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIfSystemProperty;
import org.junit.jupiter.api.io.TempDir;
import org.apache.uniffle.common.KerberizedHadoopBase;
+@DisabledIfSystemProperty(
+ named = "uniffle.test.skip.kerberos",
+ matches = "true",
+ disabledReason = "MiniKdc is not compatible with JDK 17+")
public class ShuffleKerberizedHadoopStorageUtilsTest extends KerberizedHadoopBase {
@BeforeAll