Search in sources :

Example 1 with RestoreOffsets

use of org.apache.samza.storage.TransactionalStateTaskRestoreManager.RestoreOffsets 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);
}
Also used : ArrayListMultimap(com.google.common.collect.ArrayListMultimap) ListMultimap(com.google.common.collect.ListMultimap) SSPMetadataCache(org.apache.samza.system.SSPMetadataCache) HashMap(java.util.HashMap) TaskModel(org.apache.samza.job.model.TaskModel) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Matchers.anyString(org.mockito.Matchers.anyString) FileUtil(org.apache.samza.util.FileUtil) Answer(org.mockito.stubbing.Answer) ImmutableList(com.google.common.collect.ImmutableList) SystemConsumer(org.apache.samza.system.SystemConsumer) Mockito.verifyNoMoreInteractions(org.mockito.Mockito.verifyNoMoreInteractions) SystemStream(org.apache.samza.system.SystemStream) Matchers.eq(org.mockito.Matchers.eq) Map(java.util.Map) Assert.fail(org.junit.Assert.fail) Path(java.nio.file.Path) MapConfig(org.apache.samza.config.MapConfig) ImmutableSet(com.google.common.collect.ImmutableSet) TaskName(org.apache.samza.container.TaskName) ImmutableMap(com.google.common.collect.ImmutableMap) TaskConfig(org.apache.samza.config.TaskConfig) Assert.assertNotNull(org.junit.Assert.assertNotNull) Partition(org.apache.samza.Partition) Set(java.util.Set) Assert.assertTrue(org.junit.Assert.assertTrue) Clock(org.apache.samza.util.Clock) Test(org.junit.Test) Mockito.times(org.mockito.Mockito.times) Mockito.when(org.mockito.Mockito.when) StoreActions(org.apache.samza.storage.TransactionalStateTaskRestoreManager.StoreActions) File(java.io.File) CheckpointId(org.apache.samza.checkpoint.CheckpointId) Mockito.verify(org.mockito.Mockito.verify) Matchers.any(org.mockito.Matchers.any) Mockito(org.mockito.Mockito) TaskMode(org.apache.samza.job.model.TaskMode) Mockito.never(org.mockito.Mockito.never) Assert.assertNull(org.junit.Assert.assertNull) RestoreOffsets(org.apache.samza.storage.TransactionalStateTaskRestoreManager.RestoreOffsets) SystemAdmin(org.apache.samza.system.SystemAdmin) SystemStreamPartitionMetadata(org.apache.samza.system.SystemStreamMetadata.SystemStreamPartitionMetadata) KafkaStateCheckpointMarker(org.apache.samza.checkpoint.kafka.KafkaStateCheckpointMarker) Config(org.apache.samza.config.Config) Collections(java.util.Collections) Assert.assertEquals(org.junit.Assert.assertEquals) SystemAdmins(org.apache.samza.system.SystemAdmins) Mockito.mock(org.mockito.Mockito.mock) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Partition(org.apache.samza.Partition) SystemConsumer(org.apache.samza.system.SystemConsumer) StoreActions(org.apache.samza.storage.TransactionalStateTaskRestoreManager.StoreActions) SystemStream(org.apache.samza.system.SystemStream) Matchers.anyString(org.mockito.Matchers.anyString) SystemStreamPartitionMetadata(org.apache.samza.system.SystemStreamMetadata.SystemStreamPartitionMetadata) RestoreOffsets(org.apache.samza.storage.TransactionalStateTaskRestoreManager.RestoreOffsets) TaskName(org.apache.samza.container.TaskName) SystemAdmin(org.apache.samza.system.SystemAdmin) SystemAdmins(org.apache.samza.system.SystemAdmins) HashMap(java.util.HashMap) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) TaskModel(org.apache.samza.job.model.TaskModel) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Test(org.junit.Test)

Example 2 with RestoreOffsets

