io.scalecube.services.transport.rsocket.RSocketClientTransport Maven / Gradle / Ivy

Go to download
package io.scalecube.services.transport.rsocket;

import io.rsocket.RSocket;
import io.rsocket.RSocketFactory;
import io.rsocket.transport.netty.client.TcpClientTransport;
import io.rsocket.util.ByteBufPayload;
import io.scalecube.services.transport.api.Address;
import io.scalecube.services.transport.api.ClientChannel;
import io.scalecube.services.transport.api.ClientTransport;
import io.scalecube.services.transport.api.ServiceMessageCodec;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
import reactor.netty.resources.LoopResources;
import reactor.netty.tcp.TcpClient;

public class RSocketClientTransport implements ClientTransport {

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

  private final ThreadLocal>> rsockets =
      ThreadLocal.withInitial(ConcurrentHashMap::new);

  private final ServiceMessageCodec codec;
  private final LoopResources loopResources;

  /**
   * Constructor for this transport.
   *
   * @param codec message codec
   * @param loopResources client loop resources
   */
  public RSocketClientTransport(ServiceMessageCodec codec, LoopResources loopResources) {
    this.codec = codec;
    this.loopResources = loopResources;
  }

  @Override
  public ClientChannel create(Address address) {
    final Map> monoMap = rsockets.get(); // keep reference for threadsafety
    Mono rsocket =
        monoMap.computeIfAbsent(address, address1 -> connect(address1, monoMap));
    return new RSocketClientChannel(rsocket, codec);
  }

  private Mono connect(Address address, Map> monoMap) {
    TcpClient tcpClient =
        TcpClient.newConnection() // create non-pooled
            .runOn(loopResources)
            .host(address.host())
            .port(address.port());

    Mono rsocketMono =
        RSocketFactory.connect()
            .frameDecoder(
                frame ->
                    ByteBufPayload.create(
                        frame.sliceData().retain(), frame.sliceMetadata().retain()))
            .transport(() -> TcpClientTransport.create(tcpClient))
            .start();

    return rsocketMono
        .doOnSuccess(
            rsocket -> {
              LOGGER.info("Connected successfully on {}", address);
              // setup shutdown hook
              rsocket
                  .onClose()
                  .doOnTerminate(
                      () -> {
                        monoMap.remove(address);
                        LOGGER.info("Connection closed on {}", address);
                      })
                  .subscribe(null, th -> LOGGER.warn("Exception on closing rsocket: {}", th));
            })
        .doOnError(
            throwable -> {
              LOGGER.warn("Connect failed on {}, cause: {}", address, throwable);
              monoMap.remove(address);
            })
        .cache();
  }
}