com.spotify.helios.client.RetryingRequestDispatcher.java Source code

Java tutorial

Introduction

Here is the source code for com.spotify.helios.client.RetryingRequestDispatcher.java

Source

/*
 * Copyright (c) 2015 Spotify AB.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.spotify.helios.client;

import com.spotify.helios.common.Clock;
import com.spotify.helios.common.SystemClock;

import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.base.Supplier;
import com.google.common.base.Throwables;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningScheduledExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.SettableFuture;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * A {@link RequestDispatcher} that retries.
 */
class RetryingRequestDispatcher implements RequestDispatcher {

    private static final Logger log = LoggerFactory.getLogger(RetryingRequestDispatcher.class);

    private final ListeningScheduledExecutorService executorService;
    private final RequestDispatcher delegate;
    private final Clock clock;
    private final long retryTimeoutMillis;
    private final long delayMillis;

    private RetryingRequestDispatcher(final RequestDispatcher delegate,
            final ListeningScheduledExecutorService executorService, final Clock clock,
            final long retryTimeoutMillis, final long delayMillis) {
        this.delegate = delegate;
        this.executorService = executorService;
        this.clock = clock;
        this.retryTimeoutMillis = retryTimeoutMillis;
        this.delayMillis = delayMillis;
    }

    @Override
    public ListenableFuture<Response> request(final URI uri, final String method, final byte[] entityBytes,
            final Map<String, List<String>> headers) {
        final long deadline = clock.now().getMillis() + retryTimeoutMillis;
        final SettableFuture<Response> future = SettableFuture.create();
        final Supplier<ListenableFuture<Response>> code = new Supplier<ListenableFuture<Response>>() {
            @Override
            public ListenableFuture<Response> get() {
                return delegate.request(uri, method, entityBytes, headers);
            }
        };
        startRetry(future, code, deadline, delayMillis, uri);
        return future;
    }

    @Override
    public void close() throws IOException {
        delegate.close();
    }

    private void startRetry(final SettableFuture<Response> future, final Supplier<ListenableFuture<Response>> code,
            final long deadline, final long delayMillis, final URI uri) {

        ListenableFuture<Response> codeFuture;
        try {
            codeFuture = code.get();
        } catch (Exception e) {
            log.debug("Failed to connect to {}, retrying in {} seconds.", uri.toString(),
                    TimeUnit.MILLISECONDS.toSeconds(delayMillis));
            log.debug("Specific reason for connection failure follows", e);
            handleFailure(future, code, deadline, delayMillis, e, uri);
            return;
        }

        Futures.addCallback(codeFuture, new FutureCallback<Response>() {
            @Override
            public void onSuccess(Response result) {
                future.set(result);
            }

            @Override
            public void onFailure(@NotNull Throwable t) {
                log.warn("Failed to connect to {}, retrying in {} seconds. Exception chain was: {} ",
                        uri.toString(), TimeUnit.MILLISECONDS.toSeconds(delayMillis), getChainAsString(t));
                log.debug("Specific reason for connection failure follows", t);
                handleFailure(future, code, deadline, delayMillis, t, uri);
            }
        });
    }

    private static String getChainAsString(final Throwable t) {
        final List<Throwable> causalChain = Throwables.getCausalChain(t);
        final List<String> messages = Lists.transform(causalChain, new Function<Throwable, String>() {
            @Override
            public String apply(final Throwable input) {
                return input.toString();
            }
        });
        return Joiner.on(", ").join(messages);
    }

    private void handleFailure(final SettableFuture<Response> future,
            final Supplier<ListenableFuture<Response>> code, final long deadline, final long delayMillis,
            final Throwable t, final URI uri) {
        if (clock.now().getMillis() < deadline) {
            if (delayMillis > 0) {
                executorService.schedule(new Runnable() {
                    @Override
                    public void run() {
                        startRetry(future, code, deadline - 1, delayMillis, uri);
                    }
                }, delayMillis, TimeUnit.MILLISECONDS);
            } else {
                startRetry(future, code, deadline - 1, delayMillis, uri);
            }
        } else {
            future.setException(t);
        }
    }

    static Builder forDispatcher(RequestDispatcher delegate) {
        return new Builder(delegate);
    }

    public static final class Builder {

        private final RequestDispatcher delegate;
        private ListeningScheduledExecutorService executor;
        private Clock clock = new SystemClock();
        private long retryTimeoutMillis = 60000;
        private long delayMillis = 5000;

        private Builder(final RequestDispatcher delegate) {
            this.delegate = delegate;
        }

        public Builder setExecutor(ScheduledExecutorService executorService) {
            this.executor = MoreExecutors.listeningDecorator(executorService);
            return this;
        }

        /** Defaults to SystemClock. */
        public Builder setClock(Clock clock) {
            this.clock = clock;
            return this;
        }

        /**
         * Set the total amount of time that the RetryingRequestDispatcher allows before giving up on
         * the request. Defaults to 60 seconds.
         */
        public Builder setRetryTimeout(long timeout, TimeUnit unit) {
            this.retryTimeoutMillis = unit.toMillis(timeout);
            return this;
        }

        /**
         * Set how much time the RetryingRequestDispatcher waits to retry the request. Defaults to 5
         * seconds.
         */
        public Builder setDelayOnFailure(long delay, TimeUnit unit) {
            this.delayMillis = unit.toMillis(delay);
            return this;
        }

        public RetryingRequestDispatcher build() {
            return new RetryingRequestDispatcher(delegate, executor, clock, retryTimeoutMillis, delayMillis);
        }
    }
}