org.apache.beam.sdk.metrics.MetricsFilter Java Examples

@Override
public MetricResults metrics() {
  return new MetricResults() {
    @Override
    public MetricQueryResults queryMetrics(@Nullable MetricsFilter filter) {
      return new MetricQueryResults() {
        @Override
        public Iterable<MetricResult<Long>> getCounters() {
          return Collections.emptyList();
        }

        @Override
        public Iterable<MetricResult<DistributionResult>> getDistributions() {
          return Collections.emptyList();
        }

        @Override
        public Iterable<MetricResult<GaugeResult>> getGauges() {
          return Collections.emptyList();
        }
      };
    }
  };
}
 

@Override
public MetricQueryResults queryMetrics(@Nullable MetricsFilter filter) {
  ImmutableList.Builder<MetricResult<Long>> counterResults = ImmutableList.builder();
  for (Entry<MetricKey, DirectMetric<Long, Long>> counter : counters.entries()) {
    maybeExtractResult(filter, counterResults, counter);
  }
  ImmutableList.Builder<MetricResult<DistributionResult>> distributionResults =
      ImmutableList.builder();
  for (Entry<MetricKey, DirectMetric<DistributionData, DistributionResult>> distribution :
      distributions.entries()) {
    maybeExtractResult(filter, distributionResults, distribution);
  }
  ImmutableList.Builder<MetricResult<GaugeResult>> gaugeResults = ImmutableList.builder();
  for (Entry<MetricKey, DirectMetric<GaugeData, GaugeResult>> gauge : gauges.entries()) {
    maybeExtractResult(filter, gaugeResults, gauge);
  }

  return MetricQueryResults.create(
      counterResults.build(), distributionResults.build(), gaugeResults.build());
}
 

@Test
@Category(NeedsRunner.class)
public void testWordCountSimple() {
  PCollection<KV<String, Long>> pc =
      pipeline.apply(Create.of(INPUT_STRS)).apply(new CountWords());
  PAssert.that(pc).containsInAnyOrder(KV.of("hello", 2L), KV.of(("world"), 1L));
  PipelineResult result = pipeline.run();
  result.waitUntilFinish();

  Map<String, Long> expectedCounters = new HashMap<>();
  expectedCounters.put("emptyLines", 2L);
  for (MetricResult c :
      result.metrics().queryMetrics(MetricsFilter.builder().build()).getCounters()) {
    String name = c.getName().getName();
    if (expectedCounters.containsKey(name)) {
      assertEquals(expectedCounters.get(name), c.getCommitted());
      expectedCounters.remove(name);
    }
  }
  assertTrue(expectedCounters.isEmpty());
}
 

/**
 * Return the current value for a long counter, or -1 if can't be retrieved. Note this uses only
 * attempted metrics because some runners don't support committed metrics.
 */
public long getCounterMetric(String name) {
  MetricQueryResults metrics =
      result
          .metrics()
          .queryMetrics(
              MetricsFilter.builder()
                  .addNameFilter(MetricNameFilter.named(namespace, name))
                  .build());
  Iterable<MetricResult<Long>> counters = metrics.getCounters();

  checkIfMetricResultIsUnique(name, counters);

  try {
    MetricResult<Long> metricResult = counters.iterator().next();
    return metricResult.getAttempted();
  } catch (NoSuchElementException e) {
    LOG.error("Failed to get metric {}, from namespace {}", name, namespace);
  }
  return ERRONEOUS_METRIC_VALUE;
}
 

@Override
public MetricQueryResults queryMetrics(@Nullable MetricsFilter filter) {
  List<MetricUpdate> metricUpdates;
  ImmutableList<MetricResult<Long>> counters = ImmutableList.of();
  ImmutableList<MetricResult<DistributionResult>> distributions = ImmutableList.of();
  ImmutableList<MetricResult<GaugeResult>> gauges = ImmutableList.of();
  JobMetrics jobMetrics;
  try {
    jobMetrics = getJobMetrics();
  } catch (IOException e) {
    LOG.warn("Unable to query job metrics.\n");
    return MetricQueryResults.create(counters, distributions, gauges);
  }
  metricUpdates = firstNonNull(jobMetrics.getMetrics(), Collections.emptyList());
  return populateMetricQueryResults(metricUpdates, filter);
}
 

