Examples with MockSegmentStreamFactory io.pravega.client.stream.mock.MockSegmentStreamFactory used on opensource projects

Example 51 with MockSegmentStreamFactory

use of io.pravega.client.stream.mock.MockSegmentStreamFactory in project pravega by pravega.

the class ReaderGroupStateManagerTest method testRemoveReader.

@Test(timeout = 10000)
public void testRemoveReader() throws ReaderNotInReaderGroupException {
    String scope = "scope";
    String stream = "stream";
    PravegaNodeUri endpoint = new PravegaNodeUri("localhost", SERVICE_PORT);
    MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
    MockController controller = new MockController(endpoint.getEndpoint(), endpoint.getPort(), connectionFactory, false);
    createScopeAndStream(scope, stream, controller);
    MockSegmentStreamFactory streamFactory = new MockSegmentStreamFactory();
    @Cleanup SynchronizerClientFactory clientFactory = new ClientFactoryImpl(scope, controller, connectionFactory, streamFactory, streamFactory, streamFactory, streamFactory);
    SynchronizerConfig config = SynchronizerConfig.builder().build();
    @Cleanup StateSynchronizer<ReaderGroupState> stateSynchronizer = createState(stream, clientFactory, config);
    AtomicLong clock = new AtomicLong();
    Map<SegmentWithRange, Long> segments = new HashMap<>();
    SegmentWithRange segment0 = new SegmentWithRange(new Segment(scope, stream, 0), 0.0, 0.5);
    SegmentWithRange segment1 = new SegmentWithRange(new Segment(scope, stream, 1), 0.5, 1.0);
    segments.put(segment0, 123L);
    segments.put(segment1, 456L);
    stateSynchronizer.initialize(new ReaderGroupState.ReaderGroupStateInit(ReaderGroupConfig.builder().stream(Stream.of(scope, stream)).build(), segments, Collections.emptyMap(), false));
    ReaderGroupStateManager readerState1 = new ReaderGroupStateManager(scope, stream, "testReader", stateSynchronizer, controller, clock::get);
    readerState1.initializeReader(0);
    Segment toRelease = readerState1.findSegmentToReleaseIfRequired();
    assertNull(toRelease);
    Map<SegmentWithRange, Long> newSegments = readerState1.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap()));
    assertFalse(newSegments.isEmpty());
    assertEquals(2, newSegments.size());
    ReaderGroupStateManager readerState2 = new ReaderGroupStateManager(scope, stream, "testReader2", stateSynchronizer, controller, clock::get);
    readerState2.initializeReader(0);
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    assertNotNull(readerState1.findSegmentToReleaseIfRequired());
    boolean released = readerState1.releaseSegment(new Segment(scope, stream, 0), 789L, 0L, new PositionImpl(segments));
    assertTrue(released);
    newSegments = readerState2.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap()));
    assertEquals(1, newSegments.size());
    assertEquals(Long.valueOf(789L), newSegments.get(segment0));
    StateSynchronizer<ReaderGroupState> spied = spy(stateSynchronizer);
    ReaderGroupStateManager.readerShutdown("testReader2", null, spied);
    // verify that fetch updates is called once on the spied state synchronizer.
    verify(spied, times(1)).fetchUpdates();
    AssertExtensions.assertThrows(ReaderNotInReaderGroupException.class, () -> readerState2.releaseSegment(new Segment(scope, stream, 0), 711L, 0L, new PositionImpl(Collections.emptyMap())));
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    newSegments = readerState1.acquireNewSegmentsIfNeeded(0, new PositionImpl(segments));
    assertEquals(1, newSegments.size());
    assertEquals(Long.valueOf(789L), newSegments.get(segment0));
    AssertExtensions.assertThrows(ReaderNotInReaderGroupException.class, () -> readerState2.acquireNewSegmentsIfNeeded(0L, new PositionImpl(Collections.emptyMap())));
}

Example 52 with MockSegmentStreamFactory

use of io.pravega.client.stream.mock.MockSegmentStreamFactory in project pravega by pravega.

the class ReaderGroupStateManagerTest method testSegmentSplit.

