[hotfix] handle in the same way in StreamExecutionEnvironment

apache · Jun 4, 2019 · d8bab9a · d8bab9a
1 parent 83b7293
commit d8bab9a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/flink-java/src/main/java/org/apache/flink/api/java/ExecutionEnvironment.java b/flink-java/src/main/java/org/apache/flink/api/java/ExecutionEnvironment.java
@@ -1065,11 +1065,10 @@ private static String getDefaultName() {
  */
  public static ExecutionEnvironment getExecutionEnvironment() {
 
- return contextEnvironmentFactory == null ?
- (contextEnvironmentFactoryThreadLocal.get() == null ?
- createLocalEnvironment() :
- contextEnvironmentFactoryThreadLocal.get().createExecutionEnvironment()) :
- contextEnvironmentFactory.createExecutionEnvironment();
+ return contextEnvironmentFactoryThreadLocal.get() == null ?
+ (contextEnvironmentFactory == null ?
+ createLocalEnvironment() : contextEnvironmentFactory.createExecutionEnvironment()) :
+ contextEnvironmentFactoryThreadLocal.get().createExecutionEnvironment();
  }
 
  /**

diff --git a/.../src/main/java/org/apache/flink/streaming/api/environment/StreamExecutionEnvironment.java b/.../src/main/java/org/apache/flink/streaming/api/environment/StreamExecutionEnvironment.java
@@ -117,6 +117,9 @@ public abstract class StreamExecutionEnvironment {
  */
  private static StreamExecutionEnvironmentFactory contextEnvironmentFactory;
 
+ /** The ThreadLocal used to store {@link StreamExecutionEnvironmentFactory}. */
+ private static ThreadLocal<StreamExecutionEnvironmentFactory> contextEnvironmentFactoryThreadLocal = new ThreadLocal<>();
+
  /** The default parallelism used when creating a local environment. */
  private static int defaultLocalParallelism = Runtime.getRuntime().availableProcessors();
 
@@ -1568,6 +1571,9 @@ public void addOperator(StreamTransformation<?> transformation) {
  * executed.
  */
  public static StreamExecutionEnvironment getExecutionEnvironment() {
+ if (contextEnvironmentFactoryThreadLocal.get() != null) {
+ return contextEnvironmentFactoryThreadLocal.get().createExecutionEnvironment();
+ }
  if (contextEnvironmentFactory != null) {
  return contextEnvironmentFactory.createExecutionEnvironment();
  }
@@ -1766,10 +1772,12 @@ public static void setDefaultLocalParallelism(int parallelism) {
 
  protected static void initializeContextEnvironment(StreamExecutionEnvironmentFactory ctx) {
  contextEnvironmentFactory = ctx;
+ contextEnvironmentFactoryThreadLocal.set(contextEnvironmentFactory);
  }
 
  protected static void resetContextEnvironment() {
  contextEnvironmentFactory = null;
+ contextEnvironmentFactoryThreadLocal.remove();
  }
 
  /**