@Test
public void testCounterCommittedUnsupportedInAttemptedAccumulatedMetricResults() {
  MetricsContainerStepMap attemptedMetrics = new MetricsContainerStepMap();
  attemptedMetrics.update(STEP1, metricsContainer);
  MetricResults metricResults = asAttemptedOnlyMetricResults(attemptedMetrics);

  MetricQueryResults step1res =
      metricResults.queryMetrics(MetricsFilter.builder().addStep(STEP1).build());

  thrown.expect(UnsupportedOperationException.class);
  thrown.expectMessage("This runner does not currently support committed metrics results.");

  assertCounter(COUNTER_NAME, step1res, STEP1, VALUE, true);
}
 

@Test
public void testDistributionCommittedUnsupportedInAttemptedAccumulatedMetricResults() {
  MetricsContainerStepMap attemptedMetrics = new MetricsContainerStepMap();
  attemptedMetrics.update(STEP1, metricsContainer);
  MetricResults metricResults = asAttemptedOnlyMetricResults(attemptedMetrics);

  MetricQueryResults step1res =
      metricResults.queryMetrics(MetricsFilter.builder().addStep(STEP1).build());

  thrown.expect(UnsupportedOperationException.class);
  thrown.expectMessage("This runner does not currently support committed metrics results.");

  assertDistribution(
      DISTRIBUTION_NAME, step1res, STEP1, DistributionResult.IDENTITY_ELEMENT, true);
}
 

@Test
public void testGaugeCommittedUnsupportedInAttemptedAccumulatedMetricResults() {
  MetricsContainerStepMap attemptedMetrics = new MetricsContainerStepMap();
  attemptedMetrics.update(STEP1, metricsContainer);
  MetricResults metricResults = asAttemptedOnlyMetricResults(attemptedMetrics);

  MetricQueryResults step1res =
      metricResults.queryMetrics(MetricsFilter.builder().addStep(STEP1).build());

  thrown.expect(UnsupportedOperationException.class);
  thrown.expectMessage("This runner does not currently support committed metrics results.");

  assertGauge(GAUGE_NAME, step1res, STEP1, GaugeResult.empty(), true);
}
 

@Override
public MetricQueryResults queryMetrics(MetricsFilter filter) {
  return MetricQueryResults.create(
      Iterables.filter(
          this.counters, (counter) -> MetricFiltering.matches(filter, counter.getKey())),
      Iterables.filter(
          this.distributions,
          (distribution) -> MetricFiltering.matches(filter, distribution.getKey())),
      Iterables.filter(this.gauges, (gauge) -> MetricFiltering.matches(filter, gauge.getKey())));
}
 

@Override
public synchronized MetricQueryResults queryMetrics(@Nullable MetricsFilter filter) {
  if (metricsAccumulator != null) {
    updateLocalMetrics(metricsAccumulator);
  }
  return new QueryResults(
      counters.filter(filter), distributions.filter(filter), gauges.filter(filter));
}
 

@Test
@Category(StreamingTest.class)
public void testUnboundedSourceMetrics() {
  final long minElements = 1000;

  // Use a GenerateSequence for the UnboundedSequence, but push the watermark to infinity at
  // minElements to let the test pipeline cleanly shut it down.  Shutdown will occur shortly
  // afterwards, but at least minElements will be reported in the metrics.
  PCollection<Long> pc =
      pipeline.apply(
          GenerateSequence.from(1)
              .withRate(minElements / 10, Duration.millis(500L))
              .withTimestampFn(
                  t -> t < minElements ? Instant.now() : BoundedWindow.TIMESTAMP_MAX_VALUE));
  assertThat(pc.isBounded(), is(PCollection.IsBounded.UNBOUNDED));

  PipelineResult pipelineResult = pipeline.run();

  MetricQueryResults metrics =
      pipelineResult
          .metrics()
          .queryMetrics(
              MetricsFilter.builder()
                  .addNameFilter(
                      MetricNameFilter.named(
                          ELEMENTS_READ.getNamespace(), ELEMENTS_READ.getName()))
                  .build());

  assertThat(
      metrics.getCounters(),
      hasItem(
          metricsResult(
              ELEMENTS_READ.getNamespace(),
              ELEMENTS_READ.getName(),
              "GenerateSequence/Read(UnboundedCountingSource)",
              greaterThanOrEqualTo(minElements),
              false)));
}
 

/**
 * Update Flink's internal metrics ({@link this#flinkCounterCache}) with the latest metrics for a
 * given step.
 */
void updateMetrics(String stepName) {
  MetricResults metricResults = asAttemptedOnlyMetricResults(metricsContainers);
  MetricQueryResults metricQueryResults =
      metricResults.queryMetrics(MetricsFilter.builder().addStep(stepName).build());
  updateCounters(metricQueryResults.getCounters());
  updateDistributions(metricQueryResults.getDistributions());
  updateGauge(metricQueryResults.getGauges());
}
 

