pl.allegro.tech.hermes.client.ReactiveHermesClient Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hermes-client Show documentation
Fast and reliable message broker built on top of Kafka.
The newest version!
package pl.allegro.tech.hermes.client;

import static pl.allegro.tech.hermes.client.HermesMessage.hermesMessage;
import static pl.allegro.tech.hermes.client.HermesResponseBuilder.hermesFailureResponse;
import static reactor.core.Exceptions.isRetryExhausted;

import java.net.URI;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Predicate;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Scheduler;
import reactor.util.retry.Retry;

public class ReactiveHermesClient {
  private static final Logger logger = LoggerFactory.getLogger(ReactiveHermesClient.class);
  private static final String RETRY_CONTEXT_KEY = "hermes-retry-context-key";

  private final ReactiveHermesSender sender;
  private final String uri;
  private final Map defaultHeaders;
  private final AtomicInteger currentlySending = new AtomicInteger(0);
  private final int maxRetries;
  private final Predicate retryCondition;
  private final Duration retrySleep;
  private final Duration maxRetrySleep;
  private final Double jitterFactor;
  private final Scheduler scheduler;
  private volatile boolean shutdown = false;
  private final List messageDeliveryListeners;

  ReactiveHermesClient(
      ReactiveHermesSender sender,
      URI uri,
      Map defaultHeaders,
      int maxRetries,
      Predicate retryCondition,
      long retrySleepInMillis,
      long maxRetrySleepInMillis,
      double jitterFactor,
      Scheduler scheduler) {
    this.sender = sender;
    this.uri = createUri(uri);
    this.defaultHeaders = Collections.unmodifiableMap(new HashMap<>(defaultHeaders));
    this.maxRetries = maxRetries;
    this.retryCondition = retryCondition;
    this.retrySleep = Duration.ofMillis(retrySleepInMillis);
    this.maxRetrySleep = Duration.ofMillis(maxRetrySleepInMillis);
    this.jitterFactor = jitterFactor;
    this.scheduler = scheduler;
    this.messageDeliveryListeners = new ArrayList<>();
  }

  private String createUri(URI uri) {
    String uriString = uri.toString();
    return uriString + (uriString.endsWith("/") ? "" : "/") + "topics/";
  }

  public Mono publishJSON(String topic, byte[] message) {
    return publish(hermesMessage(topic, message).json().build());
  }

  public Mono publishJSON(String topic, String message) {
    return publish(hermesMessage(topic, message).json().build());
  }

  public Mono publishAvro(String topic, int schemaVersion, byte[] message) {
    return publish(hermesMessage(topic, message).avro(schemaVersion).build());
  }

  public Mono publish(String topic, String message) {
    return publish(hermesMessage(topic, message).build());
  }

  public Mono publish(String topic, String contentType, byte[] message) {
    return publish(hermesMessage(topic, message).withContentType(contentType).build());
  }

  public Mono publish(String topic, String contentType, String message) {
    return publish(hermesMessage(topic, message).withContentType(contentType).build());
  }

  public Mono publish(
      String topic, String contentType, int schemaVersion, byte[] message) {
    return publish(
        hermesMessage(topic, message)
            .withContentType(contentType)
            .withSchemaVersion(schemaVersion)
            .build());
  }

  public Mono publish(HermesMessage message) {
    if (shutdown) {
      return completedWithShutdownException();
    }
    HermesMessage.appendDefaults(message, defaultHeaders);
    return publishWithRetries(message);
  }

  public boolean addMessageDeliveryListener(MessageDeliveryListener listener) {
    return messageDeliveryListeners.add(listener);
  }

  private Mono publishWithRetries(HermesMessage message) {
    currentlySending.incrementAndGet();
    Mono sendOnceResult =
        sendOnce(message)
            .flatMap(this::testRetryCondition)
            .onErrorResume(exception -> mapExceptionToFailedAttempt(exception, message))
            .subscribeOn(scheduler);

    return retry(sendOnceResult)
        .doOnSuccess(
            hr -> {
              if (hr.response.isSuccess() || !hr.matchesRetryPolicy) {
                handleSuccessfulRetry(hr.response, hr.attempt);
              } else {
                handleFailure(hr.response, hr.attempt);
              }
            })
        .map(result -> result.response)
        .doFinally(s -> currentlySending.decrementAndGet());
  }

  private Mono retry(Mono sendOnceResult) {
    Retry retrySpec = prepareRetrySpec();
    return sendOnceResult
        .flatMap(this::unwrapFailedAttemptAsException)
        .retryWhen(retrySpec)
        .contextWrite(ctx -> ctx.put(RETRY_CONTEXT_KEY, HermesRetryContext.emptyRetryContext()))
        .onErrorResume(Exception.class, this::unwrapRetryExhaustedException);
  }

  private Mono unwrapRetryExhaustedException(Exception exception) {
    if (isRetryExhausted(exception)) {
      RetryFailedException rfe = (RetryFailedException) (exception.getCause());
      Failed failedAttempt = rfe.failed;
      HermesResponse hermesResponse = failedAttempt.hermesResponse;
      handleMaxRetriesExceeded(hermesResponse, failedAttempt.attempt);
      Throwable cause = rfe.getCause();
      if (cause instanceof ShouldRetryException) {
        ShouldRetryException sre = (ShouldRetryException) cause;
        return Mono.just((Attempt) Result.attempt(sre.hermesResponse, failedAttempt.attempt, true));
      }
      if (hermesResponse.isFailure()) {
        handleFailure(hermesResponse, failedAttempt.attempt);
      }
    }
    return Mono.error(exception);
  }

