Add status code label to server latency histogram

src/main/java/me/dinowernli/grpc/prometheus/Configuration.java

@@ -19,14 +19,16 @@ public class Configuration {
   private final CollectorRegistry collectorRegistry;
   private final double[] latencyBuckets;
   private final List<String> labelHeaders;
+  private final boolean addCodeLabelToHistograms;
 
   /** Returns a {@link Configuration} for recording all cheap metrics about the rpcs. */
   public static Configuration cheapMetricsOnly() {
     return new Configuration(
         false /* isIncludeLatencyHistograms */,
         CollectorRegistry.defaultRegistry,
         DEFAULT_LATENCY_BUCKETS,
-        new ArrayList<>());
+        new ArrayList<>(),
+        false);
   }
 
   /**
@@ -38,7 +40,8 @@ public static Configuration allMetrics() {
         true /* isIncludeLatencyHistograms */,
         CollectorRegistry.defaultRegistry,
         DEFAULT_LATENCY_BUCKETS,
-        new ArrayList<>());
+        new ArrayList<>(),
+        false);
   }
 
   /**
@@ -47,15 +50,24 @@ public static Configuration allMetrics() {
    */
   public Configuration withCollectorRegistry(CollectorRegistry collectorRegistry) {
     return new Configuration(
-        isIncludeLatencyHistograms, collectorRegistry, latencyBuckets, labelHeaders);
+        isIncludeLatencyHistograms,
+        collectorRegistry,
+        latencyBuckets,
+        labelHeaders,
+        addCodeLabelToHistograms);
   }
 
   /**
    * Returns a copy {@link Configuration} with the difference that the latency histogram values are
    * recorded with the specified set of buckets.
    */
   public Configuration withLatencyBuckets(double[] buckets) {
-    return new Configuration(isIncludeLatencyHistograms, collectorRegistry, buckets, labelHeaders);
+    return new Configuration(
+        isIncludeLatencyHistograms,
+        collectorRegistry,
+        buckets,
+        labelHeaders,
+        addCodeLabelToHistograms);
   }
 
   /**
@@ -77,7 +89,25 @@ public Configuration withLabelHeaders(List<String> headers) {
     List<String> newHeaders = new ArrayList<>(labelHeaders);
     newHeaders.addAll(headers);
     return new Configuration(
-        isIncludeLatencyHistograms, collectorRegistry, latencyBuckets, newHeaders);
+        isIncludeLatencyHistograms,
+        collectorRegistry,
+        latencyBuckets,
+        newHeaders,
+        addCodeLabelToHistograms);
+  }
+
+  /**
+   * Returns a copy {@link Configuration} with the difference that status code label will be added
+   * to latency histogram. If latency histogram itself is disabled, this takes no effect.
+   */
+  public Configuration withCodeLabelInLatencyHistogram() {
+    return new Configuration(
+        isIncludeLatencyHistograms,
+        collectorRegistry,
+        latencyBuckets,
+        labelHeaders,
+        // addCodeLabelToHistograms: set true only if isIncludeLatencyHistograms is also true
+        isIncludeLatencyHistograms);
   }
 
   /** Returns whether or not latency histograms for calls should be included. */
@@ -100,6 +130,11 @@ public List<String> getLabelHeaders() {
     return labelHeaders;
   }
 
+  /** Returns whether or not status code label should be added to latency histogram. */
+  public boolean addCodeLabelToHistograms() {
+    return addCodeLabelToHistograms;
+  }
+
   /**
    * Returns the sanitized version of the label headers, after turning all hyphens to underscores.
    */
@@ -111,10 +146,12 @@ private Configuration(
       boolean isIncludeLatencyHistograms,
       CollectorRegistry collectorRegistry,
       double[] latencyBuckets,
-      List<String> labelHeaders) {
+      List<String> labelHeaders,
+      boolean addCodeLabelToHistograms) {
     this.isIncludeLatencyHistograms = isIncludeLatencyHistograms;
     this.collectorRegistry = collectorRegistry;
     this.latencyBuckets = latencyBuckets;
     this.labelHeaders = labelHeaders;
+    this.addCodeLabelToHistograms = addCodeLabelToHistograms;
   }
 }

src/main/java/me/dinowernli/grpc/prometheus/MonitoringServerCall.java

@@ -8,6 +8,7 @@
 import io.grpc.Status;
 import java.time.Clock;
 import java.time.Instant;