private DataflowMetricQueryResultsFactory(
    DataflowPipelineJob dataflowPipelineJob,
    Iterable<MetricUpdate> metricUpdates,
    MetricsFilter filter) {
  this.dataflowPipelineJob = dataflowPipelineJob;
  this.metricUpdates = metricUpdates;
  this.filter = filter;

  tentativeByName = new HashMap<>();
  committedByName = new HashMap<>();
  metricHashKeys = new HashSet<>();
}
 

private <ResultT> void maybeExtractResult(
    MetricsFilter filter,
    ImmutableList.Builder<MetricResult<ResultT>> resultsBuilder,
    Map.Entry<MetricKey, ? extends DirectMetric<?, ResultT>> entry) {
  if (MetricFiltering.matches(filter, entry.getKey())) {
    resultsBuilder.add(
        MetricResult.create(
            entry.getKey(),
            entry.getValue().extractCommitted(),
            entry.getValue().extractLatestAttempted()));
  }
}
 

@SuppressWarnings("unchecked")
@Test
public void testApplyAttemptedCountersQueryOneNamespace() {
  metrics.updatePhysical(
      bundle1,
      MetricUpdates.create(
          ImmutableList.of(
              MetricUpdate.create(MetricKey.create("step1", NAME1), 5L),
              MetricUpdate.create(MetricKey.create("step1", NAME3), 8L)),
          ImmutableList.of(),
          ImmutableList.of()));
  metrics.updatePhysical(
      bundle1,
      MetricUpdates.create(
          ImmutableList.of(
              MetricUpdate.create(MetricKey.create("step2", NAME1), 7L),
              MetricUpdate.create(MetricKey.create("step1", NAME3), 4L)),
          ImmutableList.of(),
          ImmutableList.of()));

  MetricQueryResults results =
      metrics.queryMetrics(MetricsFilter.builder().addNameFilter(inNamespace("ns1")).build());

  assertThat(
      results.getCounters(),
      containsInAnyOrder(
          attemptedMetricsResult("ns1", "name1", "step1", 5L),
          attemptedMetricsResult("ns1", "name1", "step2", 7L)));

  assertThat(
      results.getCounters(),
      containsInAnyOrder(
          committedMetricsResult("ns1", "name1", "step1", 0L),
          committedMetricsResult("ns1", "name1", "step2", 0L)));
}
 

@SuppressWarnings("unchecked")
@Test
public void testApplyAttemptedQueryCompositeScope() {
  metrics.updatePhysical(
      bundle1,
      MetricUpdates.create(
          ImmutableList.of(
              MetricUpdate.create(MetricKey.create("Outer1/Inner1", NAME1), 5L),
              MetricUpdate.create(MetricKey.create("Outer1/Inner2", NAME1), 8L)),
          ImmutableList.of(),
          ImmutableList.of()));
  metrics.updatePhysical(
      bundle1,
      MetricUpdates.create(
          ImmutableList.of(
              MetricUpdate.create(MetricKey.create("Outer1/Inner1", NAME1), 12L),
              MetricUpdate.create(MetricKey.create("Outer2/Inner2", NAME1), 18L)),
          ImmutableList.of(),
          ImmutableList.of()));

  MetricQueryResults results =
      metrics.queryMetrics(MetricsFilter.builder().addStep("Outer1").build());

  assertThat(
      results.getCounters(),
      containsInAnyOrder(
          attemptedMetricsResult("ns1", "name1", "Outer1/Inner1", 12L),
          attemptedMetricsResult("ns1", "name1", "Outer1/Inner2", 8L)));

  assertThat(
      results.getCounters(),
      containsInAnyOrder(
          committedMetricsResult("ns1", "name1", "Outer1/Inner1", 0L),
          committedMetricsResult("ns1", "name1", "Outer1/Inner2", 0L)));
}
 

