io.rsocket.RSocket Java Examples

/**
 * Creates a new {@link RSocket}.
 * @param delegate the {@link RSocket} to delegate to
 * @param meterRegistry the {@link MeterRegistry} to use
 * @param tags additional tags to attach to {@link Meter}s
 * @throws IllegalArgumentException if {@code delegate} or {@code meterRegistry} is
 * {@code null}
 */
public MicrometerResponderRSocket(RSocket delegate, MeterRegistry meterRegistry,
		Tag... tags) {
	Assert.notNull(delegate, "delegate must not be null");
	Assert.notNull(meterRegistry, "meterRegistry must not be null");

	this.delegate = delegate;
	this.metadataPush = new InteractionCounters(meterRegistry, "metadata.push", tags);
	this.requestChannel = new InteractionCounters(meterRegistry, "request.channel",
			tags);
	this.requestFireAndForget = new InteractionCounters(meterRegistry, "request.fnf",
			tags);
	this.requestResponse = new InteractionTimers(meterRegistry, "request.response",
			tags);
	this.requestStream = new InteractionCounters(meterRegistry, "request.stream",
			tags);
}
 

@Test
void timeoutOnPermanentDisconnect() {
  CloseableChannel closeable = newServerRSocket().block();

  DisconnectableClientTransport clientTransport =
      new DisconnectableClientTransport(clientTransport(closeable.address()));

  int sessionDurationSeconds = 5;
  RSocket rSocket = newClientRSocket(clientTransport, sessionDurationSeconds).block();

  Mono.delay(Duration.ofSeconds(1)).subscribe(v -> clientTransport.disconnectPermanently());

  StepVerifier.create(
          rSocket.requestChannel(testRequest()).then().doFinally(s -> closeable.dispose()))
      .expectError(ClosedChannelException.class)
      .verify(Duration.ofSeconds(7));
}
 

@Test(timeout = 2000)
public void testRequestPropagatesCorrectlyForRequestChannel() {
  rule.setRequestAcceptor(
      new RSocket() {
        @Override
        public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
          return Flux.from(payloads)
              // specifically limits request to 3 in order to prevent 256 request from limitRate
              // hidden on the responder side
              .limitRequest(3);
        }
      });

  Flux.range(0, 3)
      .map(i -> DefaultPayload.create("" + i))
      .as(rule.crs::requestChannel)
      .as(publisher -> StepVerifier.create(publisher, 3))
      .expectSubscription()
      .expectNextCount(3)
      .expectComplete()
      .verify(Duration.ofMillis(5000));
}
 

public Flux<Payload> requestChannel(Payload signal, Publisher<Payload> payloads) {
    BinaryRoutingMetadata binaryRoutingMetadata = binaryRoutingMetadata(signal.metadata());
    GSVRoutingMetadata gsvRoutingMetadata;
    if (binaryRoutingMetadata != null) {
        gsvRoutingMetadata = GSVRoutingMetadata.from(new String(binaryRoutingMetadata.getRoutingText(), StandardCharsets.UTF_8));
    } else {
        RSocketCompositeMetadata compositeMetadata = RSocketCompositeMetadata.from(signal.metadata());
        gsvRoutingMetadata = compositeMetadata.getRoutingMetaData();
        if (gsvRoutingMetadata == null) {
            return Flux.error(new InvalidException(RsocketErrorCode.message("RST-600404")));
        }
    }
    Mono<RSocket> destination = findDestination(gsvRoutingMetadata);
    return destination.flatMapMany(rsocket -> {
        recordServiceInvoke(principal.getName(), gsvRoutingMetadata.gsv());
        metrics(gsvRoutingMetadata, "0x07");
        return rsocket.requestChannel(payloads);
    });
}
 

