/**
 * Requests that a cluster controller sets all nodes in the cluster to the requested state.
 *
 * @throws IOException if there was a problem communicating with the cluster controller
 */
@Override
public ClusterControllerStateResponse setApplicationState(
        OrchestratorContext context,
        ClusterControllerNodeState wantedState) throws IOException {
    ClusterControllerStateRequest.State state = new ClusterControllerStateRequest.State(wantedState, REQUEST_REASON);
    ClusterControllerStateRequest stateRequest = new ClusterControllerStateRequest(
            state, ClusterControllerStateRequest.Condition.FORCE, null);
    ClusterControllerClientTimeouts timeouts = context.getClusterControllerTimeouts();

    try {
        return clusterControllerApi.apply(api -> api.setClusterState(
                clusterName,
                timeouts.getServerTimeoutOrThrow().toMillis() / 1000.0f,
                stateRequest),
                timeouts);
    } catch (IOException | UncheckedTimeoutException e) {
        final String message = String.format(
                "Giving up setting %s for cluster %s",
                stateRequest,
                clusterName);

        throw new IOException(message, e);
    }
}
 

@Test
public void testBasics() {
    ManualClock clock = new ManualClock();
    clock.setInstant(Instant.ofEpochSecond(0));
    TimeBudget timeBudget = TimeBudget.fromNow(clock, Duration.ofSeconds(10));

    clock.advance(Duration.ofSeconds(7));
    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().get().toMillis());
    clock.setInstant(Instant.ofEpochSecond(9, 999000001));
    try {
        timeBudget.timeLeftOrThrow();
        fail();
    } catch (UncheckedTimeoutException e) {
        // OK
    }
}
 

private <T> T waitUntilStable(Supplier<T> computation, Duration timeout)
{
    T lastValue = computation.get();
    long start = System.nanoTime();
    while (!currentThread().isInterrupted() && nanosSince(start).compareTo(timeout) < 0) {
        sleepUninterruptibly(100, MILLISECONDS);
        T currentValue = computation.get();
        if (currentValue.equals(lastValue)) {
            return currentValue;
        }
        lastValue = currentValue;
    }
    throw new UncheckedTimeoutException();
}
 

@Override
public void backupShard(UUID uuid, File source)
{
    try {
        store.backupShard(uuid, source);
    }
    catch (UncheckedTimeoutException e) {
        timeoutException(uuid, "Shard backup timed out");
    }
}
 

@Override
public void restoreShard(UUID uuid, File target)
{
    try {
        store.restoreShard(uuid, target);
    }
    catch (UncheckedTimeoutException e) {
        timeoutException(uuid, "Shard restore timed out");
    }
}
 

@Override
public boolean deleteShard(UUID uuid)
{
    try {
        return store.deleteShard(uuid);
    }
    catch (UncheckedTimeoutException e) {
        throw timeoutException(uuid, "Shard delete timed out");
    }
}
 

@Override
public boolean shardExists(UUID uuid)
{
    try {
        return store.shardExists(uuid);
    }
    catch (UncheckedTimeoutException e) {
        throw timeoutException(uuid, "Shard existence check timed out");
    }
}
 

/**
 * Get the implementation of {@link Binding} that is registered for the given <code>request</code>.
 *
 * @param request URL pattern from request URL
 *
 * @return The implementation of {@link Binding} that is registered for the given <code>urlPattern</code>.
 *
 *
 * @throws HTTPException If the URL pattern or ContentType is not supported by this service.
 */
private Binding getBinding(HttpServletRequest request) throws HTTPException {
    final String requestURI = request.getPathInfo();
    if (requestURI == null || requestURI.isEmpty() || requestURI.equals("/")) {
        MediaType contentType = getContentType(request);
        // strip of the parameters to get rid of things like encoding
        Binding binding = this.bindingRepository.getBinding(contentType.withoutParameters());
        if (binding == null) {
            if (contentType.equals(MediaTypes.APPLICATION_KVP)) {
                throw new HTTPException(HTTPStatus.METHOD_NOT_ALLOWED);
            } else {
                throw new HTTPException(HTTPStatus.UNSUPPORTED_MEDIA_TYPE);
            }
        } else {
            if (requestTimeout > 0) {
                try {
                    return TIME_LIMITER.newProxy(binding, Binding.class, requestTimeout, TimeUnit.SECONDS);
                } catch (UncheckedTimeoutException ute) {
                    HTTPException httpException = new HTTPException(HTTPStatus.GATEWAY_TIME_OUT);
                    httpException.addSuppressed(ute);
                    throw httpException;
                }
            }
            return binding;
        }
    }
    throw new HTTPException(HTTPStatus.NOT_FOUND);
}
 

