From 2a21414bd0376274f3fa8e6b2ae1e9bcb5ebf1a7 Mon Sep 17 00:00:00 2001
From: tangcent <pentatangcent@gmail.com>
Date: Fri, 26 Apr 2024 15:51:39 +0800
Subject: [PATCH] Add parameter to customize status codes mapped as 'NOT_FOUND'
 in OkHttpMetricsEventListener

---
 .../okhttp3/OkHttpMetricsEventListener.java   | 49 ++++++++++++++++---
 .../OkHttpMetricsEventListenerTest.java       | 23 +++++++++
 2 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/micrometer-core/src/main/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpMetricsEventListener.java b/micrometer-core/src/main/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpMetricsEventListener.java
index 8b1809303d..f8326339b8 100644
--- a/micrometer-core/src/main/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpMetricsEventListener.java
+++ b/micrometer-core/src/main/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpMetricsEventListener.java
@@ -95,6 +95,8 @@ private static Method getMethod(Class<?>... parameterTypes) {
 
     private final Function<Request, String> urlMapper;
 
+    private final Set<Integer> statusCodesMappedAsNotFound;
+
     private final Iterable<Tag> extraTags;
 
     private final Iterable<BiFunction<Request, Response, Tag>> contextSpecificTags;
@@ -107,17 +109,20 @@ private static Method getMethod(Class<?>... parameterTypes) {
     final ConcurrentMap<Call, CallState> callState = new ConcurrentHashMap<>();
 
     protected OkHttpMetricsEventListener(MeterRegistry registry, String requestsMetricName,
-            Function<Request, String> urlMapper, Iterable<Tag> extraTags,
+            Function<Request, String> urlMapper, Set<Integer> statusCodesMappedAsNotFound, Iterable<Tag> extraTags,
             Iterable<BiFunction<Request, Response, Tag>> contextSpecificTags) {
-        this(registry, requestsMetricName, urlMapper, extraTags, contextSpecificTags, emptyList(), true);
+        this(registry, requestsMetricName, urlMapper, statusCodesMappedAsNotFound, extraTags, contextSpecificTags,
+                emptyList(), true);
     }
 
     OkHttpMetricsEventListener(MeterRegistry registry, String requestsMetricName, Function<Request, String> urlMapper,
-            Iterable<Tag> extraTags, Iterable<BiFunction<Request, Response, Tag>> contextSpecificTags,
-            Iterable<String> requestTagKeys, boolean includeHostTag) {
+            Set<Integer> statusCodesMappedAsNotFound, Iterable<Tag> extraTags,
+            Iterable<BiFunction<Request, Response, Tag>> contextSpecificTags, Iterable<String> requestTagKeys,
+            boolean includeHostTag) {
         this.registry = registry;
         this.requestsMetricName = requestsMetricName;
         this.urlMapper = urlMapper;
+        this.statusCodesMappedAsNotFound = statusCodesMappedAsNotFound;
         this.extraTags = extraTags;
         this.contextSpecificTags = contextSpecificTags;
         this.includeHostTag = includeHostTag;
@@ -200,7 +205,7 @@ private String getUriTag(CallState state, @Nullable Request request) {
         if (request == null) {
             return TAG_VALUE_UNKNOWN;
         }
-        return state.response != null && (state.response.code() == 404 || state.response.code() == 301) ? "NOT_FOUND"
+        return state.response != null && statusCodesMappedAsNotFound.contains(state.response.code()) ? "NOT_FOUND"
                 : urlMapper.apply(request);
     }
 
@@ -271,6 +276,8 @@ public static class Builder {
         private Function<Request, String> uriMapper = (request) -> Optional.ofNullable(request.header(URI_PATTERN))
             .orElse("none");
 
+        private Set<Integer> statusCodesMappedAsNotFound = new HashSet<>(Arrays.asList(301, 404));
+
         private Tags tags = Tags.empty();
 
         private Collection<BiFunction<Request, Response, Tag>> contextSpecificTags = new ArrayList<>();
@@ -316,6 +323,34 @@ public Builder uriMapper(Function<Request, String> uriMapper) {
             return this;
         }
 
+        /**
+         * Sets specific status codes to be mapped to a generic "NOT_FOUND" category for
+         * the 'uri' tag when recording metrics. This mapping helps in avoiding tag
+         * explosion and reducing metric dimensionality by treating these status codes
+         * identically.
+         * @param statusCodesMappedAsNotFound One or more status codes to be treated as
+         * "NOT_FOUND".
+         * @return this builder for chaining
+         */
+        public Builder statusCodesMappedAsNotFound(Integer... statusCodesMappedAsNotFound) {
+            return statusCodesMappedAsNotFound(Arrays.asList(statusCodesMappedAsNotFound));
+        }
+
+        /**
+         * Sets specific status codes to be mapped to a generic "NOT_FOUND" category for
+         * the 'uri' tag when recording metrics. This mapping helps in avoiding tag
+         * explosion and reducing metric dimensionality by treating these status codes
+         * identically.
+         * @param statusCodesMappedAsNotFound One or more status codes to be treated as
+         * "NOT_FOUND".
+         * @return this builder for chaining
+         */
+        public Builder statusCodesMappedAsNotFound(Iterable<Integer> statusCodesMappedAsNotFound) {
+            this.statusCodesMappedAsNotFound = new HashSet<>();
+            statusCodesMappedAsNotFound.forEach(this.statusCodesMappedAsNotFound::add);
+            return this;
+        }
+
         /**
          * Historically, OkHttp Metrics provided by {@link OkHttpMetricsEventListener}
          * included a {@code host} tag for the target host being called. To align with
@@ -361,8 +396,8 @@ public Builder requestTagKeys(Iterable<String> requestTagKeys) {
         }
 
         public OkHttpMetricsEventListener build() {
-            return new OkHttpMetricsEventListener(registry, name, uriMapper, tags, contextSpecificTags, requestTagKeys,
-                    includeHostTag);
+            return new OkHttpMetricsEventListener(registry, name, uriMapper, statusCodesMappedAsNotFound, tags,
+                    contextSpecificTags, requestTagKeys, includeHostTag);
         }
 
     }
diff --git a/micrometer-core/src/test/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpMetricsEventListenerTest.java b/micrometer-core/src/test/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpMetricsEventListenerTest.java
index 4a14a91588..3901509ed0 100644
--- a/micrometer-core/src/test/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpMetricsEventListenerTest.java
+++ b/micrometer-core/src/test/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpMetricsEventListenerTest.java
@@ -162,6 +162,29 @@ void uriTagWorksWithUriMapper(@WiremockResolver.Wiremock WireMockServer server)
             .count()).isEqualTo(1L);
     }
 
+    @Test
+    void timeWithStatusCodesMappedAsNotFound(@WiremockResolver.Wiremock WireMockServer server) throws IOException {
+        server.stubFor(any(anyUrl()).willReturn(aResponse().withStatus(404)));
+        Request request = new Request.Builder().url(server.baseUrl()).build();
+
+        OkHttpClient client = new OkHttpClient.Builder()
+            .eventListener(OkHttpMetricsEventListener.builder(registry, "okhttp.requests")
+                .statusCodesMappedAsNotFound(404) // Specifying that 404 should be treated
+                                                  // as 'NOT_FOUND'
+                .uriMapper(URI_MAPPER)
+                .tags(Tags.of("foo", "bar"))
+                .build())
+            .build();
+
+        client.newCall(request).execute().close();
+
+        assertThat(registry.get("okhttp.requests")
+            .tags("foo", "bar", "status", "404", "outcome", "CLIENT_ERROR", "uri", "NOT_FOUND", "target.host",
+                    "localhost", "target.port", String.valueOf(server.port()), "target.scheme", "http")
+            .timer()
+            .count()).isEqualTo(1L);
+    }
+
     @Test
     void contextSpecificTags(@WiremockResolver.Wiremock WireMockServer server) throws IOException {
         server.stubFor(any(anyUrl()));