Add blocking and poll_interval options

to `wait` operator

digdag-docs/src/operators/wait.md

@@ -2,7 +2,7 @@
 
 **wait>** operator waits a specific duration in the workflow.
 
-This operator is similar to `sh>: sleep 5`, but this is a non-blocking operator and should be always available even in security-restricted environment.
+This operator seems similar to `sh>: sleep 5`, but this works in both blocking and non-blocking modes and should be always available even in security-restricted environment.
 
     +wait_10s:
       wait>: 10s
@@ -13,3 +13,23 @@ This operator is similar to `sh>: sleep 5`, but this is a non-blocking operator
 
   Duration to wait.
 
+* **blocking**: BOOLEAN
+
+  Digdag agent internally executes this operator in blocking mode and the agent keeps waiting if this option is set to true (default: false)
+
+  Examples:
+
+  ```
+  blocking: true
+  ```
+
+* **poll_interval**: DURATION
+
+  This option is used only with non-blocking mode. If it's set, digdag agent internally gets awake and checks at a specific interval if the duration has passed. If not set, digdag agent gets awake only when a specific duration passes.
+
+  Examples:
+
+  ```
+  poll_interval: 5s
+  ```
+

digdag-standards/src/main/java/io/digdag/standards/operator/WaitOperatorFactory.java

@@ -1,5 +1,6 @@
 package io.digdag.standards.operator;
 
+import com.google.common.base.Optional;
 import com.google.inject.Inject;
 import io.digdag.client.config.Config;
 import io.digdag.client.config.ConfigElement;
@@ -18,6 +19,7 @@
 
 import java.time.Duration;
 import java.time.Instant;
+import java.util.concurrent.TimeUnit;
 
 
 public class WaitOperatorFactory
@@ -57,21 +59,53 @@ private WaitOperator(OperatorContext context)
             this.workspace = Workspace.ofTaskRequest(context.getProjectPath(), request);
         }
 
-        public TaskResult run()
+        private Duration duration(Config config)
         {
             Duration duration;
             try {
-                Config config = request.getConfig();
                 duration = Durations.parseDuration(config.get("_command", String.class));
-                if (duration == null) {
-                    throw new ConfigException("Wait duration should be specified");
+            }
+            catch (RuntimeException re) {
+                throw new ConfigException("Invalid configuration", re);
+            }
+            logger.debug("wait duration: {}", duration);
+            return duration;
+        }
+
+        private boolean blocking(Config config)
+        {
+            boolean blocking = config.get("blocking", boolean.class, false);
+            logger.debug("wait blocking mode: {}", blocking);
+            return blocking;
+        }
+
+        private Duration pollInterval(Config config)
+        {
+            Duration pollInterval;
+            try {
+                Optional<String> pollIntervalStr = config.getOptional("poll_interval", String.class);
+                if (!pollIntervalStr.isPresent()) {
+                    return null;
                 }
+                pollInterval = Durations.parseDuration(pollIntervalStr.get());
             }
             catch (RuntimeException re) {
                 throw new ConfigException("Invalid configuration", re);
             }
+            logger.debug("wait poll_interval: {}", pollInterval);
+            return pollInterval;
+        }
 
-            logger.debug("wait duration: {}", duration);
+        public TaskResult run()
+        {
+            Config config = request.getConfig();
+
+            Duration duration = duration(config);
+            boolean blocking = blocking(config);
+            Duration pollInterval = pollInterval(config);
+            if (blocking && pollInterval != null) {
+                throw new ConfigException("poll_interval can't be specified with blocking:true");
+            }
 
             Instant now = Instant.now();
             Instant start = request.getLastStateParams()
@@ -81,19 +115,39 @@ public TaskResult run()
 
             if (now.isAfter(start.plusMillis(duration.toMillis()))) {
                 logger.info("wait finished. start:{}", start);
+                return TaskResult.empty(request);
+            }
+
+            // Wait at least 1 second
+            long waitDurationSeconds = Math.max(
+                    Duration.between(now, start.plusMillis(duration.toMillis())).getSeconds(),
+                    1);
+
+            if (blocking) {
+                logger.debug("waiting for {}s", waitDurationSeconds);
+                try {
+                    TimeUnit.SECONDS.sleep(waitDurationSeconds);
+                    return TaskResult.empty(request);
+                }
+                catch (InterruptedException e) {
+                    // The blocking wait will be restart from the beginning when interrupted.
+                    //
+                    // There is room to improve this by making the task resume from when interrupted.
+                    // But this operator, especially blocking mode, is for development use,
+                    // so we'll go with this simple implementation for now.
+                    throw new RuntimeException("`wait` operator with blocking mode is interrupted and this will be restart from the beginning of the wait");
+                }
             }
             else {
-                // Wait at least 1 second
-                long nextPollSecs = Math.max(
-                        Duration.between(now, start.plusMillis(duration.toMillis())).getSeconds(),
-                        1);
-                logger.debug("polling after {}s", nextPollSecs);
+                if (pollInterval != null) {
+                    waitDurationSeconds = pollInterval.getSeconds();
+                }
+                logger.debug("polling after {}s", waitDurationSeconds);
                 throw TaskExecutionException.ofNextPolling(
-                        (int)nextPollSecs,
+                        (int) waitDurationSeconds,
                         ConfigElement.copyOf(
                                 request.getLastStateParams().set(WAIT_START_TIME_PARAM, start.toEpochMilli())));
             }
-            return TaskResult.empty(request);
         }
     }
 }