@Test(timeout = 20000)
public void testSegmentSplit() throws ReaderNotInReaderGroupException {
    String scope = "scope";
    String stream = "stream";
    PravegaNodeUri endpoint = new PravegaNodeUri("localhost", SERVICE_PORT);
    MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
    SegmentWithRange initialSegment = new SegmentWithRange(new Segment(scope, stream, 0), 0, 1);
    SegmentWithRange successorA = new SegmentWithRange(new Segment(scope, stream, 1), 0, 0.5);
    SegmentWithRange successorB = new SegmentWithRange(new Segment(scope, stream, 2), 0.5, 1.0);
    MockController controller = new MockControllerWithSuccessors(endpoint.getEndpoint(), endpoint.getPort(), connectionFactory, new StreamSegmentsWithPredecessors(ImmutableMap.of(successorA, singletonList(0L), successorB, singletonList(0L)), ""));
    createScopeAndStream(scope, stream, controller);
    MockSegmentStreamFactory streamFactory = new MockSegmentStreamFactory();
    @Cleanup SynchronizerClientFactory clientFactory = new ClientFactoryImpl(scope, controller, connectionFactory, streamFactory, streamFactory, streamFactory, streamFactory);
    SynchronizerConfig config = SynchronizerConfig.builder().build();
    @Cleanup StateSynchronizer<ReaderGroupState> stateSynchronizer = createState(stream, clientFactory, config);
    Map<SegmentWithRange, Long> segments = new HashMap<>();
    segments.put(initialSegment, 1L);
    stateSynchronizer.initialize(new ReaderGroupState.ReaderGroupStateInit(ReaderGroupConfig.builder().stream(Stream.of(scope, stream)).build(), segments, Collections.emptyMap(), false));
    val readerState = new ReaderGroupStateManager(scope, stream, "testReader", stateSynchronizer, controller, null);
    readerState.initializeReader(0);
    Map<SegmentWithRange, Long> newSegments = readerState.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap()));
    assertEquals(1, newSegments.size());
    assertEquals(Long.valueOf(1), newSegments.get(initialSegment));
    readerState.handleEndOfSegment(initialSegment);
    newSegments = readerState.acquireNewSegmentsIfNeeded(0, new PositionImpl(ImmutableMap.of(initialSegment, 1L)));
    assertEquals(2, newSegments.size());
    assertEquals(Long.valueOf(0), newSegments.get(successorA));
    assertEquals(Long.valueOf(0), newSegments.get(successorB));
    newSegments = readerState.acquireNewSegmentsIfNeeded(0, new PositionImpl(newSegments));
    assertTrue(newSegments.isEmpty());
}

Example 53 with MockSegmentStreamFactory

use of io.pravega.client.stream.mock.MockSegmentStreamFactory in project pravega by pravega.

the class ReaderGroupStateManagerTest method testReachEndInvalidState.

@Test(timeout = 10000)
public void testReachEndInvalidState() throws ReaderNotInReaderGroupException {
    String scope = "scope";
    String stream = "stream";
    PravegaNodeUri endpoint = new PravegaNodeUri("localhost", SERVICE_PORT);
    MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
    MockController controller = new MockController(endpoint.getEndpoint(), endpoint.getPort(), connectionFactory, false);
    createScopeAndStream(scope, stream, controller);
    MockSegmentStreamFactory streamFactory = new MockSegmentStreamFactory();
    @Cleanup SynchronizerClientFactory clientFactory = new ClientFactoryImpl(scope, controller, connectionFactory, streamFactory, streamFactory, streamFactory, streamFactory);
    @Cleanup StateSynchronizer<ReaderGroupState> stateSynchronizer = createState(stream, clientFactory, SynchronizerConfig.builder().build());
    Map<SegmentWithRange, Long> segments = new HashMap<>();
    SegmentWithRange segment = new SegmentWithRange(new Segment(scope, stream, 0), 0.0, 1.0);
    segments.put(segment, 1L);
    StreamCutImpl start = new StreamCutImpl(Stream.of(scope, stream), ImmutableMap.of(segment.getSegment(), 0L));
    StreamCutImpl end = new StreamCutImpl(Stream.of(scope, stream), ImmutableMap.of(segment.getSegment(), 100L));
    ReaderGroupConfig config = ReaderGroupConfig.builder().stream(Stream.of(scope, stream), start, end).build();
    stateSynchronizer.initialize(new ReaderGroupState.ReaderGroupStateInit(config, segments, ReaderGroupImpl.getEndSegmentsForStreams(config), false));
    ReaderGroupStateManager readerState = new ReaderGroupStateManager(scope, stream, "testReader", stateSynchronizer, controller, null);
    readerState.initializeReader(0);
    readerState.readerShutdown(new PositionImpl(segments));
    assertThrows(ReaderNotInReaderGroupException.class, () -> readerState.handleEndOfSegment(segment));
    // restore reader without segment
    readerState.initializeReader(0);
    assertThrows(ReaderNotInReaderGroupException.class, () -> readerState.handleEndOfSegment(segment));
    // Test it can release if it has the segment
    readerState.acquireNewSegmentsIfNeeded(0, new PositionImpl(segments));
    readerState.handleEndOfSegment(segment);
}

