Allow DNSToSwitchMapping to access BookieAddressResolver

In order to allow Pulsar to leverage BookieId with custom Bookie-To-Rack Mapping (ZkBookieRackAffinityMapping) we need to inject the BookieAddressResolver into it.

There is not way to do it currently because it is created by the BookKeeper client, and we have to inject the same BookieAddressResolver used by the client itself.

*Changes*
- add new setBookieAddressResolver method (use 'default' methods in order to preserve backward compatibility
- call the new method in every point in which we instantiate a DNSToSwitchMapping and when we receive an instance at construction time

*Compatibility*
- this change is 100% compatible with previous versions

Reviewers: Rajan Dhabalia <rdhabalia@apache.org>, Anup Ghatage <ghatage@apache.org>

This closes #2519 from eolivelli/fix/dnsswitch-access-bookieaddress-resolver
diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java
index 0a594d3..c34defa 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java
@@ -499,7 +499,9 @@
 
         BookieAddressResolver bookieAddressResolver =
                 new DefaultBookieAddressResolver(metadataDriver.getRegistrationClient());
-
+        if (dnsResolver != null) {
+            dnsResolver.setBookieAddressResolver(bookieAddressResolver);
+        }
         // initialize the ensemble placement
         this.placementPolicy = initializeEnsemblePlacementPolicy(conf,
                 dnsResolver, this.requestTimer, this.featureProvider, this.statsLogger, bookieAddressResolver);
diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImpl.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImpl.java
index 5568f0a..8ee3f69 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImpl.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImpl.java
@@ -282,6 +282,9 @@
                 }
             }
         }
