[FLINK-33986][runtime] Extend ShuffleMaster to support snapshot and r…

…estore state.

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/DefaultShuffleMasterSnapshotContext.java

@@ -0,0 +1,22 @@
+/*
+ * 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.flink.runtime.shuffle;
+
+/** Default {@link ShuffleMasterSnapshotContext} implementation. */
+public class DefaultShuffleMasterSnapshotContext implements ShuffleMasterSnapshotContext {}

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/EmptyShuffleMasterSnapshot.java

@@ -0,0 +1,37 @@
+/*
+ * 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.flink.runtime.shuffle;
+
+/**
+ * A singleton implementation of {@link ShuffleMasterSnapshot} that represents an empty snapshot of
+ * shuffle master.
+ */
+public class EmptyShuffleMasterSnapshot implements ShuffleMasterSnapshot {
+
+    private static final EmptyShuffleMasterSnapshot INSTANCE = new EmptyShuffleMasterSnapshot();
+
+    @Override
+    public boolean isIncremental() {
+        return false;
+    }
+
+    public static EmptyShuffleMasterSnapshot getInstance() {
+        return INSTANCE;
+    }
+}

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/NettyShuffleMaster.java

@@ -187,4 +187,16 @@ public void registerJob(JobShuffleContext context) {
     public void unregisterJob(JobID jobId) {
         jobShuffleContexts.remove(jobId);
     }
+
+    @Override
+    public boolean supportsBatchSnapshot() {
+        return true;
+    }
+
+    @Override
+    public void snapshotState(
+            CompletableFuture<ShuffleMasterSnapshot> snapshotFuture,
+            ShuffleMasterSnapshotContext context) {
+        snapshotFuture.complete(EmptyShuffleMasterSnapshot.getInstance());
+    }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/ShuffleMaster.java

@@ -23,6 +23,7 @@
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -121,4 +122,22 @@ default CompletableFuture<Collection<PartitionWithMetrics>> getAllPartitionWithM
             JobID jobId) {
         return CompletableFuture.completedFuture(Collections.emptyList());
     }
+
+    /**
+     * Whether the shuffle master supports taking snapshot in batch scenarios if {@link
+     * org.apache.flink.configuration.BatchExecutionOptions#JOB_RECOVERY_ENABLED} is true. If it
+     * returns true, Flink will call {@link #snapshotState} to take snapshot, and call {@link
+     * #restoreState} to restore the state of shuffle master.
+     */
+    default boolean supportsBatchSnapshot() {
+        return false;
+    }
+
+    /** Triggers a snapshot of the shuffle master's state. */
+    default void snapshotState(
+            CompletableFuture<ShuffleMasterSnapshot> snapshotFuture,
+            ShuffleMasterSnapshotContext context) {}
+
+    /** Restores the state of the shuffle master from the provided snapshots. */
+    default void restoreState(List<ShuffleMasterSnapshot> snapshots) {}
 }

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/ShuffleMasterSnapshot.java

@@ -0,0 +1,31 @@
+/*
+ * 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.flink.runtime.shuffle;
+
+import java.io.Serializable;
+
+/**
+ * This class represents a snapshot of shuffle master, which can be used to restore the internal
+ * state of the shuffle master. It is full or incremental.
+ */
+public interface ShuffleMasterSnapshot extends Serializable {
+
+    /** Returns whether the snapshot is incremental. */
+    boolean isIncremental();
+}

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/ShuffleMasterSnapshotContext.java

@@ -0,0 +1,22 @@
+/*
+ * 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.flink.runtime.shuffle;
+
+/** Snapshot context used to create {@link ShuffleMasterSnapshot}. */
+public interface ShuffleMasterSnapshotContext {}