Example 54 with MockSegmentStreamFactory

use of io.pravega.client.stream.mock.MockSegmentStreamFactory in project pravega by pravega.

the class ReaderGroupStateManagerTest method testCheckpointContainsAllShards.

@Test(timeout = 10000)
public void testCheckpointContainsAllShards() throws ReaderNotInReaderGroupException {
    String scope = "scope";
    String stream = "stream";
    PravegaNodeUri endpoint = new PravegaNodeUri("localhost", SERVICE_PORT);
    MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
    SegmentWithRange segment0 = new SegmentWithRange(new Segment(scope, stream, 0), 0.0, 0.33);
    SegmentWithRange segment1 = new SegmentWithRange(new Segment(scope, stream, 1), 0.33, 0.66);
    SegmentWithRange segment2 = new SegmentWithRange(new Segment(scope, stream, 2), 0.66, 1.0);
    MockController controller = new MockControllerWithSuccessors(endpoint.getEndpoint(), endpoint.getPort(), connectionFactory, new StreamSegmentsWithPredecessors(ImmutableMap.of(), ""));
    createScopeAndStream(scope, stream, controller);
    MockSegmentStreamFactory streamFactory = new MockSegmentStreamFactory();
    @Cleanup SynchronizerClientFactory clientFactory = new ClientFactoryImpl(scope, controller, connectionFactory, streamFactory, streamFactory, streamFactory, streamFactory);
    SynchronizerConfig config = SynchronizerConfig.builder().build();
    @Cleanup StateSynchronizer<ReaderGroupState> stateSynchronizer = createState(stream, clientFactory, config);
    Map<SegmentWithRange, Long> segments = ImmutableMap.of(segment0, 0L, segment1, 1L, segment2, 2L);
    stateSynchronizer.initialize(new ReaderGroupState.ReaderGroupStateInit(ReaderGroupConfig.builder().stream(Stream.of(scope, stream)).build(), segments, Collections.emptyMap(), false));
    val readerState1 = new ReaderGroupStateManager(scope, stream, "reader1", stateSynchronizer, controller, null);
    readerState1.initializeReader(0);
    val readerState2 = new ReaderGroupStateManager(scope, stream, "reader2", stateSynchronizer, controller, null);
    readerState2.initializeReader(0);
    assertEquals(segments, stateSynchronizer.getState().getUnassignedSegments());
    stateSynchronizer.updateStateUnconditionally(new CreateCheckpoint("CP1"));
    stateSynchronizer.fetchUpdates();
    assertEquals("CP1", readerState1.getCheckpoint());
    assertEquals(Collections.emptyMap(), readerState1.acquireNewSegmentsIfNeeded(1, new PositionImpl(Collections.emptyMap())));
    assertEquals(Collections.emptyMap(), readerState2.acquireNewSegmentsIfNeeded(2, new PositionImpl(Collections.emptyMap())));
    assertEquals("CP1", readerState2.getCheckpoint());
    readerState1.checkpoint("CP1", new PositionImpl(Collections.emptyMap()));
    readerState2.checkpoint("CP1", new PositionImpl(Collections.emptyMap()));
    assertTrue(stateSynchronizer.getState().isCheckpointComplete("CP1"));
    assertEquals(ImmutableMap.of(segment0.getSegment(), 0L, segment1.getSegment(), 1L, segment2.getSegment(), 2L), stateSynchronizer.getState().getPositionsForCompletedCheckpoint("CP1"));
}

