Apache flink强制着色依赖的源代码

import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
import org.apache.flink.runtime.io.network.partition.consumer.InputChannel;
import org.apache.flink.runtime.io.network.partition.producer.SingleInputGate;
import org.apache.flink.runtime.taskmanager.TaskManager;

public class ForcedColoringExample {
    public static void main(String[] args) {
        // 配置 TaskManager
        TaskManager taskManager = new TaskManager(
            new Configuration(),
            new NetworkBufferPool(1024, 64 * 1024),
            new SingleInputGate<>(),
            new InputChannel<>(),
            true // 启用强制着色
        );

        // 启动 TaskManager
        taskManager.start();
    }
}

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

public class ForcedColoringKeyExample {
    public static void main(String[] args) throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStream<Tuple2<String, Integer>> dataStream = env.fromElements(
            new Tuple2<>("key1", 1),
            new Tuple2<>("key2", 2),
            new Tuple2<>("key1", 3),
            new Tuple2<>("key2", 4)
        );

        dataStream.keyBy(0)
                  .map(new RichMapFunction<Tuple2<String, Integer>, Tuple2<String, Integer>>() {
                      @Override
                      public void open(Configuration parameters) throws Exception {
                          // 启用强制着色
                          getRuntimeContext().setParallelism(2);
                      }

                      @Override
                      public Tuple2<String, Integer> map(Tuple2<String, Integer> value) throws Exception {
                          return value;
                      }
                  })
                  .print();

        env.execute("Forced Coloring Key Example");
    }
}