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

Example 56 with MockConnectionFactoryImpl

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

the class ReaderGroupStateManagerTest method testReleaseCompletedSegment.

@Test(timeout = 10000)
public void testReleaseCompletedSegment() 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);
    controller.createScope(scope);
    controller.createStream(scope, stream, StreamConfiguration.builder().scalingPolicy(ScalingPolicy.fixed(2)).build());
    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<>();
    segments.put(new SegmentWithRange(new Segment(scope, stream, 0), 0.0, 0.5), 123L);
    segments.put(new SegmentWithRange(new Segment(scope, stream, 1), 0.5, 1.0), 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());
    Position pos = new PositionImpl(ImmutableMap.of(new SegmentWithRange(new Segment(scope, stream, 0), 0.0, 0.5), -1L, new SegmentWithRange(new Segment(scope, stream, 1), 0.5, 1.0), 789L));
    ReaderGroupStateManager.readerShutdown("testReader", pos, stateSynchronizer);
    newSegments = readerState2.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap()));
    assertEquals(2, newSegments.size());
    assertEquals(Long.valueOf(789L), newSegments.get(new SegmentWithRange(new Segment(scope, stream, 1), 0.5, 1.0)));
    assertEquals(0, stateSynchronizer.getState().getNumberOfUnassignedSegments());
    AssertExtensions.assertThrows(ReaderNotInReaderGroupException.class, () -> readerState1.acquireNewSegmentsIfNeeded(0L, new PositionImpl(Collections.emptyMap())));
}

Example 57 with MockConnectionFactoryImpl

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

the class ReaderGroupStateManagerTest method testSegmentsCannotBeReleasedWithoutCheckpoint.

@Test(timeout = 10000)
public void testSegmentsCannotBeReleasedWithoutCheckpoint() 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);
    AtomicLong clock = new AtomicLong();
    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, clock::get);
    readerState1.initializeReader(0);
    val readerState2 = new ReaderGroupStateManager(scope, stream, "reader2", stateSynchronizer, controller, clock::get);
    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()));
    assertEquals(ImmutableMap.of(segment0.getSegment(), 0L, segment1.getSegment(), 1L, segment2.getSegment(), 2L), stateSynchronizer.getState().getPositionsForCompletedCheckpoint("CP1"));
    Map<SegmentWithRange, Long> segments1 = readerState1.acquireNewSegmentsIfNeeded(1, new PositionImpl(Collections.emptyMap()));
    Map<SegmentWithRange, Long> segments2 = readerState2.acquireNewSegmentsIfNeeded(2, new PositionImpl(Collections.emptyMap()));
    assertFalse(segments1.isEmpty());
    assertFalse(segments2.isEmpty());
    assertEquals(0, stateSynchronizer.getState().getNumberOfUnassignedSegments());
    // Induce imbalance
    for (Entry<SegmentWithRange, Long> entry : segments1.entrySet()) {
        stateSynchronizer.updateStateUnconditionally(new ReaderGroupState.ReleaseSegment("reader1", entry.getKey().getSegment(), entry.getValue()));
        stateSynchronizer.updateStateUnconditionally(new ReaderGroupState.AcquireSegment("reader2", entry.getKey().getSegment()));
    }
    stateSynchronizer.updateStateUnconditionally(new CreateCheckpoint("CP2"));
    stateSynchronizer.fetchUpdates();
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    assertNull(readerState1.findSegmentToReleaseIfRequired());
    assertNull(readerState2.findSegmentToReleaseIfRequired());
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    assertFalse(readerState2.releaseSegment(segments2.keySet().iterator().next().getSegment(), 20, 2, new PositionImpl(segments)));
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    readerState1.checkpoint("CP2", new PositionImpl(Collections.emptyMap()));
    readerState2.checkpoint("CP2", new PositionImpl(segments));
    assertEquals(ImmutableMap.of(segment0.getSegment(), 0L, segment1.getSegment(), 1L, segment2.getSegment(), 2L), stateSynchronizer.getState().getPositionsForCompletedCheckpoint("CP2"));
    Segment toRelease = readerState2.findSegmentToReleaseIfRequired();
    assertNotNull(toRelease);
    assertTrue(readerState2.releaseSegment(toRelease, 10, 1, new PositionImpl(segments)));
    assertEquals(1, stateSynchronizer.getState().getNumberOfUnassignedSegments());
}

Example 58 with MockConnectionFactoryImpl

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

the class ReaderGroupStateManagerTest method testUpdateConfigIfNeeded.