@Test(expected = UncheckedTimeoutException.class)
public void testAttemptTimeLimit02() throws Throwable {
    try {
    	SleepyOut sleepyOut = new SleepyOut(10 * 1000L);
        r.call(sleepyOut);
    } catch (Exception e) {
    	throw e.getCause();
    }
}
 

@Test(expected = UncheckedTimeoutException.class)
public void testAttemptTimeLimit02() throws Throwable {
    try {
    	SleepyOut sleepyOut = new SleepyOut(10 * 1000L);
        r.call(sleepyOut);
    } catch (Exception e) {
    	throw e.getCause();
    }
}
 

@Override
public boolean acquire(long timeout, TimeUnit unit) {
    if (throwExceptionOnLock)
        throw new CuratorLockException("Thrown by mock");
    if (timeoutOnLock) return false;
    
    try {
        lock = locks.lock(path, timeout, unit);
        return true;
    }
    catch (UncheckedTimeoutException e) {
        return false;
    }
}
 

/**
 *  1) locks the status service for an application instance.
 *  2) fails all operations in this thread when the session is lost,
 *     since session loss might cause the lock to be lost.
 *     Since it only fails operations in this thread,
 *     all operations depending on a lock, including the locking itself, must be done in this thread.
 *     Note that since it is the thread that fails, all status operations in this thread will fail
 *     even if they're not supposed to be guarded by this lock
 *     (i.e. the request is for another applicationInstanceReference)
 */
@Override
public ApplicationLock lockApplication(OrchestratorContext context, ApplicationInstanceReference reference)
        throws UncheckedTimeoutException {

    Runnable onRegistryClose;

    // A multi-application operation, aka batch suspension, will first issue a probe
    // then a non-probe. With "large locks", the lock is not release in between -
    // no lock is taken on the non-probe. Instead, the release is done on the multi-application
    // context close.
    if (context.hasLock(reference)) {
        onRegistryClose = () -> {};
    } else {
        Runnable unlock = acquireLock(context, reference);
        if (context.registerLockAcquisition(reference, unlock)) {
            onRegistryClose = () -> {};
        } else {
            onRegistryClose = unlock;
        }
    }

    try {
        return new ZkApplicationLock(
                this,
                curator,
                onRegistryClose,
                reference,
                context.isProbe(),
                hostInfosCache);
    } catch (Throwable t) {
        // In case the constructor throws an exception.
        onRegistryClose.run();
        throw t;
    }
}
 

private static WebApplicationException webExceptionFromTimeout(String operationDescription,
                                                               HostName hostName,
                                                               UncheckedTimeoutException e) {
    // Return timeouts as 409 Conflict instead of 504 Gateway Timeout to reduce noise in 5xx graphs.
    return createWebException(operationDescription, hostName, e,
            HostedVespaPolicy.DEADLINE_CONSTRAINT, e.getMessage(), Response.Status.CONFLICT);
}
 

@Override
public Duration getReadTimeoutOrThrow() {
    Duration timeLeft = timeBudget.timeLeft().get();

    // timeLeft = CONNECT_TIMEOUT + readTimeout
    Duration readTimeout = timeLeft.minus(CONNECT_TIMEOUT);

    if (readTimeout.toMillis() <= 0) {
        throw new UncheckedTimeoutException("Timed out after " + timeBudget.originalTimeout().get());
    }

    return readTimeout;
}
 

public Duration getServerTimeoutOrThrow() {
    // readTimeout = DOWNSTREAM_OVERHEAD + serverTimeout
    Duration serverTimeout = getReadTimeoutOrThrow().minus(DOWNSTREAM_OVERHEAD);

    if (serverTimeout.toMillis() < MIN_SERVER_TIMEOUT.toMillis()) {
        throw new UncheckedTimeoutException("Timed out after " + timeBudget.originalTimeout().get());
    }

    return serverTimeout;
}
 

/**
 * Requests that a cluster controller sets the requested node to the requested state.
 *
 * @throws IOException if there was a problem communicating with the cluster controller
 */