public void onRSocketClosed(String rsocketUri, RSocket rsocket, @Nullable Throwable cause) {
    //in last rsocket uris or not
    if (this.lastRSocketUris.contains(rsocketUri)) {
        this.unHealthyUriSet.add(rsocketUri);
        if (activeSockets.containsKey(rsocketUri)) {
            activeSockets.remove(rsocketUri);
            this.randomSelector = new RandomSelector<>(this.serviceId, new ArrayList<>(activeSockets.values()));
            log.error(RsocketErrorCode.message("RST-500407", rsocketUri));
            tryToReconnect(rsocketUri, cause);
        }
        if (!rsocket.isDisposed()) {
            try {
                rsocket.dispose();
            } catch (Exception ignore) {

            }
        }
    }
}
 

@ParameterizedTest
@MethodSource("allInteractions")
void singleSubscriber(Function<RSocket, Publisher<?>> interaction) {
  Flux<?> response = Flux.from(interaction.apply(rSocketRequester));

  AssertSubscriber assertSubscriberA = AssertSubscriber.create();
  AssertSubscriber assertSubscriberB = AssertSubscriber.create();

  response.subscribe(assertSubscriberA);
  response.subscribe(assertSubscriberB);

  connection.addToReceivedBuffer(PayloadFrameCodec.encodeComplete(connection.alloc(), 1));

  assertSubscriberA.assertTerminated();
  assertSubscriberB.assertTerminated();

  Assertions.assertThat(requestFramesCount(connection.getSent())).isEqualTo(1);
}
 

@Test
public void multipleFilters() {
	TestFilter filter1 = new TestFilter();
	TestFilter filter2 = new TestFilter();
	TestFilter filter3 = new TestFilter();

	RSocket socket = new GatewaySocketAcceptor(this.factory,
			Arrays.asList(filter1, filter2, filter3), this.meterRegistry,
			this.properties, this.metadataExtractor)
					.accept(this.setupPayload, this.sendingSocket)
					.block(Duration.ZERO);

	assertThat(filter1.invoked()).isTrue();
	assertThat(filter2.invoked()).isTrue();
	assertThat(filter3.invoked()).isTrue();
	assertThat(socket).isNotNull();
}
 

BackgroundWorker(
    Flux<Task> taskProducer,
    ConcurrentMap<String, BlockingQueue<Task>> idToCompletedTasksMap,
    ConcurrentMap<String, RSocket> idToRSocketMap) {

  this.idToCompletedTasksMap = idToCompletedTasksMap;
  this.idToRSocketMap = idToRSocketMap;

  // mimic a long running task processing
  taskProducer
      .concatMap(
          t ->
              Mono.delay(Duration.ofMillis(ThreadLocalRandom.current().nextInt(200, 2000)))
                  .thenReturn(t))
      .subscribe(this);
}
 

@Test
public void shortcircuitFilter() {

	TestFilter filter1 = new TestFilter();
	ShortcircuitingFilter filter2 = new ShortcircuitingFilter();
	TestFilter filter3 = new TestFilter();

	RSocket socket = new GatewaySocketAcceptor(this.factory,
			Arrays.asList(filter1, filter2, filter3), this.meterRegistry,
			this.properties, this.metadataExtractor)
					.accept(this.setupPayload, this.sendingSocket)
					.block(Duration.ZERO);

	assertThat(filter1.invoked()).isTrue();
	assertThat(filter2.invoked()).isTrue();
	assertThat(filter3.invoked()).isFalse();
	assertThat(socket).isNull();
}
 

@Test(timeout = 2000)
public void testHandlerEmitsCustomError() {
  rule.setRequestAcceptor(
      new RSocket() {
        @Override
        public Mono<Payload> requestResponse(Payload payload) {
          return Mono.error(
              new CustomRSocketException(0x00000501, "Deliberate Custom exception."));
        }
      });
  rule.crs
      .requestResponse(EmptyPayload.INSTANCE)
      .as(StepVerifier::create)
      .expectErrorSatisfies(
          t ->
              Assertions.assertThat(t)
                  .isInstanceOf(CustomRSocketException.class)
                  .hasMessage("Deliberate Custom exception.")
                  .hasFieldOrPropertyWithValue("errorCode", 0x00000501))
      .verify();
}
 

