io.scalecube.services.gateway.transport.websocket.WebsocketGatewayClient Maven / Gradle / Ivy

Go to download
package io.scalecube.services.gateway.transport.websocket;

import io.netty.buffer.ByteBuf;
import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
import io.scalecube.services.api.ServiceMessage;
import io.scalecube.services.gateway.transport.GatewayClient;
import io.scalecube.services.gateway.transport.GatewayClientCodec;
import io.scalecube.services.gateway.transport.GatewayClientSettings;
import java.time.Duration;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.MonoProcessor;
import reactor.netty.Connection;
import reactor.netty.http.client.HttpClient;
import reactor.netty.resources.LoopResources;

public final class WebsocketGatewayClient implements GatewayClient {

  private static final Logger LOGGER = LoggerFactory.getLogger(WebsocketGatewayClient.class);

  private static final String STREAM_ID = "sid";
  private static final String SIGNAL = "sig";

  @SuppressWarnings("rawtypes")
  private static final AtomicReferenceFieldUpdater
      websocketMonoUpdater =
          AtomicReferenceFieldUpdater.newUpdater(
              WebsocketGatewayClient.class, Mono.class, "websocketMono");

  private final GatewayClientCodec codec;
  private final GatewayClientSettings settings;
  private final HttpClient httpClient;
  private final AtomicLong sidCounter = new AtomicLong();
  private final LoopResources loopResources;
  private final MonoProcessor close = MonoProcessor.create();
  private final MonoProcessor onClose = MonoProcessor.create();

  @SuppressWarnings("unused")
  private volatile Mono websocketMono;

  /**
   * Creates instance of websocket client transport.
   *
   * @param settings client settings
   * @param codec client codec.
   */
  public WebsocketGatewayClient(GatewayClientSettings settings, GatewayClientCodec codec) {
    this.settings = settings;
    this.codec = codec;
    this.loopResources = LoopResources.create("websocket-gateway-client");

    httpClient =
        HttpClient.newConnection()
            .followRedirect(settings.followRedirect())
            .tcpConfiguration(
                tcpClient -> {
                  if (settings.sslProvider() != null) {
                    tcpClient = tcpClient.secure(settings.sslProvider());
                  }
                  return tcpClient
                      .wiretap(settings.wiretap())
                      .runOn(loopResources)
                      .host(settings.host())
                      .port(settings.port());
                });

    // Setup cleanup
    close
        .then(doClose())
        .doFinally(s -> onClose.onComplete())
        .doOnTerminate(() -> LOGGER.info("Closed WebsocketGatewayClient resources"))
        .subscribe(
            null, ex -> LOGGER.warn("Exception occurred on WebsocketGatewayClient close: " + ex));
  }

  @Override
  public Mono requestResponse(ServiceMessage request) {
    return Mono.defer(
        () -> {
          long sid = sidCounter.incrementAndGet();
          return getOrConnect()
              .flatMap(
                  session ->
                      session
                          .send(encodeRequest(request, sid), sid)
                          .doOnSubscribe(s -> LOGGER.debug("Sending request {}", request))
                          .then(session.newMonoProcessor(sid))
                          .doOnCancel(() -> handleCancel(sid, session))
                          .doFinally(s -> session.removeProcessor(sid)));
        });
  }

  @Override
  public Flux requestStream(ServiceMessage request) {
    return Flux.defer(
        () -> {
          long sid = sidCounter.incrementAndGet();
          return getOrConnect()
              .flatMapMany(
                  session ->
                      session
                          .send(encodeRequest(request, sid), sid)
                          .doOnSubscribe(s -> LOGGER.debug("Sending request {}", request))
                          .thenMany(session.newUnicastProcessor(sid))
                          .doOnCancel(() -> handleCancel(sid, session))
                          .doFinally(s -> session.removeProcessor(sid)));
        });
  }

  @Override
  public Flux requestChannel(Flux requests) {
    return Flux.error(
        new UnsupportedOperationException(
            "requestChannel is not supported by WebSocket transport implementation"));
  }

  @Override
  public void close() {
    close.onComplete();
  }

  @Override
  public Mono onClose() {
    return onClose;
  }

  private Mono doClose() {
    return Mono.defer(loopResources::disposeLater);
  }

  public GatewayClientCodec getCodec() {
    return codec;
  }

  private Mono getOrConnect() {
    // noinspection unchecked
    return Mono.defer(() -> websocketMonoUpdater.updateAndGet(this, this::getOrConnect0));
  }

  private Mono getOrConnect0(Mono prev) {
    if (prev != null) {
      return prev;
    }

    Duration keepAliveInterval = settings.keepAliveInterval();

    return httpClient
        .websocket()
        .uri("/")
        .connect()
        .map(
            connection ->
                keepAliveInterval != Duration.ZERO
                    ? connection
                        .onReadIdle(keepAliveInterval.toMillis(), () -> onReadIdle(connection))
                        .onWriteIdle(keepAliveInterval.toMillis(), () -> onWriteIdle(connection))
                    : connection)
        .map(
            connection -> {
              WebsocketSession session = new WebsocketSession(codec, connection);
              LOGGER.info("Created {} on {}:{}", session, settings.host(), settings.port());
              // setup shutdown hook
              session
                  .onClose()
                  .doOnTerminate(
                      () -> {
                        websocketMonoUpdater.getAndSet(this, null); // clear reference
                        LOGGER.info(
                            "Closed {} on {}:{}", session, settings.host(), settings.port());
                      })
                  .subscribe(
                      null,
                      th ->
                          LOGGER.warn(
                              "Exception on closing session={}, cause: {}",
                              session.id(),
                              th.toString()));
              return session;
            })
        .doOnError(
            ex -> {
              LOGGER.warn(
                  "Failed to connect on {}:{}, cause: {}", settings.host(), settings.port(), ex);
              websocketMonoUpdater.getAndSet(this, null); // clear reference
            })
        .cache();
  }

  private void onWriteIdle(Connection connection) {
    LOGGER.debug("Sending keepalive on writeIdle");
    connection
        .outbound()
        .sendObject(new PingWebSocketFrame())
        .then()
        .subscribe(null, ex -> LOGGER.warn("Can't send keepalive on writeIdle: " + ex));
  }

  private void onReadIdle(Connection connection) {
    LOGGER.debug("Sending keepalive on readIdle");
    connection
        .outbound()
        .sendObject(new PingWebSocketFrame())
        .then()
        .subscribe(null, ex -> LOGGER.warn("Can't send keepalive on readIdle: " + ex));
  }

  private void handleCancel(long sid, WebsocketSession session) {
    ByteBuf byteBuf =
        codec.encode(
            ServiceMessage.builder()
                .header(STREAM_ID, sid)
                .header(SIGNAL, Signal.CANCEL.codeAsString())
                .build());
    session
        .send(byteBuf, sid)
        .subscribe(
            null,
            th ->
                LOGGER.error(
                    "Exception on sending CANCEL signal for session={}", session.id(), th));
  }

  private ByteBuf encodeRequest(ServiceMessage message, long sid) {
    return codec.encode(ServiceMessage.from(message).header(STREAM_ID, sid).build());
  }
}