@Test(timeout = 10000)
public void testUpdateConfigIfNeeded() {
    String scope = "scope";
    String stream = "stream";
    String stream2 = "stream2";
    String groupName = "group";
    String rgStream = NameUtils.getStreamForReaderGroup(groupName);
    ReaderGroupConfig clientConfig = ReaderGroupConfig.builder().stream(Stream.of(scope, stream)).build();
    ReaderGroupConfig controllerConfig = ReaderGroupConfig.builder().stream(Stream.of(scope, stream2)).build();
    controllerConfig = ReaderGroupConfig.cloneConfig(controllerConfig, clientConfig.getReaderGroupId(), 1L);
    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);
    createScopeAndStream(scope, stream2, controller);
    MockSegmentStreamFactory streamFactory = new MockSegmentStreamFactory();
    @Cleanup SynchronizerClientFactory clientFactory = new ClientFactoryImpl(scope, controller, connectionFactory, streamFactory, streamFactory, streamFactory, streamFactory);
    // Initialize RG state with updatingConfig as true.
    SynchronizerConfig config = SynchronizerConfig.builder().build();
    createScopeAndStream(scope, rgStream, controller);
    @Cleanup StateSynchronizer<ReaderGroupState> state = createState(rgStream, clientFactory, config);
    AtomicLong clock = new AtomicLong();
    Map<SegmentWithRange, Long> segments = new HashMap<>();
    segments.put(new SegmentWithRange(new Segment(scope, stream, 0), 0.0, 0.25), 0L);
    segments.put(new SegmentWithRange(new Segment(scope, stream, 1), 0.25, 0.5), 1L);
    segments.put(new SegmentWithRange(new Segment(scope, stream, 2), 0.5, 0.75), 2L);
    segments.put(new SegmentWithRange(new Segment(scope, stream, 3), 0.65, 1.0), 3L);
    state.initialize(new ReaderGroupState.ReaderGroupStateInit(clientConfig, segments, Collections.emptyMap(), true));
    controller.createReaderGroup(scope, groupName, controllerConfig);
    ReaderGroupStateManager reader1 = new ReaderGroupStateManager(scope, groupName, "reader1", state, controller, clock::get);
    reader1.initializeReader(100);
    assertTrue(state.getState().isUpdatingConfig());
    assertEquals(clientConfig, state.getState().getConfig());
    clock.addAndGet(ReaderGroupStateManager.UPDATE_CONFIG_WINDOW.toNanos());
    reader1.updateConfigIfNeeded();
    assertFalse(state.getState().isUpdatingConfig());
    assertEquals(controllerConfig, state.getState().getConfig());
}

Example 59 with MockConnectionFactoryImpl

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

the class ReaderGroupStateManagerTest method testReleaseWhenReadersAdded.

