set-node-state timeout in CC

2 files changed, 42 insertions, 18 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/time/TimeBudget.java b/vespajlib/src/main/java/com/yahoo/time/TimeBudget.java
index fa18cb5e467..a7963df0208 100644
--- a/vespajlib/src/main/java/com/yahoo/time/TimeBudget.java
+++ b/vespajlib/src/main/java/com/yahoo/time/TimeBudget.java
@@ -6,26 +6,31 @@ import com.google.common.util.concurrent.UncheckedTimeoutException;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Optional;
 
 /**
- * A TimeBudget can be used to track the time of an ongoing operation with a timeout.
+ * A TimeBudget can be used to track the time of an ongoing operation, possibly with a timeout.
  *
  * @author hakon
  */
 public class TimeBudget {
     private final Clock clock;
     private final Instant start;
-    private final Duration timeout;
+    private final Optional<Duration> timeout;
 
     /** Returns a TimeBudget with a start time of now, and with the given timeout. */
     public static TimeBudget fromNow(Clock clock, Duration timeout) {
-        return new TimeBudget(clock, clock.instant(), timeout);
+        return new TimeBudget(clock, clock.instant(), Optional.of(timeout));
     }
 
-    private TimeBudget(Clock clock, Instant start, Duration timeout) {
+    public static TimeBudget from(Clock clock, Instant start, Optional<Duration> timeout) {
+        return new TimeBudget(clock, start, timeout);
+    }
+
+    private TimeBudget(Clock clock, Instant start, Optional<Duration> timeout) {
         this.clock = clock;
         this.start = start;
-        this.timeout = makeNonNegative(timeout);
+        this.timeout = timeout.map(TimeBudget::makeNonNegative);
     }
 
     /** Returns time since start. */
@@ -33,26 +38,32 @@ public class TimeBudget {
         return nonNegativeBetween(start, clock.instant());
     }
 
-    /** Returns the original timeout. */
-    public Duration originalTimeout() {
+    /** Returns the original timeout, if any. */
+    public Optional<Duration> originalTimeout() {
         return timeout;
     }
 
+    /** Returns the deadline, if present. */
+    public Optional<Instant> deadline() {
+        return timeout.map(start::plus);
+    }
+
     /**
-     * Returns the time until deadline.
+     * Returns the time until deadline, if there is one.
      *
      * @return time until deadline. It's toMillis() is guaranteed to be positive.
      * @throws UncheckedTimeoutException if the deadline has been reached or passed.
      */
-    public Duration timeLeftOrThrow() {
-        Instant now = clock.instant();
-        Duration left = Duration.between(now, start.plus(timeout));
-        if (left.toMillis() <= 0) {
-            throw new UncheckedTimeoutException("Time since start " + nonNegativeBetween(start, now) +
-                    " exceeds timeout " + timeout);
-        }
+    public Optional<Duration> timeLeftOrThrow() {
+        return timeout.map(timeout -> {
+            Duration passed = timePassed();
+            Duration left = timeout.minus(passed);
+            if (left.toMillis() <= 0) {
+                throw new UncheckedTimeoutException("Time since start " + passed + " exceeds timeout " + this.timeout);
+            }
 
-        return left;
+            return left;
+        });
     }
 
     private static Duration nonNegativeBetween(Instant start, Instant end) {
diff --git a/vespajlib/src/test/java/com/yahoo/time/TimeBudgetTest.java b/vespajlib/src/test/java/com/yahoo/time/TimeBudgetTest.java
index a51081067b7..f197cc34b32 100644
--- a/vespajlib/src/test/java/com/yahoo/time/TimeBudgetTest.java
+++ b/vespajlib/src/test/java/com/yahoo/time/TimeBudgetTest.java
@@ -8,8 +8,10 @@ import org.junit.Test;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 
@@ -23,12 +25,12 @@ public class TimeBudgetTest {
         TimeBudget timeBudget = TimeBudget.fromNow(clock, Duration.ofSeconds(10));
 
         clock.advance(Duration.ofSeconds(7));
-        assertEquals(Duration.ofSeconds(3), timeBudget.timeLeftOrThrow());
+        assertEquals(Duration.ofSeconds(3), timeBudget.timeLeftOrThrow().get());
 
         // Verify that toMillis() of >=1 is fine, but 0 is not.
 
         clock.setInstant(Instant.ofEpochSecond(9, 999000000));
-        assertEquals(1, timeBudget.timeLeftOrThrow().toMillis());
+        assertEquals(1, timeBudget.timeLeftOrThrow().get().toMillis());
         clock.setInstant(Instant.ofEpochSecond(9, 999000001));
         try {
             timeBudget.timeLeftOrThrow();
@@ -37,4 +39,15 @@ public class TimeBudgetTest {
             // OK
         }
     }
+
+    @Test
+    public void noDeadline() {
+        ManualClock clock = new ManualClock();
+        clock.setInstant(Instant.ofEpochSecond(0));
+        TimeBudget timeBudget = TimeBudget.from(clock, clock.instant(), Optional.empty());
+
+        assertFalse(timeBudget.originalTimeout().isPresent());
+        assertFalse(timeBudget.timeLeftOrThrow().isPresent());
+        assertFalse(timeBudget.deadline().isPresent());
+    }
 }
\ No newline at end of file