Examples with PushConnectionsHandle io.confluent.ksql.physical.scalablepush.PushRouting.PushConnectionsHandle used on opensource projects

Example 11 with PushConnectionsHandle

use of io.confluent.ksql.physical.scalablepush.PushRouting.PushConnectionsHandle in project ksql by confluentinc.

the class PushRoutingTest method shouldSucceed_removeRemoteNode.

@Test
public void shouldSucceed_removeRemoteNode() throws ExecutionException, InterruptedException {
    // Given:
    final AtomicReference<Set<KsqlNode>> nodes = new AtomicReference<>(ImmutableSet.of(ksqlNodeLocal, ksqlNodeRemote));
    final PushRouting routing = new PushRouting(sqr -> nodes.get(), 50, true);
    // When:
    final PushConnectionsHandle handle = handlePushRouting(routing);
    context.runOnContext(v -> {
        localPublisher.accept(LOCAL_ROW1);
        localPublisher.accept(LOCAL_ROW2);
        remotePublisher.accept(REMOTE_ROW1);
        remotePublisher.accept(REMOTE_ROW2);
    });
    Set<List<?>> rows = waitOnRows(4);
    final RoutingResult result = handle.get(ksqlNodeRemote).get();
    nodes.set(ImmutableSet.of(ksqlNodeLocal));
    while (handle.get(ksqlNodeRemote).isPresent()) {
        Thread.sleep(100);
        continue;
    }
    handle.close();
    // Then:
    assertThat(rows.contains(LOCAL_ROW1.value().values()), is(true));
    assertThat(rows.contains(LOCAL_ROW2.value().values()), is(true));
    assertThat(rows.contains(REMOTE_ROW1.getRow().get().getColumns()), is(true));
    assertThat(rows.contains(REMOTE_ROW2.getRow().get().getColumns()), is(true));
    assertThat(result.getStatus(), is(RoutingResultStatus.REMOVED));
}

Example 12 with PushConnectionsHandle

use of io.confluent.ksql.physical.scalablepush.PushRouting.PushConnectionsHandle in project ksql by confluentinc.

the class PushRoutingTest method shouldSucceed_gapDetectedLocal_retry.

@Test
public void shouldSucceed_gapDetectedLocal_retry() throws ExecutionException, InterruptedException {
    // Given:
    final AtomicReference<Set<KsqlNode>> nodes = new AtomicReference<>(ImmutableSet.of(ksqlNodeLocal, ksqlNodeRemote));
    final PushRouting routing = new PushRouting(sqr -> nodes.get(), 50, true);
    AtomicReference<TestLocalPublisher> localPublisher = new AtomicReference<>();
    AtomicInteger localCount = new AtomicInteger(0);
    when(pushPhysicalPlanManager.execute()).thenAnswer(a -> {
        localPublisher.set(new TestLocalPublisher(context));
        localCount.incrementAndGet();
        if (localCount.get() == 2) {
            localPublisher.get().accept(LOCAL_ROW2);
        }
        return localPublisher.get();
    });
    doAnswer(a -> {
        final Optional<PushOffsetRange> newOffsetRange = a.getArgument(0);
        assertThat(newOffsetRange.isPresent(), is(true));
        assertThat(newOffsetRange.get().getEndOffsets(), is(ImmutableList.of(0L, 3L)));
        return null;
    }).when(pushPhysicalPlanManager).reset(any());
    // When:
    final PushConnectionsHandle handle = handlePushRouting(routing);
    context.runOnContext(v -> {
        localPublisher.get().accept(LOCAL_CONTINUATION_TOKEN1);
        localPublisher.get().accept(LOCAL_ROW1);
        localPublisher.get().accept(LOCAL_CONTINUATION_TOKEN_GAP);
    });
    Set<List<?>> rows = waitOnRows(2);
    handle.close();
    // Then:
    verify(pushPhysicalPlanManager, times(2)).execute();
    assertThat(rows.contains(LOCAL_ROW1.value().values()), is(true));
    assertThat(rows.contains(LOCAL_ROW2.value().values()), is(true));
}

Example 13 with PushConnectionsHandle

use of io.confluent.ksql.physical.scalablepush.PushRouting.PushConnectionsHandle in project ksql by confluentinc.

the class PushRoutingTest method shouldSucceed_remoteNodeExceptionWithRetry.

