[FLINK-5799] [rpc] Let RpcService.scheduleRunnable return a Scheduled…

…Future

The returned ScheduledFuture instance allows to cancel a scheduled runnable and obtain
the delay until the runnable will be executed. Furthermore, it allows to wait on the
completion of the runnable.

This closes #3311.

flink-runtime/src/main/java/org/apache/flink/runtime/rpc/RpcService.java

@@ -24,6 +24,7 @@
 
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -119,7 +120,7 @@ public interface RpcService {
 	 * @param runnable Runnable to be executed
 	 * @param delay    The delay after which the runnable will be executed
 	 */
-	void scheduleRunnable(Runnable runnable, long delay, TimeUnit unit);
+	ScheduledFuture<?> scheduleRunnable(Runnable runnable, long delay, TimeUnit unit);
 
 	/**
 	 * Execute the given runnable in the executor of the RPC service. This method can be used to run

flink-runtime/src/main/java/org/apache/flink/runtime/rpc/akka/AkkaRpcService.java

@@ -276,12 +276,12 @@ public ScheduledExecutor getScheduledExecutor() {
 	}
 
 	@Override
-	public void scheduleRunnable(Runnable runnable, long delay, TimeUnit unit) {
+	public ScheduledFuture<?> scheduleRunnable(Runnable runnable, long delay, TimeUnit unit) {
 		checkNotNull(runnable, "runnable");
 		checkNotNull(unit, "unit");
-		checkArgument(delay >= 0, "delay must be zero or larger");
+		checkArgument(delay >= 0L, "delay must be zero or larger");
 
-		actorSystem.scheduler().scheduleOnce(new FiniteDuration(delay, unit), runnable, actorSystem.dispatcher());
+		return internalScheduledExecutor.schedule(runnable, delay, unit);
 	}
 
 	@Override

flink-runtime/src/test/java/org/apache/flink/runtime/rpc/TestingSerialRpcService.java

@@ -37,10 +37,14 @@
 import java.util.UUID;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Delayed;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
@@ -67,10 +71,12 @@ public TestingSerialRpcService() {
 	}
 
 	@Override
-	public void scheduleRunnable(final Runnable runnable, final long delay, final TimeUnit unit) {
+	public ScheduledFuture<?> scheduleRunnable(final Runnable runnable, final long delay, final TimeUnit unit) {
 		try {
 			unit.sleep(delay);
 			runnable.run();
+
+			return new DoneScheduledFuture<Void>(null);
 		} catch (Throwable e) {
 			throw new RuntimeException(e);
 		}
@@ -434,4 +440,49 @@ private static boolean isRpcTimeout(Annotation[] annotations) {
 		}
 
 	}
+
+	private static class DoneScheduledFuture<V> implements ScheduledFuture<V> {
+
+		private final V value;
+
+		private DoneScheduledFuture(V value) {
+			this.value = value;
+		}
+
+		@Override
+		public long getDelay(TimeUnit unit) {
+			return 0L;
+		}
+
+		@Override
+		public int compareTo(Delayed o) {
+			return 0;
+		}
+
+		@Override
+		public boolean cancel(boolean mayInterruptIfRunning) {
+			return false;
+		}
+
+		@Override
+		public boolean isCancelled() {
+			return false;
+		}
+
+		@Override
+		public boolean isDone() {
+			return true;
+		}
+
+		@Override
+		public V get() throws InterruptedException, ExecutionException {
+			return value;
+		}
+
+		@Override
+		public V get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
+			return value;
+		}
+	}
+
 }

flink-runtime/src/test/java/org/apache/flink/runtime/rpc/akka/AkkaRpcServiceTest.java

@@ -69,14 +69,16 @@ public void testScheduleRunnable() throws Exception {
 		final long delay = 100;
 		final long start = System.nanoTime();
 
-		akkaRpcService.scheduleRunnable(new Runnable() {
+		ScheduledFuture<?> scheduledFuture = akkaRpcService.scheduleRunnable(new Runnable() {
 			@Override
 			public void run() {
 				latch.trigger();
 			}
 		}, delay, TimeUnit.MILLISECONDS);
 
-		latch.await();
+		scheduledFuture.get();
+
+		assertTrue(latch.isTriggered());
 		final long stop = System.nanoTime();
 
 		assertTrue("call was not properly delayed", ((stop - start) / 1000000) >= delay);