org.apache.samza.context.Context Java Examples

/**
 * initializes the ProjectMapFunction before any message is processed
 * @param context the {@link Context} for this task
 */
@Override
public void init(Context context) {
  this.context = context;
  this.translatorContext =
      ((SamzaSqlApplicationContext) context.getApplicationTaskContext()).getTranslatorContexts().get(queryId);
  this.project = (Project) this.translatorContext.getRelNode(projectId);
  this.expr = this.translatorContext.getExpressionCompiler().compile(project.getInputs(), project.getProjects());
  ContainerContext containerContext = context.getContainerContext();
  metricsRegistry = containerContext.getContainerMetricsRegistry();
  processingTime = new SamzaHistogram(metricsRegistry, logicalOpId, TranslatorConstants.PROCESSING_TIME_NAME);
  inputEvents = metricsRegistry.newCounter(logicalOpId, TranslatorConstants.INPUT_EVENTS_NAME);
  inputEvents.clear();
  outputEvents = metricsRegistry.newCounter(logicalOpId, TranslatorConstants.OUTPUT_EVENTS_NAME);
  outputEvents.clear();
}
 

@Test
public void testCloseDuringInitializationErrors() throws Exception {
  Context context = mock(Context.class);
  JobContext jobContext = mock(JobContext.class);
  when(context.getJobContext()).thenReturn(jobContext);
  doThrow(new RuntimeException("Failed to get config")).when(jobContext).getConfig();
  StreamOperatorTask operatorTask = new StreamOperatorTask(mock(OperatorSpecGraph.class), mock(Clock.class));
  try {
    operatorTask.init(context);
  } catch (RuntimeException e) {
    if (e instanceof NullPointerException) {
      fail("Unexpected null pointer exception");
    }
  }
  operatorTask.close();
}
 

FilterOperatorSpec(FilterFunction<M> filterFn, String opId) {
  super(new FlatMapFunction<M, M>() {
    @Override
    public Collection<M> apply(M message) {
      return new ArrayList<M>() {
        {
          if (filterFn.apply(message)) {
            this.add(message);
          }
        }
      };
    }

    @Override
    public void init(Context context) {
      filterFn.init(context);
    }

    @Override
    public void close() {
      filterFn.close();
    }
  }, OpCode.FILTER, opId);
  this.filterFn = filterFn;
}
 

MapOperatorSpec(MapFunction<M, OM> mapFn, String opId) {
  super(new FlatMapFunction<M, OM>() {
    @Override
    public Collection<OM> apply(M message) {
      return new ArrayList<OM>() {
        {
          OM r = mapFn.apply(message);
          if (r != null) {
            this.add(r);
          }
        }
      };
    }

    @Override
    public void init(Context context) {
      mapFn.init(context);
    }

    @Override
    public void close() {
      mapFn.close();
    }
  }, OpCode.MAP, opId);
  this.mapFn = mapFn;
}
 

private LocalTable createTable(boolean isTimerDisabled) {
  Map<String, String> config = new HashMap<>();
  if (isTimerDisabled) {
    config.put(MetricsConfig.METRICS_TIMER_ENABLED, "false");
  }
  Context context = mock(Context.class);
  JobContext jobContext = mock(JobContext.class);
  when(context.getJobContext()).thenReturn(jobContext);
  when(jobContext.getConfig()).thenReturn(new MapConfig(config));
  ContainerContext containerContext = mock(ContainerContext.class);
  when(context.getContainerContext()).thenReturn(containerContext);
  when(containerContext.getContainerMetricsRegistry()).thenReturn(metricsRegistry);

  LocalTable table =  new LocalTable("t1", kvStore);
  table.init(context);

  return table;
}
 

@Override
protected void handleInit(Context context) {

  KeyValueStore<TimeSeriesKey<K>, Object> store =
      (KeyValueStore<TimeSeriesKey<K>, Object>) context.getTaskContext().getStore(windowOpSpec.getOpId());

  if (initializer != null) {
    initializer.init(context);
  }

  if (keyFn != null) {
    keyFn.init(context);
  }

  // For aggregating windows, we use the store in over-write mode since we only retain the aggregated
  // value. Else, we use the store in append-mode.
  if (foldLeftFn != null) {
    foldLeftFn.init(context);
    timeSeriesStore = new TimeSeriesStoreImpl(store, false);
  } else {
    timeSeriesStore = new TimeSeriesStoreImpl(store, true);
  }
}
 