@Test(timeout = 20000)
public void testReleaseWhenReadersAdded() 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();
    SegmentWithRange s0 = new SegmentWithRange(new Segment(scope, stream, 0), 0.0, 0.1);
    SegmentWithRange s1 = new SegmentWithRange(new Segment(scope, stream, 1), 0.1, 0.2);
    SegmentWithRange s2 = new SegmentWithRange(new Segment(scope, stream, 2), 0.2, 0.3);
    SegmentWithRange s3 = new SegmentWithRange(new Segment(scope, stream, 3), 0.3, 0.4);
    SegmentWithRange s4 = new SegmentWithRange(new Segment(scope, stream, 4), 0.4, 0.5);
    SegmentWithRange s5 = new SegmentWithRange(new Segment(scope, stream, 5), 0.5, 1.0);
    Map<SegmentWithRange, Long> segments = new HashMap<>();
    segments.put(s0, 0L);
    segments.put(s1, 1L);
    segments.put(s2, 2L);
    segments.put(s3, 3L);
    segments.put(s4, 4L);
    segments.put(s5, 5L);
    stateSynchronizer.initialize(new ReaderGroupState.ReaderGroupStateInit(ReaderGroupConfig.builder().stream(Stream.of(scope, stream)).build(), segments, Collections.emptyMap(), false));
    ReaderGroupStateManager reader1 = new ReaderGroupStateManager(scope, stream, "reader1", stateSynchronizer, controller, clock::get);
    reader1.initializeReader(0);
    Map<SegmentWithRange, Long> segments1 = reader1.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap()));
    assertEquals(6, segments1.size());
    ReaderGroupStateManager reader2 = new ReaderGroupStateManager(scope, stream, "reader2", stateSynchronizer, controller, clock::get);
    reader2.initializeReader(0);
    assertTrue(reader2.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap())).isEmpty());
    assertNull(reader1.findSegmentToReleaseIfRequired());
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    assertNotNull(reader1.findSegmentToReleaseIfRequired());
    reader1.releaseSegment(new Segment(scope, stream, 3), 3, 0, new PositionImpl(segments));
    assertNull(reader1.findSegmentToReleaseIfRequired());
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    assertNotNull(reader1.findSegmentToReleaseIfRequired());
    reader1.releaseSegment(new Segment(scope, stream, 4), 4, 0, new PositionImpl(segments));
    assertNull(reader1.findSegmentToReleaseIfRequired());
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    assertNotNull(reader1.findSegmentToReleaseIfRequired());
    reader1.releaseSegment(new Segment(scope, stream, 5), 5, 0, new PositionImpl(segments));
    assertNull(reader1.findSegmentToReleaseIfRequired());
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    assertNull(reader1.findSegmentToReleaseIfRequired());
    Map<SegmentWithRange, Long> segments2 = reader2.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap()));
    assertEquals(3, segments2.size());
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    ReaderGroupStateManager reader3 = new ReaderGroupStateManager(scope, stream, "reader3", stateSynchronizer, controller, clock::get);
    reader3.initializeReader(0);
    assertTrue(reader3.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap())).isEmpty());
    assertNotNull(reader1.findSegmentToReleaseIfRequired());
    reader1.releaseSegment(new Segment(scope, stream, 0), 0, 0, new PositionImpl(ImmutableMap.of(s0, 10L, s1, 11L, s2, 12L)));
    assertNull(reader1.findSegmentToReleaseIfRequired());
    assertNotNull(reader2.findSegmentToReleaseIfRequired());
    reader2.releaseSegment(new Segment(scope, stream, 3), 3, 0, new PositionImpl(ImmutableMap.of(s3, 13L, s4, 14L, s5, 15L)));
    assertNull(reader2.findSegmentToReleaseIfRequired());
    clock.addAndGet(ReaderGroupStateManager.TIME_UNIT.toNanos());
    Map<SegmentWithRange, Long> segments3 = reader3.acquireNewSegmentsIfNeeded(0, new PositionImpl(Collections.emptyMap()));
    assertEquals(2, segments3.size());
    clock.addAndGet(ReaderGroupStateManager.UPDATE_WINDOW.toNanos());
    reader3.updateLagIfNeeded(0, new PositionImpl(ImmutableMap.of(s0, 20L, s3, 23L)));
    assertNull(reader1.findSegmentToReleaseIfRequired());
    assertNull(reader2.findSegmentToReleaseIfRequired());
    assertNull(reader3.findSegmentToReleaseIfRequired());
    Map<SegmentWithRange, Long> expected = new HashMap<>();
    expected.put(s0, 20L);
    expected.put(s1, 11L);
    expected.put(s2, 12L);
    expected.put(s3, 23L);
    expected.put(s4, 14L);
    expected.put(s5, 15L);
    assertEquals(expected, stateSynchronizer.getState().getLastReadPositions(Stream.of(scope, stream)));
}

Example 60 with MockConnectionFactoryImpl

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

the class BatchClientImplTest method testSegmentIterator.

@Test(timeout = 5000)
public void testSegmentIterator() throws ConnectionFailedException {
    MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
    ClientConnection connection = Mockito.mock(ClientConnection.class);
    PravegaNodeUri location = new PravegaNodeUri("localhost", 0);
    Mockito.doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            CreateSegment request = (CreateSegment) invocation.getArgument(0);
            connectionFactory.getProcessor(location).process(new SegmentCreated(request.getRequestId(), request.getSegment()));
            return null;
        }
    }).when(connection).send(Mockito.any(CreateSegment.class));
    Mockito.doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            GetStreamSegmentInfo request = (GetStreamSegmentInfo) invocation.getArgument(0);
            connectionFactory.getProcessor(location).process(new StreamSegmentInfo(request.getRequestId(), request.getSegmentName(), true, false, false, 0, 0, 0));
            return null;
        }
    }).when(connection).send(Mockito.any(GetStreamSegmentInfo.class));
    connectionFactory.provideConnection(location, connection);
    MockController mockController = new MockController(location.getEndpoint(), location.getPort(), connectionFactory);
    BatchClientImpl client = new BatchClientImpl(mockController, connectionFactory);
    Stream stream = new StreamImpl("scope", "stream");
    mockController.createScope("scope");
    mockController.createStream(StreamConfiguration.builder().scope("scope").streamName("stream").scalingPolicy(ScalingPolicy.fixed(3)).build()).join();
    Iterator<SegmentRange> segments = client.getSegments(stream, null, null).getIterator();
    assertTrue(segments.hasNext());
    assertEquals(0, segments.next().asImpl().getSegment().getSegmentNumber());
    assertTrue(segments.hasNext());
    assertEquals(1, segments.next().asImpl().getSegment().getSegmentNumber());
    assertTrue(segments.hasNext());
    assertEquals(2, segments.next().asImpl().getSegment().getSegmentNumber());
    assertFalse(segments.hasNext());
}