Examples with Controller io.pravega.client.control.impl.Controller used on opensource projects

Example 21 with Controller

use of io.pravega.client.control.impl.Controller in project pravega by pravega.

the class SegmentSelectorTest method testUsesAllSegments.

@Test
public void testUsesAllSegments() {
    Controller controller = Mockito.mock(Controller.class);
    SegmentOutputStreamFactory factory = Mockito.mock(SegmentOutputStreamFactory.class);
    SegmentSelector selector = new SegmentSelector(new StreamImpl(scope, streamName), controller, factory, config, DelegationTokenProviderFactory.createWithEmptyToken());
    TreeMap<Double, SegmentWithRange> segments = new TreeMap<>();
    addNewSegment(segments, 0, 0.0, 0.25);
    addNewSegment(segments, 1, 0.25, 0.5);
    addNewSegment(segments, 2, 0.5, 0.75);
    addNewSegment(segments, 3, 0.75, 1.0);
    StreamSegments streamSegments = new StreamSegments(segments);
    when(controller.getCurrentSegments(scope, streamName)).thenReturn(CompletableFuture.completedFuture(streamSegments));
    selector.refreshSegmentEventWriters(segmentSealedCallback);
    int[] counts = new int[4];
    Arrays.fill(counts, 0);
    for (int i = 0; i < 20; i++) {
        Segment segment = selector.getSegmentForEvent("" + i);
        assertNotNull(segment);
        counts[NameUtils.getSegmentNumber(segment.getSegmentId())]++;
    }
    for (int count : counts) {
        assertTrue(count > 1);
    }
}

Example 22 with Controller

use of io.pravega.client.control.impl.Controller in project pravega by pravega.

the class SegmentSelectorTest method testNullRoutingKey.

@Test
public void testNullRoutingKey() {
    Controller controller = Mockito.mock(Controller.class);
    SegmentOutputStreamFactory factory = Mockito.mock(SegmentOutputStreamFactory.class);
    SegmentSelector selector = new SegmentSelector(new StreamImpl(scope, streamName), controller, factory, config, DelegationTokenProviderFactory.createWithEmptyToken());
    TreeMap<Double, SegmentWithRange> segments = new TreeMap<>();
    addNewSegment(segments, 0, 0.0, 0.25);
    addNewSegment(segments, 1, 0.25, 0.5);
    addNewSegment(segments, 2, 0.5, 0.75);
    addNewSegment(segments, 3, 0.75, 1.0);
    StreamSegments streamSegments = new StreamSegments(segments);
    when(controller.getCurrentSegments(scope, streamName)).thenReturn(CompletableFuture.completedFuture(streamSegments));
    selector.refreshSegmentEventWriters(segmentSealedCallback);
    int[] counts = new int[4];
    Arrays.fill(counts, 0);
    for (int i = 0; i < 100; i++) {
        Segment segment = selector.getSegmentForEvent(null);
        assertNotNull(segment);
        counts[NameUtils.getSegmentNumber(segment.getSegmentId())]++;
    }
    for (int count : counts) {
        assertTrue(count > 1);
    }
}

Example 23 with Controller

use of io.pravega.client.control.impl.Controller in project pravega by pravega.

the class SegmentSelectorTest method testControllerNotReachable.

