Examples with SystemConsumer org.apache.samza.system.SystemConsumer used on opensource projects

Example 1 with SystemConsumer

use of org.apache.samza.system.SystemConsumer in project samza by apache.

the class ITestEventHubSystemConsumer method testSinglePartitionConsumptionHappyPath.

@Test
public void testSinglePartitionConsumptionHappyPath() throws Exception {
    int partitionId = 0;
    TestMetricsRegistry testMetrics = new TestMetricsRegistry();
    SystemStreamPartition ssp = new SystemStreamPartition(SYSTEM_NAME, STREAM_NAME1, new Partition(partitionId));
    Config eventHubConfig = createEventHubConfig();
    EventHubSystemFactory factory = new EventHubSystemFactory();
    SystemConsumer consumer = factory.getConsumer(SYSTEM_NAME, eventHubConfig, testMetrics);
    consumer.register(ssp, EventHubSystemConsumer.START_OF_STREAM);
    consumer.start();
    int numEvents = 0;
    int numRetries = 20;
    while (numRetries-- > 0) {
        List<IncomingMessageEnvelope> result = consumer.poll(Collections.singleton(ssp), 2000).get(ssp);
        numEvents = result == null ? 0 : result.size();
        if (numEvents > 0) {
            EventHubIncomingMessageEnvelope eventData = (EventHubIncomingMessageEnvelope) result.get(0);
            System.out.println("System properties: " + eventData.getEventData().getSystemProperties());
            System.out.println("Key: " + new String((byte[]) eventData.getKey()));
            System.out.println("Message: " + new String((byte[]) eventData.getMessage()));
        }
        System.out.println("Retries left: " + numRetries);
    }
    Assert.assertTrue(numEvents > 0);
}

Example 2 with SystemConsumer

use of org.apache.samza.system.SystemConsumer in project samza by apache.

the class TestTransactionalStateTaskRestoreManager method testRegisterStartingOffsets.

