This is an automated email from the ASF dual-hosted git repository. guoweijie pushed a commit to branch FLIP-410-update3 in repository https://gitbox.apache.org/repos/asf/flink.git
commit 60488e69e7aa2f772d9090a92447a876b1bec474 Author: Weijie Guo <res...@163.com> AuthorDate: Mon Apr 8 17:25:39 2024 +0800 [FLINK-34549][API] Implement JobInfo and expose it via RuntimeContext --- .../datastream/api/context/RuntimeContext.java | 3 +- ...PartitionedContext.java => DefaultJobInfo.java} | 25 ++++++-- .../impl/context/DefaultNonPartitionedContext.java | 11 ++++ .../impl/context/DefaultRuntimeContext.java | 13 +++- .../DefaultTwoOutputNonPartitionedContext.java | 12 ++++ .../datastream/impl/operators/ProcessOperator.java | 4 +- .../TwoInputBroadcastProcessOperator.java | 4 +- .../TwoInputNonBroadcastProcessOperator.java | 4 +- .../impl/operators/TwoOutputProcessOperator.java | 4 +- .../impl/context/DefaultJobInfoTest.java | 73 ++++++++++++++++++++++ .../state/api/runtime/SavepointEnvironment.java | 9 +++ .../flink/runtime/execution/Environment.java | 3 + .../DefaultExecutionGraphBuilder.java | 3 + .../runtime/executiongraph/JobInformation.java | 27 ++++++++ .../runtime/taskmanager/RuntimeEnvironment.java | 11 ++++ .../org/apache/flink/runtime/taskmanager/Task.java | 6 ++ .../operators/testutils/DummyEnvironment.java | 7 +++ .../operators/testutils/MockEnvironment.java | 14 ++++- .../testutils/MockEnvironmentBuilder.java | 15 +++++ .../api/operators/StreamingRuntimeContext.java | 5 ++ .../runtime/tasks/StreamMockEnvironment.java | 6 ++ 21 files changed, 241 insertions(+), 18 deletions(-) diff --git a/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/RuntimeContext.java b/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/RuntimeContext.java index 6a1af18b183..40476469b6b 100644 --- a/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/RuntimeContext.java +++ b/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/RuntimeContext.java @@ -28,5 +28,6 @@ import org.apache.flink.annotation.Experimental; */ @Experimental public interface RuntimeContext { - // TODO Introduce related methods in the subsequent RP. + /** Get the {@link JobInfo} of this process function. */ + JobInfo getJobInfo(); } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultJobInfo.java similarity index 53% copy from flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java copy to flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultJobInfo.java index 017bdcb8e25..faac1a9c3f5 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultJobInfo.java @@ -18,14 +18,27 @@ package org.apache.flink.datastream.impl.context; -import org.apache.flink.datastream.api.context.NonPartitionedContext; -import org.apache.flink.datastream.api.function.ApplyPartitionFunction; +import org.apache.flink.datastream.api.context.JobInfo; +import org.apache.flink.runtime.jobgraph.JobType; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; -/** The default implementation of {@link NonPartitionedContext}. */ -public class DefaultNonPartitionedContext<OUT> implements NonPartitionedContext<OUT> { +/** Default implementation of {@link JobInfo}. */ +public class DefaultJobInfo implements JobInfo { + private final StreamingRuntimeContext operatorContext; + + public DefaultJobInfo(StreamingRuntimeContext streamingRuntimeContext) { + this.operatorContext = streamingRuntimeContext; + } + + @Override + public String getJobName() { + return operatorContext.getJobInfo().getJobName(); + } @Override - public void applyToAllPartitions(ApplyPartitionFunction<OUT> applyPartitionFunction) { - // TODO implements this method. + public ExecutionMode getExecutionMode() { + return operatorContext.getJobType() == JobType.STREAMING + ? ExecutionMode.STREAMING + : ExecutionMode.BATCH; } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java index 017bdcb8e25..a099272bd6d 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java @@ -18,14 +18,25 @@ package org.apache.flink.datastream.impl.context; +import org.apache.flink.datastream.api.context.JobInfo; import org.apache.flink.datastream.api.context.NonPartitionedContext; import org.apache.flink.datastream.api.function.ApplyPartitionFunction; /** The default implementation of {@link NonPartitionedContext}. */ public class DefaultNonPartitionedContext<OUT> implements NonPartitionedContext<OUT> { + private final DefaultRuntimeContext context; + + public DefaultNonPartitionedContext(DefaultRuntimeContext context) { + this.context = context; + } @Override public void applyToAllPartitions(ApplyPartitionFunction<OUT> applyPartitionFunction) { // TODO implements this method. } + + @Override + public JobInfo getJobInfo() { + return context.getJobInfo(); + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultRuntimeContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultRuntimeContext.java index 22f61223ca1..492c8fa5c2c 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultRuntimeContext.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultRuntimeContext.java @@ -18,9 +18,20 @@ package org.apache.flink.datastream.impl.context; +import org.apache.flink.datastream.api.context.JobInfo; import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; /** The default implementation of {@link RuntimeContext}. */ public class DefaultRuntimeContext implements RuntimeContext { - // TODO implements this class + private final DefaultJobInfo jobInfo; + + public DefaultRuntimeContext(StreamingRuntimeContext operatorContext) { + this.jobInfo = new DefaultJobInfo(operatorContext); + } + + @Override + public JobInfo getJobInfo() { + return jobInfo; + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContext.java index 7a76e923191..a6043f31a06 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContext.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContext.java @@ -18,15 +18,27 @@ package org.apache.flink.datastream.impl.context; +import org.apache.flink.datastream.api.context.JobInfo; import org.apache.flink.datastream.api.context.TwoOutputNonPartitionedContext; import org.apache.flink.datastream.api.function.TwoOutputApplyPartitionFunction; /** The default implementation of {@link TwoOutputNonPartitionedContext}. */ public class DefaultTwoOutputNonPartitionedContext<OUT1, OUT2> implements TwoOutputNonPartitionedContext<OUT1, OUT2> { + private final DefaultRuntimeContext context; + + public DefaultTwoOutputNonPartitionedContext(DefaultRuntimeContext context) { + this.context = context; + } + @Override public void applyToAllPartitions( TwoOutputApplyPartitionFunction<OUT1, OUT2> applyPartitionFunction) { // TODO implements this method. } + + @Override + public JobInfo getJobInfo() { + return context.getJobInfo(); + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java index 13204b27b02..40cc595ce53 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java @@ -52,9 +52,9 @@ public class ProcessOperator<IN, OUT> @Override public void open() throws Exception { super.open(); - context = new DefaultRuntimeContext(); + context = new DefaultRuntimeContext(getRuntimeContext()); partitionedContext = new DefaultPartitionedContext(context); - nonPartitionedContext = new DefaultNonPartitionedContext<>(); + nonPartitionedContext = new DefaultNonPartitionedContext<>(context); outputCollector = getOutputCollector(); } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java index 7c768715375..f8fdca63bd5 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java @@ -56,9 +56,9 @@ public class TwoInputBroadcastProcessOperator<IN1, IN2, OUT> public void open() throws Exception { super.open(); this.collector = getOutputCollector(); - this.context = new DefaultRuntimeContext(); + this.context = new DefaultRuntimeContext(getRuntimeContext()); this.partitionedContext = new DefaultPartitionedContext(context); - this.nonPartitionedContext = new DefaultNonPartitionedContext<>(); + this.nonPartitionedContext = new DefaultNonPartitionedContext<>(context); } @Override diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java index 36e403afc35..dff0284ae28 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java @@ -56,9 +56,9 @@ public class TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> public void open() throws Exception { super.open(); this.collector = getOutputCollector(); - this.context = new DefaultRuntimeContext(); + this.context = new DefaultRuntimeContext(getRuntimeContext()); this.partitionedContext = new DefaultPartitionedContext(context); - this.nonPartitionedContext = new DefaultNonPartitionedContext<>(); + this.nonPartitionedContext = new DefaultNonPartitionedContext<>(context); } @Override diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java index 671f0ee15cb..42e39ee7965 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java @@ -67,9 +67,9 @@ public class TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> public void open() throws Exception { this.mainCollector = getMainCollector(); this.sideCollector = getSideCollector(); - this.context = new DefaultRuntimeContext(); + this.context = new DefaultRuntimeContext(getRuntimeContext()); this.partitionedContext = new DefaultPartitionedContext(context); - this.nonPartitionedContext = new DefaultTwoOutputNonPartitionedContext<>(); + this.nonPartitionedContext = new DefaultTwoOutputNonPartitionedContext<>(context); } @Override diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultJobInfoTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultJobInfoTest.java new file mode 100644 index 00000000000..3e08793286f --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultJobInfoTest.java @@ -0,0 +1,73 @@ +/* + * 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.datastream.impl.context; + +import org.apache.flink.datastream.api.context.JobInfo; +import org.apache.flink.runtime.jobgraph.JobType; +import org.apache.flink.runtime.memory.MemoryManager; +import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder; +import org.apache.flink.streaming.util.MockStreamingRuntimeContext; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link DefaultJobInfo}. */ +class DefaultJobInfoTest { + @Test + void testGetJobName() { + String jobName = "Test-Job"; + DefaultJobInfo jobInfo = createDefaultJobInfo(jobName); + assertThat(jobInfo.getJobName()).isEqualTo(jobName); + } + + @Test + void testGetExecutionMode() { + DefaultJobInfo batchJob = createDefaultJobInfo(JobType.BATCH); + assertThat(batchJob.getExecutionMode()).isEqualTo(JobInfo.ExecutionMode.BATCH); + + DefaultJobInfo streamingJob = createDefaultJobInfo(JobType.STREAMING); + assertThat(streamingJob.getExecutionMode()).isEqualTo(JobInfo.ExecutionMode.STREAMING); + } + + private static DefaultJobInfo createDefaultJobInfo(String jobName) { + return createDefaultJobInfo(JobType.STREAMING, jobName); + } + + private static DefaultJobInfo createDefaultJobInfo(JobType jobType) { + return createDefaultJobInfo(jobType, "mock-job"); + } + + private static DefaultJobInfo createDefaultJobInfo(JobType jobType, String jobName) { + return new DefaultJobInfo( + new MockStreamingRuntimeContext( + false, + 2, + 1, + new MockEnvironmentBuilder() + .setTaskName("mockTask") + .setManagedMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE) + .setParallelism(2) + .setMaxParallelism(2) + .setSubtaskIndex(1) + .setJobType(jobType) + .setJobName(jobName) + .build())); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java index 350d0fc9875..00f471f1a7b 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java @@ -45,6 +45,7 @@ import org.apache.flink.runtime.io.disk.iomanager.IOManagerAsync; import org.apache.flink.runtime.io.network.TaskEventDispatcher; import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter; import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertexID; import org.apache.flink.runtime.jobgraph.OperatorID; import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider; @@ -85,6 +86,8 @@ public class SavepointEnvironment implements Environment { private final JobID jobID; + private final JobType jobType; + private final JobVertexID vertexID; private final ExecutionAttemptID attemptID; @@ -124,6 +127,7 @@ public class SavepointEnvironment implements Environment { PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) { this.jobID = new JobID(); this.vertexID = new JobVertexID(); + this.jobType = JobType.STREAMING; this.attemptID = new ExecutionAttemptID( new ExecutionGraphID(), new ExecutionVertexID(vertexID, indexOfSubtask), 0); @@ -156,6 +160,11 @@ public class SavepointEnvironment implements Environment { return jobID; } + @Override + public JobType getJobType() { + return jobType; + } + @Override public JobVertexID getJobVertexId() { return vertexID; diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java index 467f698045b..80199f490c4 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java @@ -37,6 +37,7 @@ import org.apache.flink.runtime.io.disk.iomanager.IOManager; import org.apache.flink.runtime.io.network.TaskEventDispatcher; import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter; import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertexID; import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider; import org.apache.flink.runtime.jobgraph.tasks.TaskOperatorEventGateway; @@ -77,6 +78,8 @@ public interface Environment { */ JobID getJobID(); + JobType getJobType(); + /** * Gets the ID of the JobVertex for which this task executes a parallel subtask. * diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphBuilder.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphBuilder.java index 7e41af896df..57be7f49abf 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphBuilder.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphBuilder.java @@ -38,6 +38,7 @@ import org.apache.flink.runtime.executiongraph.failover.partitionrelease.Partiti import org.apache.flink.runtime.executiongraph.failover.partitionrelease.PartitionGroupReleaseStrategyFactoryLoader; import org.apache.flink.runtime.io.network.partition.JobMasterPartitionTracker; import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertex; import org.apache.flink.runtime.jobgraph.jsonplan.JsonPlanGenerator; import org.apache.flink.runtime.jobgraph.tasks.CheckpointCoordinatorConfiguration; @@ -103,10 +104,12 @@ public class DefaultExecutionGraphBuilder { final String jobName = jobGraph.getName(); final JobID jobId = jobGraph.getJobID(); + final JobType jobType = jobGraph.getJobType(); final JobInformation jobInformation = new JobInformation( jobId, + jobType, jobName, jobGraph.getSerializedExecutionConfig(), jobGraph.getJobConfiguration(), diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/JobInformation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/JobInformation.java index 5792caa2711..94a9fc9b34b 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/JobInformation.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/JobInformation.java @@ -23,6 +23,7 @@ import org.apache.flink.api.common.JobID; import org.apache.flink.configuration.Configuration; import org.apache.flink.configuration.UnmodifiableConfiguration; import org.apache.flink.runtime.blob.PermanentBlobKey; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.util.Preconditions; import org.apache.flink.util.SerializedValue; @@ -42,6 +43,9 @@ public class JobInformation implements Serializable { /** Id of the job. */ private final JobID jobId; + /** Type of the job. */ + private final JobType jobType; + /** Job name. */ private final String jobName; @@ -64,7 +68,26 @@ public class JobInformation implements Serializable { Configuration jobConfiguration, Collection<PermanentBlobKey> requiredJarFileBlobKeys, Collection<URL> requiredClasspathURLs) { + this( + jobId, + JobType.STREAMING, + jobName, + serializedExecutionConfig, + jobConfiguration, + requiredJarFileBlobKeys, + requiredClasspathURLs); + } + + public JobInformation( + JobID jobId, + JobType jobType, + String jobName, + SerializedValue<ExecutionConfig> serializedExecutionConfig, + Configuration jobConfiguration, + Collection<PermanentBlobKey> requiredJarFileBlobKeys, + Collection<URL> requiredClasspathURLs) { this.jobId = Preconditions.checkNotNull(jobId); + this.jobType = Preconditions.checkNotNull(jobType); this.jobName = Preconditions.checkNotNull(jobName); this.serializedExecutionConfig = Preconditions.checkNotNull(serializedExecutionConfig); this.jobConfiguration = @@ -83,6 +106,10 @@ public class JobInformation implements Serializable { return jobName; } + public JobType getJobType() { + return jobType; + } + public SerializedValue<ExecutionConfig> getSerializedExecutionConfig() { return serializedExecutionConfig; } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java index 104b42c2b1a..d5e297c54b0 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java @@ -38,6 +38,7 @@ import org.apache.flink.runtime.io.disk.iomanager.IOManager; import org.apache.flink.runtime.io.network.TaskEventDispatcher; import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter; import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertexID; import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider; import org.apache.flink.runtime.jobgraph.tasks.TaskOperatorEventGateway; @@ -64,6 +65,9 @@ public class RuntimeEnvironment implements Environment { private final JobID jobId; private final JobVertexID jobVertexId; + + private final JobType jobType; + private final ExecutionAttemptID executionId; private final JobInfo jobInfo; @@ -118,6 +122,7 @@ public class RuntimeEnvironment implements Environment { public RuntimeEnvironment( JobID jobId, + JobType jobType, JobVertexID jobVertexId, ExecutionAttemptID executionId, ExecutionConfig executionConfig, @@ -149,6 +154,7 @@ public class RuntimeEnvironment implements Environment { TaskManagerActions taskManagerActions) { this.jobId = checkNotNull(jobId); + this.jobType = checkNotNull(jobType); this.jobVertexId = checkNotNull(jobVertexId); this.executionId = checkNotNull(executionId); this.jobInfo = checkNotNull(jobInfo); @@ -192,6 +198,11 @@ public class RuntimeEnvironment implements Environment { return jobId; } + @Override + public JobType getJobType() { + return jobType; + } + @Override public JobVertexID getJobVertexId() { return jobVertexId; diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java index e813de5c398..afc74f9a98d 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java @@ -62,6 +62,7 @@ import org.apache.flink.runtime.io.network.partition.ResultPartitionID; import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate; import org.apache.flink.runtime.io.network.partition.consumer.InputGate; import org.apache.flink.runtime.jobgraph.IntermediateDataSetID; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertexID; import org.apache.flink.runtime.jobgraph.OperatorID; import org.apache.flink.runtime.jobgraph.tasks.CheckpointableTask; @@ -162,6 +163,9 @@ public class Task /** The job that the task belongs to. */ private final JobID jobId; + /** The type of this job. */ + private final JobType jobType; + /** The vertex in the JobGraph whose code the task executes. */ private final JobVertexID vertexId; @@ -352,6 +356,7 @@ public class Task String.valueOf(slotAllocationId)); this.jobId = jobInformation.getJobId(); + this.jobType = jobInformation.getJobType(); this.vertexId = taskInformation.getJobVertexId(); this.executionId = Preconditions.checkNotNull(executionAttemptID); this.allocationId = Preconditions.checkNotNull(slotAllocationId); @@ -695,6 +700,7 @@ public class Task Environment env = new RuntimeEnvironment( jobId, + jobType, vertexId, executionId, executionConfig, diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java index edb85cf19c5..0a6a425c4ac 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java @@ -39,6 +39,7 @@ import org.apache.flink.runtime.io.disk.iomanager.IOManager; import org.apache.flink.runtime.io.network.TaskEventDispatcher; import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter; import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertexID; import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider; import org.apache.flink.runtime.jobgraph.tasks.TaskOperatorEventGateway; @@ -72,6 +73,7 @@ public class DummyEnvironment implements Environment { private final JobInfo jobInfo = new JobInfoImpl(new JobID(), "DummyJob"); private final JobVertexID jobVertexId = new JobVertexID(); + private final JobType jobType = JobType.STREAMING; private final ExecutionAttemptID executionId; private final ExecutionConfig executionConfig = new ExecutionConfig(); private final TaskInfo taskInfo; @@ -128,6 +130,11 @@ public class DummyEnvironment implements Environment { return jobInfo.getJobId(); } + @Override + public JobType getJobType() { + return jobType; + } + @Override public JobVertexID getJobVertexId() { return jobVertexId; diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java index 9d0249c1257..ec9bfcb917d 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java @@ -42,6 +42,7 @@ import org.apache.flink.runtime.io.network.api.writer.RecordCollectingResultPart import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter; import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate; import org.apache.flink.runtime.io.network.partition.consumer.IteratorWrappingTestSingleInputGate; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertexID; import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider; import org.apache.flink.runtime.jobgraph.tasks.TaskOperatorEventGateway; @@ -85,6 +86,8 @@ public class MockEnvironment implements Environment, AutoCloseable { private final TaskInfo taskInfo; + private final JobType jobType; + private final ExecutionConfig executionConfig; private final MemoryManager memManager; @@ -149,7 +152,9 @@ public class MockEnvironment implements Environment, AutoCloseable { protected MockEnvironment( JobID jobID, + String jobName, JobVertexID jobVertexID, + JobType jobType, String taskName, MockInputSplitProvider inputSplitProvider, int bufferSize, @@ -168,9 +173,9 @@ public class MockEnvironment implements Environment, AutoCloseable { ExternalResourceInfoProvider externalResourceInfoProvider, ChannelStateWriteRequestExecutorFactory channelStateExecutorFactory) { - this.jobInfo = new JobInfoImpl(jobID, "MockJob"); + this.jobInfo = new JobInfoImpl(jobID, jobName); this.jobVertexID = jobVertexID; - + this.jobType = jobType; this.taskInfo = new TaskInfoImpl(taskName, maxParallelism, subtaskIndex, parallelism, 0); this.jobConfiguration = new Configuration(); this.taskConfiguration = taskConfiguration; @@ -269,6 +274,11 @@ public class MockEnvironment implements Environment, AutoCloseable { return this.jobInfo.getJobId(); } + @Override + public JobType getJobType() { + return jobType; + } + @Override public Configuration getJobConfiguration() { return this.jobConfiguration; diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironmentBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironmentBuilder.java index ec61b9a6fa5..2a8c7c273a8 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironmentBuilder.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironmentBuilder.java @@ -25,6 +25,7 @@ import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequestExecu import org.apache.flink.runtime.externalresource.ExternalResourceInfoProvider; import org.apache.flink.runtime.io.disk.iomanager.IOManager; import org.apache.flink.runtime.io.disk.iomanager.IOManagerAsync; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertexID; import org.apache.flink.runtime.memory.MemoryManager; import org.apache.flink.runtime.memory.MemoryManagerBuilder; @@ -41,6 +42,7 @@ import org.apache.flink.util.UserCodeClassLoader; public class MockEnvironmentBuilder { private String taskName = "mock-task"; + private String jobName = "mock-job"; private MockInputSplitProvider inputSplitProvider = null; private int bufferSize = 16; private TaskStateManager taskStateManager = new TestTaskStateManager(); @@ -54,6 +56,7 @@ public class MockEnvironmentBuilder { TestingUserCodeClassLoader.newBuilder().build(); private JobID jobID = new JobID(); private JobVertexID jobVertexID = new JobVertexID(); + private JobType jobType = JobType.STREAMING; private TaskMetricGroup taskMetricGroup = UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(); private TaskManagerRuntimeInfo taskManagerRuntimeInfo = new TestingTaskManagerRuntimeInfo(); @@ -136,6 +139,11 @@ public class MockEnvironmentBuilder { return this; } + public MockEnvironmentBuilder setJobName(String jobName) { + this.jobName = jobName; + return this; + } + public MockEnvironmentBuilder setJobVertexID(JobVertexID jobVertexID) { this.jobVertexID = jobVertexID; return this; @@ -168,13 +176,20 @@ public class MockEnvironmentBuilder { return this; } + public MockEnvironmentBuilder setJobType(JobType jobType) { + this.jobType = jobType; + return this; + } + public MockEnvironment build() { if (ioManager == null) { ioManager = new IOManagerAsync(); } return new MockEnvironment( jobID, + jobName, jobVertexID, + jobType, taskName, inputSplitProvider, bufferSize, diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContext.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContext.java index 2e82d1c2b08..d7bb2bf57f6 100644 --- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContext.java +++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContext.java @@ -40,6 +40,7 @@ import org.apache.flink.configuration.Configuration; import org.apache.flink.metrics.groups.OperatorMetricGroup; import org.apache.flink.runtime.execution.Environment; import org.apache.flink.runtime.externalresource.ExternalResourceInfoProvider; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.OperatorID; import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider; import org.apache.flink.runtime.state.v2.KeyedStateStoreV2; @@ -169,6 +170,10 @@ public class StreamingRuntimeContext extends AbstractRuntimeUDFContext { return taskEnvironment.getJobConfiguration(); } + public JobType getJobType() { + return taskEnvironment.getJobType(); + } + @Override public Set<ExternalResourceInfo> getExternalResourceInfos(String resourceName) { return externalResourceInfoProvider.getExternalResourceInfos(resourceName); diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java index 7e9f04be81c..9fb6b733762 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java @@ -43,6 +43,7 @@ import org.apache.flink.runtime.io.network.TaskEventDispatcher; import org.apache.flink.runtime.io.network.api.writer.RecordOrEventCollectingResultPartitionWriter; import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter; import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate; +import org.apache.flink.runtime.jobgraph.JobType; import org.apache.flink.runtime.jobgraph.JobVertexID; import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider; import org.apache.flink.runtime.jobgraph.tasks.TaskOperatorEventGateway; @@ -285,6 +286,11 @@ public class StreamMockEnvironment implements Environment { return this.jobInfo.getJobId(); } + @Override + public JobType getJobType() { + return JobType.STREAMING; + } + @Override public Configuration getJobConfiguration() { return this.jobConfiguration;