[FLINK-30798][runtime] Refactor OutputFormat#open to expose attempt number

Author: ifndef-SleePy

flink-core/src/main/java/org/apache/flink/api/common/io/OutputFormat.java

@@ -62,9 +62,22 @@ public interface OutputFormat<IT> extends Serializable {
      * @param taskNumber The number of the parallel instance.
      * @param numTasks The number of parallel tasks.
      * @throws IOException Thrown, if the output could not be opened due to an I/O problem.
+     * @deprecated Use {@link #open(InitializationContext)} instead
      */
-    void open(int taskNumber, int numTasks) throws IOException;
+    @Deprecated
+    default void open(int taskNumber, int numTasks) throws IOException {}
 
+    /**
+     * Opens a parallel instance of the output format to store the result of its parallel instance.
+     *
+     * <p>When this method is called, the output format it guaranteed to be configured.
+     *
+     * @param context The context to get task parallel infos.
+     * @throws IOException Thrown, if the output could not be opened due to an I/O problem.
+     */
+    default void open(InitializationContext context) throws IOException {
+        open(context.getTaskNumber(), context.getNumTasks());
+    }
     /**
      * Adds a record to the output.
      *
@@ -85,4 +98,30 @@ public interface OutputFormat<IT> extends Serializable {
      * @throws IOException Thrown, if the input could not be closed properly.
      */
     void close() throws IOException;
+
+    /** The context exposes some runtime info for initializing output format. */
+    @Public
+    interface InitializationContext {
+        /**
+         * Gets the parallelism with which the parallel task runs.
+         *
+         * @return The parallelism with which the parallel task runs.
+         */
+        int getNumTasks();
+
+        /**
+         * Gets the number of this parallel subtask. The numbering starts from 0 and goes up to
+         * parallelism-1 (parallelism as returned by {@link #getNumTasks()}).
+         *
+         * @return The index of the parallel subtask.
+         */
+        int getTaskNumber();
+
+        /**
+         * Gets the attempt number of this parallel subtask. First attempt is numbered 0.
+         *
+         * @return Attempt number of the subtask.
+         */
+        int getAttemptNumber();
+    }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/operators/DataSinkTask.java

@@ -22,6 +22,7 @@
 import org.apache.flink.api.common.functions.RuntimeContext;
 import org.apache.flink.api.common.io.CleanupWhenUnsuccessful;
 import org.apache.flink.api.common.io.OutputFormat;
+import org.apache.flink.api.common.io.OutputFormat.InitializationContext;
 import org.apache.flink.api.common.io.RichOutputFormat;
 import org.apache.flink.api.common.typeutils.TypeComparatorFactory;
 import org.apache.flink.api.common.typeutils.TypeSerializer;
@@ -215,8 +216,22 @@ public void invoke() throws Exception {
 
             // open
             format.open(
-                    this.getEnvironment().getTaskInfo().getIndexOfThisSubtask(),
-                    this.getEnvironment().getTaskInfo().getNumberOfParallelSubtasks());
+                    new InitializationContext() {
+                        @Override
+                        public int getNumTasks() {
+                            return getEnvironment().getTaskInfo().getNumberOfParallelSubtasks();
+                        }
+
+                        @Override
+                        public int getTaskNumber() {
+                            return getEnvironment().getTaskInfo().getIndexOfThisSubtask();
+                        }
+
+                        @Override
+                        public int getAttemptNumber() {
+                            return getEnvironment().getTaskInfo().getAttemptNumber();
+                        }
+                    });
 
             if (objectReuseEnabled) {
                 IT record = serializer.createInstance();

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/sink/OutputFormatSinkFunction.java

@@ -22,6 +22,7 @@
 import org.apache.flink.api.common.functions.RuntimeContext;
 import org.apache.flink.api.common.io.CleanupWhenUnsuccessful;
 import org.apache.flink.api.common.io.OutputFormat;
+import org.apache.flink.api.common.io.OutputFormat.InitializationContext;
 import org.apache.flink.api.common.io.RichOutputFormat;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.typeutils.InputTypeConfigurable;
@@ -62,7 +63,23 @@ public void open(Configuration parameters) throws Exception {
         format.configure(parameters);
         int indexInSubtaskGroup = context.getIndexOfThisSubtask();
         int currentNumberOfSubtasks = context.getNumberOfParallelSubtasks();
-        format.open(indexInSubtaskGroup, currentNumberOfSubtasks);
+        format.open(
+                new InitializationContext() {
+                    @Override
+                    public int getNumTasks() {
+                        return currentNumberOfSubtasks;
+                    }
+
+                    @Override
+                    public int getTaskNumber() {
+                        return indexInSubtaskGroup;
+                    }
+
+                    @Override
+                    public int getAttemptNumber() {
+                        return context.getAttemptNumber();
+                    }
+                });
     }
 
     @Override

flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/runtime/utils/StreamTestSink.scala

@@ -386,7 +386,7 @@ class TestingOutputFormat[T](tz: TimeZone) extends OutputFormat[T] {
 
   def configure(var1: Configuration): Unit = {}
 
-  def open(taskNumber: Int, numTasks: Int): Unit = {
+  override def open(taskNumber: Int, numTasks: Int): Unit = {
     localRetractResults = mutable.ArrayBuffer.empty[String]
     StreamTestSink.synchronized {
       StreamTestSink.globalResults(index) += (taskNumber -> localRetractResults)

pom.xml

@@ -2185,6 +2185,8 @@ under the License.
 								<!-- UnionSerializerConfigSnapshot was a PublicEvolving and Deprecated class that has been removed, embedded inside a Public CoGroupedStreams class, triggering this false failure -->
 								<exclude>org.apache.flink.streaming.api.datastream.CoGroupedStreams$UnionSerializerConfigSnapshot</exclude>
 								<exclude>org.apache.flink.api.connector.source.SourceReaderContext#currentParallelism()</exclude>
+								<exclude>org.apache.flink.api.common.io.OutputFormat#open(int,int)</exclude>
+								<exclude>org.apache.flink.api.common.io.OutputFormat#open(org.apache.flink.api.common.io.OutputFormat$InitializationContext)</exclude>
 								<!-- MARKER: end exclusions -->
 							</excludes>
 							<accessModifier>public</accessModifier>