digdag-tests/src/test/java/acceptance/WaitIT.java

@@ -39,7 +39,6 @@ public ExecResult(CommandStatus commandStatus, Duration duration)
         }
     }
 
-
     private ExecResult runAndMonitorDuration(Supplier<CommandStatus> task)
     {
         Instant start = Instant.now();
@@ -48,30 +47,67 @@ private ExecResult runAndMonitorDuration(Supplier<CommandStatus> task)
         return new ExecResult(commandStatus, duration);
     }
 
-    @Test
-    public void testRun()
+    private void testWorkflow(String workflowName, int expectedDuration)
             throws Exception
     {
+        String nowaitResourcePath = "acceptance/wait/nowait.dig";
+        String targetResourcePath = "acceptance/wait/" + workflowName;
+
         Duration baselineDuration;
         {
-            copyResource("acceptance/wait/nowait.dig", root().resolve("nowait.dig"));
+            copyResource(nowaitResourcePath, root().resolve("wait.dig"));
             ExecResult result = runAndMonitorDuration(() ->
-                    main("run", "-o", root().toString(), "--project", root().toString(), "nowait.dig"));
+                    main("run", "-o", root().toString(), "--project", root().toString(), "wait.dig"));
             CommandStatus status = result.commandStatus;
             assertThat(status.errUtf8(), status.code(), is(0));
             baselineDuration = result.duration;
         }
 
         {
-            copyResource("acceptance/wait/wait_10s.dig", root().resolve("wait_10s.dig"));
+            copyResource(targetResourcePath, root().resolve("wait.dig"));
             ExecResult result = runAndMonitorDuration(() ->
-                    main("run", "-o", root().toString(), "--project", root().toString(), "wait_10s.dig"));
+                    main("run", "-o", root().toString(), "--project", root().toString(), "wait.dig"));
             CommandStatus status = result.commandStatus;
             assertThat(status.errUtf8(), status.code(), is(0));
             assertThat(result.duration, greaterThan(baselineDuration));
             assertThat(result.duration, lessThan(
                     // Actual wait duration can be longer than the specified 10 seconds for some reason
-                    baselineDuration.plusSeconds((long) (10 * 1.5))));
+                    baselineDuration.plusSeconds((long) (expectedDuration * 1.5))));
         }
     }
+
+    @Test
+    public void testSimpleVersion()
+            throws Exception
+    {
+        testWorkflow("wait.dig", 10);
+    }
+
+    @Test
+    public void testBlockingMode()
+            throws Exception
+    {
+        testWorkflow("wait_blocking.dig", 10);
+    }
+
+    @Test
+    public void testPollInterval()
+            throws Exception
+    {
+        testWorkflow("wait_poll_interval.dig", 10);
+    }
+
+    @Test
+    public void testInvalidConfig()
+            throws Exception
+    {
+        String targetResourcePath = "acceptance/wait/wait_invalid_config.dig";
+
+        copyResource(targetResourcePath, root().resolve("wait.dig"));
+        ExecResult result = runAndMonitorDuration(() ->
+                main("run", "-o", root().toString(), "--project", root().toString(), "wait.dig"));
+        CommandStatus status = result.commandStatus;
+        // The workflow contains a conflict configuration and it should fail.
+        assertThat(status.errUtf8(), status.code(), is(1));
+    }
 }

digdag-tests/src/test/resources/acceptance/wait/wait_blocking.dig

@@ -0,0 +1,3 @@
++wait:
+  wait>: 10s
+  blocking: true

digdag-tests/src/test/resources/acceptance/wait/wait_invalid_config.dig

@@ -0,0 +1,4 @@
++wait:
+  wait>: 10s
+  blocking: true
+  poll_interval: 5s

digdag-tests/src/test/resources/acceptance/wait/wait_nonblocking.dig

@@ -0,0 +1,3 @@
++wait:
+  wait>: 10s
+  blocking: false

digdag-tests/src/test/resources/acceptance/wait/wait_poll_interval.dig

@@ -0,0 +1,3 @@
++wait:
+  wait>: 10s
+  poll_interval: 5s