@Override
public void init(Context context) {
  this.tagToRateLimiterMap = Collections.unmodifiableMap(tagToTargetRateMap.entrySet().stream()
      .map(e -> {
        String tag = e.getKey();
        JobModel jobModel = ((TaskContextImpl) context.getTaskContext()).getJobModel();
        int numTasks = jobModel.getContainers().values().stream()
            .mapToInt(cm -> cm.getTasks().size())
            .sum();
        double effectiveRate = (double) e.getValue() / numTasks;
        TaskName taskName = context.getTaskContext().getTaskModel().getTaskName();
        LOGGER.info(String.format("Effective rate limit for task %s and tag %s is %f", taskName, tag,
            effectiveRate));
        if (effectiveRate < 1.0) {
          LOGGER.warn(String.format("Effective limit rate (%f) is very low. "
                          + "Total rate limit is %d while number of tasks is %d. Consider increasing the rate limit.",
                    effectiveRate, e.getValue(), numTasks));
        }
        return new ImmutablePair<>(tag, com.google.common.util.concurrent.RateLimiter.create(effectiveRate));
      })
      .collect(Collectors.toMap(ImmutablePair::getKey, ImmutablePair::getValue))
  );
  initialized = true;
}
 

@Override
public void init(Context context) {
  this.context = context;
  this.translatorContext = ((SamzaSqlApplicationContext) context.getApplicationTaskContext()).getTranslatorContexts().get(queryId);
  this.filter = (LogicalFilter) this.translatorContext.getRelNode(filterId);
  this.expr = this.translatorContext.getExpressionCompiler().compile(filter.getInputs(), Collections.singletonList(filter.getCondition()));
  ContainerContext containerContext = context.getContainerContext();
  metricsRegistry = containerContext.getContainerMetricsRegistry();
  processingTime = new SamzaHistogram(metricsRegistry, logicalOpId, TranslatorConstants.PROCESSING_TIME_NAME);
  inputEvents = metricsRegistry.newCounter(logicalOpId, TranslatorConstants.INPUT_EVENTS_NAME);
  inputEvents.clear();
  filteredOutEvents = metricsRegistry.newCounter(logicalOpId, TranslatorConstants.FILTERED_EVENTS_NAME);
  filteredOutEvents.clear();
  outputEvents = metricsRegistry.newCounter(logicalOpId, TranslatorConstants.OUTPUT_EVENTS_NAME);
  outputEvents.clear();
}
 

@SuppressWarnings("unchecked")
public void init(Context context) {
  this.store = (KeyValueStore<String, String>) context.getTaskContext().getStore("mystore");
  System.out.println("Contents of store: ");
  KeyValueIterator<String, String> iter = store.all();
  while (iter.hasNext()) {
    Entry<String, String> entry = iter.next();
    System.out.println(entry.getKey() + " => " + entry.getValue());
  }
  iter.close();
}
 

/** Create DoFnRunner for portable runner. */
public static <InT, FnOutT> DoFnRunner<InT, FnOutT> createPortable(
    SamzaPipelineOptions pipelineOptions,
    BagState<WindowedValue<InT>> bundledEventsBag,
    DoFnRunners.OutputManager outputManager,
    StageBundleFactory stageBundleFactory,
    TupleTag<FnOutT> mainOutputTag,
    Map<String, TupleTag<?>> idToTupleTagMap,
    Context context,
    String transformFullName) {
  final SamzaExecutionContext executionContext =
      (SamzaExecutionContext) context.getApplicationContainerContext();
  final DoFnRunner<InT, FnOutT> sdkHarnessDoFnRunner =
      new SdkHarnessDoFnRunner<>(
          outputManager, stageBundleFactory, mainOutputTag, idToTupleTagMap, bundledEventsBag);
  return DoFnRunnerWithMetrics.wrap(
      sdkHarnessDoFnRunner, executionContext.getMetricsContainer(), transformFullName);
}
 

@Override
public void init(Context context) throws Exception {
  maxMessages = context.getJobContext().getConfig().getInt("task.max.messages", 50);
  String outputSystemStreamString = context.getJobContext().getConfig().get("task.outputs", null);
  if (outputSystemStreamString == null) {
    throw new ConfigException("Missing required configuration: task.outputs");
  }
  outputSystemStream = StreamUtil.getSystemStreamFromNames(outputSystemStreamString);
}
 

public ScalarUdf createInstance(String clazz, String udfName, Context context) {
  // Configs should be same for all the UDF methods within a UDF. Hence taking the first one.
  Config udfConfig = udfMetadata.get(udfName).get(0).getUdfConfig();
  ScalarUdf scalarUdf = ReflectionUtil.getObj(clazz, ScalarUdf.class);
  scalarUdf.init(udfConfig, context);
  return scalarUdf;
}
 