@Override
public Mono<Object> resolveArgument(MethodParameter parameter, Message<?> message) {
	Object headerValue = message.getHeaders().get(RSOCKET_REQUESTER_HEADER);
	Assert.notNull(headerValue, "Missing '" + RSOCKET_REQUESTER_HEADER + "'");

	Assert.isInstanceOf(RSocketRequester.class, headerValue, "Expected header value of type RSocketRequester");
	RSocketRequester requester = (RSocketRequester) headerValue;

	Class<?> type = parameter.getParameterType();
	if (RSocketRequester.class.equals(type)) {
		return Mono.just(requester);
	}
	else if (RSocket.class.isAssignableFrom(type)) {
		return Mono.just(requester.rsocket());
	}
	else {
		return Mono.error(new IllegalArgumentException("Unexpected parameter type: " + parameter));
	}
}
 

@Before
public void init() {
	this.factory = mock(GatewayRSocketFactory.class);
	this.setupPayload = mock(ConnectionSetupPayload.class);
	this.sendingSocket = mock(RSocket.class);
	this.meterRegistry = new SimpleMeterRegistry();

	this.metadataExtractor.metadataToExtract(ROUTE_SETUP_MIME_TYPE, RouteSetup.class,
			"routesetup");

	when(this.factory.create(any(TagsMetadata.class)))
			.thenReturn(mock(GatewayRSocket.class));

	when(this.setupPayload.metadataMimeType())
			.thenReturn(Metadata.COMPOSITE_MIME_TYPE.toString());

	when(this.setupPayload.hasMetadata()).thenReturn(true);

	MetadataEncoder encoder = new MetadataEncoder(Metadata.COMPOSITE_MIME_TYPE,
			this.rSocketStrategies);
	encoder.metadata(RouteSetup.of(1L, "myservice").build(), ROUTE_SETUP_MIME_TYPE);
	DataBuffer dataBuffer = encoder.encode();
	DataBuffer data = MetadataEncoder.emptyDataBuffer(rSocketStrategies);
	Payload payload = PayloadUtils.createPayload(data, dataBuffer);
	when(setupPayload.metadata()).thenReturn(payload.metadata());
}
 

@Test
void serverMissingResume() {
  CloseableChannel closeableChannel =
      RSocketServer.create(SocketAcceptor.with(new TestResponderRSocket()))
          .bind(serverTransport(SERVER_HOST, SERVER_PORT))
          .block();

  RSocket rSocket =
      RSocketConnector.create()
          .resume(new Resume())
          .connect(clientTransport(closeableChannel.address()))
          .block();

  StepVerifier.create(rSocket.onClose().doFinally(s -> closeableChannel.dispose()))
      .expectErrorMatches(
          err ->
              err instanceof UnsupportedSetupException
                  && "resume not supported".equals(err.getMessage()))
      .verify(Duration.ofSeconds(5));

  Assertions.assertThat(rSocket.isDisposed()).isTrue();
}
 

PendingRequestRSocket(MetadataExtractor metadataExtractor,
		Function<RegisteredEvent, Mono<Route>> routeFinder,
		Consumer<TagsMetadata> metadataCallback,
		MonoProcessor<RSocket> rSocketProcessor) {
	this.routeFinder = routeFinder;
	this.metadataExtractor = metadataExtractor;
	this.metadataCallback = metadataCallback;
	this.rSocketProcessor = rSocketProcessor;
}
 

