Go VR Flink test on Flink 2.0

Author: Abacn

runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkExecutionEnvironments.java

@@ -36,6 +36,7 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.RuntimeExecutionMode;
+import org.apache.flink.api.common.serialization.SerializerConfigImpl;
 import org.apache.flink.configuration.CheckpointingOptions;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.CoreOptions;
@@ -270,6 +271,7 @@ public static StreamExecutionEnvironment createStreamExecutionEnvironment(
       flinkStreamEnv.getConfig().setAutoWatermarkInterval(options.getAutoWatermarkInterval());
     }
     configureWebUIOptions(flinkStreamEnv.getConfig(), options.as(PipelineOptions.class));
+    configureCustomKryoSerializers(flinkStreamEnv.getConfig());
 
     return flinkStreamEnv;
   }
@@ -294,6 +296,14 @@ private static void configureWebUIOptions(
     }
   }
 
+  private static void configureCustomKryoSerializers(ExecutionConfig config) {
+    SerializerConfigImpl serializerConfig = (SerializerConfigImpl) config.getSerializerConfig();
+    // Force Beam schema to use JavaSerializer to fix serialization involving ImmutableMap
+    serializerConfig.registerTypeWithKryoSerializer(
+        org.apache.beam.sdk.schemas.Schema.class,
+        com.esotericsoftware.kryo.serializers.JavaSerializer.class);
+  }
+
   private static class GlobalJobParametersImpl extends ExecutionConfig.GlobalJobParameters {
     private final Map<String, String> jobOptions;
 

runners/flink/flink_runner.gradle

@@ -230,6 +230,13 @@ dependencies {
   // configuration (https://issues.apache.org/jira/browse/BEAM-11732).
   permitUnusedDeclared "org.apache.flink:flink-clients:$flink_version"
 
+  // align with Flink's kryo version
+  def kryo_version = '2.24.0'
+  if (flink_major.startsWith('2')) {
+    kryo_version = '5.6.2'
+  }
+  implementation "com.esotericsoftware.kryo:kryo:$kryo_version"
+
   implementation "org.apache.flink:flink-streaming-java:$flink_version"
   testImplementation "org.apache.flink:flink-statebackend-rocksdb:$flink_version"
   testImplementation "org.apache.flink:flink-streaming-java:$flink_version:tests"

runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkExecutionEnvironments.java

@@ -298,6 +298,7 @@ public static StreamExecutionEnvironment createStreamExecutionEnvironment(
     configureStateBackend(options, flinkStreamEnv);
 
     configureWebUIOptions(flinkStreamEnv.getConfig(), options.as(PipelineOptions.class));
+    configureCustomKryoSerializers(flinkStreamEnv.getConfig());
 
     return flinkStreamEnv;
   }
@@ -322,6 +323,13 @@ private static void configureWebUIOptions(
     }
   }
 
+  private static void configureCustomKryoSerializers(ExecutionConfig config) {
+    // Force Beam schema to use JavaSerializer to fix serialization involving ImmutableMap
+    config.registerTypeWithKryoSerializer(
+        org.apache.beam.sdk.schemas.Schema.class,
+        com.esotericsoftware.kryo.serializers.JavaSerializer.class);
+  }
+
   private static class GlobalJobParametersImpl extends ExecutionConfig.GlobalJobParameters {
     private final Map<String, String> jobOptions;
 

sdks/go/test/build.gradle

@@ -79,8 +79,7 @@ task dataflowValidatesRunnerARM64() {
 task flinkValidatesRunner {
   group = "Verification"
 
-  // TODO(https://github.com/apache/beam/issues/37600) use project.ext.latestFlinkVersion after resolved
-  def flinkVersion = '1.20'
+  def flinkVersion = project.ext.latestFlinkVersion
 
   dependsOn ":sdks:go:test:goBuild"
   dependsOn ":sdks:go:container:docker"