Search in sources :

Example 6 with ContainerContext

use of org.apache.samza.context.ContainerContext in project samza by apache.

the class TestFilterTranslator method testTranslate.

@Test
public void testTranslate() throws IOException, ClassNotFoundException {
    // setup mock values to the constructor of FilterTranslator
    LogicalFilter mockFilter = PowerMockito.mock(LogicalFilter.class);
    Context mockContext = mock(Context.class);
    ContainerContext mockContainerContext = mock(ContainerContext.class);
    TranslatorContext mockTranslatorContext = mock(TranslatorContext.class);
    TestMetricsRegistryImpl metricsRegistry = new TestMetricsRegistryImpl();
    RelNode mockInput = mock(RelNode.class);
    when(mockFilter.getInput()).thenReturn(mockInput);
    when(mockInput.getId()).thenReturn(1);
    when(mockFilter.getId()).thenReturn(2);
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec<Object, SamzaSqlRelMessage> mockInputOp = mock(OperatorSpec.class);
    MessageStream<SamzaSqlRelMessage> mockStream = new MessageStreamImpl<>(mockGraph, mockInputOp);
    when(mockTranslatorContext.getMessageStream(eq(1))).thenReturn(mockStream);
    doAnswer(this.getRegisterMessageStreamAnswer()).when(mockTranslatorContext).registerMessageStream(eq(2), any(MessageStream.class));
    RexToJavaCompiler mockCompiler = mock(RexToJavaCompiler.class);
    when(mockTranslatorContext.getExpressionCompiler()).thenReturn(mockCompiler);
    Expression mockExpr = mock(Expression.class);
    when(mockCompiler.compile(any(), any())).thenReturn(mockExpr);
    when(mockContext.getContainerContext()).thenReturn(mockContainerContext);
    when(mockContainerContext.getContainerMetricsRegistry()).thenReturn(metricsRegistry);
    // Apply translate() method to verify that we are getting the correct filter operator constructed
    FilterTranslator filterTranslator = new FilterTranslator(1);
    filterTranslator.translate(mockFilter, LOGICAL_OP_ID, mockTranslatorContext);
    // make sure that context has been registered with LogicFilter and output message streams
    verify(mockTranslatorContext, times(1)).registerRelNode(2, mockFilter);
    verify(mockTranslatorContext, times(1)).registerMessageStream(2, this.getRegisteredMessageStream(2));
    when(mockTranslatorContext.getRelNode(2)).thenReturn(mockFilter);
    when(mockTranslatorContext.getMessageStream(2)).thenReturn(this.getRegisteredMessageStream(2));
    StreamOperatorSpec filterSpec = (StreamOperatorSpec) Whitebox.getInternalState(this.getRegisteredMessageStream(2), "operatorSpec");
    assertNotNull(filterSpec);
    assertEquals(filterSpec.getOpCode(), OperatorSpec.OpCode.FILTER);
    // Verify that the describe() method will establish the context for the filter function
    Map<Integer, TranslatorContext> mockContexts = new HashMap<>();
    mockContexts.put(1, mockTranslatorContext);
    when(mockContext.getApplicationTaskContext()).thenReturn(new SamzaSqlApplicationContext(mockContexts));
    filterSpec.getTransformFn().init(mockContext);
    FilterFunction filterFn = (FilterFunction) Whitebox.getInternalState(filterSpec, "filterFn");
    assertNotNull(filterFn);
    assertEquals(mockTranslatorContext, Whitebox.getInternalState(filterFn, "translatorContext"));
    assertEquals(mockFilter, Whitebox.getInternalState(filterFn, "filter"));
    assertEquals(mockExpr, Whitebox.getInternalState(filterFn, "expr"));
    // Verify MetricsRegistry works with Project
    assertEquals(1, metricsRegistry.getGauges().size());
    assertTrue(metricsRegistry.getGauges().get(LOGICAL_OP_ID).size() > 0);
    assertEquals(1, metricsRegistry.getCounters().size());
    assertEquals(3, metricsRegistry.getCounters().get(LOGICAL_OP_ID).size());
    assertEquals(0, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(0, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
    // Calling filterFn.apply() to verify the filter function is correctly applied to the input message
    SamzaSqlRelMessage mockInputMsg = new SamzaSqlRelMessage(new ArrayList<>(), new ArrayList<>(), new SamzaSqlRelMsgMetadata(0L, 0L));
    SamzaSqlExecutionContext executionContext = mock(SamzaSqlExecutionContext.class);
    DataContext dataContext = mock(DataContext.class);
    when(mockTranslatorContext.getExecutionContext()).thenReturn(executionContext);
    when(mockTranslatorContext.getDataContext()).thenReturn(dataContext);
    Object[] result = new Object[1];
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = new Boolean(true);
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    assertTrue(filterFn.apply(mockInputMsg));
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = new Boolean(false);
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    assertFalse(filterFn.apply(mockInputMsg));
    // Verify filterFn.apply() updates the MetricsRegistry metrics
    assertEquals(2, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(1, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
}
Also used : MessageStreamImpl(org.apache.samza.operators.MessageStreamImpl) FilterFunction(org.apache.samza.operators.functions.FilterFunction) HashMap(java.util.HashMap) ContainerContext(org.apache.samza.context.ContainerContext) StreamOperatorSpec(org.apache.samza.operators.spec.StreamOperatorSpec) DataContext(org.apache.calcite.DataContext) StreamApplicationDescriptorImpl(org.apache.samza.application.descriptors.StreamApplicationDescriptorImpl) SamzaSqlApplicationContext(org.apache.samza.sql.runner.SamzaSqlApplicationContext) MessageStream(org.apache.samza.operators.MessageStream) ContainerContext(org.apache.samza.context.ContainerContext) SamzaSqlExecutionContext(org.apache.samza.sql.data.SamzaSqlExecutionContext) DataContext(org.apache.calcite.DataContext) Context(org.apache.samza.context.Context) SamzaSqlApplicationContext(org.apache.samza.sql.runner.SamzaSqlApplicationContext) SamzaSqlRelMsgMetadata(org.apache.samza.sql.data.SamzaSqlRelMsgMetadata) LogicalFilter(org.apache.calcite.rel.logical.LogicalFilter) RexToJavaCompiler(org.apache.samza.sql.data.RexToJavaCompiler) TestMetricsRegistryImpl(org.apache.samza.sql.util.TestMetricsRegistryImpl) RelNode(org.apache.calcite.rel.RelNode) Expression(org.apache.samza.sql.data.Expression) SamzaSqlExecutionContext(org.apache.samza.sql.data.SamzaSqlExecutionContext) SamzaSqlRelMessage(org.apache.samza.sql.data.SamzaSqlRelMessage) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest) Test(org.junit.Test)

Example 7 with ContainerContext

use of org.apache.samza.context.ContainerContext in project samza by apache.

the class TestJoinTranslator method testTranslateStreamToTableJoin.

private void testTranslateStreamToTableJoin(boolean isRemoteTable) throws IOException, ClassNotFoundException {
    // setup mock values to the constructor of JoinTranslator
    final String logicalOpId = "sql0_join3";
    final int queryId = 0;
    LogicalJoin mockJoin = PowerMockito.mock(LogicalJoin.class);
    TranslatorContext mockTranslatorContext = mock(TranslatorContext.class);
    RelNode mockLeftInput = PowerMockito.mock(LogicalTableScan.class);
    RelNode mockRightInput = mock(RelNode.class);
    List<RelNode> inputs = new ArrayList<>();
    inputs.add(mockLeftInput);
    inputs.add(mockRightInput);
    RelOptTable mockLeftTable = mock(RelOptTable.class);
    when(mockLeftInput.getTable()).thenReturn(mockLeftTable);
    List<String> qualifiedTableName = Arrays.asList("test", "LeftTable");
    when(mockLeftTable.getQualifiedName()).thenReturn(qualifiedTableName);
    when(mockLeftInput.getId()).thenReturn(1);
    when(mockRightInput.getId()).thenReturn(2);
    when(mockJoin.getId()).thenReturn(3);
    when(mockJoin.getInputs()).thenReturn(inputs);
    when(mockJoin.getLeft()).thenReturn(mockLeftInput);
    when(mockJoin.getRight()).thenReturn(mockRightInput);
    RexCall mockJoinCondition = mock(RexCall.class);
    when(mockJoinCondition.isAlwaysTrue()).thenReturn(false);
    when(mockJoinCondition.getKind()).thenReturn(SqlKind.EQUALS);
    when(mockJoin.getCondition()).thenReturn(mockJoinCondition);
    RexInputRef mockLeftConditionInput = mock(RexInputRef.class);
    RexInputRef mockRightConditionInput = mock(RexInputRef.class);
    when(mockLeftConditionInput.getIndex()).thenReturn(0);
    when(mockRightConditionInput.getIndex()).thenReturn(0);
    List<RexNode> condOperands = new ArrayList<>();
    condOperands.add(mockLeftConditionInput);
    condOperands.add(mockRightConditionInput);
    when(mockJoinCondition.getOperands()).thenReturn(condOperands);
    RelDataType mockLeftCondDataType = mock(RelDataType.class);
    RelDataType mockRightCondDataType = mock(RelDataType.class);
    when(mockLeftCondDataType.getSqlTypeName()).thenReturn(SqlTypeName.BOOLEAN);
    when(mockRightCondDataType.getSqlTypeName()).thenReturn(SqlTypeName.BOOLEAN);
    when(mockLeftConditionInput.getType()).thenReturn(mockLeftCondDataType);
    when(mockRightConditionInput.getType()).thenReturn(mockRightCondDataType);
    RelDataType mockLeftRowType = mock(RelDataType.class);
    // ?? why ??
    when(mockLeftRowType.getFieldCount()).thenReturn(0);
    when(mockLeftInput.getRowType()).thenReturn(mockLeftRowType);
    List<String> leftFieldNames = Collections.singletonList("test_table_field1");
    List<String> rightStreamFieldNames = Collections.singletonList("test_stream_field1");
    when(mockLeftRowType.getFieldNames()).thenReturn(leftFieldNames);
    RelDataType mockRightRowType = mock(RelDataType.class);
    when(mockRightInput.getRowType()).thenReturn(mockRightRowType);
    when(mockRightRowType.getFieldNames()).thenReturn(rightStreamFieldNames);
    StreamApplicationDescriptorImpl mockAppDesc = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec<Object, SamzaSqlRelMessage> mockLeftInputOp = mock(OperatorSpec.class);
    MessageStream<SamzaSqlRelMessage> mockLeftInputStream = new MessageStreamImpl<>(mockAppDesc, mockLeftInputOp);
    when(mockTranslatorContext.getMessageStream(eq(mockLeftInput.getId()))).thenReturn(mockLeftInputStream);
    OperatorSpec<Object, SamzaSqlRelMessage> mockRightInputOp = mock(OperatorSpec.class);
    MessageStream<SamzaSqlRelMessage> mockRightInputStream = new MessageStreamImpl<>(mockAppDesc, mockRightInputOp);
    when(mockTranslatorContext.getMessageStream(eq(mockRightInput.getId()))).thenReturn(mockRightInputStream);
    when(mockTranslatorContext.getStreamAppDescriptor()).thenReturn(mockAppDesc);
    InputOperatorSpec mockInputOp = mock(InputOperatorSpec.class);
    OutputStreamImpl mockOutputStream = mock(OutputStreamImpl.class);
    when(mockInputOp.isKeyed()).thenReturn(true);
    when(mockOutputStream.isKeyed()).thenReturn(true);
    doAnswer(this.getRegisterMessageStreamAnswer()).when(mockTranslatorContext).registerMessageStream(eq(3), any(MessageStream.class));
    RexToJavaCompiler mockCompiler = mock(RexToJavaCompiler.class);
    when(mockTranslatorContext.getExpressionCompiler()).thenReturn(mockCompiler);
    Expression mockExpr = mock(Expression.class);
    when(mockCompiler.compile(any(), any())).thenReturn(mockExpr);
    if (isRemoteTable) {
        doAnswer(this.getRegisteredTableAnswer()).when(mockAppDesc).getTable(any(RemoteTableDescriptor.class));
    } else {
        IntermediateMessageStreamImpl mockPartitionedStream = new IntermediateMessageStreamImpl(mockAppDesc, mockInputOp, mockOutputStream);
        when(mockAppDesc.getIntermediateStream(any(String.class), any(Serde.class), eq(false))).thenReturn(mockPartitionedStream);
        doAnswer(this.getRegisteredTableAnswer()).when(mockAppDesc).getTable(any(RocksDbTableDescriptor.class));
    }
    when(mockJoin.getJoinType()).thenReturn(JoinRelType.INNER);
    SamzaSqlExecutionContext mockExecutionContext = mock(SamzaSqlExecutionContext.class);
    when(mockTranslatorContext.getExecutionContext()).thenReturn(mockExecutionContext);
    SamzaSqlApplicationConfig mockAppConfig = mock(SamzaSqlApplicationConfig.class);
    when(mockExecutionContext.getSamzaSqlApplicationConfig()).thenReturn(mockAppConfig);
    Map<String, SqlIOConfig> ssConfigBySource = mock(HashMap.class);
    when(mockAppConfig.getInputSystemStreamConfigBySource()).thenReturn(ssConfigBySource);
    SqlIOConfig mockIOConfig = mock(SqlIOConfig.class);
    TableDescriptor mockTableDesc;
    if (isRemoteTable) {
        mockTableDesc = mock(RemoteTableDescriptor.class);
    } else {
        mockTableDesc = mock(RocksDbTableDescriptor.class);
    }
    when(ssConfigBySource.get(String.join(".", qualifiedTableName))).thenReturn(mockIOConfig);
    when(mockIOConfig.getTableDescriptor()).thenReturn(Optional.of(mockTableDesc));
    JoinTranslator joinTranslator = new JoinTranslator(logicalOpId, "", queryId);
    // Verify Metrics Works with Join
    Context mockContext = mock(Context.class);
    ContainerContext mockContainerContext = mock(ContainerContext.class);
    TestMetricsRegistryImpl testMetricsRegistryImpl = new TestMetricsRegistryImpl();
    when(mockContext.getContainerContext()).thenReturn(mockContainerContext);
    when(mockContainerContext.getContainerMetricsRegistry()).thenReturn(testMetricsRegistryImpl);
    TranslatorInputMetricsMapFunction inputMetricsMF = joinTranslator.getInputMetricsMF();
    assertNotNull(inputMetricsMF);
    inputMetricsMF.init(mockContext);
    TranslatorOutputMetricsMapFunction outputMetricsMF = joinTranslator.getOutputMetricsMF();
    assertNotNull(outputMetricsMF);
    outputMetricsMF.init(mockContext);
    assertEquals(1, testMetricsRegistryImpl.getCounters().size());
    assertEquals(2, testMetricsRegistryImpl.getCounters().get(logicalOpId).size());
    assertEquals(0, testMetricsRegistryImpl.getCounters().get(logicalOpId).get(0).getCount());
    assertEquals(0, testMetricsRegistryImpl.getCounters().get(logicalOpId).get(1).getCount());
    assertEquals(1, testMetricsRegistryImpl.getGauges().size());
    // Apply translate() method to verify that we are getting the correct map operator constructed
    joinTranslator.translate(mockJoin, mockTranslatorContext);
    // make sure that context has been registered with LogicFilter and output message streams
    verify(mockTranslatorContext, times(1)).registerMessageStream(3, this.getRegisteredMessageStream(3));
    when(mockTranslatorContext.getRelNode(3)).thenReturn(mockJoin);
    when(mockTranslatorContext.getMessageStream(3)).thenReturn(this.getRegisteredMessageStream(3));
    StreamTableJoinOperatorSpec joinSpec = (StreamTableJoinOperatorSpec) Whitebox.getInternalState(this.getRegisteredMessageStream(3), "operatorSpec");
    assertNotNull(joinSpec);
    assertEquals(joinSpec.getOpCode(), OperatorSpec.OpCode.JOIN);
    // Verify joinSpec has the corresponding setup
    StreamTableJoinFunction joinFn = joinSpec.getJoinFn();
    assertNotNull(joinFn);
    if (isRemoteTable) {
        assertTrue(joinFn instanceof SamzaSqlRemoteTableJoinFunction);
    } else {
        assertTrue(joinFn instanceof SamzaSqlLocalTableJoinFunction);
    }
    assertTrue(Whitebox.getInternalState(joinFn, "isTablePosOnRight").equals(false));
    assertEquals(Collections.singletonList(0), Whitebox.getInternalState(joinFn, "streamFieldIds"));
    assertEquals(leftFieldNames, Whitebox.getInternalState(joinFn, "tableFieldNames"));
    List<String> outputFieldNames = new ArrayList<>();
    outputFieldNames.addAll(leftFieldNames);
    outputFieldNames.addAll(rightStreamFieldNames);
    assertEquals(outputFieldNames, Whitebox.getInternalState(joinFn, "outFieldNames"));
}
Also used : Serde(org.apache.samza.serializers.Serde) IntermediateMessageStreamImpl(org.apache.samza.operators.stream.IntermediateMessageStreamImpl) MessageStreamImpl(org.apache.samza.operators.MessageStreamImpl) OutputStreamImpl(org.apache.samza.operators.spec.OutputStreamImpl) ArrayList(java.util.ArrayList) RelDataType(org.apache.calcite.rel.type.RelDataType) RexCall(org.apache.calcite.rex.RexCall) ContainerContext(org.apache.samza.context.ContainerContext) StreamApplicationDescriptorImpl(org.apache.samza.application.descriptors.StreamApplicationDescriptorImpl) RocksDbTableDescriptor(org.apache.samza.storage.kv.descriptors.RocksDbTableDescriptor) MessageStream(org.apache.samza.operators.MessageStream) StreamTableJoinFunction(org.apache.samza.operators.functions.StreamTableJoinFunction) SqlIOConfig(org.apache.samza.sql.interfaces.SqlIOConfig) SamzaSqlExecutionContext(org.apache.samza.sql.data.SamzaSqlExecutionContext) Context(org.apache.samza.context.Context) ContainerContext(org.apache.samza.context.ContainerContext) InputOperatorSpec(org.apache.samza.operators.spec.InputOperatorSpec) SamzaSqlApplicationConfig(org.apache.samza.sql.runner.SamzaSqlApplicationConfig) RemoteTableDescriptor(org.apache.samza.table.descriptors.RemoteTableDescriptor) IntermediateMessageStreamImpl(org.apache.samza.operators.stream.IntermediateMessageStreamImpl) RexToJavaCompiler(org.apache.samza.sql.data.RexToJavaCompiler) TableDescriptor(org.apache.samza.table.descriptors.TableDescriptor) RocksDbTableDescriptor(org.apache.samza.storage.kv.descriptors.RocksDbTableDescriptor) RemoteTableDescriptor(org.apache.samza.table.descriptors.RemoteTableDescriptor) TestMetricsRegistryImpl(org.apache.samza.sql.util.TestMetricsRegistryImpl) RelNode(org.apache.calcite.rel.RelNode) Expression(org.apache.samza.sql.data.Expression) LogicalJoin(org.apache.calcite.rel.logical.LogicalJoin) RexInputRef(org.apache.calcite.rex.RexInputRef) SamzaSqlExecutionContext(org.apache.samza.sql.data.SamzaSqlExecutionContext) RelOptTable(org.apache.calcite.plan.RelOptTable) StreamTableJoinOperatorSpec(org.apache.samza.operators.spec.StreamTableJoinOperatorSpec) SamzaSqlRelMessage(org.apache.samza.sql.data.SamzaSqlRelMessage) RexNode(org.apache.calcite.rex.RexNode)

Example 8 with ContainerContext

use of org.apache.samza.context.ContainerContext in project samza by apache.

the class StorageRecovery method getContainerStorageManagers.

/**
 * create one TaskStorageManager for each task. Add all of them to the
 * List<TaskStorageManager>
 */
@SuppressWarnings("rawtypes")
private void getContainerStorageManagers() {
    Set<String> factoryClasses = new StorageConfig(jobConfig).getRestoreFactories();
    Map<String, StateBackendFactory> stateBackendFactories = factoryClasses.stream().collect(Collectors.toMap(factoryClass -> factoryClass, factoryClass -> ReflectionUtil.getObj(factoryClass, StateBackendFactory.class)));
    Clock clock = SystemClock.instance();
    StreamMetadataCache streamMetadataCache = new StreamMetadataCache(systemAdmins, 5000, clock);
    // don't worry about prefetching for this; looks like the tool doesn't flush to offset files anyways
    Map<String, SystemFactory> systemFactories = new SystemConfig(jobConfig).getSystemFactories();
    CheckpointManager checkpointManager = new TaskConfig(jobConfig).getCheckpointManager(new MetricsRegistryMap()).orElse(null);
    for (ContainerModel containerModel : containers.values()) {
        ContainerContext containerContext = new ContainerContextImpl(containerModel, new MetricsRegistryMap());
        ContainerStorageManager containerStorageManager = new ContainerStorageManager(checkpointManager, containerModel, streamMetadataCache, systemAdmins, changeLogSystemStreams, new HashMap<>(), storageEngineFactories, systemFactories, this.getSerdes(), jobConfig, new HashMap<>(), new SamzaContainerMetrics(containerModel.getId(), new MetricsRegistryMap(), ""), JobContextImpl.fromConfigWithDefaults(jobConfig, jobModel), containerContext, stateBackendFactories, new HashMap<>(), storeBaseDir, storeBaseDir, null, new SystemClock());
        this.containerStorageManagers.put(containerModel.getId(), containerStorageManager);
    }
}
Also used : StreamMetadataCache(org.apache.samza.system.StreamMetadataCache) CoordinatorStreamStore(org.apache.samza.coordinator.metadatastore.CoordinatorStreamStore) JobContextImpl(org.apache.samza.context.JobContextImpl) LoggerFactory(org.slf4j.LoggerFactory) HashMap(java.util.HashMap) TaskModel(org.apache.samza.job.model.TaskModel) Serde(org.apache.samza.serializers.Serde) ContainerContextImpl(org.apache.samza.context.ContainerContextImpl) JobModelManager(org.apache.samza.coordinator.JobModelManager) SerdeFactory(org.apache.samza.serializers.SerdeFactory) CheckpointManager(org.apache.samza.checkpoint.CheckpointManager) SamzaContainerMetrics(org.apache.samza.container.SamzaContainerMetrics) SystemStream(org.apache.samza.system.SystemStream) Map(java.util.Map) SystemConfig(org.apache.samza.config.SystemConfig) StreamUtil(org.apache.samza.util.StreamUtil) JobModel(org.apache.samza.job.model.JobModel) StorageConfig(org.apache.samza.config.StorageConfig) Logger(org.slf4j.Logger) SerializerConfig(org.apache.samza.config.SerializerConfig) TaskConfig(org.apache.samza.config.TaskConfig) ContainerContext(org.apache.samza.context.ContainerContext) Set(java.util.Set) SystemFactory(org.apache.samza.system.SystemFactory) Clock(org.apache.samza.util.Clock) Collectors(java.util.stream.Collectors) File(java.io.File) SamzaException(org.apache.samza.SamzaException) List(java.util.List) SystemClock(org.apache.samza.util.SystemClock) ReflectionUtil(org.apache.samza.util.ReflectionUtil) ContainerModel(org.apache.samza.job.model.ContainerModel) Optional(java.util.Optional) Config(org.apache.samza.config.Config) CoordinatorStreamUtil(org.apache.samza.util.CoordinatorStreamUtil) MetricsRegistryMap(org.apache.samza.metrics.MetricsRegistryMap) SystemAdmins(org.apache.samza.system.SystemAdmins) StreamMetadataCache(org.apache.samza.system.StreamMetadataCache) SystemFactory(org.apache.samza.system.SystemFactory) SystemConfig(org.apache.samza.config.SystemConfig) SystemClock(org.apache.samza.util.SystemClock) StorageConfig(org.apache.samza.config.StorageConfig) CheckpointManager(org.apache.samza.checkpoint.CheckpointManager) TaskConfig(org.apache.samza.config.TaskConfig) ContainerContextImpl(org.apache.samza.context.ContainerContextImpl) Clock(org.apache.samza.util.Clock) SystemClock(org.apache.samza.util.SystemClock) ContainerModel(org.apache.samza.job.model.ContainerModel) ContainerContext(org.apache.samza.context.ContainerContext) MetricsRegistryMap(org.apache.samza.metrics.MetricsRegistryMap) SamzaContainerMetrics(org.apache.samza.container.SamzaContainerMetrics)

Example 9 with ContainerContext

use of org.apache.samza.context.ContainerContext in project samza by apache.

the class TestContainerStorageManager method testNoConfiguredDurableStores.

@Test
public void testNoConfiguredDurableStores() 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 mocked configs for specifying serdes
    Map<String, String> configMap = new HashMap<>();
    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));
    CheckpointManager checkpointManager = mock(CheckpointManager.class);
    when(checkpointManager.readLastCheckpoint(any(TaskName.class))).thenReturn(new CheckpointV1(new HashMap<>()));
    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);
    when(backendFactory.getRestoreManager(any(), any(), any(), any(), any(), any(), any(), any(), any(), any(), any())).thenReturn(restoreManager);
    doAnswer(invocation -> {
        storeRestoreCallCount++;
        return CompletableFuture.completedFuture(null);
    }).when(restoreManager).restore();
    // Create the container storage manager
    ContainerStorageManager containerStorageManager = new ContainerStorageManager(checkpointManager, mockContainerModel, mock(StreamMetadataCache.class), mock(SystemAdmins.class), new HashMap<>(), new HashMap<>(), new HashMap<>(), new HashMap<>(), serdes, config, taskInstanceMetrics, samzaContainerMetrics, mock(JobContext.class), mockContainerContext, new HashMap<>(), mock(Map.class), DEFAULT_LOGGED_STORE_BASE_DIR, DEFAULT_STORE_BASE_DIR, null, new SystemClock());
    containerStorageManager.start();
    containerStorageManager.shutdown();
    for (Gauge gauge : taskRestoreMetricGauges.values()) {
        Assert.assertTrue("Restoration time gauge value should never be invoked", mockingDetails(gauge).getInvocations().size() == 0);
    }
    Assert.assertEquals("Store restore count should be 2 because there are 0 stores", 0, this.storeRestoreCallCount);
    Assert.assertEquals(0, this.systemConsumerCreationCount);
    Assert.assertEquals(0, this.systemConsumerStopCount);
    Assert.assertEquals(0, this.systemConsumerStartCount);
}
Also used : Serde(org.apache.samza.serializers.Serde) StreamMetadataCache(org.apache.samza.system.StreamMetadataCache) StringSerdeFactory(org.apache.samza.serializers.StringSerdeFactory) SystemClock(org.apache.samza.util.SystemClock) HashMap(java.util.HashMap) MapConfig(org.apache.samza.config.MapConfig) StorageConfig(org.apache.samza.config.StorageConfig) Config(org.apache.samza.config.Config) TaskConfig(org.apache.samza.config.TaskConfig) CheckpointManager(org.apache.samza.checkpoint.CheckpointManager) ContainerModel(org.apache.samza.job.model.ContainerModel) Gauge(org.apache.samza.metrics.Gauge) ContainerContext(org.apache.samza.context.ContainerContext) TaskName(org.apache.samza.container.TaskName) CheckpointV1(org.apache.samza.checkpoint.CheckpointV1) MapConfig(org.apache.samza.config.MapConfig) JobContext(org.apache.samza.context.JobContext) SystemAdmins(org.apache.samza.system.SystemAdmins) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) HashMap(java.util.HashMap) SamzaContainerMetrics(org.apache.samza.container.SamzaContainerMetrics) Test(org.junit.Test)

