Throw UncheckedTimeoutException on timeout

4 files changed, 20 insertions, 12 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/TimeoutBudget.java b/configserver/src/main/java/com/yahoo/vespa/config/server/TimeoutBudget.java
index fa30fa3561a..80bde90fac2 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/TimeoutBudget.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/TimeoutBudget.java
@@ -1,11 +1,14 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.config.server;
 
+import com.google.common.util.concurrent.UncheckedTimeoutException;
+
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Supplier;
 
 /**
  * Handles a timeout logic by providing higher level abstraction for asking if there is time left.
@@ -31,13 +34,12 @@ public class TimeoutBudget {
 
     public Duration timeLeft() {
         Instant now = clock.instant();
-        measurements.add(new Measurement(now));
         Duration duration = Duration.between(now, endTime);
-        return duration.isNegative() ? Duration.ofMillis(0) : duration;
+        return duration.isNegative() ? Duration.ZERO : duration;
     }
 
     public boolean hasTimeLeft() {
-        return hasTimeLeft("");
+        return clock.instant().isBefore(endTime);
     }
 
     public boolean hasTimeLeft(String step) {
@@ -66,15 +68,20 @@ public class TimeoutBudget {
         return buf.toString();
     }
 
+    /**
+     * @param exceptionMessage exception message for the exception that will be thrown if there is no time left
+     * @throws UncheckedTimeoutException if this has no time left
+     */
+    public void assertNotTimedOut(Supplier<String> exceptionMessage) {
+        if (hasTimeLeft()) return;
+        throw new UncheckedTimeoutException(exceptionMessage.get());
+    }
+
     private static class Measurement {
 
         private final Instant timestamp;
         private final String label;
 
-        Measurement(Instant timestamp) {
-            this(timestamp, "");
-        }
-
         Measurement(Instant timestamp, String label) {
             this.timestamp = timestamp;
             this.label = label;
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/deploy/Deployment.java b/configserver/src/main/java/com/yahoo/vespa/config/server/deploy/Deployment.java
index 3726ea97fcc..fc7bd70679a 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/deploy/Deployment.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/deploy/Deployment.java
@@ -117,7 +117,7 @@ public class Deployment implements com.yahoo.config.provision.Deployment {
         ApplicationId applicationId = session.getApplicationId();
         try (ActionTimer timer = applicationRepository.timerFor(applicationId, "deployment.activateMillis")) {
             TimeoutBudget timeoutBudget = params.getTimeoutBudget();
-            if ( ! timeoutBudget.hasTimeLeft()) throw new RuntimeException("Timeout exceeded when trying to activate '" + applicationId + "'");
+            timeoutBudget.assertNotTimedOut(() -> "Timeout exceeded when trying to activate '" + applicationId + "'");
 
             RemoteSession previousActiveSession;
             CompletionWaiter waiter;
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
index a16d3d0c221..6aed054486e 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.config.server.session;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.UncheckedTimeoutException;
 import com.google.inject.Inject;
 import com.yahoo.cloud.config.ConfigserverConfig;
 import com.yahoo.component.Version;
@@ -238,8 +239,8 @@ public class SessionPreparer {
             TimeoutBudget timeoutBudget = params.getTimeoutBudget();
             if (! timeoutBudget.hasTimeLeft(step)) {
                 String used = timeoutBudget.timesUsed();
-                throw new RuntimeException("prepare timed out " + used + " after " + step +
-                                           " step (timeout " + timeoutBudget.timeout() + "): " + applicationId);
+                throw new UncheckedTimeoutException("prepare timed out " + used + " after " + step +
+                        " step (timeout " + timeoutBudget.timeout() + "): " + applicationId);
             }
         }
 
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/TimeoutBudgetTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/TimeoutBudgetTest.java
index 6059b77061a..ac43964a279 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/TimeoutBudgetTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/TimeoutBudgetTest.java
@@ -35,7 +35,7 @@ public class TimeoutBudgetTest {
         assertThat(budget.timeLeft().toMillis(), is(0L));
 
         clock.advance(Duration.ofMillis(1));
-        assertThat(budget.timesUsed(), is("[0 ms, 1 ms, 5 ms, 0 ms, 1 ms, 5 ms, total: 13 ms]"));
+        assertThat(budget.timesUsed(), is("[total: 13 ms]"));
     }
 
     @Test
@@ -55,7 +55,7 @@ public class TimeoutBudgetTest {
         assertThat(budget.hasTimeLeft(), is(false));
 
         clock.advance(Duration.ofMillis(1));
-        assertThat(budget.timesUsed(), is("[0 ms, 1 ms, 5 ms, 0 ms, 1 ms, 5 ms, total: 13 ms]"));
+        assertThat(budget.timesUsed(), is("[total: 13 ms]"));
     }
 
     @Test