@Test
public void shouldSucceed_remoteNodeExceptionWithRetry() throws ExecutionException, InterruptedException {
    // Given:
    final AtomicReference<Set<KsqlNode>> nodes = new AtomicReference<>(ImmutableSet.of(ksqlNodeLocal, ksqlNodeRemote));
    final PushRouting routing = new PushRouting(sqr -> nodes.get(), 50, true);
    AtomicReference<TestRemotePublisher> remotePublisher = new AtomicReference<>();
    AtomicInteger remoteCount = new AtomicInteger(0);
    when(simpleKsqlClient.makeQueryRequestStreamed(any(), any(), any(), any())).thenAnswer(a -> {
        remotePublisher.set(new TestRemotePublisher(context));
        if (remoteCount.incrementAndGet() == 2) {
            remotePublisher.get().accept(REMOTE_ROW1);
            remotePublisher.get().accept(REMOTE_ROW2);
        }
        return createFuture(RestResponse.successful(200, remotePublisher.get()));
    });
    // When:
    final PushConnectionsHandle handle = handlePushRouting(routing);
    final AtomicReference<Throwable> exception = new AtomicReference<>(null);
    handle.onException(exception::set);
    context.runOnContext(v -> {
        localPublisher.accept(LOCAL_ROW1);
        localPublisher.accept(LOCAL_ROW2);
        remotePublisher.get().error(new RuntimeException("Random error"));
    });
    Set<List<?>> rows = waitOnRows(4);
    handle.close();
    // Then:
    assertThat(rows.contains(LOCAL_ROW1.value().values()), is(true));
    assertThat(rows.contains(LOCAL_ROW2.value().values()), is(true));
    assertThat(rows.contains(REMOTE_ROW1.getRow().get().getColumns()), is(true));
    assertThat(rows.contains(REMOTE_ROW2.getRow().get().getColumns()), is(true));
}

Example 14 with PushConnectionsHandle

use of io.confluent.ksql.physical.scalablepush.PushRouting.PushConnectionsHandle in project ksql by confluentinc.

the class PushRoutingTest method shouldFail_errorRemoteCall.

@Test
public void shouldFail_errorRemoteCall() throws ExecutionException, InterruptedException {
    // Given:
    when(locator.locate()).thenReturn(ImmutableList.of(ksqlNodeRemote, ksqlNodeRemote2));
    final PushRouting routing = new PushRouting();
    TestRemotePublisher remotePublisher = new TestRemotePublisher(context);
    when(simpleKsqlClient.makeQueryRequestStreamed(eq(ksqlNodeRemote.location()), any(), any(), any())).thenReturn(createErrorFuture(new RuntimeException("Error remote!")));
    when(simpleKsqlClient.makeQueryRequestStreamed(eq(ksqlNodeRemote2.location()), any(), any(), any())).thenReturn(createFuture(RestResponse.successful(200, remotePublisher)));
    // When:
    final PushConnectionsHandle handle = handlePushRouting(routing);
    // Then:
    assertThat(handle.getError().getMessage(), containsString("Error remote!"));
    assertThat(remotePublisher.isClosed(), is(true));
}

Example 15 with PushConnectionsHandle

use of io.confluent.ksql.physical.scalablepush.PushRouting.PushConnectionsHandle in project ksql by confluentinc.

the class PushRoutingTest method shouldSucceed_gapDetectedRemote_disableAlos.

@Test
public void shouldSucceed_gapDetectedRemote_disableAlos() throws ExecutionException, InterruptedException {
    // Given:
    when(pushRoutingOptions.alosEnabled()).thenReturn(false);
    final AtomicReference<Set<KsqlNode>> nodes = new AtomicReference<>(ImmutableSet.of(ksqlNodeLocal, ksqlNodeRemote));
    final PushRouting routing = new PushRouting(sqr -> nodes.get(), 50, true);
    AtomicReference<TestRemotePublisher> remotePublisher = new AtomicReference<>();
    AtomicInteger remoteCount = new AtomicInteger(0);
    when(simpleKsqlClient.makeQueryRequestStreamed(any(), any(), any(), any())).thenAnswer(a -> {
        remotePublisher.set(new TestRemotePublisher(context));
        remoteCount.incrementAndGet();
        final Map<String, ?> requestProperties = a.getArgument(3);
        String continuationToken = (String) requestProperties.get(KsqlRequestConfig.KSQL_REQUEST_QUERY_PUSH_CONTINUATION_TOKEN);
        if (remoteCount.get() == 1) {
            assertThat(continuationToken, nullValue());
            context.runOnContext(v -> {
                remotePublisher.get().accept(REMOTE_ROW2);
            });
        } else if (remoteCount.get() == 2) {
            fail();
        }
        return createFuture(RestResponse.successful(200, remotePublisher.get()));
    });
    // When:
    final PushConnectionsHandle handle = handlePushRouting(routing);
    final AtomicReference<Throwable> exception = new AtomicReference<>(null);
    handle.onException(exception::set);
    context.runOnContext(v -> {
        remotePublisher.get().accept(REMOTE_CONTINUATION_TOKEN1);
        remotePublisher.get().accept(REMOTE_ROW1);
        remotePublisher.get().accept(REMOTE_CONTINUATION_TOKEN_GAP);
    });
    Set<List<?>> rows = waitOnRows(2);
    handle.close();
    // Then:
    verify(simpleKsqlClient, times(1)).makeQueryRequestStreamed(any(), any(), any(), any());
    assertThat(rows.contains(REMOTE_ROW1.getRow().get().getColumns()), is(true));
    assertThat(rows.contains(REMOTE_ROW2.getRow().get().getColumns()), is(true));
}