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

Example 6 with PushConnectionsHandle

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

the class PushRoutingTest method shouldSucceed_gapDetectedRemote_retry.

@Test
public void shouldSucceed_gapDetectedRemote_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<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());
        } else if (remoteCount.get() == 2) {
            assertThat(continuationToken, notNullValue());
            final PushOffsetRange range = PushOffsetRange.deserialize(continuationToken);
            assertThat(range.getEndOffsets().getDenseRepresentation(), is(ImmutableList.of(0L, 3L)));
            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 -> {
        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(2)).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));
}

Example 7 with PushConnectionsHandle

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

the class PushRoutingTest method shouldFail_hitRequestLimitRemote.

@Test
public void shouldFail_hitRequestLimitRemote() throws ExecutionException, InterruptedException {
    // Given:
    when(locator.locate()).thenReturn(ImmutableList.of(ksqlNodeRemote));
    transientQueryQueue = new TransientQueryQueue(OptionalInt.empty(), 1, 100);
    final PushRouting routing = new PushRouting();
    // When:
    final PushConnectionsHandle handle = handlePushRouting(routing);
    context.runOnContext(v -> {
        remotePublisher.accept(REMOTE_ROW1);
        remotePublisher.accept(REMOTE_ROW2);
    });
    // Then:
    Set<List<?>> rows = waitOnRows(1);
    handle.close();
    assertThat(rows.contains(REMOTE_ROW1.getRow().get().getColumns()), is(true));
    assertThat(handle.getError().getMessage(), containsString("Hit limit of request queue"));
}

Example 8 with PushConnectionsHandle

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

the class PushRoutingTest method shouldSucceed_gapDetectedLocal_disableAlos.

@Test
public void shouldSucceed_gapDetectedLocal_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<TestLocalPublisher> localPublisher = new AtomicReference<>();
    AtomicInteger localCount = new AtomicInteger(0);
    when(pushPhysicalPlanManager.execute()).thenAnswer(a -> {
        localPublisher.set(new TestLocalPublisher(context));
        localCount.incrementAndGet();
        context.runOnContext(v -> {
            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(1)).execute();
    assertThat(rows.contains(LOCAL_ROW1.value().values()), is(true));
    assertThat(rows.contains(LOCAL_ROW2.value().values()), is(true));
}

Example 9 with PushConnectionsHandle

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

the class PushRoutingTest method shouldSucceed_gapDetectedRemote_noRetry.

@Test
public void shouldSucceed_gapDetectedRemote_noRetry() throws ExecutionException, InterruptedException {
    // Given:
    final AtomicReference<Set<KsqlNode>> nodes = new AtomicReference<>(ImmutableSet.of(ksqlNodeLocal, ksqlNodeRemote));
    final PushRouting routing = new PushRouting(sqr -> nodes.get(), 50, false);
    // When:
    final PushConnectionsHandle handle = handlePushRouting(routing);
    context.runOnContext(v -> {
        remotePublisher.accept(REMOTE_ROW1);
        remotePublisher.accept(REMOTE_CONTINUATION_TOKEN1);
        remotePublisher.accept(REMOTE_ROW2);
        remotePublisher.accept(REMOTE_CONTINUATION_TOKEN_GAP);
    });
    Set<List<?>> rows = waitOnRows(2);
    waitOnNodeStatus(handle, ksqlNodeRemote, RoutingResultStatus.OFFSET_GAP_FOUND);
    handle.close();
    // Then:
    assertThat(rows.contains(REMOTE_ROW1.getRow().get().getColumns()), is(true));
    assertThat(rows.contains(REMOTE_ROW2.getRow().get().getColumns()), is(true));
    assertThat(handle.get(ksqlNodeRemote).get().getStatus(), is(RoutingResultStatus.OFFSET_GAP_FOUND));
}

Example 10 with PushConnectionsHandle

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

the class PushRoutingTest method shouldSucceed_forward.

@Test
public void shouldSucceed_forward() throws ExecutionException, InterruptedException {
    // Given:
    final PushRouting routing = new PushRouting();
    // 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);
    });
    // Then:
    Set<List<?>> rows = waitOnRows(4);
    handle.close();
    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));
}