@Test
// see https://github.com/rsocket/rsocket-java/issues/858
public void testWorkaround858() {
  ByteBuf buffer = rule.alloc().buffer();
  buffer.writeCharSequence("test", CharsetUtil.UTF_8);

  TestPublisher<Payload> testPublisher = TestPublisher.create();

  rule.setAcceptingSocket(
      new RSocket() {
        @Override
        public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
          Flux.from(payloads).doOnNext(ReferenceCounted::release).subscribe();

          return testPublisher.flux();
        }
      });

  rule.connection.addToReceivedBuffer(
      RequestChannelFrameCodec.encodeReleasingPayload(
          rule.alloc(), 1, false, 1, ByteBufPayload.create(buffer)));
  rule.connection.addToReceivedBuffer(
      ErrorFrameCodec.encode(rule.alloc(), 1, new RuntimeException("test")));

  Assertions.assertThat(rule.connection.getSent())
      .hasSize(1)
      .first()
      .matches(bb -> FrameHeaderCodec.frameType(bb) == REQUEST_N)
      .matches(ReferenceCounted::release);

  Assertions.assertThat(rule.socket.isDisposed()).isFalse();
  testPublisher.assertWasCancelled();

  rule.assertHasNoLeaks();
}
 

void rejectSetupTcp(
    Function<InetSocketAddress, ServerTransport<CloseableChannel>> serverTransport,
    Function<InetSocketAddress, ClientTransport> clientTransport) {

  String errorMessage = "error";
  RejectingAcceptor acceptor = new RejectingAcceptor(errorMessage);
  Mono<RSocket> serverRequester = acceptor.requesterRSocket();

  CloseableChannel channel =
      RSocketServer.create(acceptor)
          .bind(serverTransport.apply(new InetSocketAddress("localhost", 0)))
          .block(Duration.ofSeconds(5));

  ErrorConsumer errorConsumer = new ErrorConsumer();

  RSocket clientRequester =
      RSocketConnector.connectWith(clientTransport.apply(channel.address()))
          .doOnError(errorConsumer)
          .block(Duration.ofSeconds(5));

  StepVerifier.create(errorConsumer.errors().next())
      .expectNextMatches(
          err -> err instanceof RejectedSetupException && errorMessage.equals(err.getMessage()))
      .expectComplete()
      .verify(Duration.ofSeconds(5));

  StepVerifier.create(clientRequester.onClose()).expectComplete().verify(Duration.ofSeconds(5));

  StepVerifier.create(serverRequester.flatMap(socket -> socket.onClose()))
      .expectComplete()
      .verify(Duration.ofSeconds(5));

  StepVerifier.create(clientRequester.requestResponse(DefaultPayload.create("test")))
      .expectErrorMatches(
          err -> err instanceof RejectedSetupException && errorMessage.equals(err.getMessage()))
      .verify(Duration.ofSeconds(5));

  channel.dispose();
}
 

@Override
public Flux<Payload> requestChannel(Payload payload, Publisher<Payload> payloads) {
	return processor("pending-request-rc", payload).flatMapMany(tuple -> {
		RSocket rSocket = tuple.getT1();
		if (rSocket instanceof ResponderRSocket) {
			ResponderRSocket socket = (ResponderRSocket) rSocket;
			return socket.requestChannel(payload, payloads);
		}
		return rSocket.requestChannel(payloads);
	});
}
 

@Test
void rSocketDisposedOnMissingKeepAlives() {
  RSocket rSocket = requesterState.rSocket();

  Mono.delay(Duration.ofMillis(2000)).block();

  Assertions.assertThat(rSocket.isDisposed()).isTrue();
  rSocket
      .onClose()
      .as(StepVerifier::create)
      .expectError(ConnectionErrorException.class)
      .verify(Duration.ofMillis(100));
}
 

@Override
protected void init() {
  acceptingSocket =
      new RSocket() {
        @Override
        public Mono<Payload> requestResponse(Payload payload) {
          return Mono.just(payload);
        }
      };
  super.init();
}
 

