com.jauntsdn.rsocket.LeaseRSocketResponder Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of rsocket-core Show documentation
RSocket core library
There is a newer version: 0.9.8
package com.jauntsdn.rsocket;

import static com.jauntsdn.rsocket.StreamErrorMappers.*;

import com.jauntsdn.rsocket.frame.FrameType;
import com.jauntsdn.rsocket.frame.decoder.PayloadDecoder;
import com.jauntsdn.rsocket.lease.*;
import io.netty.buffer.ByteBufAllocator;
import java.util.function.Consumer;
import org.reactivestreams.Publisher;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Signal;

class LeaseRSocketResponder extends RSocketResponder {
  private final ResponderLeaseHandler leaseHandler;
  private final Disposable leaseDisposable;

  LeaseRSocketResponder(
      ByteBufAllocator allocator,
      DuplexConnection connection,
      RSocket requestHandler,
      PayloadDecoder payloadDecoder,
      Consumer errorConsumer,
      ErrorFrameMapper errorFrameMapper,
      ResponderLeaseHandler leaseHandler) {
    super(allocator, connection, requestHandler, payloadDecoder, errorConsumer, errorFrameMapper);
    this.leaseHandler = leaseHandler;
    this.leaseDisposable = leaseHandler.send(this::sendFrame);
  }

  @Override
  public Mono fireAndForget(Payload payload) {
    try {
      ResponderLeaseHandler leaseHandler = this.leaseHandler;
      FrameType requestType = FrameType.REQUEST_FNF;
      Signal requestOrReject = leaseHandler.allowRequest(requestType, payload.metadata());
      Throwable reject = requestOrReject.getThrowable();
      if (reject != null) {
        payload.release();
        return Mono.error(reject);
      } else {
        Mono response = super.fireAndForget(payload);
        if (leaseHandler.requireStats()) {
          MonoSignalStatsRecorder recorder =
              new MonoSignalStatsRecorder(leaseHandler, requestType, requestOrReject.get());
          return response.doOnEach(recorder).doOnCancel(recorder);
        }
        return response;
      }
    } catch (Throwable t) {
      return Mono.error(t);
    }
  }

  @Override
  public Mono requestResponse(Payload payload) {
    try {
      ResponderLeaseHandler leaseHandler = this.leaseHandler;
      FrameType requestType = FrameType.REQUEST_RESPONSE;
      Signal requestOrReject = leaseHandler.allowRequest(requestType, payload.metadata());
      Throwable reject = requestOrReject.getThrowable();
      if (reject != null) {
        payload.release();
        return Mono.error(reject);
      } else {
        Mono response = super.requestResponse(payload);
        if (leaseHandler.requireStats()) {
          MonoSignalStatsRecorder recorder =
              new MonoSignalStatsRecorder(leaseHandler, requestType, requestOrReject.get());
          return response.doOnEach(recorder).doOnCancel(recorder);
        }
        return response;
      }

    } catch (Throwable t) {
      return Mono.error(t);
    }
  }

  @Override
  public Flux requestStream(Payload payload) {
    try {
      ResponderLeaseHandler leaseHandler = this.leaseHandler;
      FrameType requestType = FrameType.REQUEST_STREAM;
      Signal requestOrReject = leaseHandler.allowRequest(requestType, payload.metadata());
      Throwable reject = requestOrReject.getThrowable();
      if (reject != null) {
        payload.release();
        return Flux.error(reject);
      } else {
        Flux response = super.requestStream(payload);
        if (leaseHandler.requireStats()) {
          FluxSignalStatsRecorder recorder =
              new FluxSignalStatsRecorder(leaseHandler, requestType, requestOrReject.get());
          return response.doOnEach(recorder).doOnCancel(recorder);
        }
        return response;
      }
    } catch (Throwable t) {
      return Flux.error(t);
    }
  }

  @Override
  public Flux requestChannel(Payload payload, Publisher payloads) {
    try {
      ResponderLeaseHandler leaseHandler = this.leaseHandler;
      FrameType requestType = FrameType.REQUEST_CHANNEL;
      Signal requestOrReject = leaseHandler.allowRequest(requestType, payload.metadata());
      Throwable reject = requestOrReject.getThrowable();
      if (reject != null) {
        payload.release();
        return Flux.error(reject);
      } else {
        Flux response = super.requestChannel(payload, payloads);
        if (leaseHandler.requireStats()) {
          FluxSignalStatsRecorder recorder =
              new FluxSignalStatsRecorder(leaseHandler, requestType, requestOrReject.get());
          return response.doOnEach(recorder).doOnCancel(recorder);
        }
        return response;
      }
    } catch (Throwable t) {
      return Flux.error(t);
    }
  }

  @Override
  void terminate() {
    super.terminate();
    leaseDisposable.dispose();
  }
}