@SuppressWarnings("unchecked")
@Test
public void testPartialScopeMatchingInMetricsQuery() {
  metrics.updatePhysical(
      bundle1,
      MetricUpdates.create(
          ImmutableList.of(
              MetricUpdate.create(MetricKey.create("Top1/Outer1/Inner1", NAME1), 5L),
              MetricUpdate.create(MetricKey.create("Top1/Outer1/Inner2", NAME1), 8L)),
          ImmutableList.of(),
          ImmutableList.of()));
  metrics.updatePhysical(
      bundle1,
      MetricUpdates.create(
          ImmutableList.of(
              MetricUpdate.create(MetricKey.create("Top2/Outer1/Inner1", NAME1), 12L),
              MetricUpdate.create(MetricKey.create("Top1/Outer2/Inner2", NAME1), 18L)),
          ImmutableList.of(),
          ImmutableList.of()));

  MetricQueryResults results =
      metrics.queryMetrics(MetricsFilter.builder().addStep("Top1/Outer1").build());

  assertThat(
      results.getCounters(),
      containsInAnyOrder(
          attemptedMetricsResult("ns1", "name1", "Top1/Outer1/Inner1", 5L),
          attemptedMetricsResult("ns1", "name1", "Top1/Outer1/Inner2", 8L)));

  results = metrics.queryMetrics(MetricsFilter.builder().addStep("Inner2").build());

  assertThat(
      results.getCounters(),
      containsInAnyOrder(
          attemptedMetricsResult("ns1", "name1", "Top1/Outer1/Inner2", 8L),
          attemptedMetricsResult("ns1", "name1", "Top1/Outer2/Inner2", 18L)));
}
 

/**
 * Update Flink's internal metrics ({@link #flinkCounterCache}) with the latest metrics for
 * a given step.
 */
private void updateMetrics(String stepName) {
	MetricResults metricResults = asAttemptedOnlyMetricResults(metricsContainers);
	MetricQueryResults metricQueryResults =
		metricResults.queryMetrics(MetricsFilter.builder().addStep(stepName).build());
	updateCounterOrMeter(metricQueryResults.getCounters());
	updateDistributions(metricQueryResults.getDistributions());
	updateGauge(metricQueryResults.getGauges());
}
 

@Ignore
public void createRunPipeline( PipelineMeta pipelineMeta ) throws Exception {

  /*
  FileOutputStream fos = new FileOutputStream( "/tmp/"+pipelineMeta.getName()+".ktr" );
  fos.write( pipelineMeta.getXML().getBytes() );
  fos.close();
  */

  PipelineOptions pipelineOptions = PipelineOptionsFactory.create();

  pipelineOptions.setJobName( pipelineMeta.getName() );
  pipelineOptions.setUserAgent( BeamConst.STRING_HOP_BEAM );

  BeamDirectPipelineRunConfiguration beamRunConfig = new BeamDirectPipelineRunConfiguration();
  beamRunConfig.setTempLocation( System.getProperty( "java.io.tmpdir" ) );

  // No extra plugins to load : null option
  HopPipelineMetaToBeamPipelineConverter converter = new HopPipelineMetaToBeamPipelineConverter( pipelineMeta, metadataProvider, beamRunConfig );
  Pipeline pipeline = converter.createPipeline();

  PipelineResult pipelineResult = pipeline.run();
  pipelineResult.waitUntilFinish();

  MetricResults metricResults = pipelineResult.metrics();

  MetricQueryResults allResults = metricResults.queryMetrics( MetricsFilter.builder().build() );
  for ( MetricResult<Long> result : allResults.getCounters() ) {
    System.out.println( "Name: " + result.getName() + " Attempted: " + result.getAttempted() );
  }
}
 

@Override
public MetricQueryResults queryMetrics(@Nullable MetricsFilter filter) {
  return MetricQueryResults.create(
      Iterables.filter(counters, counter -> MetricFiltering.matches(filter, counter.getKey())),
      Iterables.filter(
          distributions, distribution -> MetricFiltering.matches(filter, distribution.getKey())),
      Iterables.filter(gauges, gauge -> MetricFiltering.matches(filter, gauge.getKey())));
}
 

private void logMetrics( PipelineResult pipelineResult ) {
  MetricResults metricResults = pipelineResult.metrics();

  logChannel.logBasic( "  ----------------- Metrics refresh @ " + new SimpleDateFormat( "yyyy/MM/dd HH:mm:ss" ).format( new Date() ) + " -----------------------" );

  MetricQueryResults allResults = metricResults.queryMetrics( MetricsFilter.builder().build() );
  for ( MetricResult<Long> result : allResults.getCounters() ) {
    logChannel.logBasic( "Name: " + result.getName() + " Attempted: " + result.getAttempted() );
  }
}
 