@Override
public ClusterControllerStateResponse setNodeState(OrchestratorContext context,
                                                   int storageNodeIndex,
                                                   ClusterControllerNodeState wantedState) throws IOException {
    ClusterControllerStateRequest.State state = new ClusterControllerStateRequest.State(wantedState, REQUEST_REASON);
    ClusterControllerStateRequest stateRequest = new ClusterControllerStateRequest(
            state,
            ClusterControllerStateRequest.Condition.SAFE,
            context.isProbe() ? true : null);
    ClusterControllerClientTimeouts timeouts = context.getClusterControllerTimeouts();

    try {
        return clusterControllerApi.apply(api -> api.setNodeState(
                clusterName,
                storageNodeIndex,
                timeouts.getServerTimeoutOrThrow().toMillis() / 1000.0f,
                stateRequest),
                timeouts);
    } catch (IOException | UncheckedTimeoutException e) {
        String message = String.format(
                "Giving up setting %s for storage node with index %d in cluster %s: %s",
                stateRequest,
                storageNodeIndex,
                clusterName,
                e.getMessage());

        throw new IOException(message, e);
    }
}
 

@Test
public void throws_409_on_timeout() throws HostNameNotFoundException, HostStateChangeDeniedException {
    Orchestrator orchestrator = mock(Orchestrator.class);
    doThrow(new UncheckedTimeoutException("Timeout Message")).when(orchestrator).resume(any(HostName.class));

    try {
        HostResource hostResource = new HostResource(orchestrator, uriInfo);
        hostResource.resume("hostname");
        fail();
    } catch (WebApplicationException w) {
        assertEquals(409, w.getResponse().getStatus());
        assertEquals("resume failed: Timeout Message [deadline]", w.getMessage());
    }
}
 

@Test
public void throws_409_on_suspendAll_timeout() throws BatchHostStateChangeDeniedException, BatchHostNameNotFoundException, BatchInternalErrorException {
    Orchestrator orchestrator = mock(Orchestrator.class);
    doThrow(new UncheckedTimeoutException("Timeout Message")).when(orchestrator).suspendAll(any(), any());

    try {
        HostSuspensionResource resource = new HostSuspensionResource(orchestrator);
        resource.suspendAll("parenthost", Arrays.asList("h1", "h2", "h3"));
        fail();
    } catch (WebApplicationException w) {
        assertEquals(409, w.getResponse().getStatus());
    }
}
 

@Test
public void alreadyTimedOut() {
    clock.advance(Duration.ofSeconds(4));

    try {
        timeouts.getServerTimeoutOrThrow();
        fail();
    } catch (UncheckedTimeoutException e) {
        assertEquals("Timed out after PT3S", e.getMessage());
    }
}
 

@Test
public void justTooLittleTime() {
    clock.advance(originalTimeout.minus(MINIMUM_TIME_LEFT).plus(Duration.ofMillis(1)));
    try {
        timeouts.getServerTimeoutOrThrow();
        fail();
    } catch (UncheckedTimeoutException e) {
        assertEquals("Timed out after PT3S", e.getMessage());
    }
}
 

@Override
public Lock lockMaintenanceJob(String jobName) {
    try {
        return tryLock(lockRoot.append("maintenanceJobLocks").append(jobName));
    } catch (TimeoutException e) {
        throw new UncheckedTimeoutException(e);
    }
}
 

/** Try locking with a low timeout, meaning it is OK to fail lock acquisition.
 *
 * Useful for maintenance jobs, where there is no point in running the jobs back to back.
 */
private Lock tryLock(Path path) throws TimeoutException {
    try {
        return curator.lock(path, tryLockTimeout);
    }
    catch (UncheckedTimeoutException e) {
        throw new TimeoutException(e.getMessage());
    }
}
 

@Test
public void test_application_lock_failure() throws IOException {
    String exceptionMessage = "Timed out after waiting PT1M to acquire lock '/provision/v1/locks/foo/bar/default'";
    long sessionId = applicationRepository.createSession(applicationId(), timeoutBudget, app);
    FailingSessionPrepareHandler handler = new FailingSessionPrepareHandler(SessionPrepareHandler.testOnlyContext(),
                                                                            applicationRepository,
                                                                            componentRegistry.getConfigserverConfig(),
                                                                            new ApplicationLockException(new UncheckedTimeoutException(exceptionMessage)));
    HttpResponse response = handler.handle(createTestRequest(pathPrefix, HttpRequest.Method.PUT, Cmd.PREPARED, sessionId));
    assertEquals(500, response.getStatus());
    Slime data = getData(response);
    assertThat(data.get().field("error-code").asString(), is(HttpErrorResponse.errorCodes.APPLICATION_LOCK_FAILURE.name()));
    assertThat(data.get().field("message").asString(), is(exceptionMessage));
}
 

