Repository: beam Updated Branches: refs/heads/master 7903e59c4 -> 3bffe0e00
[BEAM-1862] SplittableDoFnOperator should close the ScheduledExecutorService Project: http://git-wip-us.apache.org/repos/asf/beam/repo Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/fc4534cd Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/fc4534cd Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/fc4534cd Branch: refs/heads/master Commit: fc4534cd6e5366a5f12cefebcd52ac1fe7cdde41 Parents: 7903e59 Author: JingsongLi <lzljs3620...@aliyun.com> Authored: Tue Apr 4 18:28:15 2017 +0800 Committer: Aljoscha Krettek <aljoscha.kret...@gmail.com> Committed: Fri May 5 12:01:49 2017 +0200 ---------------------------------------------------------------------- .../streaming/SplittableDoFnOperator.java | 28 +++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/beam/blob/fc4534cd/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java ---------------------------------------------------------------------- diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java index 7d54cfa..968fc0a 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java @@ -24,6 +24,8 @@ import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import org.apache.beam.runners.core.ElementAndRestriction; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.KeyedWorkItems; @@ -57,6 +59,8 @@ public class SplittableDoFnOperator< extends DoFnOperator< KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, FnOutputT, OutputT> { + private transient ScheduledExecutorService executorService; + public SplittableDoFnOperator( DoFn<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, FnOutputT> doFn, String stepName, @@ -108,6 +112,8 @@ public class SplittableDoFnOperator< } }; + executorService = Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory()); + ((SplittableParDo.ProcessFn) doFn).setStateInternalsFactory(stateInternalsFactory); ((SplittableParDo.ProcessFn) doFn).setTimerInternalsFactory(timerInternalsFactory); ((SplittableParDo.ProcessFn) doFn).setProcessElementInvoker( @@ -137,7 +143,7 @@ public class SplittableDoFnOperator< } }, sideInputReader, - Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory()), + executorService, 10000, Duration.standardSeconds(10))); } @@ -149,4 +155,24 @@ public class SplittableDoFnOperator< (String) stateInternals.getKey(), Collections.singletonList(timer.getNamespace())))); } + + @Override + public void close() throws Exception { + super.close(); + + executorService.shutdown(); + + long shutdownTimeout = Duration.standardSeconds(10).getMillis(); + try { + if (!executorService.awaitTermination(shutdownTimeout, TimeUnit.MILLISECONDS)) { + LOG.debug("The scheduled executor service did not properly terminate. Shutting " + + "it down now."); + executorService.shutdownNow(); + } + } catch (InterruptedException e) { + LOG.debug("Could not properly await the termination of the scheduled executor service.", e); + executorService.shutdownNow(); + } + } + }