Rewrites Prometheus filter to neither double-register or double-count

Adrian Cole · adriancole · commit 74336c80c65e · 2017-11-23T20:41:32.000+08:00
diff --git a/zipkin-autoconfigure/metrics-prometheus/src/main/java/zipkin/autoconfigure/metrics/PrometheusMetricsAutoConfiguration.java b/zipkin-autoconfigure/metrics-prometheus/src/main/java/zipkin/autoconfigure/metrics/PrometheusMetricsAutoConfiguration.java
@@ -14,12 +14,22 @@
 
 package zipkin.autoconfigure.metrics;
 
+import io.prometheus.client.Histogram;
 import io.prometheus.client.hotspot.DefaultExports;
 import io.prometheus.client.spring.boot.EnablePrometheusEndpoint;
 import io.prometheus.client.spring.boot.EnableSpringBootMetricsCollector;
 import io.prometheus.client.spring.web.EnablePrometheusTiming;
+import java.io.IOException;
+import javax.servlet.AsyncContext;
+import javax.servlet.AsyncEvent;
+import javax.servlet.AsyncListener;
 import javax.servlet.Filter;
+import javax.servlet.FilterChain;
+import javax.servlet.FilterConfig;
 import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletRequest;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
@@ -32,11 +42,108 @@ public class PrometheusMetricsAutoConfiguration {
     DefaultExports.initialize();
   }
 
-  @Bean
-  public Filter prometheusMetricsFilter() throws ServletException {
-    return new io.prometheus.client.filter.MetricsFilter("http_request_duration_seconds",
-      "Response time histogram",
-      0,
-      null);
+  // Obviates the state bug in MetricsFilter which implicitly registers and hides something you
+  // can't create twice
+  static final Histogram http_request_duration_seconds = Histogram.build()
+    .labelNames("path", "method")
+    .help("Response time histogram")
+    .name("http_request_duration_seconds")
+    .register();
+
+  @Bean("http_request_duration_seconds") Histogram http_request_duration_seconds() {
+    return http_request_duration_seconds;
+  }
+
+  @Bean public Filter prometheusMetricsFilter() {
+    return new PrometheusDurationFilter();
+  }
+
+  /**
+   * The normal prometheus metrics filter implicitly registers a histogram which is hidden in a
+   * field and not deregistered on destroy. A registration of any second instance of that filter
+   * fails trying to re-register the same collector by design (by brian-brazil). The rationale is
+   * that you are not supposed to recreate the same histogram. However, this design prevents us from
+   * doing that. brian-bazil's hard stance on this makes the filter unusable for applications who
+   * run tests.
+   *
+   * <p>This filter replaces the normal prometheus filter, correcting the design flaw by allowing us
+   * to re-use the JVM singleton. It also corrects a major flaw in the upstream filter which results
+   * in double-counting of requests when they are performed asynchronously. When the culture changes
+   * in the prometheus project such that bugs are fixable, please submit this so that it can help
+   * others. For more info, see "brian's bomb" https://github.com/openzipkin/zipkin/issues/1811
+   */
+  static final class PrometheusDurationFilter implements Filter {
+    @Override public void init(FilterConfig filterConfig) {
+    }
+
+    /**
+     * Note that upstream also has a problem which is that it doesn't handle async properly.
+     * MetricsFilter results in double-counting, which this implementation avoids.
+     */
+    @Override public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse,
+      FilterChain filterChain) throws IOException, ServletException {
+
+      HttpServletRequest request = (HttpServletRequest) servletRequest;
+
+      // async servlets will enter the filter twice
+      if (request.getAttribute("PrometheusDurationFilter") != null) {
+        filterChain.doFilter(request, servletResponse);
+        return;
+      }
+
+      request.setAttribute("PrometheusDurationFilter", "true");
+
+      Histogram.Timer timer = http_request_duration_seconds
+        .labels(request.getRequestURI(), request.getMethod())
+        .startTimer();
+
+      try {
+        filterChain.doFilter(servletRequest, servletResponse);
+      } finally {
+        if (request.isAsyncStarted()) { // we don't have the actual response, handle later
+          request.getAsyncContext().addListener(new CompleteTimer(timer));
+        } else { // we have a synchronous response, so we can finish the recording
+          timer.observeDuration();
+        }
+      }
+    }
+
+    @Override public void destroy() {
+    }
+  }
+
+  /** Inspired by WingtipsRequestSpanCompletionAsyncListener */
+  static final class CompleteTimer implements AsyncListener {
+    final Histogram.Timer timer;
+    volatile boolean completed = false;
+
+    CompleteTimer(Histogram.Timer timer) {
+      this.timer = timer;
+    }
+
+    @Override public void onComplete(AsyncEvent e) {
+      tryComplete();
+    }
+
+    @Override public void onTimeout(AsyncEvent e) {
+      tryComplete();
+    }
+
+    @Override public void onError(AsyncEvent e) {
+      tryComplete();
+    }
+
+    /** Only observes the first completion event */
+    void tryComplete() {
+      if (completed) return;
+      timer.observeDuration();
+      completed = true;
+    }
+
+    /** If another async is created (ex via asyncContext.dispatch), this needs to be re-attached */
+    @Override public void onStartAsync(AsyncEvent event) {
+      AsyncContext eventAsyncContext = event.getAsyncContext();
+      if (eventAsyncContext != null) eventAsyncContext.addListener(this);
+    }
   }
 }