  private Mono unwrapFailedAttemptAsException(Result result) {
    if (result instanceof Failed) {
      Failed failed = (Failed) result;
      return Mono.error(new RetryFailedException(failed));
    } else {
      return Mono.just((Attempt) result);
    }
  }

  private Mono mapExceptionToFailedAttempt(
      Throwable throwable, HermesMessage hermesMessage) {
    return getNextAttempt()
        .map(attempt -> Result.failureByException(throwable, hermesMessage, attempt));
  }

  private Mono testRetryCondition(HermesResponse response) {
    return getNextAttempt()
        .map(
            attempt -> {
              if (retryCondition.test(response)) {
                return Result.retryableFailure(response, attempt);
              } else {
                return Result.attempt(response, attempt, false);
              }
            });
  }

  private Retry prepareRetrySpec() {
    if (retrySleep.isZero()) {
      return Retry.max(maxRetries).doAfterRetry(this::handleRetryAttempt);
    } else {
      return Retry.backoff(maxRetries, retrySleep)
          .maxBackoff(maxRetrySleep)
          .jitter(jitterFactor)
          .doAfterRetry(this::handleRetryAttempt);
    }
  }

  private void handleRetryAttempt(Retry.RetrySignal retrySignal) {
    RetryFailedException failedException = (RetryFailedException) retrySignal.failure();
    handleFailedAttempt(failedException.failed.hermesResponse, retrySignal.totalRetries() + 1);
  }

  private Mono getNextAttempt() {
    return Mono.deferContextual(Mono::just)
        .map(
            ctx ->
                ctx.getOrDefault(RETRY_CONTEXT_KEY, HermesRetryContext.emptyRetryContext())
                    .getAndIncrementAttempt());
  }

  private Mono sendOnce(HermesMessage message) {
    return Mono.defer(
        () -> {
          long startTime = System.nanoTime();
          try {
            return sender
                .sendReactively(URI.create(uri + message.getTopic()), message)
                .onErrorResume(e -> Mono.just(hermesFailureResponse(e, message)))
                .doOnNext(
                    resp -> {
                      long latency = System.nanoTime() - startTime;
                      messageDeliveryListeners.forEach(l -> l.onSend(resp, latency));
                    });

          } catch (Exception e) {
            return Mono.error(e);
          }
        });
  }

  private Mono completedWithShutdownException() {
    return Mono.error(new HermesClientShutdownException());
  }

  public Mono closeAsync(long pollInterval) {
    shutdown = true;
    CompletableFuture voidCompletableFuture =
        new HermesClientTermination(pollInterval)
            .observe(() -> currentlySending.get() == 0)
            .whenComplete((response, ex) -> scheduler.dispose());
    return Mono.fromFuture(voidCompletableFuture);
  }

  public void close(long pollInterval, long timeout) {
    closeAsync(pollInterval).block(Duration.ofMillis(timeout));
  }

  private void handleMaxRetriesExceeded(HermesResponse response, int attemptCount) {
    messageDeliveryListeners.forEach(l -> l.onMaxRetriesExceeded(response, attemptCount));
    logger.error(
        "Failed to send message to topic {} after {} attempts",
        response.getHermesMessage().getTopic(),
        attemptCount);
  }

  private void handleFailedAttempt(HermesResponse response, long attemptCount) {
    messageDeliveryListeners.forEach(l -> l.onFailedRetry(response, (int) attemptCount));
  }

  private void handleFailure(HermesResponse response, long attemptCount) {
    messageDeliveryListeners.forEach(l -> l.onFailure(response, (int) attemptCount));
  }

  private void handleSuccessfulRetry(HermesResponse response, long attemptCount) {
    messageDeliveryListeners.forEach(l -> l.onSuccessfulRetry(response, (int) attemptCount));
  }

  private static class ShouldRetryException extends Exception {
    private final HermesResponse hermesResponse;

    public ShouldRetryException(HermesResponse hermesResponse) {
      this.hermesResponse = hermesResponse;
    }

    public HermesResponse getHermesResponse() {
      return hermesResponse;
    }
  }

  private static class RetryFailedException extends Exception {
    private final Failed failed;

    public RetryFailedException(Failed failed) {
      super(failed.cause);
      this.failed = failed;
    }
  }

  private interface Result {
    static Result attempt(HermesResponse response, int attempt, boolean qualifiedForRetry) {
      return new Attempt(response, attempt, qualifiedForRetry);
    }

    static Result retryableFailure(HermesResponse response, int attempt) {
      return new Failed(response, attempt, new ShouldRetryException(response));
    }

    static Result failureByException(
        Throwable throwable, HermesMessage hermesMessage, int attempt) {
      return new Failed(hermesFailureResponse(throwable, hermesMessage), attempt, throwable);
    }
  }

  private static class Attempt implements Result {
    private final HermesResponse response;
    private final int attempt;
    private final boolean matchesRetryPolicy;

    private Attempt(HermesResponse response, int attempt, boolean matchesRetryPolicy) {
      this.response = response;
      this.attempt = attempt;
      this.matchesRetryPolicy = matchesRetryPolicy;
    }
  }

  private static class Failed implements Result {
    private final HermesResponse hermesResponse;
    private final int attempt;
    private final Throwable cause;

    private Failed(HermesResponse hermesResponse, int attempt, Throwable cause) {
      this.attempt = attempt;
      this.cause = cause;
      this.hermesResponse = hermesResponse;
    }
  }

  private static class HermesRetryContext {
    static HermesRetryContext emptyRetryContext() {
      return new HermesRetryContext();
    }

    private int attempt;

    HermesRetryContext() {
      attempt = 1;
    }

    int getAndIncrementAttempt() {
      return attempt++;
    }
  }
}