diff --git a/core/src/main/java/com/linecorp/armeria/client/Endpoint.java b/core/src/main/java/com/linecorp/armeria/client/Endpoint.java index cfd6ceac464..500e8622ab9 100644 --- a/core/src/main/java/com/linecorp/armeria/client/Endpoint.java +++ b/core/src/main/java/com/linecorp/armeria/client/Endpoint.java @@ -696,7 +696,7 @@ public Endpoint withAttr(AttributeKey key, @Nullable T value) { if (value == null) { return this; } - return withAttrs(Attributes.of(key, value)); + return replaceAttrs(Attributes.of(key, value)); } if (attributes.attr(key) == value) { @@ -704,17 +704,41 @@ public Endpoint withAttr(AttributeKey key, @Nullable T value) { } else { final AttributesBuilder attributesBuilder = attributes.toBuilder(); attributesBuilder.set(key, value); - return withAttrs(attributesBuilder.build()); + return replaceAttrs(attributesBuilder.build()); } } + /** + * Returns a new {@link Endpoint} with the specified {@link Attributes}. + * Note that the {@link #attrs()} of this {@link Endpoint} is merged with the specified + * {@link Attributes}. For attributes with the same {@link AttributeKey}, the attribute + * in {@param newAttributes} has higher precedence. + */ + @UnstableApi + @SuppressWarnings("unchecked") + public Endpoint withAttrs(Attributes newAttributes) { + requireNonNull(newAttributes, "newAttributes"); + if (newAttributes.isEmpty()) { + return this; + } + if (attrs().isEmpty()) { + return replaceAttrs(newAttributes); + } + final AttributesBuilder builder = attrs().toBuilder(); + newAttributes.attrs().forEachRemaining(entry -> { + final AttributeKey key = (AttributeKey) entry.getKey(); + builder.set(key, entry.getValue()); + }); + return new Endpoint(type, host, ipAddr, port, weight, builder.build()); + } + /** * Returns a new {@link Endpoint} with the specified {@link Attributes}. * Note that the {@link #attrs()} of this {@link Endpoint} is replaced with the specified * {@link Attributes}. */ @UnstableApi - public Endpoint withAttrs(Attributes newAttributes) { + public Endpoint replaceAttrs(Attributes newAttributes) { requireNonNull(newAttributes, "newAttributes"); if (attrs().isEmpty() && newAttributes.isEmpty()) { return this; diff --git a/core/src/test/java/com/linecorp/armeria/client/EndpointTest.java b/core/src/test/java/com/linecorp/armeria/client/EndpointTest.java index 4502bcad6c1..e1b6a04633f 100644 --- a/core/src/test/java/com/linecorp/armeria/client/EndpointTest.java +++ b/core/src/test/java/com/linecorp/armeria/client/EndpointTest.java @@ -745,8 +745,8 @@ void attrs() { attrs2.set(key1, "value1-2"); attrs2.set(key3, "value3"); - final Endpoint endpointB = endpoint.withAttrs(attrs.build()); - final Endpoint endpointC = endpointB.withAttrs(attrs2.build()); + final Endpoint endpointB = endpoint.replaceAttrs(attrs.build()); + final Endpoint endpointC = endpointB.replaceAttrs(attrs2.build()); assertThat(endpointB.attr(key1)) .isEqualTo("value1"); @@ -767,10 +767,40 @@ void attrs() { Maps.immutableEntry(key3, "value3")); // Reset attrs with an empty attributes. - final Endpoint newEndpointB = endpointB.withAttrs(Attributes.of()); + final Endpoint newEndpointB = endpointB.replaceAttrs(Attributes.of()); assertThat(newEndpointB.attrs().isEmpty()).isTrue(); - final Endpoint sameEndpoint = endpoint.withAttrs(Attributes.of()); + final Endpoint sameEndpoint = endpoint.replaceAttrs(Attributes.of()); assertThat(sameEndpoint).isSameAs(endpoint); } + + @Test + void withAttrs() { + final AttributeKey key1 = AttributeKey.valueOf("key1"); + final AttributeKey key2 = AttributeKey.valueOf("key2"); + + final Endpoint endpoint = Endpoint.parse("a").withAttrs(Attributes.of(key1, "val1")); + assertThat(endpoint.attrs().attrs()) + .toIterable() + .containsExactlyInAnyOrder(Maps.immutableEntry(key1, "val1")); + + assertThat(endpoint.withAttrs(Attributes.of(key2, "val2")).attrs().attrs()) + .toIterable() + .containsExactlyInAnyOrder(Maps.immutableEntry(key1, "val1"), + Maps.immutableEntry(key2, "val2")); + + assertThat(endpoint.withAttrs(Attributes.of(key1, "val1")).attrs().attrs()) + .toIterable() + .containsExactlyInAnyOrder(Maps.immutableEntry(key1, "val1")); + + assertThat(endpoint.withAttrs(Attributes.of(key1, "val2")).attrs().attrs()) + .toIterable() + .containsExactlyInAnyOrder(Maps.immutableEntry(key1, "val2")); + + assertThat(endpoint.withAttrs(Attributes.of()).attrs().attrs()) + .toIterable() + .containsExactlyInAnyOrder(Maps.immutableEntry(key1, "val1")); + + assertThat(endpoint.withAttrs(Attributes.of())).isSameAs(endpoint); + } }