Example 10 with ContainerContext

use of org.apache.samza.context.ContainerContext 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());
}
Also used : Serde(org.apache.samza.serializers.Serde) StreamMetadataCache(org.apache.samza.system.StreamMetadataCache) SystemConsumer(org.apache.samza.system.SystemConsumer) SystemFactory(org.apache.samza.system.SystemFactory) StringSerdeFactory(org.apache.samza.serializers.StringSerdeFactory) HashMap(java.util.HashMap) MapConfig(org.apache.samza.config.MapConfig) StorageConfig(org.apache.samza.config.StorageConfig) Config(org.apache.samza.config.Config) TaskConfig(org.apache.samza.config.TaskConfig) ContainerModel(org.apache.samza.job.model.ContainerModel) ContainerContext(org.apache.samza.context.ContainerContext) CheckpointV1(org.apache.samza.checkpoint.CheckpointV1) MapConfig(org.apache.samza.config.MapConfig) SystemAdmins(org.apache.samza.system.SystemAdmins) HashSet(java.util.HashSet) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Partition(org.apache.samza.Partition) SystemClock(org.apache.samza.util.SystemClock) SystemStream(org.apache.samza.system.SystemStream) CheckpointManager(org.apache.samza.checkpoint.CheckpointManager) SystemStreamMetadata(org.apache.samza.system.SystemStreamMetadata) TaskName(org.apache.samza.container.TaskName) InvocationOnMock(org.mockito.invocation.InvocationOnMock) SSPMetadataCache(org.apache.samza.system.SSPMetadataCache) ExecutorService(java.util.concurrent.ExecutorService) SystemAdmin(org.apache.samza.system.SystemAdmin) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) HashMap(java.util.HashMap) SamzaContainerMetrics(org.apache.samza.container.SamzaContainerMetrics) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Before(org.junit.Before)

Aggregations

ContainerContext (org.apache.samza.context.ContainerContext)17 HashMap (java.util.HashMap)11 JobContext (org.apache.samza.context.JobContext)9 Context (org.apache.samza.context.Context)8 Config (org.apache.samza.config.Config)6 MapConfig (org.apache.samza.config.MapConfig)6 ContainerModel (org.apache.samza.job.model.ContainerModel)6 Map (java.util.Map)5 SamzaSqlExecutionContext (org.apache.samza.sql.data.SamzaSqlExecutionContext)5 CheckpointManager (org.apache.samza.checkpoint.CheckpointManager)4 StorageConfig (org.apache.samza.config.StorageConfig)4 TaskConfig (org.apache.samza.config.TaskConfig)4 SamzaContainerMetrics (org.apache.samza.container.SamzaContainerMetrics)4 TaskContext (org.apache.samza.context.TaskContext)4 TaskModel (org.apache.samza.job.model.TaskModel)4 MetricsRegistry (org.apache.samza.metrics.MetricsRegistry)4 Serde (org.apache.samza.serializers.Serde)4 StreamMetadataCache (org.apache.samza.system.StreamMetadataCache)4 SystemAdmins (org.apache.samza.system.SystemAdmins)4 File (java.io.File)3