Examples with PushOffsetRange io.confluent.ksql.util.PushOffsetRange used on opensource projects

Example 6 with PushOffsetRange

use of io.confluent.ksql.util.PushOffsetRange 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 7 with PushOffsetRange

use of io.confluent.ksql.util.PushOffsetRange in project ksql by confluentinc.

the class ScalablePushConsumer method handleProgressToken.

private void handleProgressToken(final PushOffsetVector startOffsetVector, final PushOffsetVector endOffsetVector) {
    final PushOffsetRange range = new PushOffsetRange(Optional.of(startOffsetVector), endOffsetVector);
    for (ProcessingQueue queue : processingQueues.values()) {
        final QueryRow row = OffsetsRow.of(clock.millis(), range);
        queue.offer(row);
    }
}

Example 8 with PushOffsetRange

use of io.confluent.ksql.util.PushOffsetRange in project ksql by confluentinc.

the class ScalablePushRegistryTest method setUp.

@Before
public void setUp() {
    when(ksqlTopic.getKafkaTopicName()).thenReturn(TOPIC);
    when(kafkaConsumerFactory.create(any(), any(), any(), any(), any(), any())).thenReturn(kafkaConsumer);
    catchupCoordinator = new TestCatchupCoordinator();
    latestConsumer = new TestLatestConsumer(TOPIC, false, SCHEMA, kafkaConsumer, catchupCoordinator, assignment -> {
    }, ksqlConfig, Clock.systemUTC());
    latestConsumer2 = new TestLatestConsumer(TOPIC, false, SCHEMA, kafkaConsumer, catchupCoordinator, assignment -> {
    }, ksqlConfig, Clock.systemUTC());
    catchupConsumer = new TestCatchupConsumer(TOPIC, false, SCHEMA, kafkaConsumer, () -> latestConsumer, catchupCoordinator, pushOffsetRange, Clock.systemUTC(), pq -> {
    });
    when(latestConsumerFactory.create(any(), anyBoolean(), any(), any(), any(), any(), any(), any())).thenReturn(latestConsumer, latestConsumer2);
    when(catchupConsumerFactory.create(any(), anyBoolean(), any(), any(), any(), any(), any(), any(), anyLong(), any())).thenReturn(catchupConsumer);
    when(ksqlTopic.getKeyFormat()).thenReturn(keyFormat);
    when(keyFormat.isWindowed()).thenReturn(false);
    realExecutorService = Executors.newFixedThreadPool(2);
    doAnswer(a -> {
        final Runnable runnable = a.getArgument(0);
        startLatestRunnable.set(runnable);
        realExecutorService.submit(runnable);
        return null;
    }).when(executorService).submit(any(Runnable.class));
    doAnswer(a -> {
        final Runnable runnable = a.getArgument(0);
        realExecutorService.submit(runnable);
        return null;
    }).when(catchupService).submit(any(Runnable.class));
    when(processingQueue.getQueryId()).thenReturn(new QueryId("q1"));
    when(processingQueue2.getQueryId()).thenReturn(new QueryId("q2"));
    registry = new ScalablePushRegistry(locator, SCHEMA, false, ImmutableMap.of(), ksqlTopic, serviceContext, ksqlConfig, SOURCE_APP_ID, kafkaConsumerFactory, latestConsumerFactory, catchupConsumerFactory, executorService, catchupService);
    when(ksqlConfig.getInt(KsqlConfig.KSQL_QUERY_PUSH_V2_MAX_CATCHUP_CONSUMERS)).thenReturn(10);
}

Example 9 with PushOffsetRange

use of io.confluent.ksql.util.PushOffsetRange in project ksql by confluentinc.

the class CatchupConsumerTest method shouldRunConsumer_timeoutWaitingForAssignment.

@Test
public void shouldRunConsumer_timeoutWaitingForAssignment() {
    // Given:
    PushOffsetRange offsetRange = new PushOffsetRange(Optional.empty(), new PushOffsetVector(ImmutableList.of(1L, 2L)));
    when(latestConsumer.getAssignment()).thenReturn(null);
    when(clock.millis()).thenReturn(CURRENT_TIME_MS, CURRENT_TIME_MS + WAIT_FOR_ASSIGNMENT_MS - 1, CURRENT_TIME_MS + WAIT_FOR_ASSIGNMENT_MS + 1);
    try (CatchupConsumer consumer = new CatchupConsumer(TOPIC, false, SCHEMA, kafkaConsumer, () -> latestConsumer, catchupCoordinator, offsetRange, clock, sleepFn, waitFn, 0, pq -> caughtUp = true)) {
        // When:
        consumer.register(queue);
        final Exception e = assertThrows(KsqlException.class, consumer::run);
        // Then:
        assertThat(e.getMessage(), containsString("Timed out waiting for assignment from Latest"));
    }
}

Example 10 with PushOffsetRange

use of io.confluent.ksql.util.PushOffsetRange 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));
}