@Ignore
public void createRunPipeline( TransMeta transMeta ) throws Exception {

  /*
  FileOutputStream fos = new FileOutputStream( "/tmp/"+transMeta.getName()+".ktr" );
  fos.write( transMeta.getXML().getBytes() );
  fos.close();
  */

  PipelineOptions pipelineOptions = PipelineOptionsFactory.create();

  pipelineOptions.setJobName( transMeta.getName() );
  pipelineOptions.setUserAgent( BeamConst.STRING_KETTLE_BEAM );

  BeamJobConfig jobConfig = new BeamJobConfig();
  jobConfig.setName("Direct runner test");
  jobConfig.setRunnerTypeName( RunnerType.Direct.name() );

  // No extra plugins to load : null option
  TransMetaPipelineConverter converter = new TransMetaPipelineConverter( transMeta, metaStore, (String) null, jobConfig );
  Pipeline pipeline = converter.createPipeline( pipelineOptions );

  PipelineResult pipelineResult = pipeline.run();
  pipelineResult.waitUntilFinish();

  MetricResults metricResults = pipelineResult.metrics();

  MetricQueryResults allResults = metricResults.queryMetrics( MetricsFilter.builder().build() );
  for ( MetricResult<Long> result : allResults.getCounters() ) {
    System.out.println( "Name: " + result.getName() + " Attempted: " + result.getAttempted() );
  }
}
 

@Test
@Category({NeedsRunner.class, UsesCommittedMetrics.class, UsesCounterMetrics.class})
public void testRunPTransform() {
  final String namespace = PipelineRunnerTest.class.getName();
  final Counter counter = Metrics.counter(namespace, "count");
  final PipelineResult result =
      PipelineRunner.fromOptions(p.getOptions())
          .run(
              new PTransform<PBegin, POutput>() {
                @Override
                public POutput expand(PBegin input) {
                  PCollection<Double> output =
                      input
                          .apply(Create.of(1, 2, 3, 4))
                          .apply("ScaleByTwo", MapElements.via(new ScaleFn<>(2.0, counter)));
                  PAssert.that(output).containsInAnyOrder(2.0, 4.0, 6.0, 8.0);
                  return output;
                }
              });

  // Checking counters to verify the pipeline actually ran.
  assertThat(
      result
          .metrics()
          .queryMetrics(
              MetricsFilter.builder()
                  .addNameFilter(MetricNameFilter.inNamespace(namespace))
                  .build())
          .getCounters(),
      hasItem(metricsResult(namespace, "count", "ScaleByTwo", 4L, true)));
}
 

/**
 * Verifies all {{@link PAssert PAsserts}} in the pipeline have been executed and were successful.
 *
 * <p>Note this only runs for runners which support Metrics. Runners which do not should verify
 * this in some other way. See: https://issues.apache.org/jira/browse/BEAM-2001
 */
public static void verifyPAssertsSucceeded(Pipeline pipeline, PipelineResult pipelineResult) {
  if (MetricsEnvironment.isMetricsSupported()) {
    long expectedNumberOfAssertions = (long) PAssert.countAsserts(pipeline);

    long successfulAssertions = 0;
    Iterable<MetricResult<Long>> successCounterResults =
        pipelineResult
            .metrics()
            .queryMetrics(
                MetricsFilter.builder()
                    .addNameFilter(MetricNameFilter.named(PAssert.class, PAssert.SUCCESS_COUNTER))
                    .build())
            .getCounters();
    for (MetricResult<Long> counter : successCounterResults) {
      if (counter.getAttempted() > 0) {
        successfulAssertions++;
      }
    }

    assertThat(
        String.format(
            "Expected %d successful assertions, but found %d.",
            expectedNumberOfAssertions, successfulAssertions),
        successfulAssertions,
        is(expectedNumberOfAssertions));
  }
}
 

public static Map<String, Long> getMetrics(final PipelineResult result) {
  final MetricQueryResults metricQueryResults =
      result.metrics().queryMetrics(MetricsFilter.builder().build());

  final Map<String, Long> gauges =
      StreamSupport.stream(metricQueryResults.getGauges().spliterator(), false)
          .collect(
              Collectors.groupingBy(
                  MetricResult::getName,
                  Collectors.reducing(
                      GaugeResult.empty(),
                      GET_COMMITTED_GAUGE,
                      BinaryOperator.maxBy(Comparator.comparing(GaugeResult::getTimestamp)))))
          .entrySet()
          .stream()
          .collect(Collectors.toMap(e -> e.getKey().getName(), e -> e.getValue().getValue()));

  final Map<String, Long> counters =
      StreamSupport.stream(metricQueryResults.getCounters().spliterator(), false)
          .collect(
              Collectors.groupingBy(
                  m -> m.getName().getName(), Collectors.summingLong(GET_COMMITTED_COUNTER)));
  Map<String, Long> ret = new HashMap<>();
  ret.putAll(gauges);
  ret.putAll(counters);
  addCalculatedMetrics(counters, ret);
  return Collections.unmodifiableMap(ret);
}
 