Example 55 with MockSegmentStreamFactory

use of io.pravega.client.stream.mock.MockSegmentStreamFactory in project pravega by pravega.

the class ReaderGroupStateManagerTest method testCheckpointWithCBR.

@Test(timeout = 10000)
public void testCheckpointWithCBR() throws ReaderNotInReaderGroupException {
    String scope = "scope";
    String stream = "stream";
    PravegaNodeUri endpoint = new PravegaNodeUri("localhost", SERVICE_PORT);
    MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
    SegmentWithRange initialSegment = new SegmentWithRange(new Segment(scope, stream, 0), 0, 1);
    SegmentWithRange successorA = new SegmentWithRange(new Segment(scope, stream, 1), 0.0, 0.5);
    SegmentWithRange successorB = new SegmentWithRange(new Segment(scope, stream, 2), 0.5, 1.0);
    MockController controller = new MockControllerWithSuccessors(endpoint.getEndpoint(), endpoint.getPort(), connectionFactory, new StreamSegmentsWithPredecessors(ImmutableMap.of(successorA, singletonList(0L), successorB, singletonList(0L)), ""));
    createScopeAndStream(scope, stream, controller);
    MockSegmentStreamFactory streamFactory = new MockSegmentStreamFactory();
    @Cleanup SynchronizerClientFactory clientFactory = new ClientFactoryImpl(scope, controller, connectionFactory, streamFactory, streamFactory, streamFactory, streamFactory);
    SynchronizerConfig config = SynchronizerConfig.builder().build();
    @Cleanup StateSynchronizer<ReaderGroupState> stateSynchronizer = createState(stream, clientFactory, config);
    Map<SegmentWithRange, Long> segments = new HashMap<>();
    segments.put(initialSegment, 1L);
    ReaderGroupConfig rgConfig = ReaderGroupConfig.builder().stream(Stream.of(scope, stream)).disableAutomaticCheckpoints().retentionType(ReaderGroupConfig.StreamDataRetention.AUTOMATIC_RELEASE_AT_LAST_CHECKPOINT).build();
    stateSynchronizer.initialize(new ReaderGroupState.ReaderGroupStateInit(rgConfig, segments, Collections.emptyMap(), false));
    val readerState = new ReaderGroupStateManager(scope, stream, "testReader", stateSynchronizer, controller, null);
    readerState.initializeReader(0);
    assertNull(readerState.getCheckpoint());
    stateSynchronizer.updateStateUnconditionally(new CreateCheckpoint("CP1"));
    stateSynchronizer.fetchUpdates();
    assertEquals("CP1", readerState.getCheckpoint());
    assertEquals("CP1", readerState.getCheckpoint());
    assertTrue(stateSynchronizer.getState().getCheckpointState().isLastCheckpointPublished());
    readerState.checkpoint("CP1", new PositionImpl(Collections.emptyMap()));
    assertNull(readerState.getCheckpoint());
    stateSynchronizer.fetchUpdates();
    assertFalse(stateSynchronizer.getState().getCheckpointState().isLastCheckpointPublished());
    readerState.updateTruncationStreamCutIfNeeded();
    stateSynchronizer.fetchUpdates();
    assertTrue(stateSynchronizer.getState().isCheckpointComplete("CP1"));
    assertTrue(stateSynchronizer.getState().getCheckpointState().isLastCheckpointPublished());
    stateSynchronizer.updateStateUnconditionally(new CreateCheckpoint("CP2"));
    stateSynchronizer.updateStateUnconditionally(new CreateCheckpoint("CP3"));
    stateSynchronizer.fetchUpdates();
    assertEquals("CP2", readerState.getCheckpoint());
    readerState.checkpoint("CP2", new PositionImpl(Collections.emptyMap()));
    assertEquals("CP3", readerState.getCheckpoint());
    readerState.checkpoint("CP3", new PositionImpl(Collections.emptyMap()));
    assertNull(readerState.getCheckpoint());
}