com.wavefront.agent.preprocessor.SpanAllowAnnotationTransformer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of proxy Show documentation
Service for batching and relaying metric traffic to Wavefront
There is a newer version: 4.36
package com.wavefront.agent.preprocessor;

import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Predicate;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import wavefront.report.Annotation;
import wavefront.report.Span;

/**
 * Only allow span annotations that match the allowed list.
 *
 * @author [email protected]
 */
public class SpanAllowAnnotationTransformer implements Function {
  private static final Set SYSTEM_TAGS =
      ImmutableSet.of("service", "application", "cluster", "shard");

  private final Map allowedKeys;
  private final PreprocessorRuleMetrics ruleMetrics;
  private final Predicate v2Predicate;

  SpanAllowAnnotationTransformer(
      final Map keys,
      @Nullable final Predicate v2Predicate,
      final PreprocessorRuleMetrics ruleMetrics) {
    this.allowedKeys = new HashMap<>(keys.size() + SYSTEM_TAGS.size());
    SYSTEM_TAGS.forEach(x -> allowedKeys.put(x, null));
    keys.forEach((k, v) -> allowedKeys.put(k, v == null ? null : Pattern.compile(v)));
    Preconditions.checkNotNull(ruleMetrics, "PreprocessorRuleMetrics can't be null");
    this.ruleMetrics = ruleMetrics;
    this.v2Predicate = v2Predicate != null ? v2Predicate : x -> true;
  }

  @Nullable
  @Override
  public Span apply(@Nullable Span span) {
    if (span == null) return null;
    long startNanos = ruleMetrics.ruleStart();
    try {
      if (!v2Predicate.test(span)) return span;

      List annotations =
          span.getAnnotations().stream()
              .filter(x -> allowedKeys.containsKey(x.getKey()))
              .filter(x -> isPatternNullOrMatches(allowedKeys.get(x.getKey()), x.getValue()))
              .collect(Collectors.toList());
      if (annotations.size() < span.getAnnotations().size()) {
        span.setAnnotations(annotations);
        ruleMetrics.incrementRuleAppliedCounter();
      }
      return span;
    } finally {
      ruleMetrics.ruleEnd(startNanos);
    }
  }

  private static boolean isPatternNullOrMatches(@Nullable Pattern pattern, String string) {
    return pattern == null || pattern.matcher(string).matches();
  }

  /**
   * Create an instance based on loaded yaml fragment.
   *
   * @param ruleMap yaml map
   * @param v2Predicate the v2 predicate
   * @param ruleMetrics metrics container
   * @return SpanAllowAnnotationTransformer instance
   */
  public static SpanAllowAnnotationTransformer create(
      Map ruleMap,
      @Nullable final Predicate v2Predicate,
      final PreprocessorRuleMetrics ruleMetrics) {
    Object keys = ruleMap.get("allow");
    if (keys instanceof Map) {
      //noinspection unchecked
      return new SpanAllowAnnotationTransformer(
          (Map) keys, v2Predicate, ruleMetrics);
    } else if (keys instanceof List) {
      Map map = new HashMap<>();
      //noinspection unchecked
      ((List) keys).forEach(x -> map.put(x, null));
      return new SpanAllowAnnotationTransformer(map, null, ruleMetrics);
    }
    throw new IllegalArgumentException("[allow] is not a list or a map");
  }
}