Skip to content

Commit 4b56a2b

Browse files
aIbrahiimaIbrahiim
committed
fix FlinkSubmissionTest
Allow SecurityManager in Flink tests on Java 21 Fix Flink runner tests on Java 21
1 parent 477c747 commit 4b56a2b

5 files changed

Lines changed: 125 additions & 32 deletions

File tree

‎.github/trigger_files/beam_PreCommit_Java.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
{
22
"comment": "Modify this file in a trivial way to cause this test suite to run.",
3-
"modification": 1
3+
"modification": 5
44
}

‎runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/FlinkSubmissionTest.java‎

Lines changed: 46 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,12 @@
1919

2020
import java.io.File;
2121
import java.lang.reflect.Field;
22-
import java.lang.reflect.Modifier;
22+
import java.lang.reflect.Method;
2323
import java.nio.charset.StandardCharsets;
2424
import java.nio.file.Files;
2525
import java.security.Permission;
2626
import java.util.Collection;
27+
import java.util.HashMap;
2728
import java.util.Map;
2829
import java.util.concurrent.TimeUnit;
2930
import java.util.stream.Collectors;
@@ -61,7 +62,8 @@ public class FlinkSubmissionTest {
6162
private static final Logger LOG = LoggerFactory.getLogger(FlinkSubmissionTest.class);
6263

6364
@ClassRule public static final TemporaryFolder TEMP_FOLDER = new TemporaryFolder();
64-
private static final Map<String, String> ENV = System.getenv();
65+
// Snapshot at class init; System.getenv() is a live view that we mutate in-place below.
66+
private static final Map<String, String> ENV = new HashMap<>(System.getenv());
6567
private static final SecurityManager SECURITY_MANAGER = System.getSecurityManager();
6668

6769
/** Flink cluster that runs over the lifespan of the tests. */
@@ -210,21 +212,52 @@ private static void restoreEnvironment() throws Exception {
210212
* We modify the JVM's environment variables here. This is necessary for the end-to-end test
211213
* because Flink's CliFrontend requires a Flink configuration file for which the location can only
212214
* be set using the {@code ConfigConstants.ENV_FLINK_CONF_DIR} environment variable.
215+
*
216+
* <p>On Unix JDKs, {@code ProcessEnvironment.theEnvironment} is a {@code Map<Variable,Value>}.
217+
* {@code System.getenv(String)} looks up {@code Variable} keys, so inserting plain {@code String}
218+
* keys (via {@code putAll}) is a no-op for lookups and leaves Flink unable to find {@code
219+
* FLINK_CONF_DIR}. On Windows the map uses {@code String} keys (plus a case-insensitive map).
213220
*/
221+
@SuppressWarnings("unchecked")
214222
private static void modifyEnv(Map<String, String> env) throws Exception {
215-
Class processEnv = Class.forName("java.lang.ProcessEnvironment");
216-
Field envField = processEnv.getDeclaredField("theUnmodifiableEnvironment");
223+
Class<?> processEnvironment = Class.forName("java.lang.ProcessEnvironment");
217224

218-
Field modifiersField = Field.class.getDeclaredField("modifiers");
219-
modifiersField.setAccessible(true);
220-
modifiersField.setInt(envField, envField.getModifiers() & ~Modifier.FINAL);
221-
222-
envField.setAccessible(true);
223-
envField.set(null, env);
224-
envField.setAccessible(false);
225+
Field theEnvironmentField = processEnvironment.getDeclaredField("theEnvironment");
226+
theEnvironmentField.setAccessible(true);
227+
Map<Object, Object> envMap = (Map<Object, Object>) theEnvironmentField.get(null);
228+
envMap.clear();
229+
try {
230+
// Unix: keys/values are ProcessEnvironment$Variable / $Value, not String.
231+
Class<?> variableClass = Class.forName("java.lang.ProcessEnvironment$Variable");
232+
Class<?> valueClass = Class.forName("java.lang.ProcessEnvironment$Value");
233+
Method valueOfVariable = variableClass.getDeclaredMethod("valueOf", String.class);
234+
Method valueOfValue = valueClass.getDeclaredMethod("valueOf", String.class);
235+
valueOfVariable.setAccessible(true);
236+
valueOfValue.setAccessible(true);
237+
for (Map.Entry<String, String> entry : env.entrySet()) {
238+
envMap.put(
239+
valueOfVariable.invoke(null, entry.getKey()),
240+
valueOfValue.invoke(null, entry.getValue()));
241+
}
242+
} catch (ClassNotFoundException e) {
243+
// Windows (and similar): theEnvironment uses String keys.
244+
envMap.putAll(env);
245+
}
225246

226-
modifiersField.setInt(envField, envField.getModifiers() & Modifier.FINAL);
227-
modifiersField.setAccessible(false);
247+
// Windows: System.getenv(String) reads the case-insensitive map when present.
248+
try {
249+
Field theCaseInsensitiveEnvironmentField =
250+
processEnvironment.getDeclaredField("theCaseInsensitiveEnvironment");
251+
theCaseInsensitiveEnvironmentField.setAccessible(true);
252+
Map<String, String> ciEnvMap =
253+
(Map<String, String>) theCaseInsensitiveEnvironmentField.get(null);
254+
if (ciEnvMap != null) {
255+
ciEnvMap.clear();
256+
ciEnvMap.putAll(env);
257+
}
258+
} catch (NoSuchFieldException e) {
259+
// Not present on Unix JDKs.
260+
}
228261
}
229262

230263
/** Prevents the CliFrontend from calling System.exit. */

‎runners/flink/flink_runner.gradle‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,13 +180,30 @@ if (use_override) {
180180
}
181181
}
182182

183+
def flinkTestJvmArgs() {
184+
def testJavaVer = project.findProperty('testJavaVersion') ? (project.property('testJavaVersion') as int) : JavaVersion.current().majorVersion.toInteger()
185+
if (testJavaVer >= 17) {
186+
return [
187+
"--add-opens=java.base/sun.nio.ch=ALL-UNNAMED",
188+
"--add-opens=java.base/java.nio=ALL-UNNAMED",
189+
"--add-opens=java.base/java.util=ALL-UNNAMED",
190+
"--add-opens=java.base/java.lang.invoke=ALL-UNNAMED",
191+
"--add-opens=java.base/java.lang=ALL-UNNAMED",
192+
"-Djava.security.manager=allow",
193+
]
194+
} else {
195+
return []
196+
}
197+
}
198+
183199
test {
184200
systemProperty "log4j.configuration", "log4j-test.properties"
185201
// Change log level to debug:
186202
// systemProperty "org.slf4j.simpleLogger.defaultLogLevel", "debug"
187203
// Change log level to debug only for the package and nested packages:
188204
// systemProperty "org.slf4j.simpleLogger.log.org.apache.beam.runners.flink.translation.wrappers.streaming", "debug"
189205
jvmArgs "-XX:-UseGCOverheadLimit"
206+
jvmArgs += flinkTestJvmArgs()
190207
if (System.getProperty("beamSurefireArgline")) {
191208
jvmArgs System.getProperty("beamSurefireArgline")
192209
}

‎runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkSubmissionTest.java‎

Lines changed: 46 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,12 @@
1919

2020
import java.io.File;
2121
import java.lang.reflect.Field;
22-
import java.lang.reflect.Modifier;
22+
import java.lang.reflect.Method;
2323
import java.nio.charset.StandardCharsets;
2424
import java.nio.file.Files;
2525
import java.security.Permission;
2626
import java.util.Collection;
27+
import java.util.HashMap;
2728
import java.util.Map;
2829
import java.util.concurrent.TimeUnit;
2930
import java.util.stream.Collectors;
@@ -62,7 +63,8 @@ public class FlinkSubmissionTest {
6263
private static final Logger LOG = LoggerFactory.getLogger(FlinkSubmissionTest.class);
6364

6465
@ClassRule public static final TemporaryFolder TEMP_FOLDER = new TemporaryFolder();
65-
private static final Map<String, String> ENV = System.getenv();
66+
// Snapshot at class init; System.getenv() is a live view that we mutate in-place below.
67+
private static final Map<String, String> ENV = new HashMap<>(System.getenv());
6668
private static final SecurityManager SECURITY_MANAGER = System.getSecurityManager();
6769

6870
/** Flink cluster that runs over the lifespan of the tests. */
@@ -230,21 +232,52 @@ private static void restoreEnvironment() throws Exception {
230232
* We modify the JVM's environment variables here. This is necessary for the end-to-end test
231233
* because Flink's CliFrontend requires a Flink configuration file for which the location can only
232234
* be set using the {@code ConfigConstants.ENV_FLINK_CONF_DIR} environment variable.
235+
*
236+
* <p>On Unix JDKs, {@code ProcessEnvironment.theEnvironment} is a {@code Map<Variable,Value>}.
237+
* {@code System.getenv(String)} looks up {@code Variable} keys, so inserting plain {@code String}
238+
* keys (via {@code putAll}) is a no-op for lookups and leaves Flink unable to find {@code
239+
* FLINK_CONF_DIR}. On Windows the map uses {@code String} keys (plus a case-insensitive map).
233240
*/
241+
@SuppressWarnings("unchecked")
234242
private static void modifyEnv(Map<String, String> env) throws Exception {
235-
Class processEnv = Class.forName("java.lang.ProcessEnvironment");
236-
Field envField = processEnv.getDeclaredField("theUnmodifiableEnvironment");
243+
Class<?> processEnvironment = Class.forName("java.lang.ProcessEnvironment");
237244

238-
Field modifiersField = Field.class.getDeclaredField("modifiers");
239-
modifiersField.setAccessible(true);
240-
modifiersField.setInt(envField, envField.getModifiers() & ~Modifier.FINAL);
241-
242-
envField.setAccessible(true);
243-
envField.set(null, env);
244-
envField.setAccessible(false);
245+
Field theEnvironmentField = processEnvironment.getDeclaredField("theEnvironment");
246+
theEnvironmentField.setAccessible(true);
247+
Map<Object, Object> envMap = (Map<Object, Object>) theEnvironmentField.get(null);
248+
envMap.clear();
249+
try {
250+
// Unix: keys/values are ProcessEnvironment$Variable / $Value, not String.
251+
Class<?> variableClass = Class.forName("java.lang.ProcessEnvironment$Variable");
252+
Class<?> valueClass = Class.forName("java.lang.ProcessEnvironment$Value");
253+
Method valueOfVariable = variableClass.getDeclaredMethod("valueOf", String.class);
254+
Method valueOfValue = valueClass.getDeclaredMethod("valueOf", String.class);
255+
valueOfVariable.setAccessible(true);
256+
valueOfValue.setAccessible(true);
257+
for (Map.Entry<String, String> entry : env.entrySet()) {
258+
envMap.put(
259+
valueOfVariable.invoke(null, entry.getKey()),
260+
valueOfValue.invoke(null, entry.getValue()));
261+
}
262+
} catch (ClassNotFoundException e) {
263+
// Windows (and similar): theEnvironment uses String keys.
264+
envMap.putAll(env);
265+
}
245266

246-
modifiersField.setInt(envField, envField.getModifiers() & Modifier.FINAL);
247-
modifiersField.setAccessible(false);
267+
// Windows: System.getenv(String) reads the case-insensitive map when present.
268+
try {
269+
Field theCaseInsensitiveEnvironmentField =
270+
processEnvironment.getDeclaredField("theCaseInsensitiveEnvironment");
271+
theCaseInsensitiveEnvironmentField.setAccessible(true);
272+
Map<String, String> ciEnvMap =
273+
(Map<String, String>) theCaseInsensitiveEnvironmentField.get(null);
274+
if (ciEnvMap != null) {
275+
ciEnvMap.clear();
276+
ciEnvMap.putAll(env);
277+
}
278+
} catch (NoSuchFieldException e) {
279+
// Not present on Unix JDKs.
280+
}
248281
}
249282

250283
/** Prevents the CliFrontend from calling System.exit. */

‎runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapperTest.java‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -665,15 +665,16 @@ private static void testSourceDoesNotShutdown(boolean shouldHaveReaders) throws
665665
// Wait to see if the wrapper shuts down immediately in case it doesn't have readers
666666
if (!shouldHaveReaders) {
667667
// The expected state is for finalizeSource to sleep instead of exiting
668-
while (true) {
669-
StackTraceElement[] callStack = thread.getStackTrace();
670-
if (callStack.length >= 2
671-
&& "sleep".equals(callStack[0].getMethodName())
672-
&& "finalizeSource".equals(callStack[1].getMethodName())) {
668+
long deadlineMs = System.currentTimeMillis() + 9_000;
669+
boolean reachedFinalizeSource = false;
670+
while (System.currentTimeMillis() < deadlineMs) {
671+
if (isInFinalizeSource(thread.getStackTrace())) {
672+
reachedFinalizeSource = true;
673673
break;
674674
}
675675
Thread.sleep(10);
676676
}
677+
assertThat(reachedFinalizeSource, is(true));
677678
}
678679
// Source should still be running even if there are no readers
679680
assertThat(sourceWrapper.isRunning(), is(true));
@@ -694,6 +695,15 @@ private static void testSourceDoesNotShutdown(boolean shouldHaveReaders) throws
694695
assertThat(thread.isAlive(), is(false));
695696
}
696697

698+
private static boolean isInFinalizeSource(StackTraceElement[] callStack) {
699+
for (StackTraceElement frame : callStack) {
700+
if ("finalizeSource".equals(frame.getMethodName())) {
701+
return true;
702+
}
703+
}
704+
return false;
705+
}
706+
697707
@Test
698708
public void testSequentialReadingFromBoundedSource() throws Exception {
699709
UnboundedReadFromBoundedSource.BoundedToUnboundedSourceAdapter<Long> source =

0 commit comments

Comments
 (0)