org.apache.beam.runners.spark.metrics.WithMetricsSupport.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.beam.runners.spark.metrics.WithMetricsSupport.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.beam.runners.spark.metrics;

import com.codahale.metrics.Counter;
import com.codahale.metrics.Gauge;
import com.codahale.metrics.Histogram;
import com.codahale.metrics.Meter;
import com.codahale.metrics.Metric;
import com.codahale.metrics.MetricFilter;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.Timer;
import com.google.common.base.Function;
import com.google.common.base.Optional;
import com.google.common.base.Predicate;
import com.google.common.base.Predicates;
import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableSortedMap;
import com.google.common.collect.Maps;
import com.google.common.collect.Ordering;
import java.util.HashMap;
import java.util.Map;
import java.util.SortedMap;
import org.apache.beam.runners.spark.aggregators.NamedAggregators;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A {@link MetricRegistry} decorator-like that supports {@link AggregatorMetric} and
 * {@link SparkBeamMetric} as {@link Gauge Gauges}.
 * <p>
 * {@link MetricRegistry} is not an interface, so this is not a by-the-book decorator.
 * That said, it delegates all metric related getters to the "decorated" instance.
 * </p>
 */
public class WithMetricsSupport extends MetricRegistry {

    private static final Logger LOG = LoggerFactory.getLogger(WithMetricsSupport.class);

    private final MetricRegistry internalMetricRegistry;

    private WithMetricsSupport(final MetricRegistry internalMetricRegistry) {
        this.internalMetricRegistry = internalMetricRegistry;
    }

    public static WithMetricsSupport forRegistry(final MetricRegistry metricRegistry) {
        return new WithMetricsSupport(metricRegistry);
    }

    @Override
    public SortedMap<String, Timer> getTimers(final MetricFilter filter) {
        return internalMetricRegistry.getTimers(filter);
    }

    @Override
    public SortedMap<String, Meter> getMeters(final MetricFilter filter) {
        return internalMetricRegistry.getMeters(filter);
    }

    @Override
    public SortedMap<String, Histogram> getHistograms(final MetricFilter filter) {
        return internalMetricRegistry.getHistograms(filter);
    }

    @Override
    public SortedMap<String, Counter> getCounters(final MetricFilter filter) {
        return internalMetricRegistry.getCounters(filter);
    }

    @Override
    public SortedMap<String, Gauge> getGauges(final MetricFilter filter) {
        return new ImmutableSortedMap.Builder<String, Gauge>(Ordering.from(String.CASE_INSENSITIVE_ORDER))
                .putAll(internalMetricRegistry.getGauges(filter))
                .putAll(extractGauges(internalMetricRegistry, filter)).build();
    }

    private Map<String, Gauge> extractGauges(final MetricRegistry metricRegistry, final MetricFilter filter) {
        Map<String, Gauge> gauges = new HashMap<>();

        // find the AggregatorMetric metrics from within all currently registered metrics
        final Optional<Map<String, Gauge>> aggregatorMetrics = FluentIterable
                .from(metricRegistry.getMetrics().entrySet()).firstMatch(isAggregatorMetric())
                .transform(aggregatorMetricToGauges());

        // find the SparkBeamMetric metrics from within all currently registered metrics
        final Optional<Map<String, Gauge>> beamMetrics = FluentIterable.from(metricRegistry.getMetrics().entrySet())
                .firstMatch(isSparkBeamMetric()).transform(beamMetricToGauges());

        if (aggregatorMetrics.isPresent()) {
            gauges.putAll(Maps.filterEntries(aggregatorMetrics.get(), matches(filter)));
        }

        if (beamMetrics.isPresent()) {
            gauges.putAll(Maps.filterEntries(beamMetrics.get(), matches(filter)));
        }

        return gauges;
    }

    private Function<Map.Entry<String, Metric>, Map<String, Gauge>> aggregatorMetricToGauges() {
        return new Function<Map.Entry<String, Metric>, Map<String, Gauge>>() {
            @Override
            public Map<String, Gauge> apply(final Map.Entry<String, Metric> entry) {
                final NamedAggregators agg = ((AggregatorMetric) entry.getValue()).getNamedAggregators();
                final String parentName = entry.getKey();
                final Map<String, Gauge> gaugeMap = Maps.transformEntries(agg.renderAll(), toGauge());
                final Map<String, Gauge> fullNameGaugeMap = Maps.newLinkedHashMap();
                for (Map.Entry<String, Gauge> gaugeEntry : gaugeMap.entrySet()) {
                    fullNameGaugeMap.put(parentName + "." + gaugeEntry.getKey(), gaugeEntry.getValue());
                }
                return Maps.filterValues(fullNameGaugeMap, Predicates.notNull());
            }
        };
    }

    private Function<Map.Entry<String, Metric>, Map<String, Gauge>> beamMetricToGauges() {
        return new Function<Map.Entry<String, Metric>, Map<String, Gauge>>() {
            @Override
            public Map<String, Gauge> apply(final Map.Entry<String, Metric> entry) {
                final Map<String, ?> metrics = ((SparkBeamMetric) entry.getValue()).renderAll();
                final String parentName = entry.getKey();
                final Map<String, Gauge> gaugeMap = Maps.transformEntries(metrics, toGauge());
                final Map<String, Gauge> fullNameGaugeMap = Maps.newLinkedHashMap();
                for (Map.Entry<String, Gauge> gaugeEntry : gaugeMap.entrySet()) {
                    fullNameGaugeMap.put(parentName + "." + gaugeEntry.getKey(), gaugeEntry.getValue());
                }
                return Maps.filterValues(fullNameGaugeMap, Predicates.notNull());
            }
        };
    }

    private Maps.EntryTransformer<String, Object, Gauge> toGauge() {
        return new Maps.EntryTransformer<String, Object, Gauge>() {

            @Override
            public Gauge transformEntry(final String name, final Object rawValue) {
                return new Gauge<Double>() {

                    @Override
                    public Double getValue() {
                        // at the moment the metric's type is assumed to be
                        // compatible with Double. While far from perfect, it seems reasonable at
                        // this point in time
                        try {
                            return Double.parseDouble(rawValue.toString());
                        } catch (final Exception e) {
                            LOG.warn("Failed reporting metric with name [{}], of type [{}], since it could not be"
                                    + " converted to double", name, rawValue.getClass().getSimpleName(), e);
                            return null;
                        }
                    }
                };
            }
        };
    }

    private Predicate<Map.Entry<String, Gauge>> matches(final MetricFilter filter) {
        return new Predicate<Map.Entry<String, Gauge>>() {
            @Override
            public boolean apply(final Map.Entry<String, Gauge> entry) {
                return filter.matches(entry.getKey(), entry.getValue());
            }
        };
    }

    private Predicate<Map.Entry<String, Metric>> isAggregatorMetric() {
        return new Predicate<Map.Entry<String, Metric>>() {
            @Override
            public boolean apply(final Map.Entry<String, Metric> metricEntry) {
                return (metricEntry.getValue() instanceof AggregatorMetric);
            }
        };
    }

    private Predicate<Map.Entry<String, Metric>> isSparkBeamMetric() {
        return new Predicate<Map.Entry<String, Metric>>() {
            @Override
            public boolean apply(final Map.Entry<String, Metric> metricEntry) {
                return (metricEntry.getValue() instanceof SparkBeamMetric);
            }
        };
    }
}