[FLINK-36112][Connector/Filesystem].Add Support for CreateFlag.NO_LOCAL_WRITE in FLINK on YARN's File Creation to Manage Disk Space and Network Load in Labeled YARN Nodes

Author: liangyu-1

flink-connectors/flink-connector-files/src/main/java/org/apache/flink/connector/file/sink/FileSink.java

@@ -582,7 +582,7 @@ public T withRollingPolicy(CheckpointRollingPolicy<IN, String> rollingPolicy) {
             return self();
         }
 
-        public T enableNoLocalWriting() {
+        public T disableLocalWriting() {
             this.noLocalWrite = true;
             return self();
         }

flink-core/src/main/java/org/apache/flink/core/fs/FileSystem.java

@@ -58,6 +58,7 @@
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.concurrent.locks.ReentrantLock;
@@ -668,11 +669,11 @@ public RecoverableWriter createRecoverableWriter() throws IOException {
 
     @PublicEvolving
     @Override
-    public RecoverableWriter createRecoverableWriter(boolean noLocalWrite) throws IOException {
-        if (noLocalWrite) {
-            return IFileSystem.super.createRecoverableWriter(noLocalWrite);
-        } else {
+    public RecoverableWriter createRecoverableWriter(Map<String, String> conf) throws IOException {
+        if (conf == null || conf.isEmpty()) {
             return createRecoverableWriter();
+        } else {
+            return IFileSystem.super.createRecoverableWriter(conf);
         }
     }
 

flink-core/src/main/java/org/apache/flink/core/fs/IFileSystem.java

@@ -30,6 +30,7 @@
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.net.URI;
+import java.util.Map;
 
 /**
  * Interface of all file systems used by Flink. This interface may be extended to implement
@@ -247,11 +248,12 @@ default RecoverableWriter createRecoverableWriter() throws IOException {
      * <p>This method is optional on file systems and various file system implementations may not
      * support this method, throwing an {@code UnsupportedOperationException}.
      *
-     * @param b Flag to indicate whether the writer should not write to local storage.
+     * @param conf Map contains a flag to indicate whether the writer should not write to local
+     *     storage. and can provide more information to instantiate the writer.
      * @return A RecoverableWriter for this file system.
      * @throws IOException Thrown, if the recoverable writer cannot be instantiated.
      */
-    default RecoverableWriter createRecoverableWriter(boolean b) throws IOException {
+    default RecoverableWriter createRecoverableWriter(Map<String, String> conf) throws IOException {
         throw new UnsupportedOperationException(
                 "This file system does not support recoverable writers that does not write on local machine.");
     }

flink-filesystems/flink-hadoop-fs/src/main/java/org/apache/flink/runtime/fs/hdfs/HadoopFileSystem.java

@@ -28,6 +28,7 @@
 import java.io.IOException;
 import java.net.URI;
 import java.util.Locale;
+import java.util.Map;
 
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
@@ -216,11 +217,18 @@ public RecoverableWriter createRecoverableWriter() throws IOException {
     }
 
     @Override
-    public RecoverableWriter createRecoverableWriter(boolean noLocalWrite) throws IOException {
+    public RecoverableWriter createRecoverableWriter(Map<String, String> conf) throws IOException {
         // This writer is only supported on a subset of file systems, and on
         // specific versions. We check these schemes and versions eagerly for better error
         // messages in the constructor of the writer.
-        return new HadoopRecoverableWriter(fs, noLocalWrite);
+        if (conf == null || conf.isEmpty()) {
+            return createRecoverableWriter();
+        } else if (conf.containsKey("fs.hdfs.no-local-write")) {
+            return new HadoopRecoverableWriter(
+                    fs, Boolean.parseBoolean(conf.get("fs.hdfs.no-local-write")));
+        } else {
+            return new HadoopRecoverableWriter(fs);
+        }
     }
 
     // ------------------------------------------------------------------------

flink-filesystems/flink-hadoop-fs/src/main/java/org/apache/flink/runtime/fs/hdfs/HadoopRecoverableWriter.java

@@ -46,7 +46,7 @@ public class HadoopRecoverableWriter implements RecoverableWriter {
     /** The Hadoop file system on which the writer operates. */
     protected final org.apache.hadoop.fs.FileSystem fs;
 
-    protected final boolean noLocalWrite;
+    private final boolean noLocalWrite;
 
     /**
      * Creates a new Recoverable writer.

flink-filesystems/flink-hadoop-fs/src/test/java/org/apache/flink/runtime/fs/hdfs/HadoopRecoverableWriterTest.java

@@ -21,22 +21,37 @@
 import org.apache.flink.core.fs.AbstractRecoverableWriterTest;
 import org.apache.flink.core.fs.FileSystem;
 import org.apache.flink.core.fs.Path;
+import org.apache.flink.core.fs.RecoverableFsDataOutputStream;
 import org.apache.flink.core.fs.RecoverableWriter;
+import org.apache.flink.mock.Whitebox;
 import org.apache.flink.runtime.util.HadoopUtils;
 import org.apache.flink.testutils.junit.utils.TempDirUtils;
 import org.apache.flink.util.OperatingSystem;
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hdfs.MiniDFSCluster;
+import org.apache.hadoop.hdfs.protocol.BlockListAsLongs;
+import org.apache.hadoop.hdfs.protocol.HdfsConstants;
+import org.apache.hadoop.hdfs.server.blockmanagement.BlockManager;
+import org.apache.hadoop.hdfs.server.blockmanagement.DatanodeDescriptor;
+import org.apache.hadoop.hdfs.server.blockmanagement.DatanodeManager;
+import org.apache.hadoop.hdfs.server.namenode.Namesystem;
+import org.apache.hadoop.hdfs.server.protocol.DatanodeStorage;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
 import java.io.File;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
 
 import static org.assertj.core.api.Assumptions.assumeThat;
+import static org.mockito.Mockito.spy;
+//import static org.mockito.Mockito.doReturn;
+//import static org.mockito.Mockito.spy;
 
 /** Tests for the {@link HadoopRecoverableWriter}. */
 class HadoopRecoverableWriterTest extends AbstractRecoverableWriterTest {
@@ -66,8 +81,11 @@ static void createHDFS() throws Exception {
 
         final Configuration hdConf = new Configuration();
         hdConf.set(MiniDFSCluster.HDFS_MINIDFS_BASEDIR, baseDir.getAbsolutePath());
+        hdConf.set("dfs.replication", "2");
+        hdConf.set("dfs.blocksize", String.valueOf(512));
+        hdConf.set("dfs.namenode.fs-limits.min-block-size", String.valueOf(512));
 
-        final MiniDFSCluster.Builder builder = new MiniDFSCluster.Builder(hdConf);
+        final MiniDFSCluster.Builder builder = new MiniDFSCluster.Builder(hdConf).numDataNodes(3);
         hdfsCluster = builder.build();
 
         final org.apache.hadoop.fs.FileSystem hdfs = hdfsCluster.getFileSystem();
@@ -87,15 +105,49 @@ static void destroyHDFS() throws Exception {
     }
 
     private RecoverableWriter getNoLocalWriteFileSystemWriter() throws Exception {
-        return fileSystem.createRecoverableWriter(true);
+        Map<String, String> conf = new HashMap<>();
+        conf.put("fs.hdfs.no-local-write", "true");
+        return fileSystem.createRecoverableWriter(conf);
     }
 
     @Test
-    void testNoLocalWrite() throws Exception {
+    void testNoLocalWriteFlag() throws Exception {
+        createHDFS();
         final HadoopRecoverableWriter writer =
                 (HadoopRecoverableWriter) getNoLocalWriteFileSystemWriter();
-
-        Assertions.assertTrue(writer.noLocalWrite);
+        BlockManager bm = hdfsCluster.getNameNode().getNamesystem().getBlockManager();
+        DatanodeManager dm = bm.getDatanodeManager();
+        try (RecoverableFsDataOutputStream os = writer.open(new Path("/tests/test-no-local"))) {
+            // Inject a DatanodeManager that returns one DataNode as local node for
+            // the client.
+            DatanodeManager spyDm = spy(dm);
+            DatanodeDescriptor dn1 =
+                    dm.getDatanodeListForReport(HdfsConstants.DatanodeReportType.LIVE).get(0);
+            Whitebox.setInternalState(bm, "datanodeManager", spyDm);
+            byte[] buf = new byte[512 * 16];
+            new Random().nextBytes(buf);
+            os.write(buf);
+        } finally {
+            Whitebox.setInternalState(bm, "datanodeManager", dm);
+        }
+        hdfsCluster.triggerBlockReports();
+        final String bpid = hdfsCluster.getNamesystem().getBlockPoolId();
+        // Total number of DataNodes is 3.
+        Assertions.assertEquals(3, hdfsCluster.getAllBlockReports(bpid).size());
+        int numDataNodesWithData = 0;
+        for (Map<DatanodeStorage, BlockListAsLongs> dnBlocks :
+                hdfsCluster.getAllBlockReports(bpid)) {
+            System.out.println(dnBlocks);
+            for (BlockListAsLongs blocks : dnBlocks.values()) {
+                if (blocks.getNumberOfBlocks() > 0) {
+                    System.out.println(dnBlocks);
+                    numDataNodesWithData++;
+                    break;
+                }
+            }
+        }
+        // Verify that only one DN has no data.
+        Assertions.assertEquals(1, 3 - numDataNodesWithData);
     }
 
     @Override