org.zalando.riptide.failsafe.FailsafePlugin Maven / Gradle / Ivy

Go to download
package org.zalando.riptide.failsafe;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import lombok.AllArgsConstructor;
import net.jodah.failsafe.Failsafe;
import net.jodah.failsafe.Policy;
import net.jodah.failsafe.RetryPolicy;
import net.jodah.failsafe.event.ExecutionAttemptedEvent;
import net.jodah.failsafe.function.CheckedConsumer;
import org.apiguardian.api.API;
import org.springframework.http.client.ClientHttpResponse;
import org.zalando.riptide.Attribute;
import org.zalando.riptide.Plugin;
import org.zalando.riptide.RequestArguments;
import org.zalando.riptide.RequestExecution;
import org.zalando.riptide.idempotency.IdempotencyPredicate;

import java.util.concurrent.ScheduledExecutorService;
import java.util.function.Predicate;
import java.util.function.UnaryOperator;
import java.util.stream.Stream;

import static lombok.AccessLevel.PRIVATE;
import static org.apiguardian.api.API.Status.MAINTAINED;

@API(status = MAINTAINED)
@AllArgsConstructor(access = PRIVATE)
public final class FailsafePlugin implements Plugin {

    public static final Attribute ATTEMPTS = Attribute.generate();

    private final ImmutableList> policies;
    private final ScheduledExecutorService scheduler;
    private final Predicate predicate;
    private final RetryListener listener;

    public FailsafePlugin(final ImmutableList> policies,
            final ScheduledExecutorService scheduler) {
        this(policies, scheduler, new IdempotencyPredicate(), RetryListener.DEFAULT);
    }

    public FailsafePlugin withPredicate(final Predicate predicate) {
        return new FailsafePlugin(policies, scheduler, predicate, listener);
    }

    public FailsafePlugin withListener(final RetryListener listener) {
        return new FailsafePlugin(policies, scheduler, predicate, listener);
    }

    @Override
    public RequestExecution aroundDispatch(final RequestExecution execution) {
        return arguments -> {
            final Policy[] policies = select(arguments);

            if (policies.length == 0) {
                return execution.execute(arguments);
            }

            return Failsafe.with(select(arguments))
                    .with(scheduler)
                    .getStageAsync(context -> execution
                            .execute(withAttempts(arguments, context.getAttemptCount())));
        };
    }

    private Policy[] select(final RequestArguments arguments) {
        final Stream> stream = policies.stream()
                .filter(skipRetriesIfNeeded(arguments))
                .map(withRetryListener(arguments));

        @SuppressWarnings("unchecked") final Policy[] policies = stream.toArray(Policy[]::new);

        return policies;
    }

    // TODO depends on the exception, e.g. pre-request exceptions are fine!
    private Predicate> skipRetriesIfNeeded(final RequestArguments arguments) {
        return predicate.test(arguments) ?
                policy -> true :
                policy -> !(policy instanceof RetryPolicy);
    }

    private UnaryOperator> withRetryListener(final RequestArguments arguments) {
        return policy -> {
            if (policy instanceof RetryPolicy) {
                final RetryPolicy retryPolicy = (RetryPolicy) policy;
                return retryPolicy.copy()
                        .onFailedAttempt(new RetryListenerAdapter(listener, arguments));
            } else {
                return policy;
            }
        };
    }

    private RequestArguments withAttempts(final RequestArguments arguments, final int attempts) {
        if (attempts == 0) {
            return arguments;
        }

        return arguments.withAttribute(ATTEMPTS, attempts);
    }

    @VisibleForTesting
    @AllArgsConstructor
    static final class RetryListenerAdapter implements CheckedConsumer> {
        private final RetryListener listener;
        private final RequestArguments arguments;

        @Override
        public void accept(final ExecutionAttemptedEvent event) {
            listener.onRetry(arguments, event);
        }
    }

}