@Override
public void init(Context context) {
  TranslatorContext translatorContext =
      ((SamzaSqlApplicationContext) context.getApplicationTaskContext()).getTranslatorContexts().get(queryId);
  this.msgConverter = translatorContext.getMsgConverter(tableName);
  this.relTableKeyConverter = translatorContext.getTableKeyConverter(tableName);
}
 

/**
 * {@inheritDoc}
 * Override {@link org.apache.samza.operators.functions.InitableFunction#init(Context)} to
 * get a KeyValueStore for persistence and the MetricsRegistry for metrics.
 */
@Override
public void init(Context context) {
  TaskContext taskContext = context.getTaskContext();
  store = (KeyValueStore<String, Integer>) taskContext.getStore("wikipedia-stats");
  repeatEdits = taskContext.getTaskMetricsRegistry().newCounter("edit-counters", "repeat-edits");
}
 

private Context createMockContext(TableDescriptor tableDescriptor) {
  Context context = mock(Context.class);

  ContainerContext containerContext = mock(ContainerContext.class);
  when(context.getContainerContext()).thenReturn(containerContext);

  MetricsRegistry metricsRegistry = mock(MetricsRegistry.class);
  when(metricsRegistry.newTimer(anyString(), anyString())).thenReturn(mock(Timer.class));
  when(metricsRegistry.newCounter(anyString(), anyString())).thenReturn(mock(Counter.class));
  when(containerContext.getContainerMetricsRegistry()).thenReturn(metricsRegistry);

  TaskContextImpl taskContext = mock(TaskContextImpl.class);
  when(context.getTaskContext()).thenReturn(taskContext);

  TaskName taskName = new TaskName("MyTask");
  TaskModel taskModel = mock(TaskModel.class);
  when(taskModel.getTaskName()).thenReturn(taskName);
  when(context.getTaskContext().getTaskModel()).thenReturn(taskModel);

  ContainerModel containerModel = mock(ContainerModel.class);
  when(containerModel.getTasks()).thenReturn(ImmutableMap.of(taskName, taskModel));
  when(containerContext.getContainerModel()).thenReturn(containerModel);

  String containerId = "container-1";
  JobModel jobModel = mock(JobModel.class);
  when(taskContext.getJobModel()).thenReturn(jobModel);
  when(jobModel.getContainers()).thenReturn(ImmutableMap.of(containerId, containerModel));

  JobContext jobContext = mock(JobContext.class);
  Config jobConfig = new MapConfig(tableDescriptor.toConfig(new MapConfig()));
  when(jobContext.getConfig()).thenReturn(jobConfig);
  when(context.getJobContext()).thenReturn(jobContext);

  return context;
}
 

public static Context getMockContext() {
  Context context = new MockContext();
  MetricsRegistry metricsRegistry = mock(MetricsRegistry.class);
  doAnswer(args -> new Timer((String) args.getArguments()[0])).when(metricsRegistry).newTimer(anyString(), anyString());
  doAnswer(args -> new Counter((String) args.getArguments()[0])).when(metricsRegistry).newCounter(anyString(), anyString());
  doAnswer(args -> new Gauge((String) args.getArguments()[0], 0)).when(metricsRegistry).newGauge(anyString(), any());
  doReturn(metricsRegistry).when(context.getContainerContext()).getContainerMetricsRegistry();
  return context;
}
 

@Override
public void init(Context context) {
  super.init(context);

  Preconditions.checkNotNull(this.context, "Must specify context for local tables.");

  kvStore = this.context.getTaskContext().getStore(tableId);
  Preconditions.checkNotNull(kvStore, String.format(
      "Backing store for table %s was not injected by SamzaContainer", tableId));

  logger.info("Initialized backing store for table " + tableId);
}
 

private PartialJoinFunction<Object, Object, Object, Object> createLeftJoinFn(JoinOperatorSpec joinOpSpec) {
  return new PartialJoinFunction<Object, Object, Object, Object>() {
    private final JoinFunction joinFn = joinOpSpec.getJoinFn();
    private KeyValueStore<Object, TimestampedValue<Object>> leftStreamState;

    @Override
    public Object apply(Object m, Object om) {
      return joinFn.apply(m, om);
    }

    @Override
    public Object getKey(Object message) {
      return joinFn.getFirstKey(message);
    }

    @Override
    public KeyValueStore<Object, TimestampedValue<Object>> getState() {
      return leftStreamState;
    }

    @Override
    public void init(Context context) {
      String leftStoreName = joinOpSpec.getLeftOpId();
      leftStreamState =
          (KeyValueStore<Object, TimestampedValue<Object>>) context.getTaskContext().getStore(leftStoreName);

      // user-defined joinFn should only be initialized once, so we do it only in left partial join function.
      joinFn.init(context);
    }

    @Override
    public void close() {
      // joinFn#close() must only be called once, so we do it it only in left partial join function.
      joinFn.close();
    }
  };
}
 

