[ISSUE-380] Refactor the flush process to fix fallback fail

integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageFaultToleranceTest.java → integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageFaultToleranceBase.java

@@ -18,7 +18,6 @@
 package org.apache.uniffle.test;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -27,11 +26,11 @@
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 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.roaringbitmap.longlong.Roaring64NavigableMap;
 
+import org.apache.uniffle.client.factory.ShuffleServerClientFactory;
 import org.apache.uniffle.client.impl.ShuffleReadClientImpl;
 import org.apache.uniffle.client.impl.grpc.ShuffleServerGrpcClient;
 import org.apache.uniffle.client.request.RssFinishShuffleRequest;
@@ -44,37 +43,17 @@
 import org.apache.uniffle.common.PartitionRange;
 import org.apache.uniffle.common.ShuffleBlockInfo;
 import org.apache.uniffle.common.ShuffleServerInfo;
-import org.apache.uniffle.coordinator.CoordinatorConf;
-import org.apache.uniffle.server.ShuffleServerConf;
 import org.apache.uniffle.storage.util.StorageType;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-public class MultiStorageFaultToleranceTest extends ShuffleReadWriteBase {
+public abstract class MultiStorageFaultToleranceBase extends ShuffleReadWriteBase {
   private ShuffleServerGrpcClient shuffleServerClient;
   private static String REMOTE_STORAGE = HDFS_URI + "rss/multi_storage_fault";
 
-  @BeforeAll
-  public static void setupServers() throws Exception {
-    final CoordinatorConf coordinatorConf = getCoordinatorConf();
-    ShuffleServerConf shuffleServerConf = getShuffleServerConf();
-    String basePath = generateBasePath();
-    shuffleServerConf.setDouble(ShuffleServerConf.CLEANUP_THRESHOLD, 0.0);
-    shuffleServerConf.setDouble(ShuffleServerConf.HIGH_WATER_MARK_OF_WRITE, 100.0);
-    shuffleServerConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 100);
-    shuffleServerConf.setLong(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 30L);
-    shuffleServerConf.setLong(ShuffleServerConf.SHUFFLE_EXPIRED_TIMEOUT_MS, 5000L);
-    shuffleServerConf.setLong(ShuffleServerConf.SERVER_APP_EXPIRED_WITHOUT_HEARTBEAT, 60L * 1000L * 60L);
-    shuffleServerConf.setLong(ShuffleServerConf.SERVER_COMMIT_TIMEOUT, 20L * 1000L);
-    shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS.name());
-    shuffleServerConf.set(ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(basePath));
-    shuffleServerConf.setLong(ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, 400L * 1024L * 1024L);
-    createAndStartServers(shuffleServerConf, coordinatorConf);
-  }
-
   @BeforeEach
   public void createClient() {
+    ShuffleServerClientFactory.getInstance().cleanupCache();
     shuffleServerClient = new ShuffleServerGrpcClient(LOCALHOST, SHUFFLE_SERVER_PORT);
   }
 
@@ -83,19 +62,19 @@ public void closeClient() {
     shuffleServerClient.close();
   }
 
+  abstract void makeChaos();
+
   @Test
-  public void hdfsFallbackTest() throws Exception {
-    String appId = "fallback_test";
+  public void fallbackTest() throws Exception {
+    String appId = "fallback_test_" + this.getClass().getSimpleName();
     Map<Long, byte[]> expectedData = Maps.newHashMap();
     Map<Integer, List<Integer>> map = Maps.newHashMap();
     map.put(0, Lists.newArrayList(0));
     registerShuffle(appId, map);
     Roaring64NavigableMap blockBitmap = Roaring64NavigableMap.bitmapOf();
     final List<ShuffleBlockInfo> blocks = createShuffleBlockList(
-        0, 0, 0, 40, 10 * 1024 * 1024, blockBitmap, expectedData);
-    assertEquals(1, cluster.getDataNodes().size());
-    cluster.stopDataNode(0);
-    assertEquals(0, cluster.getDataNodes().size());
+        0, 0, 0, 40, 2 * 1024 * 1024, blockBitmap, expectedData);
+    makeChaos();
     sendSinglePartitionToShuffleServer(appId, 0, 0, 0, blocks);
     validateResult(appId, 0, 0, blockBitmap, Roaring64NavigableMap.bitmapOf(0), expectedData);
   }

integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageHdfsFallbackTest.java

@@ -0,0 +1,56 @@
+/*
+ * 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.Arrays;
+
+import org.junit.jupiter.api.BeforeAll;
+
+import org.apache.uniffle.coordinator.CoordinatorConf;
+import org.apache.uniffle.server.ShuffleServerConf;
+import org.apache.uniffle.storage.util.StorageType;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class MultiStorageHdfsFallbackTest extends MultiStorageFaultToleranceBase {
+
+  @BeforeAll
+  public static void setupServers() throws Exception {
+    final CoordinatorConf coordinatorConf = getCoordinatorConf();
+    ShuffleServerConf shuffleServerConf = getShuffleServerConf();
+    String basePath = generateBasePath();
+    shuffleServerConf.setDouble(ShuffleServerConf.CLEANUP_THRESHOLD, 0.0);
+    shuffleServerConf.setDouble(ShuffleServerConf.HIGH_WATER_MARK_OF_WRITE, 100.0);
+    shuffleServerConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 100);
+    shuffleServerConf.setLong(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 30L);
+    shuffleServerConf.setLong(ShuffleServerConf.SHUFFLE_EXPIRED_TIMEOUT_MS, 5000L);
+    shuffleServerConf.setLong(ShuffleServerConf.SERVER_APP_EXPIRED_WITHOUT_HEARTBEAT, 60L * 1000L * 60L);
+    shuffleServerConf.setLong(ShuffleServerConf.SERVER_COMMIT_TIMEOUT, 20L * 1000L);
+    shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS.name());
+    shuffleServerConf.set(ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(basePath));
+    shuffleServerConf.setLong(ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, 1L * 1024L * 1024L);
+    createAndStartServers(shuffleServerConf, coordinatorConf);
+  }
+
+  @Override
+  public void makeChaos() {
+    assertEquals(1, cluster.getDataNodes().size());
+    cluster.stopDataNode(0);
+    assertEquals(0, cluster.getDataNodes().size());
+  }
+}

integration-test/common/src/test/java/org/apache/uniffle/test/MultiStorageLocalfileFallbackTest.java

@@ -0,0 +1,66 @@
+/*
+ * 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.Arrays;
+
+import org.junit.jupiter.api.BeforeAll;
+
+import org.apache.uniffle.coordinator.CoordinatorConf;
+import org.apache.uniffle.server.ShuffleServerConf;
+import org.apache.uniffle.server.storage.LocalStorageManager;
+import org.apache.uniffle.server.storage.LocalStorageManagerFallbackStrategy;
+import org.apache.uniffle.server.storage.MultiStorageManager;
+import org.apache.uniffle.storage.common.LocalStorage;
+import org.apache.uniffle.storage.common.Storage;
+import org.apache.uniffle.storage.util.StorageType;
+
+public class MultiStorageLocalfileFallbackTest extends MultiStorageFaultToleranceBase {
+
+  @BeforeAll
+  public static void setupServers() throws Exception {
+    final CoordinatorConf coordinatorConf = getCoordinatorConf();
+    ShuffleServerConf shuffleServerConf = getShuffleServerConf();
+    String basePath = generateBasePath();
+    shuffleServerConf.setDouble(ShuffleServerConf.CLEANUP_THRESHOLD, 0.0);
+    shuffleServerConf.setDouble(ShuffleServerConf.HIGH_WATER_MARK_OF_WRITE, 100.0);
+    shuffleServerConf.setLong(ShuffleServerConf.DISK_CAPACITY, 1024L * 1024L * 100);
+    shuffleServerConf.setLong(ShuffleServerConf.PENDING_EVENT_TIMEOUT_SEC, 30L);
+    shuffleServerConf.setLong(ShuffleServerConf.SHUFFLE_EXPIRED_TIMEOUT_MS, 5000L);
+    shuffleServerConf.setLong(ShuffleServerConf.SERVER_APP_EXPIRED_WITHOUT_HEARTBEAT, 60L * 1000L * 60L);
+    shuffleServerConf.setLong(ShuffleServerConf.SERVER_COMMIT_TIMEOUT, 20L * 1000L);
+    shuffleServerConf.setString(ShuffleServerConf.RSS_STORAGE_TYPE, StorageType.LOCALFILE_HDFS.name());
+    shuffleServerConf.set(ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(basePath));
+    shuffleServerConf.setLong(ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, 1000L * 1024L * 1024L);
+    shuffleServerConf.setString(
+        ShuffleServerConf.MULTISTORAGE_FALLBACK_STRATEGY_CLASS,
+        LocalStorageManagerFallbackStrategy.class.getCanonicalName()
+    );
+    createAndStartServers(shuffleServerConf, coordinatorConf);
+  }
+
+  @Override
+  public void makeChaos() {
+    LocalStorageManager warmStorageManager =
+        (LocalStorageManager) ((MultiStorageManager)shuffleServers.get(0).getStorageManager()).getWarmStorageManager();
+    for (Storage storage : warmStorageManager.getStorages()) {
+      LocalStorage localStorage = (LocalStorage) storage;
+      localStorage.markSpaceFull();
+    }
+  }
+}

internal-client/src/main/java/org/apache/uniffle/client/factory/ShuffleServerClientFactory.java

@@ -59,4 +59,10 @@ public synchronized ShuffleServerClient getShuffleServerClient(
     }
     return serverToClients.get(shuffleServerInfo);
   }
+
+  // Only for tests
+  public synchronized void cleanupCache() {
+    clients.values().stream().flatMap(x -> x.values().stream()).forEach(ShuffleServerClient::close);
+    this.clients = Maps.newConcurrentMap();
+  }
 }

server/src/main/java/org/apache/uniffle/server/ShuffleDataFlushEvent.java

@@ -17,14 +17,20 @@
 
 package org.apache.uniffle.server;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import org.apache.uniffle.common.ShufflePartitionedBlock;
 import org.apache.uniffle.server.buffer.ShuffleBuffer;
+import org.apache.uniffle.storage.common.Storage;
 
 public class ShuffleDataFlushEvent {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ShuffleDataFlushEvent.class);
 
   private final long eventId;
   private final String appId;
@@ -37,6 +43,10 @@ public class ShuffleDataFlushEvent {
   private final ShuffleBuffer shuffleBuffer;
   private final AtomicInteger retryTimes = new AtomicInteger();
 
+  private boolean isPended = false;
+  private Storage underStorage;
+  private final List<Runnable> cleanupCallbackChains;
+
   public ShuffleDataFlushEvent(
       long eventId,
       String appId,
@@ -56,6 +66,7 @@ public ShuffleDataFlushEvent(
     this.shuffleBlocks = shuffleBlocks;
     this.valid = valid;
     this.shuffleBuffer = shuffleBuffer;
+    this.cleanupCallbackChains = new ArrayList<>();
   }
 
   public List<ShufflePartitionedBlock> getShuffleBlocks() {
@@ -105,6 +116,42 @@ public void increaseRetryTimes() {
     retryTimes.incrementAndGet();
   }
 
+  public boolean isPended() {
+    return isPended;
+  }
+
+  public void markPended() {
+    isPended = true;
+  }
+
+  public Storage getUnderStorage() {
+    return underStorage;
+  }
+
+  public void setUnderStorage(Storage underStorage) {
+    this.underStorage = underStorage;
+  }
+
+  public boolean doCleanup() {
+    boolean ret = true;
+    for (Runnable cleanupCallback : cleanupCallbackChains) {
+      try {
+        cleanupCallback.run();
+      } catch (Exception e) {
+        ret = false;
+        LOGGER.error("Errors doing cleanup callback. event: {}", this, e);
+      }
+    }
+    return ret;
+  }
+
+  public void addCleanupCallback(
+      Runnable cleanupCallback) {
+    if (cleanupCallback != null) {
+      cleanupCallbackChains.add(cleanupCallback);
+    }
+  }
+
   @Override
   public String toString() {
     return "ShuffleDataFlushEvent: eventId=" + eventId