+        if (dnsResolver != null) {
+            dnsResolver.setBookieAddressResolver(bookieAddressResolver);
+        }
         slowBookies = CacheBuilder.newBuilder()
             .expireAfterWrite(conf.getBookieFailureHistoryExpirationMSec(), TimeUnit.MILLISECONDS)
             .build(new CacheLoader<BookieId, Long>() {
diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/TopologyAwareEnsemblePlacementPolicy.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/TopologyAwareEnsemblePlacementPolicy.java
index 7167c96..438053f 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/TopologyAwareEnsemblePlacementPolicy.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/TopologyAwareEnsemblePlacementPolicy.java
@@ -546,6 +546,11 @@
         }
 
         @Override
+        public void setBookieAddressResolver(BookieAddressResolver bookieAddressResolver) {
+            this.resolver.setBookieAddressResolver(bookieAddressResolver);
+        }
+
+        @Override
         public List<String> resolve(List<String> names) {
             if (names == null) {
                 return Collections.emptyList();
diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/ZoneawareEnsemblePlacementPolicyImpl.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/ZoneawareEnsemblePlacementPolicyImpl.java
index fcb3c73..990d60f 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/ZoneawareEnsemblePlacementPolicyImpl.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/ZoneawareEnsemblePlacementPolicyImpl.java
@@ -240,6 +240,7 @@
 
         this.dnsResolver = new DNSResolverDecorator(actualDNSResolver, () -> this.getDefaultFaultDomain(),
                 failedToResolveNetworkLocationCounter);
+        dnsResolver.setBookieAddressResolver(bookieAddressResolver);
         this.stabilizePeriodSeconds = conf.getNetworkTopologyStabilizePeriodSeconds();
         // create the network topology
         if (stabilizePeriodSeconds > 0) {
diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/net/AbstractDNSToSwitchMapping.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/net/AbstractDNSToSwitchMapping.java
index a0ee422..d658d8e 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/net/AbstractDNSToSwitchMapping.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/net/AbstractDNSToSwitchMapping.java
@@ -22,6 +22,7 @@
 import java.util.Set;
 
 import org.apache.bookkeeper.conf.Configurable;
+import org.apache.bookkeeper.proto.BookieAddressResolver;
 import org.apache.commons.configuration.Configuration;
 import org.apache.commons.lang.StringUtils;
 
@@ -42,6 +43,7 @@
 public abstract class AbstractDNSToSwitchMapping implements DNSToSwitchMapping, Configurable {
 
     private Configuration conf;
+    private BookieAddressResolver bookieAddressResolver;
 
     /**
      * Create an unconfigured instance.
@@ -59,6 +61,15 @@
         this.conf = conf;
     }
 
+    public BookieAddressResolver getBookieAddressResolver() {
+        return bookieAddressResolver;
+    }
+
+    @Override
+    public void setBookieAddressResolver(BookieAddressResolver bookieAddressResolver) {
+        this.bookieAddressResolver = bookieAddressResolver;
+    }
+
     @Override
     public Configuration getConf() {
         return conf;
diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/net/DNSToSwitchMapping.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/net/DNSToSwitchMapping.java
index bdf9ce5..c2a8152 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/net/DNSToSwitchMapping.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/net/DNSToSwitchMapping.java
@@ -20,6 +20,7 @@
 import com.google.common.annotations.Beta;
 
 import java.util.List;
+import org.apache.bookkeeper.proto.BookieAddressResolver;
 
 /**
  * An interface that must be implemented to allow pluggable
@@ -67,4 +68,11 @@
     default boolean useHostName() {
         return true;
     }
+
+    /**
+     * Receives the current BookieAddressResolver.
+     * @param bookieAddressResolver
+     */
+    default void setBookieAddressResolver(BookieAddressResolver bookieAddressResolver) {
+    }
 }
diff --git a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/BookKeeperTest.java b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/BookKeeperTest.java
index c19e391..d1ad5fd 100644
--- a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/BookKeeperTest.java
+++ b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/BookKeeperTest.java
@@ -27,6 +27,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -58,6 +59,7 @@
 import org.apache.bookkeeper.stats.StatsLogger;
 import org.apache.bookkeeper.test.BookKeeperClusterTestCase;
 import org.apache.bookkeeper.test.TestStatsProvider;
+import org.apache.bookkeeper.util.StaticDNSResolver;
 import org.apache.bookkeeper.zookeeper.BoundExponentialBackoffRetryPolicy;
 import org.apache.bookkeeper.zookeeper.ZooKeeperClient;
 import org.apache.bookkeeper.zookeeper.ZooKeeperWatcherBase;
@@ -85,7 +87,7 @@
     private final DigestType digestType;
 
     public BookKeeperTest() {
-        super(4);
+        super(3);
         this.digestType = DigestType.CRC32;
     }
 
@@ -1112,4 +1114,20 @@
                          1);
         }
     }
+
+    @Test
+    public void testBookieAddressResolverPassedToDNSToSwitchMapping() throws Exception {
+        ClientConfiguration conf = new ClientConfiguration();
+        conf.setMetadataServiceUri(zkUtil.getMetadataServiceUri());
+
+        StaticDNSResolver tested = new StaticDNSResolver();
+        try (BookKeeper bkc = BookKeeper
+                        .forConfig(conf)
+                        .dnsResolver(tested)
+                        .build()) {
+            bkc.createLedger(digestType, "testPasswd".getBytes()).close();
+            assertSame(bkc.getBookieAddressResolver(), tested.getBookieAddressResolver());
+        }
+    }
+
 }
diff --git a/bookkeeper-server/src/test/java/org/apache/bookkeeper/util/StaticDNSResolver.java b/bookkeeper-server/src/test/java/org/apache/bookkeeper/util/StaticDNSResolver.java
index f82310a..43a1cf6 100644
--- a/bookkeeper-server/src/test/java/org/apache/bookkeeper/util/StaticDNSResolver.java
+++ b/bookkeeper-server/src/test/java/org/apache/bookkeeper/util/StaticDNSResolver.java
@@ -73,6 +73,10 @@
 
     @Override
     public List<String> resolve(List<String> names) {
+        if (getBookieAddressResolver() == null) {
+            // test that this istance has been properly initialized
+            throw new IllegalStateException("bookieAddressResolver was not set");
+        }
         List<String> racks = new ArrayList<String>();
         for (String n : names) {
             String rack = name2Racks.get(n);