private Iterable<MetricResult<DistributionResult>> getDistributions(String name) {
  MetricQueryResults metrics =
      result
          .metrics()
          .queryMetrics(
              MetricsFilter.builder()
                  .addNameFilter(MetricNameFilter.named(namespace, name))
                  .build());
  return metrics.getDistributions();
}
 

/**
 * Grab the Beam pipeline results and convert it into engine metrics
 */
protected synchronized void populateEngineMetrics() throws HopException {
  ClassLoader oldContextClassLoader = Thread.currentThread().getContextClassLoader();
  try {
    Thread.currentThread().setContextClassLoader( this.getClass().getClassLoader() );

    EngineMetrics em = new EngineMetrics();
    evaluatePipelineStatus();

    em.setStartDate( getExecutionStartDate() );
    em.setEndDate( getExecutionEndDate() );

    if ( beamPipelineResults != null ) {
      Set<String> transformNames = new HashSet<>( Arrays.asList( pipelineMeta.getTransformNames() ) );
      Map<String, EngineComponent> componentsMap = new HashMap<>();
      MetricResults metrics = beamPipelineResults.metrics();
      MetricQueryResults allResults = metrics.queryMetrics( MetricsFilter.builder().build() );

      for ( MetricResult<Long> result : allResults.getCounters() ) {
        String metricsType = result.getName().getNamespace();
        String metricsName = result.getName().getName();
        long processed = result.getAttempted();

        // This is a transform executing in Beam
        //
        if ( transformNames.contains( metricsName ) ) {
          EngineComponent engineComponent = componentsMap.get( metricsName );
          if ( engineComponent == null ) {
            engineComponent = new EngineComponent( metricsName, 0 );
            componentsMap.put( metricsName, engineComponent );
          }
          if ( Pipeline.METRIC_NAME_READ.equalsIgnoreCase( metricsType ) ) {
            engineComponent.setLinesRead( processed );
            em.setComponentMetric( engineComponent, Pipeline.METRIC_READ, processed );
          } else if ( Pipeline.METRIC_NAME_WRITTEN.equalsIgnoreCase( metricsType ) ) {
            engineComponent.setLinesWritten( processed );
            em.setComponentMetric( engineComponent, Pipeline.METRIC_WRITTEN, processed );
          } else if ( Pipeline.METRIC_NAME_INPUT.equalsIgnoreCase( metricsType ) ) {
            engineComponent.setLinesInput( processed );
            em.setComponentMetric( engineComponent, Pipeline.METRIC_INPUT, processed );
          } else if ( Pipeline.METRIC_NAME_OUTPUT.equalsIgnoreCase( metricsType ) ) {
            engineComponent.setLinesOutput( processed );
            em.setComponentMetric( engineComponent, Pipeline.METRIC_OUTPUT, processed );
          } else if ( Pipeline.METRIC_NAME_INIT.equalsIgnoreCase( metricsType ) ) {
            em.setComponentMetric( engineComponent, Pipeline.METRIC_INIT, processed );
          } else if ( Pipeline.METRIC_NAME_FLUSH_BUFFER.equalsIgnoreCase( metricsType ) ) {
            em.setComponentMetric( engineComponent, Pipeline.METRIC_FLUSH_BUFFER, processed );
          }

          // Copy the execution start and end date from the pipeline
          //
          engineComponent.setExecutionStartDate( getExecutionStartDate() );
          engineComponent.setExecutionEndDate( getExecutionEndDate() );
          engineComponent.setExecutionDuration( calculateDuration( getExecutionStartDate(), getExecutionEndDate() ) );

          // Set the transform status to reflect the pipeline status.
          //
          switch ( beamPipelineResults.getState() ) {
            case DONE:
              engineComponent.setRunning( false );
              engineComponent.setStatus( ComponentExecutionStatus.STATUS_FINISHED );
              break;
            case CANCELLED:
            case FAILED:
            case STOPPED:
              engineComponent.setStopped( true );
              engineComponent.setRunning( false );
              engineComponent.setStatus( ComponentExecutionStatus.STATUS_STOPPED );
              break;
            case RUNNING:
              engineComponent.setRunning( true );
              engineComponent.setStopped( false );
              engineComponent.setStatus( ComponentExecutionStatus.STATUS_RUNNING );
              break;
            case UNKNOWN:
              break;
            case UPDATED:
              break;
            default:
              break;
          }
        }
      }

      em.getComponents().clear();
      em.getComponents().addAll( componentsMap.values() );
    }

    // Swap the engine metrics with the new value
    //
    synchronized ( engineMetrics ) {
      engineMetrics = em;
    }
  } finally {
    Thread.currentThread().setContextClassLoader( oldContextClassLoader );
  }
}
 

