From e8f09e45c2b4c4085c08e09970d898f1e847eae6 Mon Sep 17 00:00:00 2001 From: Alan Zhang Date: Fri, 13 Jan 2023 16:45:27 -0800 Subject: [PATCH 1/5] allow to set timeout for finishing a remote bundle processing --- .../samza/SamzaPortablePipelineOptions.java | 8 +++++ .../beam/runners/samza/runtime/DoFnOp.java | 3 +- .../samza/runtime/SamzaDoFnRunners.java | 31 +++++++++++++++++-- 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPortablePipelineOptions.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPortablePipelineOptions.java index aa8e7ceb71d7..73cb1424f47f 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPortablePipelineOptions.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPortablePipelineOptions.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.samza; +import org.apache.beam.sdk.options.Default; import org.apache.beam.sdk.options.Description; import org.apache.beam.sdk.options.PortablePipelineOptions; @@ -29,4 +30,11 @@ public interface SamzaPortablePipelineOptions String getFsTokenPath(); void setFsTokenPath(String path); + + @Description( + "Wait if necessary for completing a remote bundle processing for at most the given time (in milliseconds). if the value of timeout is negative, wait forever until the bundle processing is completed. Used only in portable mode.") + @Default.Long(-1) + long getBundleProcessingTimeout(); + + void setBundleProcessingTimeout(long timeoutMs); } diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java index 35661ae86fe1..9b33eeb58f86 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java @@ -45,6 +45,7 @@ import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.samza.SamzaExecutionContext; import org.apache.beam.runners.samza.SamzaPipelineOptions; +import org.apache.beam.runners.samza.SamzaPortablePipelineOptions; import org.apache.beam.runners.samza.util.DoFnUtils; import org.apache.beam.runners.samza.util.FutureUtils; import org.apache.beam.sdk.coders.Coder; @@ -236,7 +237,7 @@ public void open( sideInputHandler, nonKeyedStateInternalsFactory, timerInternalsFactory, - samzaPipelineOptions, + samzaPipelineOptions.as(SamzaPortablePipelineOptions.class), outputManagerFactory.create(emitter, outputFutureCollector), stageBundleFactory, samzaExecutionContext, diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java index 41fe8190dec1..26317f818965 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java @@ -21,8 +21,10 @@ import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.DoFnRunners; @@ -43,6 +45,7 @@ import org.apache.beam.runners.fnexecution.state.StateRequestHandler; import org.apache.beam.runners.samza.SamzaExecutionContext; import org.apache.beam.runners.samza.SamzaPipelineOptions; +import org.apache.beam.runners.samza.SamzaPortablePipelineOptions; import org.apache.beam.runners.samza.metrics.DoFnRunnerWithMetrics; import org.apache.beam.runners.samza.util.StateUtils; import org.apache.beam.runners.samza.util.WindowUtils; @@ -203,7 +206,7 @@ public static DoFnRunner createPortable( SideInputHandler sideInputHandler, SamzaStoreStateInternals.Factory nonKeyedStateInternalsFactory, SamzaTimerInternalsFactory timerInternalsFactory, - SamzaPipelineOptions pipelineOptions, + SamzaPortablePipelineOptions pipelineOptions, DoFnRunners.OutputManager outputManager, StageBundleFactory stageBundleFactory, SamzaExecutionContext samzaExecutionContext, @@ -232,6 +235,7 @@ public static DoFnRunner createPortable( (SamzaExecutionContext) context.getApplicationContainerContext(); final DoFnRunner underlyingRunner = new SdkHarnessDoFnRunner<>( + pipelineOptions, stepName, timerInternalsFactory, WindowUtils.getWindowStrategy( @@ -252,6 +256,7 @@ private static class SdkHarnessDoFnRunner implements DoFnRunner implements DoFnRunner timerInternalsFactory, WindowingStrategy windowingStrategy, @@ -276,6 +282,7 @@ private SdkHarnessDoFnRunner( BagState> bundledEventsBag, StateRequestHandler stateRequestHandler, SamzaExecutionContext samzaExecutionContext) { + this.pipelineOptions = pipelineOptions; this.timerInternalsFactory = timerInternalsFactory; this.windowingStrategy = windowingStrategy; this.outputManager = outputManager; @@ -426,8 +433,7 @@ public void onTimer( @Override public void finishBundle() { try { - // RemoteBundle close blocks until all results are received - remoteBundle.close(); + closeBundle(); emitResults(); emitMetrics(); bundledEventsBag.clear(); @@ -439,6 +445,25 @@ public void finishBundle() { } } + private void closeBundle() throws Exception { + long bundleProcessingTimeout = pipelineOptions.getBundleProcessingTimeout(); + if (bundleProcessingTimeout < 0) { + // RemoteBundle close blocks until all results are received + remoteBundle.close(); + } else { + CompletableFuture future = + CompletableFuture.runAsync( + () -> { + try { + remoteBundle.close(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + future.get(bundleProcessingTimeout, TimeUnit.MILLISECONDS); + } + } + @Override public void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {} From 29361580ef5da9e3f5360bf0c1aeb39359d7b41b Mon Sep 17 00:00:00 2001 From: Alan Zhang Date: Tue, 17 Jan 2023 11:15:18 -0800 Subject: [PATCH 2/5] move bundleProcessingTimeout config to SamzaPipelineOptions --- .../apache/beam/runners/samza/SamzaPipelineOptions.java | 7 +++++++ .../beam/runners/samza/SamzaPortablePipelineOptions.java | 8 -------- .../org/apache/beam/runners/samza/runtime/DoFnOp.java | 3 +-- .../beam/runners/samza/runtime/SamzaDoFnRunners.java | 7 +++---- 4 files changed, 11 insertions(+), 14 deletions(-) diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPipelineOptions.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPipelineOptions.java index 814b14f98b8a..9dd5234c35bf 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPipelineOptions.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPipelineOptions.java @@ -135,6 +135,13 @@ public interface SamzaPipelineOptions extends PipelineOptions { void setMaxBundleTimeMs(long maxBundleTimeMs); + @Description( + "Wait if necessary for completing a remote bundle processing for at most the given time (in milliseconds). if the value of timeout is negative, wait forever until the bundle processing is completed. Used only in portable mode for now.") + @Default.Long(-1) + long getBundleProcessingTimeout(); + + void setBundleProcessingTimeout(long timeoutMs); + @Description( "The number of threads to run DoFn.processElements in parallel within a bundle. Used only in non-portable mode.") @Default.Integer(1) diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPortablePipelineOptions.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPortablePipelineOptions.java index 73cb1424f47f..aa8e7ceb71d7 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPortablePipelineOptions.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPortablePipelineOptions.java @@ -17,7 +17,6 @@ */ package org.apache.beam.runners.samza; -import org.apache.beam.sdk.options.Default; import org.apache.beam.sdk.options.Description; import org.apache.beam.sdk.options.PortablePipelineOptions; @@ -30,11 +29,4 @@ public interface SamzaPortablePipelineOptions String getFsTokenPath(); void setFsTokenPath(String path); - - @Description( - "Wait if necessary for completing a remote bundle processing for at most the given time (in milliseconds). if the value of timeout is negative, wait forever until the bundle processing is completed. Used only in portable mode.") - @Default.Long(-1) - long getBundleProcessingTimeout(); - - void setBundleProcessingTimeout(long timeoutMs); } diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java index 9b33eeb58f86..35661ae86fe1 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java @@ -45,7 +45,6 @@ import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.samza.SamzaExecutionContext; import org.apache.beam.runners.samza.SamzaPipelineOptions; -import org.apache.beam.runners.samza.SamzaPortablePipelineOptions; import org.apache.beam.runners.samza.util.DoFnUtils; import org.apache.beam.runners.samza.util.FutureUtils; import org.apache.beam.sdk.coders.Coder; @@ -237,7 +236,7 @@ public void open( sideInputHandler, nonKeyedStateInternalsFactory, timerInternalsFactory, - samzaPipelineOptions.as(SamzaPortablePipelineOptions.class), + samzaPipelineOptions, outputManagerFactory.create(emitter, outputFutureCollector), stageBundleFactory, samzaExecutionContext, diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java index 26317f818965..72fcb519a380 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java @@ -45,7 +45,6 @@ import org.apache.beam.runners.fnexecution.state.StateRequestHandler; import org.apache.beam.runners.samza.SamzaExecutionContext; import org.apache.beam.runners.samza.SamzaPipelineOptions; -import org.apache.beam.runners.samza.SamzaPortablePipelineOptions; import org.apache.beam.runners.samza.metrics.DoFnRunnerWithMetrics; import org.apache.beam.runners.samza.util.StateUtils; import org.apache.beam.runners.samza.util.WindowUtils; @@ -206,7 +205,7 @@ public static DoFnRunner createPortable( SideInputHandler sideInputHandler, SamzaStoreStateInternals.Factory nonKeyedStateInternalsFactory, SamzaTimerInternalsFactory timerInternalsFactory, - SamzaPortablePipelineOptions pipelineOptions, + SamzaPipelineOptions pipelineOptions, DoFnRunners.OutputManager outputManager, StageBundleFactory stageBundleFactory, SamzaExecutionContext samzaExecutionContext, @@ -256,7 +255,7 @@ private static class SdkHarnessDoFnRunner implements DoFnRunner implements DoFnRunner timerInternalsFactory, WindowingStrategy windowingStrategy, From ace2021fa0cd4d3dc8eac9bc383ac238436b44dc Mon Sep 17 00:00:00 2001 From: Alan Zhang Date: Tue, 17 Jan 2023 13:56:31 -0800 Subject: [PATCH 3/5] add a unit test --- .../samza/runtime/SamzaDoFnRunners.java | 34 +++++-------- .../runners/samza/util/RunWithTimeout.java | 42 ++++++++++++++++ .../samza/util/RunWithTimeoutTest.java | 48 +++++++++++++++++++ 3 files changed, 102 insertions(+), 22 deletions(-) create mode 100644 runners/samza/src/main/java/org/apache/beam/runners/samza/util/RunWithTimeout.java create mode 100644 runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java index 72fcb519a380..fb461d6012b5 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java @@ -17,14 +17,14 @@ */ package org.apache.beam.runners.samza.runtime; +import static org.apache.beam.runners.samza.util.RunWithTimeout.run; + import java.util.Collections; import java.util.List; import java.util.Locale; import java.util.Map; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadLocalRandom; -import java.util.concurrent.TimeUnit; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.DoFnRunners; @@ -432,7 +432,16 @@ public void onTimer( @Override public void finishBundle() { try { - closeBundle(); + run( + pipelineOptions.getBundleProcessingTimeout(), + () -> { + // RemoteBundle close blocks until all results are received + try { + remoteBundle.close(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); emitResults(); emitMetrics(); bundledEventsBag.clear(); @@ -444,25 +453,6 @@ public void finishBundle() { } } - private void closeBundle() throws Exception { - long bundleProcessingTimeout = pipelineOptions.getBundleProcessingTimeout(); - if (bundleProcessingTimeout < 0) { - // RemoteBundle close blocks until all results are received - remoteBundle.close(); - } else { - CompletableFuture future = - CompletableFuture.runAsync( - () -> { - try { - remoteBundle.close(); - } catch (Exception e) { - throw new RuntimeException(e); - } - }); - future.get(bundleProcessingTimeout, TimeUnit.MILLISECONDS); - } - } - @Override public void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {} diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/util/RunWithTimeout.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/util/RunWithTimeout.java new file mode 100644 index 000000000000..13c2ca7f14e2 --- /dev/null +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/util/RunWithTimeout.java @@ -0,0 +1,42 @@ +/* + * 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.beam.runners.samza.util; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +public class RunWithTimeout { + + /** + * Run a function and wait for at most the given time (in milliseconds). + * + * @param timeoutInMs the time to wait for completing the function call. If the value of timeout + * is negative, wait forever until the function call is completed + * @param runnable the main function + */ + public static void run(long timeoutInMs, Runnable runnable) + throws ExecutionException, InterruptedException, TimeoutException { + if (timeoutInMs < 0) { + runnable.run(); + } else { + CompletableFuture.runAsync(runnable).get(timeoutInMs, TimeUnit.MILLISECONDS); + } + } +} diff --git a/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java b/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java new file mode 100644 index 000000000000..f263b523c8f3 --- /dev/null +++ b/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java @@ -0,0 +1,48 @@ +/* + * 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.beam.runners.samza.util; + +import java.util.concurrent.TimeoutException; +import org.junit.Test; + +public class RunWithTimeoutTest { + + @Test(expected = TimeoutException.class) + public void testRunWithTimeoutOccurred() throws Exception { + RunWithTimeout.run( + 100, + () -> { + try { + Thread.sleep(500); + } catch (InterruptedException ignored) { + } + }); + } + + @Test + public void testRunWithTimeoutDisabled() throws Exception { + RunWithTimeout.run( + -1, + () -> { + try { + Thread.sleep(500); + } catch (InterruptedException ignored) { + } + }); + } +} From dc57691b1ad97ff9ec1ff5aafdb73e09297a328b Mon Sep 17 00:00:00 2001 From: Alan Zhang Date: Tue, 17 Jan 2023 22:45:13 -0800 Subject: [PATCH 4/5] add a new unit test --- .../samza/util/RunWithTimeoutTest.java | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java b/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java index f263b523c8f3..8a9b609b8f73 100644 --- a/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java +++ b/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java @@ -17,7 +17,9 @@ */ package org.apache.beam.runners.samza.util; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeoutException; +import org.junit.Assert; import org.junit.Test; public class RunWithTimeoutTest { @@ -45,4 +47,23 @@ public void testRunWithTimeoutDisabled() throws Exception { } }); } + + @Test + public void testRunWithUserException() throws Exception { + try { + RunWithTimeout.run( + 100, + () -> { + throw new UserException(); + }); + } catch (ExecutionException e) { + Assert.assertTrue(e.getCause() instanceof UserException); + } + } + + private static class UserException extends RuntimeException { + public UserException() { + super(); + } + } } From db7073d29bb38a742014c1f4debcdb9393b0875f Mon Sep 17 00:00:00 2001 From: Alan Zhang Date: Thu, 19 Jan 2023 14:17:13 -0800 Subject: [PATCH 5/5] move runWithTimeout method inside SdkHarnessDoFnRunner --- .../samza/runtime/SamzaDoFnRunners.java | 26 ++++++++++-- .../runners/samza/util/RunWithTimeout.java | 42 ------------------- .../SdkHarnessDoFnRunnerTest.java} | 29 ++----------- 3 files changed, 26 insertions(+), 71 deletions(-) delete mode 100644 runners/samza/src/main/java/org/apache/beam/runners/samza/util/RunWithTimeout.java rename runners/samza/src/test/java/org/apache/beam/runners/samza/{util/RunWithTimeoutTest.java => runtime/SdkHarnessDoFnRunnerTest.java} (68%) diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java index fb461d6012b5..8af62059de9f 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java @@ -17,14 +17,16 @@ */ package org.apache.beam.runners.samza.runtime; -import static org.apache.beam.runners.samza.util.RunWithTimeout.run; - import java.util.Collections; import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.DoFnRunners; @@ -251,7 +253,7 @@ public static DoFnRunner createPortable( : underlyingRunner; } - private static class SdkHarnessDoFnRunner implements DoFnRunner { + static class SdkHarnessDoFnRunner implements DoFnRunner { private static final int DEFAULT_METRIC_SAMPLE_RATE = 100; @@ -432,7 +434,7 @@ public void onTimer( @Override public void finishBundle() { try { - run( + runWithTimeout( pipelineOptions.getBundleProcessingTimeout(), () -> { // RemoteBundle close blocks until all results are received @@ -453,6 +455,22 @@ public void finishBundle() { } } + /** + * Run a function and wait for at most the given time (in milliseconds). + * + * @param timeoutInMs the time to wait for completing the function call. If the value of timeout + * is negative, wait forever until the function call is completed + * @param runnable the main function + */ + static void runWithTimeout(long timeoutInMs, Runnable runnable) + throws ExecutionException, InterruptedException, TimeoutException { + if (timeoutInMs < 0) { + runnable.run(); + } else { + CompletableFuture.runAsync(runnable).get(timeoutInMs, TimeUnit.MILLISECONDS); + } + } + @Override public void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {} diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/util/RunWithTimeout.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/util/RunWithTimeout.java deleted file mode 100644 index 13c2ca7f14e2..000000000000 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/util/RunWithTimeout.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * 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.beam.runners.samza.util; - -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; - -public class RunWithTimeout { - - /** - * Run a function and wait for at most the given time (in milliseconds). - * - * @param timeoutInMs the time to wait for completing the function call. If the value of timeout - * is negative, wait forever until the function call is completed - * @param runnable the main function - */ - public static void run(long timeoutInMs, Runnable runnable) - throws ExecutionException, InterruptedException, TimeoutException { - if (timeoutInMs < 0) { - runnable.run(); - } else { - CompletableFuture.runAsync(runnable).get(timeoutInMs, TimeUnit.MILLISECONDS); - } - } -} diff --git a/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java b/runners/samza/src/test/java/org/apache/beam/runners/samza/runtime/SdkHarnessDoFnRunnerTest.java similarity index 68% rename from runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java rename to runners/samza/src/test/java/org/apache/beam/runners/samza/runtime/SdkHarnessDoFnRunnerTest.java index 8a9b609b8f73..e6029beb93b0 100644 --- a/runners/samza/src/test/java/org/apache/beam/runners/samza/util/RunWithTimeoutTest.java +++ b/runners/samza/src/test/java/org/apache/beam/runners/samza/runtime/SdkHarnessDoFnRunnerTest.java @@ -15,18 +15,16 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.runners.samza.util; +package org.apache.beam.runners.samza.runtime; -import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeoutException; -import org.junit.Assert; import org.junit.Test; -public class RunWithTimeoutTest { +public class SdkHarnessDoFnRunnerTest { @Test(expected = TimeoutException.class) public void testRunWithTimeoutOccurred() throws Exception { - RunWithTimeout.run( + SamzaDoFnRunners.SdkHarnessDoFnRunner.runWithTimeout( 100, () -> { try { @@ -38,7 +36,7 @@ public void testRunWithTimeoutOccurred() throws Exception { @Test public void testRunWithTimeoutDisabled() throws Exception { - RunWithTimeout.run( + SamzaDoFnRunners.SdkHarnessDoFnRunner.runWithTimeout( -1, () -> { try { @@ -47,23 +45,4 @@ public void testRunWithTimeoutDisabled() throws Exception { } }); } - - @Test - public void testRunWithUserException() throws Exception { - try { - RunWithTimeout.run( - 100, - () -> { - throw new UserException(); - }); - } catch (ExecutionException e) { - Assert.assertTrue(e.getCause() instanceof UserException); - } - } - - private static class UserException extends RuntimeException { - public UserException() { - super(); - } - } }