diff --git a/core/src/main/java/com/linecorp/armeria/client/DefaultDnsCache.java b/core/src/main/java/com/linecorp/armeria/client/DefaultDnsCache.java index ab4634d2197..a9b2b6e2f0f 100644 --- a/core/src/main/java/com/linecorp/armeria/client/DefaultDnsCache.java +++ b/core/src/main/java/com/linecorp/armeria/client/DefaultDnsCache.java @@ -33,6 +33,7 @@ import com.github.benmanes.caffeine.cache.Caffeine; import com.github.benmanes.caffeine.cache.RemovalCause; import com.github.benmanes.caffeine.cache.RemovalListener; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.MoreObjects; import com.google.common.base.MoreObjects.ToStringHelper; import com.google.common.base.Objects; @@ -179,6 +180,17 @@ public void addListener(DnsCacheListener listener) { listeners.add(listener); } + @Override + public void removeListener(DnsCacheListener listener) { + requireNonNull(listener, "listener"); + listeners.remove(listener); + } + + @VisibleForTesting + List listeners() { + return listeners; + } + private static final class CacheEntry { @Nullable diff --git a/core/src/main/java/com/linecorp/armeria/client/DnsCache.java b/core/src/main/java/com/linecorp/armeria/client/DnsCache.java index c3f90a20a02..3728529fd84 100644 --- a/core/src/main/java/com/linecorp/armeria/client/DnsCache.java +++ b/core/src/main/java/com/linecorp/armeria/client/DnsCache.java @@ -109,4 +109,9 @@ default void cache(DnsQuestion question, DnsRecord... records) { * occurs. */ void addListener(DnsCacheListener listener); + + /** + * Removes the specified {@link DnsCacheListener} from this {@link DnsCache}. + */ + void removeListener(DnsCacheListener listener); } diff --git a/core/src/main/java/com/linecorp/armeria/client/RefreshingAddressResolver.java b/core/src/main/java/com/linecorp/armeria/client/RefreshingAddressResolver.java index 1b31102eb1f..b27403ca9c6 100644 --- a/core/src/main/java/com/linecorp/armeria/client/RefreshingAddressResolver.java +++ b/core/src/main/java/com/linecorp/armeria/client/RefreshingAddressResolver.java @@ -67,6 +67,7 @@ final class RefreshingAddressResolver @Nullable private final ToLongFunction autoRefreshTimeoutFunction; private final boolean autoRefresh; + private final DnsCache dnsResolverCache; private volatile boolean resolverClosed; @@ -78,6 +79,7 @@ final class RefreshingAddressResolver @Nullable ToLongFunction autoRefreshTimeoutFunction) { super(eventLoop); this.addressResolverCache = addressResolverCache; + this.dnsResolverCache = dnsResolverCache; this.resolver = resolver; this.dnsRecordTypes = dnsRecordTypes; this.negativeTtl = negativeTtl; @@ -258,6 +260,7 @@ public void onEviction(DnsQuestion question, @Nullable List records, @Override public void close() { resolverClosed = true; + dnsResolverCache.removeListener(this); resolver.close(); } diff --git a/core/src/test/java/com/linecorp/armeria/client/NoopDnsCache.java b/core/src/test/java/com/linecorp/armeria/client/NoopDnsCache.java index 1fe20869198..e946e192f07 100644 --- a/core/src/test/java/com/linecorp/armeria/client/NoopDnsCache.java +++ b/core/src/test/java/com/linecorp/armeria/client/NoopDnsCache.java @@ -50,4 +50,7 @@ public void removeAll() {} @Override public void addListener(DnsCacheListener listener) {} + + @Override + public void removeListener(DnsCacheListener listener) {} } diff --git a/core/src/test/java/com/linecorp/armeria/client/RefreshingAddressResolverTest.java b/core/src/test/java/com/linecorp/armeria/client/RefreshingAddressResolverTest.java index 3a524b90d2d..9600b1748d1 100644 --- a/core/src/test/java/com/linecorp/armeria/client/RefreshingAddressResolverTest.java +++ b/core/src/test/java/com/linecorp/armeria/client/RefreshingAddressResolverTest.java @@ -905,6 +905,22 @@ public void onEviction(DnsQuestion question, @Nullable List records, } } + @Test + void shouldRemoveListenerFromDnsCacheWhenClosed() { + final EventLoop eventLoop = eventLoopExtension.get(); + final DnsCache dnsCache = DnsCache.builder() + .executor(eventLoop) + .build(); + try (RefreshingAddressResolverGroup group = new DnsResolverGroupBuilder() + .dnsCache(dnsCache) + .build(eventLoop)) { + final AddressResolver resolver = group.getResolver(eventLoop); + resolver.close(); + final DefaultDnsCache defaultDnsCache = (DefaultDnsCache) dnsCache; + assertThat(defaultDnsCache.listeners()).isEmpty(); + } + } + private static AbstractStringAssert assertIpAddress(Future staticAddr) { await().untilAsserted(() -> assertThat(staticAddr.isSuccess()).isTrue()); return assertThat(NetUtil.bytesToIpAddress(staticAddr.getNow().getAddress().getAddress()));