@Test
public void testControllerNotReachable() {
    final Segment segment0 = new Segment(scope, streamName, 0);
    final Segment segment1 = new Segment(scope, streamName, 1);
    final CompletableFuture<Void> writerFuture = new CompletableFuture<>();
    // Setup Mock.
    SegmentOutputStream s0Writer = Mockito.mock(SegmentOutputStream.class);
    SegmentOutputStream s1Writer = Mockito.mock(SegmentOutputStream.class);
    when(s0Writer.getUnackedEventsOnSeal()).thenReturn(ImmutableList.of(PendingEvent.withHeader("0", ByteBuffer.wrap("e".getBytes()), writerFuture)));
    SegmentOutputStreamFactory factory = Mockito.mock(SegmentOutputStreamFactory.class);
    when(factory.createOutputStreamForSegment(eq(segment0), ArgumentMatchers.<Consumer<Segment>>any(), any(EventWriterConfig.class), any(DelegationTokenProvider.class))).thenReturn(s0Writer);
    when(factory.createOutputStreamForSegment(eq(segment1), ArgumentMatchers.<Consumer<Segment>>any(), any(EventWriterConfig.class), any(DelegationTokenProvider.class))).thenReturn(s1Writer);
    Controller controller = Mockito.mock(Controller.class);
    SegmentSelector selector = new SegmentSelector(new StreamImpl(scope, streamName), controller, factory, config, DelegationTokenProviderFactory.createWithEmptyToken());
    TreeMap<Double, SegmentWithRange> segments = new TreeMap<>();
    addNewSegment(segments, 0, 0.0, 0.5);
    addNewSegment(segments, 1, 0.5, 1.0);
    StreamSegments streamSegments = new StreamSegments(segments);
    when(controller.getCurrentSegments(scope, streamName)).thenReturn(CompletableFuture.completedFuture(streamSegments));
    // trigger refresh.
    selector.refreshSegmentEventWriters(segmentSealedCallback);
    // simulate controller failure when controller.getSuccessors() is invoked.
    when(controller.getSuccessors(segment0)).thenAnswer(i -> {
        CompletableFuture<StreamSegmentsWithPredecessors> result = new CompletableFuture<>();
        // Controller client in case of RPC exceptions throws a StatusRuntimeException which is retried by the client.
        // If the controller client is not able to reach the controller after all the retries a RetriesExhaustedException is
        // thrown.
        result.completeExceptionally(new RetriesExhaustedException(new StatusRuntimeException(Status.DATA_LOSS)));
        return result;
    });
    assertEquals(Collections.emptyList(), selector.refreshSegmentEventWritersUponSealed(segment0, segmentSealedCallback));
    assertFutureThrows("Writer Future", writerFuture, t -> t instanceof ControllerFailureException);
}

Example 24 with Controller

use of io.pravega.client.control.impl.Controller in project pravega by pravega.

the class SegmentSelectorTest method testStreamDeletion.

@Test
public void testStreamDeletion() {
    final Segment segment0 = new Segment(scope, streamName, 0);
    final Segment segment1 = new Segment(scope, streamName, 1);
    final CompletableFuture<Void> writerFuture = new CompletableFuture<>();
    // Setup Mock.
    SegmentOutputStream s0Writer = Mockito.mock(SegmentOutputStream.class);
    SegmentOutputStream s1Writer = Mockito.mock(SegmentOutputStream.class);
    when(s0Writer.getUnackedEventsOnSeal()).thenReturn(ImmutableList.of(PendingEvent.withHeader("0", ByteBuffer.wrap("e".getBytes()), writerFuture)));
    SegmentOutputStreamFactory factory = Mockito.mock(SegmentOutputStreamFactory.class);
    when(factory.createOutputStreamForSegment(eq(segment0), ArgumentMatchers.<Consumer<Segment>>any(), any(EventWriterConfig.class), any(DelegationTokenProvider.class))).thenReturn(s0Writer);
    when(factory.createOutputStreamForSegment(eq(segment1), ArgumentMatchers.<Consumer<Segment>>any(), any(EventWriterConfig.class), any(DelegationTokenProvider.class))).thenReturn(s1Writer);
    Controller controller = Mockito.mock(Controller.class);
    SegmentSelector selector = new SegmentSelector(new StreamImpl(scope, streamName), controller, factory, config, DelegationTokenProviderFactory.createWithEmptyToken());
    TreeMap<Double, SegmentWithRange> segments = new TreeMap<>();
    addNewSegment(segments, 0, 0.0, 0.5);
    addNewSegment(segments, 1, 0.5, 1.0);
    StreamSegments streamSegments = new StreamSegments(segments);
    when(controller.getCurrentSegments(scope, streamName)).thenReturn(CompletableFuture.completedFuture(streamSegments));
    // trigger refresh.
    selector.refreshSegmentEventWriters(segmentSealedCallback);
    // simulate stream deletion where controller.getSuccessors() is completed exceptionally.
    when(controller.getSuccessors(segment0)).thenAnswer(i -> {
        CompletableFuture<StreamSegmentsWithPredecessors> result = new CompletableFuture<>();
        // Controller throws io.pravega.controller.store.stream.StoreException$DataNotFoundException which is type RuntimeException.
        // Using RunTimeException here as the controller exception is not visible.
        result.completeExceptionally(new RuntimeException());
        return result;
    });
    assertEquals(Collections.emptyList(), selector.refreshSegmentEventWritersUponSealed(segment0, segmentSealedCallback));
    assertFutureThrows("Writer Future", writerFuture, t -> t instanceof NoSuchSegmentException);
}

