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

Example 1 with SSPMetadataCache

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

the class TestTransactionalStateTaskRestoreManager method testGetCurrentChangelogOffsets.

@Test
public void testGetCurrentChangelogOffsets() {
    // test gets metadata for all and only task store changelog SSPs
    // test all changelogs have same partition
    // test does not change returned ssp metadata
    TaskModel mockTaskModel = mock(TaskModel.class);
    when(mockTaskModel.getTaskName()).thenReturn(new TaskName("Partition 0"));
    Partition taskChangelogPartition = new Partition(0);
    when(mockTaskModel.getChangelogPartition()).thenReturn(taskChangelogPartition);
    String store1Name = "store1";
    String changelog1SystemName = "system1";
    String changelog1StreamName = "store1Changelog";
    String store2Name = "store2";
    String changelog2SystemName = "system2";
    String changelog2StreamName = "store2Changelog";
    SystemStream changelog1SystemStream = new SystemStream(changelog1SystemName, changelog1StreamName);
    SystemStream changelog2SystemStream = new SystemStream(changelog2SystemName, changelog2StreamName);
    Map<String, SystemStream> mockStoreChangelogs = ImmutableMap.of(store1Name, changelog1SystemStream, store2Name, changelog2SystemStream);
    SSPMetadataCache mockSSPMetadataCache = mock(SSPMetadataCache.class);
    SystemStreamPartition changelog1SSP = new SystemStreamPartition(changelog1SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog1SSPMetadata = new SystemStreamPartitionMetadata("0", "1", "2");
    when(mockSSPMetadataCache.getMetadata(eq(changelog1SSP))).thenReturn(changelog1SSPMetadata);
    SystemStreamPartition changelog2SSP = new SystemStreamPartition(changelog2SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog2SSPMetadata = new SystemStreamPartitionMetadata("1", "2", "3");
    when(mockSSPMetadataCache.getMetadata(eq(changelog2SSP))).thenReturn(changelog2SSPMetadata);
    Map<SystemStreamPartition, SystemStreamPartitionMetadata> currentChangelogOffsets = TransactionalStateTaskRestoreManager.getCurrentChangelogOffsets(mockTaskModel, mockStoreChangelogs, mockSSPMetadataCache);
    verify(mockSSPMetadataCache, times(1)).getMetadata(changelog1SSP);
    verify(mockSSPMetadataCache, times(1)).getMetadata(changelog2SSP);
    verifyNoMoreInteractions(mockSSPMetadataCache);
    assertEquals(2, currentChangelogOffsets.size());
    assertEquals(changelog1SSPMetadata, currentChangelogOffsets.get(changelog1SSP));
    assertEquals(changelog2SSPMetadata, currentChangelogOffsets.get(changelog2SSP));
}

Example 2 with SSPMetadataCache

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

the class TestContainerStorageManager method setUp.

/**
 * Method to create a containerStorageManager with mocked dependencies
 */
@Before
public void setUp() throws InterruptedException {
    taskRestoreMetricGauges = new HashMap<>();
    this.tasks = new HashMap<>();
    this.taskInstanceMetrics = new HashMap<>();
    // Add two mocked tasks
    addMockedTask("task 0", 0);
    addMockedTask("task 1", 1);
    // Mock container metrics
    samzaContainerMetrics = mock(SamzaContainerMetrics.class);
    when(samzaContainerMetrics.taskStoreRestorationMetrics()).thenReturn(taskRestoreMetricGauges);
    // Create a map of test changeLogSSPs
    Map<String, SystemStream> changelogSystemStreams = new HashMap<>();
    changelogSystemStreams.put(STORE_NAME, new SystemStream(SYSTEM_NAME, STREAM_NAME));
    // Create mocked storage engine factories
    Map<String, StorageEngineFactory<Object, Object>> storageEngineFactories = new HashMap<>();
    StorageEngineFactory mockStorageEngineFactory = (StorageEngineFactory<Object, Object>) mock(StorageEngineFactory.class);
    StorageEngine mockStorageEngine = mock(StorageEngine.class);
    when(mockStorageEngine.getStoreProperties()).thenReturn(new StoreProperties.StorePropertiesBuilder().setLoggedStore(true).setPersistedToDisk(true).build());
    doAnswer(invocation -> {
        return mockStorageEngine;
    }).when(mockStorageEngineFactory).getStorageEngine(anyString(), any(), any(), any(), any(), any(), any(), any(), any(), any());
    storageEngineFactories.put(STORE_NAME, mockStorageEngineFactory);
    // Add instrumentation to mocked storage engine, to record the number of store.restore() calls
    doAnswer(invocation -> {
        storeRestoreCallCount++;
        return CompletableFuture.completedFuture(null);
    }).when(mockStorageEngine).restore(any());
    // Set the mocked stores' properties to be persistent
    doAnswer(invocation -> {
        return new StoreProperties.StorePropertiesBuilder().setLoggedStore(true).build();
    }).when(mockStorageEngine).getStoreProperties();
    // Mock and setup sysconsumers
    SystemConsumer mockSystemConsumer = mock(SystemConsumer.class);
    doAnswer(invocation -> {
        systemConsumerStartCount++;
        return null;
    }).when(mockSystemConsumer).start();
    doAnswer(invocation -> {
        systemConsumerStopCount++;
        return null;
    }).when(mockSystemConsumer).stop();
    // Create mocked system factories
    Map<String, SystemFactory> systemFactories = new HashMap<>();
    // Count the number of sysConsumers created
    SystemFactory mockSystemFactory = mock(SystemFactory.class);
    doAnswer(invocation -> {
        this.systemConsumerCreationCount++;
        return mockSystemConsumer;
    }).when(mockSystemFactory).getConsumer(anyString(), any(), any());
    systemFactories.put(SYSTEM_NAME, mockSystemFactory);
    // Create mocked configs for specifying serdes
    Map<String, String> configMap = new HashMap<>();
    configMap.put("stores." + STORE_NAME + ".key.serde", "stringserde");
    configMap.put("stores." + STORE_NAME + ".msg.serde", "stringserde");
    configMap.put("stores." + STORE_NAME + ".factory", mockStorageEngineFactory.getClass().getName());
    configMap.put("stores." + STORE_NAME + ".changelog", SYSTEM_NAME + "." + STREAM_NAME);
    configMap.put("serializers.registry.stringserde.class", StringSerdeFactory.class.getName());
    configMap.put(TaskConfig.TRANSACTIONAL_STATE_RETAIN_EXISTING_STATE, "true");
    Config config = new MapConfig(configMap);
    Map<String, Serde<Object>> serdes = new HashMap<>();
    serdes.put("stringserde", mock(Serde.class));
    // Create mocked system admins
    SystemAdmin mockSystemAdmin = mock(SystemAdmin.class);
    doAnswer(new Answer<Void>() {

        public Void answer(InvocationOnMock invocation) {
            Object[] args = invocation.getArguments();
            System.out.println("called with arguments: " + Arrays.toString(args));
            return null;
        }
    }).when(mockSystemAdmin).validateStream(any());
    SystemAdmins mockSystemAdmins = mock(SystemAdmins.class);
    when(mockSystemAdmins.getSystemAdmin("kafka")).thenReturn(mockSystemAdmin);
    // Create a mocked mockStreamMetadataCache
    SystemStreamMetadata.SystemStreamPartitionMetadata sspMetadata = new SystemStreamMetadata.SystemStreamPartitionMetadata("0", "50", "51");
    Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>();
    partitionMetadata.put(new Partition(0), sspMetadata);
    partitionMetadata.put(new Partition(1), sspMetadata);
    SystemStreamMetadata systemStreamMetadata = new SystemStreamMetadata(STREAM_NAME, partitionMetadata);
    StreamMetadataCache mockStreamMetadataCache = mock(StreamMetadataCache.class);
    when(mockStreamMetadataCache.getStreamMetadata(JavaConverters.asScalaSetConverter(new HashSet<SystemStream>(changelogSystemStreams.values())).asScala().toSet(), false)).thenReturn(new scala.collection.immutable.Map.Map1(new SystemStream(SYSTEM_NAME, STREAM_NAME), systemStreamMetadata));
    CheckpointManager checkpointManager = mock(CheckpointManager.class);
    when(checkpointManager.readLastCheckpoint(any(TaskName.class))).thenReturn(new CheckpointV1(new HashMap<>()));
    SSPMetadataCache mockSSPMetadataCache = mock(SSPMetadataCache.class);
    when(mockSSPMetadataCache.getMetadata(any(SystemStreamPartition.class))).thenReturn(new SystemStreamMetadata.SystemStreamPartitionMetadata("0", "10", "11"));
    ContainerContext mockContainerContext = mock(ContainerContext.class);
    ContainerModel mockContainerModel = new ContainerModel("samza-container-test", tasks);
    when(mockContainerContext.getContainerModel()).thenReturn(mockContainerModel);
    // Reset the expected number of sysConsumer create, start and stop calls, and store.restore() calls
    this.systemConsumerCreationCount = 0;
    this.systemConsumerStartCount = 0;
    this.systemConsumerStopCount = 0;
    this.storeRestoreCallCount = 0;
    StateBackendFactory backendFactory = mock(StateBackendFactory.class);
    TaskRestoreManager restoreManager = mock(TaskRestoreManager.class);
    ArgumentCaptor<ExecutorService> restoreExecutorCaptor = ArgumentCaptor.forClass(ExecutorService.class);
    when(backendFactory.getRestoreManager(any(), any(), any(), restoreExecutorCaptor.capture(), any(), any(), any(), any(), any(), any(), any())).thenReturn(restoreManager);
    doAnswer(invocation -> {
        storeRestoreCallCount++;
        return CompletableFuture.completedFuture(null);
    }).when(restoreManager).restore();
    // Create the container storage manager
    this.containerStorageManager = new ContainerStorageManager(checkpointManager, mockContainerModel, mockStreamMetadataCache, mockSystemAdmins, changelogSystemStreams, new HashMap<>(), storageEngineFactories, systemFactories, serdes, config, taskInstanceMetrics, samzaContainerMetrics, mock(JobContext.class), mockContainerContext, ImmutableMap.of(StorageConfig.KAFKA_STATE_BACKEND_FACTORY, backendFactory), mock(Map.class), DEFAULT_LOGGED_STORE_BASE_DIR, DEFAULT_STORE_BASE_DIR, null, new SystemClock());
}