@Override
public void init(Context context) {
  this.store = (KeyValueStore<String, String>) context.getTaskContext().getStore(STORE_1_NAME);
  KeyValueIterator<String, String> storeEntries = store.all();
  while (storeEntries.hasNext()) {
    actualInitialStoreContents.add(storeEntries.next().getValue());
  }
  storeEntries.close();
}
 

private PartialJoinFunction<Object, Object, Object, Object> createRightJoinFn(JoinOperatorSpec joinOpSpec) {
  return new PartialJoinFunction<Object, Object, Object, Object>() {
    private final JoinFunction joinFn = joinOpSpec.getJoinFn();
    private KeyValueStore<Object, TimestampedValue<Object>> rightStreamState;

    @Override
    public Object apply(Object m, Object om) {
      return joinFn.apply(om, m);
    }

    @Override
    public Object getKey(Object message) {
      return joinFn.getSecondKey(message);
    }

    @Override
    public void init(Context context) {
      String rightStoreName = joinOpSpec.getRightOpId();
      rightStreamState =
          (KeyValueStore<Object, TimestampedValue<Object>>) context.getTaskContext().getStore(rightStoreName);

      // user-defined joinFn should only be initialized once,
      // so we do it only in left partial join function and not here again.
    }

    @Override
    public KeyValueStore<Object, TimestampedValue<Object>> getState() {
      return rightStreamState;
    }
  };
}
 

/**
 * Initialize this {@link OperatorImpl} and its user-defined functions.
 *
 * @param internalTaskContext the {@link InternalTaskContext} for the task
 */
public final void init(InternalTaskContext internalTaskContext) {
  final Context context = internalTaskContext.getContext();

  String opId = getOpImplId();

  if (initialized) {
    throw new IllegalStateException(String.format("Attempted to initialize Operator %s more than once.", opId));
  }

  if (closed) {
    throw new IllegalStateException(String.format("Attempted to initialize Operator %s after it was closed.", opId));
  }

  this.highResClock = createHighResClock(context.getJobContext().getConfig());
  registeredOperators = new HashSet<>();
  prevOperators = new HashSet<>();
  inputStreams = new HashSet<>();

  final ContainerContext containerContext = context.getContainerContext();
  final MetricsRegistry metricsRegistry = containerContext.getContainerMetricsRegistry();
  this.numMessage = metricsRegistry.newCounter(METRICS_GROUP, opId + "-messages");
  this.handleMessageNs = metricsRegistry.newTimer(METRICS_GROUP, opId + "-handle-message-ns");
  this.handleTimerNs = metricsRegistry.newTimer(METRICS_GROUP, opId + "-handle-timer-ns");

  final TaskContext taskContext =  context.getTaskContext();
  this.taskName = taskContext.getTaskModel().getTaskName();
  this.eosStates = (EndOfStreamStates) internalTaskContext.fetchObject(EndOfStreamStates.class.getName());
  this.watermarkStates = (WatermarkStates) internalTaskContext.fetchObject(WatermarkStates.class.getName());
  this.controlMessageSender = new ControlMessageSender(internalTaskContext.getStreamMetadataCache());
  this.taskModel = taskContext.getTaskModel();
  this.callbackScheduler = taskContext.getCallbackScheduler();
  handleInit(context);

  initialized = true;
}
 

private Context createMockContext() {
  MetricsRegistry metricsRegistry = mock(MetricsRegistry.class);
  doReturn(new Counter("")).when(metricsRegistry).newCounter(anyString(), anyString());
  doReturn(new Timer("")).when(metricsRegistry).newTimer(anyString(), anyString());
  Context context = new MockContext();
  doReturn(new MapConfig()).when(context.getJobContext()).getConfig();
  doReturn(metricsRegistry).when(context.getContainerContext()).getContainerMetricsRegistry();
  return context;
}
 

private <V> CouchbaseTableWriteFunction<V> createAndInit(Class<V> valueClass, Serde<V> serde, Bucket bucket,
    AsyncBucket asyncBucket) {
  when(bucket.async()).thenReturn(asyncBucket);
  PowerMockito.stub(PowerMockito.method(CouchbaseBucketRegistry.class, "getBucket", String.class, List.class,
      CouchbaseEnvironmentConfigs.class)).toReturn(bucket);
  CouchbaseTableWriteFunction<V> writeFunction =
      new CouchbaseTableWriteFunction<>(DEFAULT_BUCKET_NAME, valueClass, DEFAULT_CLUSTER_NODE).withSerde(serde);
  writeFunction.init(mock(Context.class), mock(AsyncReadWriteTable.class));
  return writeFunction;
}
 