+import java.util.Optional;
 
 /**
  * A {@link ForwardingServerCall} which updates Prometheus metrics based on the server-side actions
@@ -61,11 +62,12 @@ private void reportStartMetrics() {
   }
 
   private void reportEndMetrics(Status status) {
-    serverMetrics.recordServerHandled(status.getCode(), requestMetadata);
+    final Status.Code code = status.getCode();
+    serverMetrics.recordServerHandled(code, requestMetadata);
     if (configuration.isIncludeLatencyHistograms()) {
       double latencySec =
           (clock.millis() - startInstant.toEpochMilli()) / (double) MILLIS_PER_SECOND;
-      serverMetrics.recordLatency(latencySec, requestMetadata);
+      serverMetrics.recordLatency(latencySec, requestMetadata, Optional.of(code));
     }
   }
 }

src/main/java/me/dinowernli/grpc/prometheus/ServerMetrics.java

@@ -28,8 +28,10 @@ class ServerMetrics {
   private static final List<String> defaultRequestLabels =
       Arrays.asList("grpc_type", "grpc_service", "grpc_method");
 
+  private static final String statusCodeLabel = "grpc_code";
+
   private static final List<String> defaultResponseLabels =
-      Arrays.asList("grpc_type", "grpc_service", "grpc_method", "code", "grpc_code");
+      Arrays.asList("grpc_type", "grpc_service", "grpc_method", "code", statusCodeLabel);
 
   private static final Counter.Builder serverStartedBuilder =
       Counter.build()
@@ -76,6 +78,7 @@ class ServerMetrics {
   private final Counter serverStreamMessagesReceived;
   private final Counter serverStreamMessagesSent;
   private final Optional<Histogram> serverHandledLatencySeconds;
+  private final boolean addCodeLabelToHistograms;
 
   private final GrpcMethod method;
 
@@ -86,14 +89,16 @@ private ServerMetrics(
       Counter serverHandled,
       Counter serverStreamMessagesReceived,
       Counter serverStreamMessagesSent,
-      Optional<Histogram> serverHandledLatencySeconds) {
+      Optional<Histogram> serverHandledLatencySeconds,
+      boolean addCodeLabelToHistograms) {
     this.labelHeaderKeys = labelHeaderKeys;
     this.method = method;
     this.serverStarted = serverStarted;
     this.serverHandled = serverHandled;
     this.serverStreamMessagesReceived = serverStreamMessagesReceived;
     this.serverStreamMessagesSent = serverStreamMessagesSent;
     this.serverHandledLatencySeconds = serverHandledLatencySeconds;
+    this.addCodeLabelToHistograms = addCodeLabelToHistograms;
   }
 
   public void recordCallStarted(Metadata metadata) {
@@ -121,13 +126,23 @@ public void recordStreamMessageReceived(Metadata metadata) {
    * Only has any effect if monitoring is configured to include latency histograms. Otherwise, this
    * does nothing.
    */
-  public void recordLatency(double latencySec, Metadata metadata) {
+  public void recordLatency(double latencySec, Metadata metadata, Optional<Code> codeOpt) {
     if (!this.serverHandledLatencySeconds.isPresent()) {
       return;
     }
-    addLabels(
-            this.serverHandledLatencySeconds.get(), customLabels(metadata, labelHeaderKeys), method)
-        .observe(latencySec);
+
+    final List<String> allLabels = new ArrayList<String>();
+    allLabels.addAll(customLabels(metadata, labelHeaderKeys));
+    if (addCodeLabelToHistograms) {
+      allLabels.add(codeOpt.map(c -> c.toString()).orElse(""));
+    }
+
+    addLabels(this.serverHandledLatencySeconds.get(), allLabels, method).observe(latencySec);
+  }
+
+  @Deprecated
+  public void recordLatency(double latencySec, Metadata metadata) {
+    recordLatency(latencySec, metadata, Optional.empty());
   }
 
   /** Knows how to produce {@link ServerMetrics} instances for individual methods. */
@@ -138,6 +153,7 @@ static class Factory {
     private final Counter serverStreamMessagesReceived;
     private final Counter serverStreamMessagesSent;
     private final Optional<Histogram> serverHandledLatencySeconds;
+    private final boolean addCodeLabelToHistograms;
 
     Factory(Configuration configuration) {
       CollectorRegistry registry = configuration.getCollectorRegistry();
@@ -160,15 +176,27 @@ static class Factory {
               .register(registry);
 
       if (configuration.isIncludeLatencyHistograms()) {
+
+        final List<String> labels = new ArrayList<String>();
+        labels.addAll(defaultRequestLabels);
+        labels.addAll(configuration.getSanitizedLabelHeaders());
+        if (configuration.addCodeLabelToHistograms()) {
+          labels.add(statusCodeLabel);
+          this.addCodeLabelToHistograms = true;
+        } else {
+          this.addCodeLabelToHistograms = false;
+        }
+
         this.serverHandledLatencySeconds =
             Optional.of(
                 serverHandledLatencySecondsBuilder
                     .buckets(configuration.getLatencyBuckets())
-                    .labelNames(
-                        asArray(defaultRequestLabels, configuration.getSanitizedLabelHeaders()))
+                    .labelNames(labels.toArray(new String[0]))
                     .register(registry));
+
       } else {
         this.serverHandledLatencySeconds = Optional.empty();
+        this.addCodeLabelToHistograms = false;
       }
     }
 
@@ -181,7 +209,8 @@ ServerMetrics createMetricsForMethod(GrpcMethod grpcMethod) {
           serverHandled,
           serverStreamMessagesReceived,
           serverStreamMessagesSent,
-          serverHandledLatencySeconds);
+          serverHandledLatencySeconds,
+          addCodeLabelToHistograms);
     }
   }
 }

src/test/java/me/dinowernli/grpc/prometheus/integration/MonitoringServerInterceptorIntegrationTest.java

@@ -248,6 +248,27 @@ public void overridesHistogramBuckets() throws Throwable {
                 "grpc_server_handled_latency_seconds",
                 "grpc_server_handled_latency_seconds_bucket"))
         .isEqualTo(expectedNum);
+
+    MetricFamilySamples.Sample sample =
+        getSample(
+            findRecordedMetricOrThrow("grpc_server_handled_latency_seconds"),
+            "grpc_server_handled_latency_seconds_bucket");
+
+    assertThat(sample.labelNames).containsExactly("grpc_type", "grpc_service", "grpc_method", "le");
+  }
+
+  @Test
+  public void addsStatusCodeLabel() throws Throwable {
+    startGrpcServer(ALL_METRICS.withCodeLabelInLatencyHistogram());
+    createGrpcBlockingStub().sayHello(REQUEST);
+
+    MetricFamilySamples.Sample sample =
+        getSample(
+            findRecordedMetricOrThrow("grpc_server_handled_latency_seconds"),
+            "grpc_server_handled_latency_seconds_bucket");
+
+    assertThat(sample.labelNames)
+        .containsExactly("grpc_type", "grpc_service", "grpc_method", "grpc_code", "le");
   }
 
   @Test