Example 25 with Controller

use of io.pravega.client.control.impl.Controller in project pravega by pravega.

the class SegmentSelectorTest method testSegmentRefreshOnSealed.

@Test
public void testSegmentRefreshOnSealed() {
    final Segment segment0 = new Segment(scope, streamName, 0);
    final Segment segment1 = new Segment(scope, streamName, computeSegmentId(1, 1));
    final Segment segment2 = new Segment(scope, streamName, computeSegmentId(2, 1));
    final CompletableFuture<Void> writerFuture = new CompletableFuture<>();
    final PendingEvent pendingEvent = PendingEvent.withHeader("0", ByteBuffer.wrap("e".getBytes()), writerFuture);
    TreeMap<Double, SegmentWithRange> segmentBeforeScale = new TreeMap<>();
    addNewSegment(segmentBeforeScale, 0, 0.0, 1.0);
    StreamSegments streamSegmentsBeforeScale = new StreamSegments(segmentBeforeScale);
    Map<SegmentWithRange, List<Long>> newRange = new HashMap<>();
    newRange.put(new SegmentWithRange(segment1, 0.0, 0.5), ImmutableList.of(segment0.getSegmentId()));
    newRange.put(new SegmentWithRange(segment2, 0.5, 1.0), ImmutableList.of(segment0.getSegmentId()));
    StreamSegmentsWithPredecessors segmentsWithPredecessors = new StreamSegmentsWithPredecessors(newRange, "");
    // Setup Mock.
    SegmentOutputStream s0Writer = Mockito.mock(SegmentOutputStream.class);
    SegmentOutputStream s1Writer = Mockito.mock(SegmentOutputStream.class);
    SegmentOutputStream s2Writer = Mockito.mock(SegmentOutputStream.class);
    SegmentOutputStreamFactory factory = Mockito.mock(SegmentOutputStreamFactory.class);
    Controller controller = Mockito.mock(Controller.class);
    when(s0Writer.getUnackedEventsOnSeal()).thenReturn(ImmutableList.of(pendingEvent));
    when(factory.createOutputStreamForSegment(eq(segment0), ArgumentMatchers.<Consumer<Segment>>any(), any(EventWriterConfig.class), any(DelegationTokenProvider.class))).thenReturn(s0Writer);
    when(factory.createOutputStreamForSegment(eq(segment1), ArgumentMatchers.<Consumer<Segment>>any(), any(EventWriterConfig.class), any(DelegationTokenProvider.class))).thenReturn(s1Writer);
    when(factory.createOutputStreamForSegment(eq(segment2), ArgumentMatchers.<Consumer<Segment>>any(), any(EventWriterConfig.class), any(DelegationTokenProvider.class))).thenReturn(s2Writer);
    // get current segments returns segment 0
    when(controller.getCurrentSegments(scope, streamName)).thenReturn(CompletableFuture.completedFuture(streamSegmentsBeforeScale));
    when(controller.getSuccessors(segment0)).thenAnswer(i -> {
        CompletableFuture<StreamSegmentsWithPredecessors> result = new CompletableFuture<>();
        result.complete(segmentsWithPredecessors);
        return result;
    });
    SegmentSelector selector = new SegmentSelector(new StreamImpl(scope, streamName), controller, factory, config, DelegationTokenProviderFactory.createWithEmptyToken());
    // trigger refresh
    selector.refreshSegmentEventWriters(segmentSealedCallback);
    // only segment 0 writer is present.
    assertEquals(s0Writer, selector.getWriters().values().stream().findAny().get());
    // trigger a referesh of writers due to segment 0 being sealed.
    List<PendingEvent> pendingEvents = selector.refreshSegmentEventWritersUponSealed(segment0, segmentSealedCallback);
    // one pending event is returned by segment0
    assertEquals(singletonList(pendingEvent), pendingEvents);
    // the current number of writers is 3, it includes the writer to segment 0.
    Map<Segment, SegmentOutputStream> writers = selector.getWriters();
    assertEquals(3, writers.size());
    assertTrue(writers.values().contains(s0Writer));
    assertTrue(writers.values().contains(s1Writer));
    assertTrue(writers.values().contains(s2Writer));
    // remove segment 0, this is done post resending the pending events.
    selector.removeSegmentWriter(segment0);
    assertFalse(selector.getWriters().values().contains(s0Writer));
}