Add "get[Low|High]CardinalityKeyValue()" on "Observation.Context"

Add two new methods `getLowCardinalityKeyValue` and
`getHighCardinalityKeyValue`.
Also, changes the internal holder from `Set` to `Map`.

Prior to this change, the usage of `LinkedHashSet` allowed multiple
`KeyValue`s with the same key. It was deduped when converted to
`KeyValues`, but it is not an ideal implementation.

This change replaces the `LinkedHashSet` with `LinkedHashMap` using the
key from `KeyValue` as a map key. This way, when a new `KeyValue` is
added, it will override the old `KeyValue` with the same key.

Author: ttddyy

micrometer-observation/src/main/java/io/micrometer/observation/Observation.java

@@ -20,7 +20,9 @@
 import io.micrometer.common.lang.NonNull;
 import io.micrometer.common.lang.Nullable;
 
-import java.util.*;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -751,9 +753,9 @@ class Context implements ContextView {
         @Nullable
         private ObservationView parentObservation;
 
-        private final Set<KeyValue> lowCardinalityKeyValues = new LinkedHashSet<>();
+        private final Map<String, KeyValue> lowCardinalityKeyValues = new LinkedHashMap<>();
 
-        private final Set<KeyValue> highCardinalityKeyValues = new LinkedHashSet<>();
+        private final Map<String, KeyValue> highCardinalityKeyValues = new LinkedHashMap<>();
 
         /**
          * The observation name.
@@ -923,7 +925,7 @@ public void clear() {
          * @return this context
          */
         public Context addLowCardinalityKeyValue(KeyValue keyValue) {
-            this.lowCardinalityKeyValues.add(keyValue);
+            this.lowCardinalityKeyValues.put(keyValue.getKey(), keyValue);
             return this;
         }
 
@@ -935,7 +937,7 @@ public Context addLowCardinalityKeyValue(KeyValue keyValue) {
          * @return this context
          */
         public Context addHighCardinalityKeyValue(KeyValue keyValue) {
-            this.highCardinalityKeyValues.add(keyValue);
+            this.highCardinalityKeyValues.put(keyValue.getKey(), keyValue);
             return this;
         }
 
@@ -962,13 +964,23 @@ public Context addHighCardinalityKeyValues(KeyValues keyValues) {
         @NonNull
         @Override
         public KeyValues getLowCardinalityKeyValues() {
-            return KeyValues.of(this.lowCardinalityKeyValues);
+            return KeyValues.of(this.lowCardinalityKeyValues.values());
         }
 
         @NonNull
         @Override
         public KeyValues getHighCardinalityKeyValues() {
-            return KeyValues.of(this.highCardinalityKeyValues);
+            return KeyValues.of(this.highCardinalityKeyValues.values());
+        }
+
+        @Override
+        public KeyValue getLowCardinalityKeyValue(String key) {
+            return this.lowCardinalityKeyValues.get(key);
+        }
+
+        @Override
+        public KeyValue getHighCardinalityKeyValue(String key) {
+            return this.highCardinalityKeyValues.get(key);
         }
 
         @NonNull
@@ -1149,6 +1161,22 @@ interface ContextView {
         @NonNull
         KeyValues getHighCardinalityKeyValues();
 
+        /**
+         * Returns a low cardinality key value or {@code null} if not present.
+         * @param key key
+         * @return a low cardinality key value or {@code null}
+         */
+        @Nullable
+        KeyValue getLowCardinalityKeyValue(String key);
+
+        /**
+         * Returns a high cardinality key value or {@code null} if not present.
+         * @param key key
+         * @return a high cardinality key value or {@code null}
+         */
+        @Nullable
+        KeyValue getHighCardinalityKeyValue(String key);
+
         /**
          * Returns all key values.
          * @return all key values

micrometer-observation/src/test/java/io/micrometer/observation/ObservationContextTest.java

@@ -15,6 +15,7 @@
  */
 package io.micrometer.observation;
 
+import io.micrometer.common.KeyValue;
 import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -109,4 +110,25 @@ void cleanEmptyContextShouldNotFail() {
         context.clear();
     }
 
+    @Test
+    void sameKeyShouldOverrideKeyValue() {
+        KeyValue low = KeyValue.of("low", "LOW");
+        KeyValue high = KeyValue.of("high", "HIGH");
+        context.addLowCardinalityKeyValue(low);
+        context.addHighCardinalityKeyValue(high);
+
+        assertThat(context.getLowCardinalityKeyValue("low")).isSameAs(low);
+        assertThat(context.getHighCardinalityKeyValue("high")).isSameAs(high);
+
+        KeyValue newLow = KeyValue.of("low", "LOW-NEW");
+        KeyValue newHigh = KeyValue.of("high", "HIGH-NEW");
+        context.addLowCardinalityKeyValue(newLow);
+        context.addHighCardinalityKeyValue(newHigh);
+
+        assertThat(context.getLowCardinalityKeyValue("low")).isSameAs(newLow);
+        assertThat(context.getHighCardinalityKeyValue("high")).isSameAs(newHigh);
+        assertThat(context.getLowCardinalityKeyValues()).containsExactly(newLow);
+        assertThat(context.getHighCardinalityKeyValues()).containsExactly(newHigh);
+    }
+
 }

micrometer-observation/src/test/java/io/micrometer/observation/ObservationTests.java

@@ -15,6 +15,7 @@
  */
 package io.micrometer.observation;
 
+import io.micrometer.common.KeyValue;
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
@@ -78,6 +79,27 @@ void havingAnObservationFilterWillMutateTheContext() {
         assertThat((String) context.get("foo")).isEqualTo("bar");
     }
 
+    @Test
+    void havingAnObservationFilterToModifyKeyValue() {
+        registry.observationConfig().observationHandler(context -> true);
+
+        Observation.Context context = new Observation.Context();
+        context.addHighCardinalityKeyValue(KeyValue.of("foo", "FOO"));
+
+        ObservationFilter filter = (ctx) -> {
+            KeyValue keyValue = ctx.getHighCardinalityKeyValue("foo");
+            assertThat(keyValue).isNotNull();
+            return ctx.addHighCardinalityKeyValue(KeyValue.of(keyValue.getKey(), keyValue.getValue() + "-modified"));
+        };
+
+        registry.observationConfig().observationFilter(filter);
+
+        Observation.start("foo", () -> context, registry).stop();
+
+        assertThat(context.getHighCardinalityKeyValue("foo")).isNotNull().extracting(KeyValue::getValue)
+                .isEqualTo("FOO-modified");
+    }
+
     @Test
     void settingParentObservationMakesAReferenceOnParentContext() {
         registry.observationConfig().observationHandler(context -> true);