private Mono<List<RSocket>> findRSocketOrCreatePending(GatewayExchange exchange,
		Route route) {
	return executeFilterChain(route.getFilters(), exchange)
			.log(GatewayRSocket.class.getName() + ".after filter chain", Level.FINEST)
			.flatMapMany(success -> {
				Map<Key, String> tags = exchange.getRoutingMetadata().getTags();
				// TODO: use frame flag
				if (tags.containsKey(new Key("multicast"))) {
					List<Tuple2<String, RSocket>> rsockets = loadBalancerFactory
							.find(exchange.getRoutingMetadata());
					return Flux.fromIterable(rsockets);
				}
				return loadBalancerFactory.choose(exchange.getRoutingMetadata())
						.flatMapMany(
								tuple -> Flux.just((Tuple2<String, RSocket>) tuple));
			}).map(tuple -> {
				// TODO: this is routeId, should it be service name?
				// Tags tags = exchange.getTags().and("responder.id", tuple.getT1());
				// exchange.setTags(tags);
				return tuple.getT2();
			}).cast(RSocket.class).map(rSocket -> {
				if (log.isDebugEnabled()) {
					log.debug("Found RSocket: " + rSocket);
				}
				return rSocket;
			}).collectList()
			.log(GatewayRSocket.class.getName() + ".find rsocket", Level.FINEST);
}
 

@Test(timeout = 15_000L)
public void testCompleteWithoutNext() {
  handler =
      new RSocket() {
        @Override
        public Flux<Payload> requestStream(Payload payload) {
          return Flux.empty();
        }
      };
  RSocket client = buildClient();
  Boolean hasElements =
      client.requestStream(DefaultPayload.create("REQUEST", "META")).log().hasElements().block();

  assertFalse(hasElements);
}
 

@ParameterizedTest
@MethodSource("interactions")
void responderDepletedAllowedLeaseRequestsAreRejected(
    BiFunction<RSocket, Payload, Publisher<?>> interaction) {
  leaseSender.onNext(Lease.create(5_000, 1));

  ByteBuf buffer = byteBufAllocator.buffer();
  buffer.writeCharSequence("test", CharsetUtil.UTF_8);
  Payload payload1 = ByteBufPayload.create(buffer);

  Flux<?> responder = Flux.from(interaction.apply(rSocketResponder, payload1));
  responder.subscribe();

  Assertions.assertThat(connection.getSent())
      .hasSize(1)
      .first()
      .matches(bb -> FrameHeaderCodec.frameType(bb) == LEASE)
      .matches(ReferenceCounted::release);

  ByteBuf buffer2 = byteBufAllocator.buffer();
  buffer2.writeCharSequence("test", CharsetUtil.UTF_8);
  Payload payload2 = ByteBufPayload.create(buffer2);

  Flux.from(interaction.apply(rSocketResponder, payload2))
      .as(StepVerifier::create)
      .expectError(MissingLeaseException.class)
      .verify(Duration.ofSeconds(5));
}
 

public void onRSocketClosed(RSocket rsocket, @Nullable Throwable cause) {
    for (Map.Entry<String, RSocket> entry : activeSockets.entrySet()) {
        if (entry.getValue() == rsocket) {
            onRSocketClosed(entry.getKey(), entry.getValue(), null);
        }
    }
    if (!rsocket.isDisposed()) {
        try {
            rsocket.dispose();
        } catch (Exception ignore) {

        }
    }
}
 

@Test
public void ensuresThatMonoFromRSocketConnectorCanBeUsedForMultipleSubscriptions() {
  Payload setupPayload = ByteBufPayload.create("TestData", "TestMetadata");

  Assertions.assertThat(setupPayload.refCnt()).isOne();

  TestClientTransport testClientTransport = new TestClientTransport();
  Mono<RSocket> connectionMono =
      RSocketConnector.create().setupPayload(setupPayload).connect(testClientTransport);

  connectionMono
      .as(StepVerifier::create)
      .expectNextCount(1)
      .expectComplete()
      .verify(Duration.ofMillis(100));

  connectionMono
      .as(StepVerifier::create)
      .expectNextCount(1)
      .expectComplete()
      .verify(Duration.ofMillis(100));

  Assertions.assertThat(testClientTransport.testConnection().getSent())
      .hasSize(2)
      .allMatch(
          bb -> {
            DefaultConnectionSetupPayload payload = new DefaultConnectionSetupPayload(bb);
            return payload.getDataUtf8().equals("TestData")
                && payload.getMetadataUtf8().equals("TestMetadata");
          })
      .allMatch(ReferenceCounted::release);
  Assertions.assertThat(setupPayload.refCnt()).isZero();
}
 

