io.github.tramchamploo.bufferslayer.RxReporter Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of bufferslayer-rx Show documentation
The newest version!
package io.github.tramchamploo.bufferslayer;

import static com.google.common.base.Preconditions.checkNotNull;
import static io.github.tramchamploo.bufferslayer.MessageDroppedException.dropped;
import static java.util.Collections.singletonList;

import io.github.tramchamploo.bufferslayer.Message.MessageKey;
import io.github.tramchamploo.bufferslayer.OverflowStrategy.Strategy;
import io.github.tramchamploo.bufferslayer.internal.Future;
import io.github.tramchamploo.bufferslayer.internal.FutureListener;
import io.github.tramchamploo.bufferslayer.internal.MessageFuture;
import io.github.tramchamploo.bufferslayer.internal.MessagePromise;
import io.reactivex.BackpressureOverflowStrategy;
import io.reactivex.BackpressureStrategy;
import io.reactivex.Flowable;
import io.reactivex.FlowableEmitter;
import io.reactivex.FlowableOnSubscribe;
import io.reactivex.Scheduler;
import io.reactivex.flowables.GroupedFlowable;
import io.reactivex.functions.Action;
import io.reactivex.functions.Consumer;
import io.reactivex.functions.Function;
import io.reactivex.schedulers.Schedulers;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class RxReporter implements Reporter, FlowableOnSubscribe> {

  static final Logger logger = LoggerFactory.getLogger(RxReporter.class);

  final long messageTimeoutNanos;
  final int bufferedMaxMessages;
  final int pendingMaxMessages;
  final Strategy overflowStrategy;
  final Sender sender;
  final ReporterMetrics metrics;
  final AtomicBoolean closed = new AtomicBoolean(false);

  private FlowableEmitter> emitter;
  private Scheduler scheduler;

  private RxReporter(Builder builder) {
    this.sender = builder.sender;
    this.metrics = builder.metrics;

    this.messageTimeoutNanos = builder.messageTimeoutNanos;
    this.bufferedMaxMessages = builder.bufferedMaxMessages;
    this.pendingMaxMessages = builder.pendingMaxMessages;
    this.overflowStrategy = builder.overflowStrategy;
    this.scheduler = builder.scheduler;

    Flowable> flowable = Flowable.create(this, BackpressureStrategy.MISSING);
    initBackpressurePolicy(flowable)
        .observeOn(Schedulers.single())
        .groupBy(new MessagePartitioner())
        .subscribe(new MessageGroupSubscriber(messageTimeoutNanos, bufferedMaxMessages, sender, scheduler));
  }

  public static  Builder builder(Sender sender) {
    return new Builder<>(sender);
  }

  private Flowable> initBackpressurePolicy(Flowable> flowable) {
    Strategy strategy = this.overflowStrategy;
    if (strategy == Strategy.DropNew) {
      return flowable.onBackpressureDrop(new Consumer>() {
        @Override
        public void accept(MessagePromise promise) {
          metricsCallback(1);
        }
      });
    } else {
      BackpressureOverflowStrategy rxStrategy = RxOverflowStrategyBridge.toRxStrategy(strategy);
      return flowable.onBackpressureBuffer(pendingMaxMessages, new Action() {
        @Override
        public void run() {
          metricsCallback(1);
        }
      }, rxStrategy);
    }
  }

  private void metricsCallback(int quantity) {
    metrics.incrementMessagesDropped(quantity);
  }

  public static final class Builder extends Reporter.Builder {

    Scheduler scheduler = Schedulers.io();

    Builder(Sender sender) {
      super(sender);
    }

    @Override
    public Builder metrics(ReporterMetrics metrics) {
      super.metrics(metrics);
      return this;
    }

    @Override
    public Builder messageTimeout(long timeout, TimeUnit unit) {
      super.messageTimeout(timeout, unit);
      return this;
    }

    @Override
    public Builder bufferedMaxMessages(int bufferedMaxMessages) {
      super.bufferedMaxMessages(bufferedMaxMessages);
      return this;
    }

    @Override
    public Builder pendingMaxMessages(int pendingMaxMessages) {
      super.pendingMaxMessages(pendingMaxMessages);
      return this;
    }

    @Override
    public Builder overflowStrategy(Strategy overflowStrategy) {
      super.overflowStrategy(overflowStrategy);
      return this;
    }

    public Builder scheduler(Scheduler scheduler) {
      this.scheduler = checkNotNull(scheduler, "scheduler");
      return this;
    }

    @Override
    public RxReporter build() {
      return new RxReporter<>(this);
    }
  }

  private class MessageGroupSubscriber implements Consumer>> {

    final long messageTimeoutNanos;
    final int bufferedMaxMessages;
    final Sender sender;
    final Scheduler scheduler;

    private MessageGroupSubscriber(long messageTimeoutNanos,
                                   int bufferedMaxMessages,
                                   Sender sender,
                                   Scheduler scheduler) {
      this.messageTimeoutNanos = messageTimeoutNanos == 0 ? Long.MAX_VALUE : messageTimeoutNanos;
      this.bufferedMaxMessages = bufferedMaxMessages;
      this.sender = sender;
      this.scheduler = scheduler;
    }

    @Override
    public void accept(GroupedFlowable> group) {
      Flowable>> bufferedMessages =
          group.buffer(messageTimeoutNanos, TimeUnit.NANOSECONDS, scheduler, bufferedMaxMessages);
      bufferedMessages.subscribeOn(scheduler).subscribe(SenderConsumerBridge.toConsumer(sender));
    }
  }

  private class MessagePartitioner implements Function, MessageKey> {

    @Override
    public MessageKey apply(MessagePromise promise) {
      return promise.message().asMessageKey();
    }
  }

  @Override
  public CheckResult check() {
    return sender.check();
  }

  @Override
  @SuppressWarnings("unchecked")
  public MessageFuture report(M message) {
    checkNotNull(message, "message");
    metrics.incrementMessages(1);

    if (closed.get()) { // Drop the message when closed.
      MessageDroppedException cause =
          dropped(new IllegalStateException("closed!"), singletonList(message));
      MessageFuture future = (MessageFuture) message.newFailedFuture(cause);
      metricsCallback(1);
      return future;
    }

    MessagePromise promise = message.newPromise();
    emitter.onNext(promise);
    setLoggingListener(promise);
    return promise;
  }

  private void setLoggingListener(MessageFuture future) {
    future.addListener(new FutureListener() {

      @Override
      public void operationComplete(Future future) {
        if (!future.isSuccess()) {
          MessageDroppedException cause = (MessageDroppedException) future.cause();
          metricsCallback(cause.dropped.size());
          logger.warn(cause.getMessage());
        }
      }
    });
  }

  @Override
  public void subscribe(FlowableEmitter> emitter) throws Exception {
    this.emitter = emitter;
  }

  @Override
  public void close() {
    if (!closed.compareAndSet(false, true)) return;
    emitter.onComplete();
  }
}