diff --git a/zipkin-server/src/test/java/zipkin/server/ZipkinServerConfigurationTest.java b/zipkin-server/src/test/java/zipkin/server/ZipkinServerConfigurationTest.java
@@ -14,6 +14,7 @@
 package zipkin.server;
 
 import com.github.kristofa.brave.Brave;
+import io.prometheus.client.Histogram;
 import org.junit.After;
 import org.junit.Test;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
@@ -27,6 +28,7 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Import;
+import zipkin.autoconfigure.metrics.PrometheusMetricsAutoConfiguration;
 import zipkin.autoconfigure.ui.ZipkinUiAutoConfiguration;
 import zipkin.server.brave.BraveConfiguration;
 
@@ -67,6 +69,18 @@ public void httpCollector_canDisable() {
     context.getBean(ZipkinHttpCollector.class);
   }
 
+  @Test public void prometheusHistogram() {
+    context.register(
+      PropertyPlaceholderAutoConfiguration.class,
+      ZipkinServerConfigurationTest.Config.class,
+      ZipkinServerConfiguration.class,
+      PrometheusMetricsAutoConfiguration.class
+    );
+    context.refresh();
+
+    assertThat(context.getBean("http_request_duration_seconds", Histogram.class)).isNotNull();
+  }
+
   @Test public void query_enabledByDefault() {
     context.register(
       PropertyPlaceholderAutoConfiguration.class,
diff --git a/zipkin-server/src/test/java/zipkin/server/ZipkinServerIntegrationTest.java b/zipkin-server/src/test/java/zipkin/server/ZipkinServerIntegrationTest.java
@@ -14,6 +14,7 @@
 package zipkin.server;
 
 import com.jayway.jsonpath.JsonPath;
+import io.prometheus.client.Histogram;
 import java.io.IOException;
 import java.util.List;
 import okhttp3.MediaType;
@@ -53,17 +54,16 @@
 @RunWith(SpringRunner.class)
 public class ZipkinServerIntegrationTest {
 
-  @Autowired
-  InMemoryStorage storage;
-  @Autowired
-  ActuateCollectorMetrics metrics;
-  @LocalServerPort
-  int zipkinPort;
+  @Autowired InMemoryStorage storage;
+  @Autowired ActuateCollectorMetrics metrics;
+  @Autowired Histogram duration;
+  @LocalServerPort int zipkinPort;
 
   OkHttpClient client = new OkHttpClient.Builder().followRedirects(false).build();
 
   @Before public void init() {
     storage.clear();
+    duration.clear();
     metrics.forTransport("http").reset();
   }
 
@@ -117,6 +117,21 @@ public class ZipkinServerIntegrationTest {
       .isEqualTo(spans.size());
   }
 
+  /** Makes sure the prometheus filter doesn't count twice */
+  @Test public void writeSpans_updatesPrometheusMetrics() throws Exception {
+    List<Span> spans = asList(LOTS_OF_SPANS[0], LOTS_OF_SPANS[1], LOTS_OF_SPANS[2]);
+    byte[] body = Codec.JSON.writeSpans(spans);
+    post("/api/v1/spans", body);
+    post("/api/v1/spans", body);
+
+    Response response = get("/prometheus");
+    assertThat(response.isSuccessful()).isTrue();
+    String prometheus = response.body().string();
+
+    assertThat(prometheus)
+      .contains("http_request_duration_seconds_count{path=\"/api/v1/spans\",method=\"POST\",} 2.0");
+  }
+
   @Test public void tracesQueryRequiresNoParameters() throws Exception {
     byte[] body = Codec.JSON.writeSpans(TRACE);
     post("/api/v1/spans", body);