use of org.apache.samza.storage.TransactionalStateTaskRestoreManager.RestoreOffsets 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");
}
Also used : ArrayListMultimap(com.google.common.collect.ArrayListMultimap) ListMultimap(com.google.common.collect.ListMultimap) SSPMetadataCache(org.apache.samza.system.SSPMetadataCache) HashMap(java.util.HashMap) TaskModel(org.apache.samza.job.model.TaskModel) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Matchers.anyString(org.mockito.Matchers.anyString) FileUtil(org.apache.samza.util.FileUtil) Answer(org.mockito.stubbing.Answer) ImmutableList(com.google.common.collect.ImmutableList) SystemConsumer(org.apache.samza.system.SystemConsumer) Mockito.verifyNoMoreInteractions(org.mockito.Mockito.verifyNoMoreInteractions) SystemStream(org.apache.samza.system.SystemStream) Matchers.eq(org.mockito.Matchers.eq) Map(java.util.Map) Assert.fail(org.junit.Assert.fail) Path(java.nio.file.Path) MapConfig(org.apache.samza.config.MapConfig) ImmutableSet(com.google.common.collect.ImmutableSet) TaskName(org.apache.samza.container.TaskName) ImmutableMap(com.google.common.collect.ImmutableMap) TaskConfig(org.apache.samza.config.TaskConfig) Assert.assertNotNull(org.junit.Assert.assertNotNull) Partition(org.apache.samza.Partition) Set(java.util.Set) Assert.assertTrue(org.junit.Assert.assertTrue) Clock(org.apache.samza.util.Clock) Test(org.junit.Test) Mockito.times(org.mockito.Mockito.times) Mockito.when(org.mockito.Mockito.when) StoreActions(org.apache.samza.storage.TransactionalStateTaskRestoreManager.StoreActions) File(java.io.File) CheckpointId(org.apache.samza.checkpoint.CheckpointId) Mockito.verify(org.mockito.Mockito.verify) Matchers.any(org.mockito.Matchers.any) Mockito(org.mockito.Mockito) TaskMode(org.apache.samza.job.model.TaskMode) Mockito.never(org.mockito.Mockito.never) Assert.assertNull(org.junit.Assert.assertNull) RestoreOffsets(org.apache.samza.storage.TransactionalStateTaskRestoreManager.RestoreOffsets) SystemAdmin(org.apache.samza.system.SystemAdmin) SystemStreamPartitionMetadata(org.apache.samza.system.SystemStreamMetadata.SystemStreamPartitionMetadata) KafkaStateCheckpointMarker(org.apache.samza.checkpoint.kafka.KafkaStateCheckpointMarker) Config(org.apache.samza.config.Config) Collections(java.util.Collections) Assert.assertEquals(org.junit.Assert.assertEquals) SystemAdmins(org.apache.samza.system.SystemAdmins) Mockito.mock(org.mockito.Mockito.mock) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Partition(org.apache.samza.Partition) SystemConsumer(org.apache.samza.system.SystemConsumer) StoreActions(org.apache.samza.storage.TransactionalStateTaskRestoreManager.StoreActions) SystemStream(org.apache.samza.system.SystemStream) Matchers.anyString(org.mockito.Matchers.anyString) SystemStreamPartitionMetadata(org.apache.samza.system.SystemStreamMetadata.SystemStreamPartitionMetadata) RestoreOffsets(org.apache.samza.storage.TransactionalStateTaskRestoreManager.RestoreOffsets) TaskName(org.apache.samza.container.TaskName) SystemAdmin(org.apache.samza.system.SystemAdmin) SystemAdmins(org.apache.samza.system.SystemAdmins) HashMap(java.util.HashMap) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) TaskModel(org.apache.samza.job.model.TaskModel) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Test(org.junit.Test)

Aggregations

ArrayListMultimap (com.google.common.collect.ArrayListMultimap)2 ImmutableList (com.google.common.collect.ImmutableList)2 ImmutableMap (com.google.common.collect.ImmutableMap)2 ImmutableSet (com.google.common.collect.ImmutableSet)2 ListMultimap (com.google.common.collect.ListMultimap)2 File (java.io.File)2 Path (java.nio.file.Path)2 Collections (java.util.Collections)2 HashMap (java.util.HashMap)2 Map (java.util.Map)2 Set (java.util.Set)2 Partition (org.apache.samza.Partition)2 CheckpointId (org.apache.samza.checkpoint.CheckpointId)2 KafkaStateCheckpointMarker (org.apache.samza.checkpoint.kafka.KafkaStateCheckpointMarker)2 Config (org.apache.samza.config.Config)2 MapConfig (org.apache.samza.config.MapConfig)2 TaskConfig (org.apache.samza.config.TaskConfig)2 TaskName (org.apache.samza.container.TaskName)2 TaskMode (org.apache.samza.job.model.TaskMode)2 TaskModel (org.apache.samza.job.model.TaskModel)2