@Test
public void testOutput() {
  final List<String> input = Arrays.asList(//
      "{\"attributeMap\":{},\"payload\":\"\"}", //
      "{\"attributeMap\":" //
          + "{\"x_pipeline_proxy\":1" //
          + "},\"payload\":\"proxied+\"}",
      "{\"attributeMap\":" //
          + "{\"x_pipeline_proxy\":\"\"" //
          + "},\"payload\":\"emptyXPP\"}",
      "{\"attributeMap\":" //
          + "{\"submission_timestamp\":\"2000-01-01T00:00:00.000000Z\"" //
          + ",\"x_forwarded_for\":\"3, 2, 1\"" //
          + "},\"payload\":\"notProxied++\"}",
      "{\"attributeMap\":" //
          + "{\"submission_timestamp\":\"2000-01-01T00:00:00.000000Z\"" //
          + ",\"x_forwarded_for\":\"4, 3, 2, 1\"" //
          + ",\"x_pipeline_proxy\":\"1999-12-31T23:59:59.999999Z\"" //
          + "},\"payload\":\"proxiedWithTimestamp\"}",
      "{\"attributeMap\":" //
          + "{\"submission_timestamp\":\"1999-12-31T23:59:59.999999Z\"" //
          + ",\"proxy_timestamp\":\"2000-01-01T00:00:00.000000Z\"" //
          + ",\"x_forwarded_for\":\"4, 3, 2\"" //
          + "},\"payload\":\"retried+\"}");

  final List<String> expected = Arrays.asList(//
      "{\"attributeMap\":{},\"payload\":\"\"}", //
      "{\"attributeMap\":{},\"payload\":\"proxied+\"}", //
      "{\"attributeMap\":{},\"payload\":\"emptyXPP\"}", //
      "{\"attributeMap\":" //
          + "{\"submission_timestamp\":\"2000-01-01T00:00:00.000000Z\"" //
          + ",\"x_forwarded_for\":\"3, 2, 1\"" //
          + "},\"payload\":\"notProxied++\"}",
      "{\"attributeMap\":" //
          + "{\"proxy_timestamp\":\"2000-01-01T00:00:00.000000Z\"" //
          + ",\"submission_timestamp\":\"1999-12-31T23:59:59.999999Z\"" //
          + ",\"x_forwarded_for\":\"4, 3, 2\"" //
          + "},\"payload\":\"proxiedWithTimestamp\"}",
      "{\"attributeMap\":" //
          + "{\"proxy_timestamp\":\"2000-01-01T00:00:00.000000Z\"" //
          + ",\"submission_timestamp\":\"1999-12-31T23:59:59.999999Z\"" //
          + ",\"x_forwarded_for\":\"4, 3, 2\"" //
          + "},\"payload\":\"retried+\"}");

  final PCollection<String> output = pipeline //
      .apply(Create.of(input)) //
      .apply(InputFileFormat.json.decode()) //
      .apply(ParseProxy.of()) //
      .apply(OutputFileFormat.json.encode());

  PAssert.that(output).containsInAnyOrder(expected);

  final PipelineResult result = pipeline.run();

  final List<MetricResult<Long>> counters = Lists.newArrayList(result.metrics()
      .queryMetrics(MetricsFilter.builder()
          .addNameFilter(MetricNameFilter.inNamespace(ParseProxy.Fn.class)).build())
      .getCounters());

  assertEquals(1, counters.size());
  counters.forEach(counter -> assertThat(counter.getCommitted(), greaterThan(0L)));
}
 