@Test
public void testRegisterStartingOffsets() {
    TaskModel mockTaskModel = mock(TaskModel.class);
    TaskName taskName = new TaskName("Partition 0");
    when(mockTaskModel.getTaskName()).thenReturn(taskName);
    Partition taskChangelogPartition = new Partition(0);
    when(mockTaskModel.getChangelogPartition()).thenReturn(taskChangelogPartition);
    String store1Name = "store1";
    String changelogSystemName = "system";
    String changelog1StreamName = "store1Changelog";
    String store2Name = "store2";
    String changelog2StreamName = "store2Changelog";
    String store3Name = "store3";
    String changelog3StreamName = "store3Changelog";
    String store4Name = "store4";
    String changelog4StreamName = "store4Changelog";
    // tests restore for store 1 and store 2 but not store 3
    Map<String, RestoreOffsets> mockRestoreOffsets = ImmutableMap.of(// tests starting offset == oldest (i.e. restore is inclusive)
    store1Name, // tests starting offset == oldest (i.e. restore is inclusive)
    new RestoreOffsets("0", "5"), // tests starting offset != oldest (i.e. restore from next offset)
    store2Name, // tests starting offset != oldest (i.e. restore from next offset)
    new RestoreOffsets("15", "20"), store4Name, // tests that null ending offsets are OK (should trim)
    new RestoreOffsets("31", null));
    StoreActions mockStoreActions = new StoreActions(ImmutableMap.of(), ArrayListMultimap.create(), mockRestoreOffsets);
    SystemStream changelog1SystemStream = new SystemStream(changelogSystemName, changelog1StreamName);
    SystemStream changelog2SystemStream = new SystemStream(changelogSystemName, changelog2StreamName);
    SystemStream changelog3SystemStream = new SystemStream(changelogSystemName, changelog3StreamName);
    SystemStream changelog4SystemStream = new SystemStream(changelogSystemName, changelog4StreamName);
    Map<String, SystemStream> mockStoreChangelogs = ImmutableMap.of(store1Name, changelog1SystemStream, store2Name, changelog2SystemStream, store3Name, changelog3SystemStream, store4Name, changelog4SystemStream);
    SystemStreamPartition changelog1SSP = new SystemStreamPartition(changelog1SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog1SSPMetadata = new SystemStreamPartitionMetadata("0", "10", "11");
    SystemStreamPartition changelog2SSP = new SystemStreamPartition(changelog2SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog2SSPMetadata = new SystemStreamPartitionMetadata("11", "20", "21");
    SystemStreamPartition changelog3SSP = new SystemStreamPartition(changelog3SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog3SSPMetadata = new SystemStreamPartitionMetadata("21", "30", "31");
    SystemStreamPartition changelog4SSP = new SystemStreamPartition(changelog4SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog4SSPMetadata = new SystemStreamPartitionMetadata("31", "40", "41");
    Map<SystemStreamPartition, SystemStreamPartitionMetadata> mockCurrentChangelogSSPMetadata = ImmutableMap.of(changelog1SSP, changelog1SSPMetadata, changelog2SSP, changelog2SSPMetadata, changelog3SSP, changelog3SSPMetadata, changelog4SSP, changelog4SSPMetadata);
    SystemAdmins mockSystemAdmins = mock(SystemAdmins.class);
    SystemAdmin mockSystemAdmin = mock(SystemAdmin.class);
    when(mockSystemAdmins.getSystemAdmin(eq(changelogSystemName))).thenReturn(mockSystemAdmin);
    Mockito.when(mockSystemAdmin.offsetComparator(anyString(), anyString())).thenAnswer((Answer<Integer>) invocation -> {
        String offset1 = (String) invocation.getArguments()[0];
        String offset2 = (String) invocation.getArguments()[1];
        return Long.valueOf(offset1).compareTo(Long.valueOf(offset2));
    });
    Mockito.when(mockSystemAdmin.getOffsetsAfter(any())).thenAnswer((Answer<Map<SystemStreamPartition, String>>) invocation -> {
        Map<SystemStreamPartition, String> offsets = (Map<SystemStreamPartition, String>) invocation.getArguments()[0];
        Map<SystemStreamPartition, String> nextOffsets = new HashMap<>();
        offsets.forEach((ssp, offset) -> nextOffsets.put(ssp, Long.toString(Long.valueOf(offset) + 1)));
        return nextOffsets;
    });
    SystemConsumer mockSystemConsumer = mock(SystemConsumer.class);
    Map<String, SystemConsumer> mockStoreConsumers = ImmutableMap.of(store1Name, mockSystemConsumer, store2Name, mockSystemConsumer, store3Name, mockSystemConsumer, store4Name, mockSystemConsumer);
    TransactionalStateTaskRestoreManager.registerStartingOffsets(mockTaskModel, mockStoreActions, mockStoreChangelogs, mockSystemAdmins, mockStoreConsumers, mockCurrentChangelogSSPMetadata);
    // verify that we first register upcoming offsets for each changelog ssp
    verify(mockSystemConsumer, times(1)).register(changelog1SSP, "11");
    verify(mockSystemConsumer, times(1)).register(changelog2SSP, "21");
    verify(mockSystemConsumer, times(1)).register(changelog3SSP, "31");
    verify(mockSystemConsumer, times(1)).register(changelog4SSP, "41");
    // then verify that we override the starting offsets for changelog 1 and 2
    // ensure that starting offset is inclusive if oldest
    verify(mockSystemConsumer, times(1)).register(changelog1SSP, "0");
    // and that it is next offset if not oldest
    verify(mockSystemConsumer, times(1)).register(changelog2SSP, "16");
    // and that null ending offset is ok
    verify(mockSystemConsumer, times(1)).register(changelog4SSP, "31");
    verifyNoMoreInteractions(mockSystemConsumer);
}

Example 3 with SystemConsumer

use of org.apache.samza.system.SystemConsumer in project samza by apache.

the class TestTransactionalStateTaskRestoreManager method testRegisterStartingOffsetsThrowsIfStartingGreaterThanEnding.

@Test(expected = IllegalStateException.class)
public void testRegisterStartingOffsetsThrowsIfStartingGreaterThanEnding() {
    TaskModel mockTaskModel = mock(TaskModel.class);
    TaskName taskName = new TaskName("Partition 0");
    when(mockTaskModel.getTaskName()).thenReturn(taskName);
    Partition taskChangelogPartition = new Partition(0);
    when(mockTaskModel.getChangelogPartition()).thenReturn(taskChangelogPartition);
    // tests starting offset > ending offset
    Map<String, RestoreOffsets> mockRestoreOffsets = ImmutableMap.of("store1", new RestoreOffsets("5", "0"));
    StoreActions mockStoreActions = new StoreActions(ImmutableMap.of(), ArrayListMultimap.create(), mockRestoreOffsets);
    String store1Name = "store1";
    String changelogSystemName = "system";
    String changelog1StreamName = "store1Changelog";
    SystemStream changelog1SystemStream = new SystemStream(changelogSystemName, changelog1StreamName);
    Map<String, SystemStream> mockStoreChangelogs = ImmutableMap.of(store1Name, changelog1SystemStream);
    SystemStreamPartition changelog1SSP = new SystemStreamPartition(changelog1SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog1SSPMetadata = new SystemStreamPartitionMetadata("0", "10", "11");
    Map<SystemStreamPartition, SystemStreamPartitionMetadata> mockCurrentChangelogSSPMetadata = ImmutableMap.of(changelog1SSP, changelog1SSPMetadata);
    SystemAdmins mockSystemAdmins = mock(SystemAdmins.class);
    SystemAdmin mockSystemAdmin = mock(SystemAdmin.class);
    when(mockSystemAdmins.getSystemAdmin(eq(changelogSystemName))).thenReturn(mockSystemAdmin);
    Mockito.when(mockSystemAdmin.offsetComparator(anyString(), anyString())).thenAnswer((Answer<Integer>) invocation -> {
        String offset1 = (String) invocation.getArguments()[0];
        String offset2 = (String) invocation.getArguments()[1];
        return Long.valueOf(offset1).compareTo(Long.valueOf(offset2));
    });
    Mockito.when(mockSystemAdmin.getOffsetsAfter(any())).thenAnswer((Answer<Map<SystemStreamPartition, String>>) invocation -> {
        Map<SystemStreamPartition, String> offsets = (Map<SystemStreamPartition, String>) invocation.getArguments()[0];
        Map<SystemStreamPartition, String> nextOffsets = new HashMap<>();
        offsets.forEach((ssp, offset) -> nextOffsets.put(ssp, Long.toString(Long.valueOf(offset) + 1)));
        return nextOffsets;
    });
    SystemConsumer mockSystemConsumer = mock(SystemConsumer.class);
    Map<String, SystemConsumer> mockStoreConsumers = ImmutableMap.of("store1", mockSystemConsumer);
    TransactionalStateTaskRestoreManager.registerStartingOffsets(mockTaskModel, mockStoreActions, mockStoreChangelogs, mockSystemAdmins, mockStoreConsumers, mockCurrentChangelogSSPMetadata);
    fail("Should have thrown an exception since starting offset > ending offset");
}

Example 4 with SystemConsumer

use of org.apache.samza.system.SystemConsumer in project samza by apache.

the class ContainerStorageManager method createConsumers.

/**
 *  Creates SystemConsumer objects for store restoration, creating one consumer per system.
 */
private static Map<String, SystemConsumer> createConsumers(Set<String> storeSystems, Map<String, SystemFactory> systemFactories, Config config, MetricsRegistry registry) {
    // Create one consumer for each system in use, map with one entry for each such system
    Map<String, SystemConsumer> consumers = new HashMap<>();
    // Iterate over the list of storeSystems and create one sysConsumer per system
    for (String storeSystemName : storeSystems) {
        SystemFactory systemFactory = systemFactories.get(storeSystemName);
        if (systemFactory == null) {
            throw new SamzaException("System " + storeSystemName + " does not exist in config");
        }
        consumers.put(storeSystemName, systemFactory.getConsumer(storeSystemName, config, registry));
    }
    return consumers;
}

Example 5 with SystemConsumer

use of org.apache.samza.system.SystemConsumer in project samza by apache.

the class TestCoordinatorStreamSystemConsumer method testOrderKeyRewrite.

/**
 * Verify that if a particular key-value is written, then another, then the original again,
 * that the original occurs last in the set.
 */
@Test
public void testOrderKeyRewrite() throws InterruptedException {
    final SystemStream systemStream = new SystemStream("system", "stream");
    final SystemStreamPartition ssp = new SystemStreamPartition(systemStream, new Partition(0));
    final SystemConsumer systemConsumer = mock(SystemConsumer.class);
    final List<IncomingMessageEnvelope> list = new ArrayList<>();
    SetConfig setConfig1 = new SetConfig("source", "key1", "value1");
    SetConfig setConfig2 = new SetConfig("source", "key1", "value2");
    SetConfig setConfig3 = new SetConfig("source", "key1", "value1");
    list.add(createIncomingMessageEnvelope(setConfig1, ssp));
    list.add(createIncomingMessageEnvelope(setConfig2, ssp));
    list.add(createIncomingMessageEnvelope(setConfig3, ssp));
    Map<SystemStreamPartition, List<IncomingMessageEnvelope>> messages = new HashMap<SystemStreamPartition, List<IncomingMessageEnvelope>>() {

        {
            put(ssp, list);
        }
    };
    when(systemConsumer.poll(anySet(), anyLong())).thenReturn(messages, Collections.<SystemStreamPartition, List<IncomingMessageEnvelope>>emptyMap());
    CoordinatorStreamSystemConsumer consumer = new CoordinatorStreamSystemConsumer(systemStream, systemConsumer, new SinglePartitionWithoutOffsetsSystemAdmin());
    consumer.bootstrap();
    Set<CoordinatorStreamMessage> bootstrappedMessages = consumer.getBootstrappedStream(SetConfig.TYPE);
    // First message should have been removed as a duplicate
    assertEquals(2, bootstrappedMessages.size());
    CoordinatorStreamMessage[] coordinatorStreamMessages = bootstrappedMessages.toArray(new CoordinatorStreamMessage[2]);
    assertEquals(setConfig2, coordinatorStreamMessages[0]);
    // Config 3 MUST be the last message, not config 2
    assertEquals(setConfig3, coordinatorStreamMessages[1]);
}