@Bean
public RSocket rSocket() {
    return RSocketFactory.connect()
                         .mimeType(MimeTypeUtils.APPLICATION_JSON_VALUE, MimeTypeUtils.APPLICATION_JSON_VALUE)
                         .frameDecoder(PayloadDecoder.ZERO_COPY)
                         .transport(TcpClientTransport.create(7000))
                         .start()
                         .block();
}
 

@Override
public @NotNull Flux<Payload> requestChannel(@NotNull Publisher<Payload> payloads) {
    RSocket next = randomSelector.next();
    if (next == null) {
        return Flux.error(new NoAvailableConnectionException(RsocketErrorCode.message("RST-200404", serviceId)));
    }
    return next.requestChannel(payloads)
            .onErrorResume(CONNECTION_ERROR_PREDICATE, error -> {
                onRSocketClosed(next, error);
                return requestChannel(payloads);
            });
}
 

@Override
public @NotNull Flux<Payload> requestStream(@NotNull Payload payload) {
    RSocket next = randomSelector.next();
    if (next == null) {
        ReferenceCountUtil.safeRelease(payload);
        return Flux.error(new NoAvailableConnectionException(RsocketErrorCode.message("RST-200404", serviceId)));
    }
    return next.requestStream(payload)
            .onErrorResume(CONNECTION_ERROR_PREDICATE, error -> {
                onRSocketClosed(next, error);
                return requestStream(payload);
            });
}
 

@Override
public @NotNull Mono<Void> fireAndForget(@NotNull Payload payload) {
    RSocket next = randomSelector.next();
    if (next == null) {
        ReferenceCountUtil.safeRelease(payload);
        return Mono.error(new NoAvailableConnectionException(RsocketErrorCode.message("RST-200404", serviceId)));
    }
    return next.fireAndForget(payload)
            .onErrorResume(CONNECTION_ERROR_PREDICATE, error -> {
                onRSocketClosed(next, error);
                return fireAndForget(payload);
            });
}
 

@Override
public @NotNull Mono<Payload> requestResponse(@NotNull Payload payload) {
    RSocket next = randomSelector.next();
    if (next == null) {
        ReferenceCountUtil.safeRelease(payload);
        return Mono.error(new NoAvailableConnectionException(RsocketErrorCode.message("RST-200404", serviceId)));
    }
    return next.requestResponse(payload)
            .onErrorResume(CONNECTION_ERROR_PREDICATE, error -> {
                onRSocketClosed(next, error);
                return requestResponse(payload);
            });
}
 

private <X, Y> Flux<Y> doRequestStream(
    final String service,
    final String route,
    final RSocket r,
    final Marshaller<X> marshaller,
    final Unmarshaller<Y> unmarshaller,
    final X o,
    final ByteBuf metadata,
    Function<? super Publisher<Y>, ? extends Publisher<Y>> metrics,
    Function<Map<String, String>, Function<? super Publisher<Y>, ? extends Publisher<Y>>>
        tracing) {
  final HashMap<String, String> map = new HashMap<>();
  return Flux.defer(
          () -> {
            try {
              ByteBuf d = marshaller.apply(o);
              ByteBuf m =
                  metadataEncoder.encode(metadata, new SimpleSpanContext(map), service, route);
              metadata.release();
              Payload payload = ByteBufPayload.create(d, m);
              return r.requestStream(payload);
            } catch (Throwable t) {
              metadata.release();
              return Flux.error(t);
            }
          })
      .map(
          p -> {
            try {
              return unmarshaller.apply(p.sliceData());
            } finally {
              p.release();
            }
          })
      .transform(metrics)
      .transform(tracing.apply(map));
}