private StreamOperatorTask createStreamOperatorTask(Clock clock, StreamApplicationDescriptorImpl graphSpec)
    throws Exception {
  Map<String, String> mapConfig = new HashMap<>();
  mapConfig.put("job.name", "jobName");
  mapConfig.put("job.id", "jobId");
  StreamTestUtils.addStreamConfigs(mapConfig, "inStream", "insystem", "instream");
  StreamTestUtils.addStreamConfigs(mapConfig, "inStream2", "insystem", "instream2");
  Context context = new MockContext(new MapConfig(mapConfig));
  TaskModel taskModel = mock(TaskModel.class);
  when(taskModel.getSystemStreamPartitions()).thenReturn(ImmutableSet
      .of(new SystemStreamPartition("insystem", "instream", new Partition(0)),
          new SystemStreamPartition("insystem", "instream2", new Partition(0))));
  when(context.getTaskContext().getTaskModel()).thenReturn(taskModel);
  when(context.getTaskContext().getTaskMetricsRegistry()).thenReturn(new MetricsRegistryMap());
  when(context.getContainerContext().getContainerMetricsRegistry()).thenReturn(new MetricsRegistryMap());
  // need to return different stores for left and right side
  IntegerSerde integerSerde = new IntegerSerde();
  TimestampedValueSerde timestampedValueSerde = new TimestampedValueSerde(new KVSerde(integerSerde, integerSerde));
  when(context.getTaskContext().getStore(eq("jobName-jobId-join-j1-L")))
      .thenReturn(new TestInMemoryStore(integerSerde, timestampedValueSerde));
  when(context.getTaskContext().getStore(eq("jobName-jobId-join-j1-R")))
      .thenReturn(new TestInMemoryStore(integerSerde, timestampedValueSerde));

  StreamOperatorTask sot = new StreamOperatorTask(graphSpec.getOperatorSpecGraph(), clock);
  sot.init(context);
  return sot;
}
 

@Override
public void init(Context context) {
  MetricsConfig metricsConfig = new MetricsConfig(context.getJobContext().getConfig());
  clock = metricsConfig.getMetricsTimerEnabled()
      ? () -> System.nanoTime()
      : () -> 0L;
  metrics = new TableMetrics(context, this, tableId);
}
 

/**
 * Constructor based on container context
 *
 * @param context {@link Context} for this task
 * @param table underlying table
 * @param tableId table Id
 */
public TableMetricsUtil(Context context, Table table, String tableId) {
  Preconditions.checkNotNull(context);
  Preconditions.checkNotNull(table);
  Preconditions.checkNotNull(tableId);

  this.metricsRegistry = context.getContainerContext().getContainerMetricsRegistry();
  this.groupName = table.getClass().getSimpleName();
  this.tableId = tableId;
}
 

@Override
public void init(Context context) {
  table.init(context);
  MetricsConfig metricsConfig = new MetricsConfig(context.getJobContext().getConfig());
  if (metricsConfig.getMetricsTimerEnabled()) {
    TableMetricsUtil tableMetricsUtil = new TableMetricsUtil(context, this, tableId);
    if (isReadRateLimited()) {
      readRateLimiter.setTimerMetric(tableMetricsUtil.newTimer("get-throttle-ns"));
    }
    if (isWriteRateLimited()) {
      writeRateLimiter.setTimerMetric(tableMetricsUtil.newTimer("put-throttle-ns"));
    }
  }
}
 

@Override
public void init(Context context) {
  super.init(context);
  asyncTable.init(context);
  if (readFn != null) {
    readFn.init(context, this);
  }
  if (writeFn != null) {
    writeFn.init(context, this);
  }
}
 

@Override
public void init(Context context) {
  table.init(context);
  final TableMetricsUtil metricsUtil = new TableMetricsUtil(context, this, tableId);

  createBatchProcessor(TableMetricsUtil.mayCreateHighResolutionClock(context.getJobContext().getConfig()),
      new BatchMetrics(metricsUtil));
}
 

@Override
public void init(Context context) {
  table.init(context);
  TableMetricsUtil metricsUtil = new TableMetricsUtil(context, this, tableId);
  if (readRetryPolicy != null) {
    readRetryMetrics = new RetryMetrics("reader", metricsUtil);
  }
  if (writeRetryPolicy != null) {
    writeRetryMetrics = new RetryMetrics("writer", metricsUtil);
  }
}