/**
 * 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 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;
    });
}
 

@Override
public void run() {
    try {
        if (jobControl.isActive(name())) {
            lockAndMaintain();
        }
    } catch (UncheckedTimeoutException ignored) {
        // Another actor is running this job
    } catch (Throwable e) {
        log.log(Level.WARNING, this + " failed. Will retry in " + interval.toMinutes() + " minutes", e);
    }
}
 

/** Acquires the single cluster-global, reentrant lock with the specified timeout for active nodes of this application */
// TODO(mpolden): Remove when all config servers take the new lock
public Lock legacyLock(ApplicationId application, Duration timeout) {
    try {
        return db.lock(legacyLockPath(application), timeout);
    }
    catch (UncheckedTimeoutException e) {
        throw new ApplicationLockException(e);
    }
}
 

private List<Column> getColumns(Connection connection, CreateTableAsSelect createTableAsSelect)
        throws SQLException
{
    io.prestosql.sql.tree.Query createSelectClause = createTableAsSelect.getQuery();

    // Rewrite the query to select zero rows, so that we can get the column names and types
    QueryBody innerQuery = createSelectClause.getQueryBody();
    io.prestosql.sql.tree.Query zeroRowsQuery;
    if (innerQuery instanceof QuerySpecification) {
        QuerySpecification querySpecification = (QuerySpecification) innerQuery;
        innerQuery = new QuerySpecification(
                querySpecification.getSelect(),
                querySpecification.getFrom(),
                querySpecification.getWhere(),
                querySpecification.getGroupBy(),
                querySpecification.getHaving(),
                querySpecification.getOrderBy(),
                querySpecification.getOffset(),
                Optional.of(new Limit("0")));

        zeroRowsQuery = new io.prestosql.sql.tree.Query(createSelectClause.getWith(), innerQuery, Optional.empty(), Optional.empty(), Optional.empty());
    }
    else {
        zeroRowsQuery = new io.prestosql.sql.tree.Query(createSelectClause.getWith(), innerQuery, Optional.empty(), Optional.empty(), Optional.of(new Limit("0")));
    }

    ImmutableList.Builder<Column> columns = ImmutableList.builder();
    try (java.sql.Statement jdbcStatement = connection.createStatement()) {
        ExecutorService executor = newSingleThreadExecutor();
        TimeLimiter limiter = SimpleTimeLimiter.create(executor);
        java.sql.Statement limitedStatement = limiter.newProxy(jdbcStatement, java.sql.Statement.class, timeout.toMillis(), TimeUnit.MILLISECONDS);
        try (ResultSet resultSet = limitedStatement.executeQuery(formatSql(zeroRowsQuery))) {
            ResultSetMetaData metaData = resultSet.getMetaData();
            for (int i = 1; i <= metaData.getColumnCount(); i++) {
                String name = metaData.getColumnName(i);
                int type = metaData.getColumnType(i);
                columns.add(new Column(name, APPROXIMATE_TYPES.contains(type)));
            }
        }
        catch (UncheckedTimeoutException e) {
            throw new SQLException("SQL statement execution timed out", e);
        }
        finally {
            executor.shutdownNow();
        }
    }

    return columns.build();
}
 

/**
 * Acquires the single cluster-global, re-entrant lock for given application. Note that this is the same lock
 * that configserver itself takes when modifying applications.
 *
 * This lock must be taken when writes to paths owned by this class may happen on both the configserver and
 * node-repository side. This behaviour is obviously wrong, but since we pass a NestedTransaction across the
 * configserver and node-repository boundary, the ownership semantics of the transaction (and its operations)
 * becomes unclear.
 *
 * Example of when to use: The config server creates a new transaction and passes the transaction to
 * {@link com.yahoo.vespa.hosted.provision.provisioning.NodeRepositoryProvisioner}, which appends operations to the
 * transaction. The config server then commits (writes) the transaction which may include operations that modify
 * data in paths owned by this class.
 */
public Lock lock(ApplicationId application, Duration timeout) {
    try {
        return db.lock(lockPath(application), timeout);
    } catch (UncheckedTimeoutException e) {
        throw new ApplicationLockException(e);
    }
}
 

/**
 * Returns a mutable host status registry for a locked application instance. All operations performed on
 * the returned registry are executed in the context of a lock, including read operations. Hence, multi-step
 * operations (e.g. read-then-write) are guaranteed to be consistent.
 */
ApplicationLock lockApplication(OrchestratorContext context, ApplicationInstanceReference reference)
        throws UncheckedTimeoutException;