@Test
public void testOutput() {
  // Some of the IPs below are chosen specifically because they are contained in the test city
  // database; see the json source for the test db in:
  // https://github.com/maxmind/MaxMind-DB/blob/664aeeb08bb50f53a1fdceac763c37f6465e44a4/source-data/GeoIP2-City-Test.json
  final List<String> input = Arrays.asList(
      "{\"attributeMap\":{\"host\":\"test\"},\"payload\":\"dGVzdA==\"}", //
      "{\"attributeMap\":{\"remote_addr\":\"202.196.224.0\"},\"payload\":\"\"}", //
      "{\"attributeMap\":" //
          + "{\"remote_addr\":\"10.0.0.2\"" //
          + ",\"x_forwarded_for\":\"192.168.1.2, 216.160.83.56, 60.1.1.1\"" //
          + "},\"payload\":\"\"}");

  final List<String> expected = Arrays.asList(//
      "{\"attributeMap\":" //
          + "{\"geo_db_version\":\"2019-01-03T21:26:19Z\"" //
          + ",\"host\":\"test\"" //
          + "},\"payload\":\"dGVzdA==\"}", //
      "{\"attributeMap\":" //
          + "{\"geo_country\":\"PH\"" //
          + ",\"geo_db_version\":\"2019-01-03T21:26:19Z\"" //
          + "},\"payload\":\"\"}", //
      "{\"attributeMap\":" //
          + "{\"geo_city\":\"Milton\"" //
          + ",\"geo_country\":\"US\"" //
          + ",\"geo_db_version\":\"2019-01-03T21:26:19Z\"" //
          + ",\"geo_subdivision1\":\"WA\"" //
          + "},\"payload\":\"\"}");

  final PCollection<String> output = pipeline //
      .apply(Create.of(input)) //
      .apply(InputFileFormat.json.decode()) //
      .apply(GeoCityLookup.of(pipeline.newProvider(MMDB), null))
      .apply(OutputFileFormat.json.encode());

  PAssert.that(output).containsInAnyOrder(expected);

  GeoCityLookup.clearSingletonsForTests();
  final PipelineResult result = pipeline.run();

  final List<MetricResult<Long>> counters = Lists.newArrayList(result.metrics()
      .queryMetrics(MetricsFilter.builder()
          .addNameFilter(MetricNameFilter.inNamespace(GeoCityLookup.Fn.class)).build())
      .getCounters());

  assertEquals(6, counters.size());
  counters.forEach(counter -> assertThat(counter.getCommitted(), greaterThan(0L)));
}
 

@Test
public void testOutput() {
  // Some of the IPs below are chosen specifically because they are contained in the test city
  // database; see the json source for the test db in:
  // https://github.com/maxmind/MaxMind-DB/blob/664aeeb08bb50f53a1fdceac763c37f6465e44a4/source-data/GeoIP2-City-Test.json
  final List<String> input = Arrays.asList(
      "{\"attributeMap\":{\"host\":\"test\"},\"payload\":\"dGVzdA==\"}", //
      "{\"attributeMap\":{\"remote_addr\":\"24.38.243.141\"},\"payload\":\"\"}", //
      "{\"attributeMap\":" //
          + "{\"remote_addr\":\"10.0.0.2\"" //
          + ",\"x_forwarded_for\":\"192.168.1.2, 23.32.32.1, 23.32.32.11\"" //
          + "},\"payload\":\"\"}");

  final List<String> expected = Arrays.asList(//
      "{\"attributeMap\":" //
          + "{\"host\":\"test\"" //
          + ",\"isp_db_version\":\"2018-01-15T22:27:16Z\"" //
          + "},\"payload\":\"dGVzdA==\"}", //
      "{\"attributeMap\":" //
          + "{\"isp_name\":\"Akamai Technologies\"" //
          + ",\"remote_addr\":\"10.0.0.2\"" //
          + ",\"isp_db_version\":\"2018-01-15T22:27:16Z\"" //
          + ",\"isp_organization\":\"Akamai Technologies\"" //
          + ",\"x_forwarded_for\":\"192.168.1.2, 23.32.32.1, 23.32.32.11\"}" //
          + ",\"payload\":\"\"}", //
      "{\"attributeMap\":" //
          + "{\"remote_addr\":\"24.38.243.141\"" //
          + ",\"isp_db_version\":\"2018-01-15T22:27:16Z\"" //
          + ",\"isp_organization\":\"LAWN MULLEN & GOOD INTERNATIONAL\"}" + ",\"payload\":\"\"}");

  final PCollection<String> output = pipeline //
      .apply(Create.of(input)) //
      .apply(InputFileFormat.json.decode()) //
      .apply(GeoIspLookup.of(pipeline.newProvider(MMDB))).apply(OutputFileFormat.json.encode());

  PAssert.that(output).containsInAnyOrder(expected);

  GeoIspLookup.clearSingletonsForTests();
  final PipelineResult result = pipeline.run();

  final List<MetricResult<Long>> counters = Lists.newArrayList(result.metrics()
      .queryMetrics(MetricsFilter.builder()
          .addNameFilter(MetricNameFilter.inNamespace(GeoIspLookup.Fn.class)).build())
      .getCounters());

  assertEquals(4, counters.size());
  counters.forEach(counter -> assertThat(counter.getCommitted(), greaterThan(0L)));
}