[client] new API to check if Bookkeeper client is connected to metadata service

bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java

@@ -1609,6 +1609,11 @@ public CompletableFuture<LedgerMetadata> getLedgerMetadata(long ledgerId) {
         });
     }
 
+    @Override
+    public CompletableFuture<Boolean> isDriverMetadataServiceAvailable() {
+        return metadataDriver.isMetadataServiceAvailable();
+    }
+
     private final ClientContext clientCtx = new ClientContext() {
             @Override
             public ClientInternalConf getConf() {

bookkeeper-server/src/main/java/org/apache/bookkeeper/client/api/BookKeeper.java

@@ -81,6 +81,13 @@ static BookKeeperBuilder newBuilder(final ClientConfiguration clientConfiguratio
      */
     CompletableFuture<LedgerMetadata> getLedgerMetadata(long ledgerId);
 
+    /**
+     * Return driver metadata service is available.
+     *
+     * @return the metadata service is available.
+     */
+    CompletableFuture<Boolean> isDriverMetadataServiceAvailable();
+
     /**
      * Close the client and release every resource.
      *

bookkeeper-server/src/main/java/org/apache/bookkeeper/meta/MetadataClientDriver.java

@@ -114,4 +114,13 @@ interface SessionStateListener {
     default CompletableFuture<Boolean> isHealthCheckEnabled() {
         return FutureUtils.value(true);
     }
+
+    /**
+     * Return driver metadata service is available.
+     *
+     * @return the metadata service is available.
+     */
+    default CompletableFuture<Boolean> isMetadataServiceAvailable() {
+        return FutureUtils.value(true);
+    }
 }

bookkeeper-server/src/main/java/org/apache/bookkeeper/meta/zk/ZKMetadataClientDriver.java

@@ -20,6 +20,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.bookkeeper.conf.ClientConfiguration;
@@ -111,4 +112,8 @@ public void setSessionStateListener(SessionStateListener sessionStateListener) {
             }
         });
     }
+
+    public CompletableFuture<Boolean> isMetadataServiceAvailable() {
+        return CompletableFuture.completedFuture(metadataServiceAvailable);
+    }
 }

bookkeeper-server/src/main/java/org/apache/bookkeeper/meta/zk/ZKMetadataDriverBase.java

@@ -27,6 +27,7 @@
 
 import java.io.IOException;
 import java.net.URI;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -53,6 +54,7 @@
 import org.apache.zookeeper.AsyncCallback;
 import org.apache.zookeeper.CreateMode;
 import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.Watcher;
 import org.apache.zookeeper.ZooKeeper;
 import org.apache.zookeeper.data.ACL;
 import org.apache.zookeeper.data.Stat;
@@ -64,6 +66,9 @@
 public class ZKMetadataDriverBase implements AutoCloseable {
 
     protected static final String SCHEME = "zk";
+
+    protected volatile boolean metadataServiceAvailable;
+
     private static final int ZK_CLIENT_WAIT_FOR_SHUTDOWN_TIMEOUT_MS = 5000;
 
     public static String getZKServersFromServiceUri(URI uri) {
@@ -179,6 +184,7 @@ protected void initialize(AbstractConfiguration<?> conf,
             // if an external zookeeper is added, use the zookeeper instance
             this.zk = (ZooKeeper) (optionalCtx.get());
             this.ownZKHandle = false;
+            this.metadataServiceAvailable = true;
         } else {
             final String metadataServiceUriStr;
             try {
@@ -212,6 +218,12 @@ protected void initialize(AbstractConfiguration<?> conf,
                     .sessionTimeoutMs(conf.getZkTimeout())
                     .operationRetryPolicy(zkRetryPolicy)
                     .requestRateLimit(conf.getZkRequestRateLimit())
+                    .watchers(Collections.singleton(watchedEvent -> {
+                        if (log.isDebugEnabled()) {
+                            log.debug("Got ZK session watch event: {}", watchedEvent);
+                        }
+                        handleState(watchedEvent.getState());
+                    }))
                     .statsLogger(statsLogger)
                     .build();
 
@@ -247,6 +259,17 @@ protected void initialize(AbstractConfiguration<?> conf,
             acls);
     }
 
+    private void handleState(Watcher.Event.KeeperState zkClientState) {
+        switch (zkClientState) {
+            case Expired:
+            case Disconnected:
+                this.metadataServiceAvailable = false;
+                break;
+            default:
+                this.metadataServiceAvailable = true;
+        }
+    }
+
     public LayoutManager getLayoutManager() {
         return layoutManager;
     }

bookkeeper-server/src/test/java/org/apache/bookkeeper/client/api/DriverMetadataServiceAvailableTest.java

@@ -0,0 +1,52 @@
+/*
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+package org.apache.bookkeeper.client.api;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import org.apache.bookkeeper.conf.ClientConfiguration;
+import org.apache.bookkeeper.test.BookKeeperClusterTestCase;
+import org.awaitility.Awaitility;
+import org.junit.Test;
+
+/**
+ * Bookkeeper Client API driver metadata service available test.
+ */
+public class DriverMetadataServiceAvailableTest extends BookKeeperClusterTestCase {
+
+    public DriverMetadataServiceAvailableTest() {
+        super(3);
+    }
+
+    @Test
+    public void testDriverMetadataServiceAvailable()
+            throws Exception {
+        ClientConfiguration conf = new ClientConfiguration();
+        conf.setMetadataServiceUri(zkUtil.getMetadataServiceUri());
+        conf.setZkTimeout(3000);
+        try (BookKeeper bkc = BookKeeper.newBuilder(conf).build()) {
+            Awaitility.await().until(() -> bkc.isDriverMetadataServiceAvailable().get());
+            zkUtil.sleepCluster(5, TimeUnit.SECONDS, new CountDownLatch(1));
+            Awaitility.await().until(() -> !bkc.isDriverMetadataServiceAvailable().get());
+            Awaitility.await().until(() -> bkc.isDriverMetadataServiceAvailable().get());
+        }
+    }
+}

bookkeeper-server/src/test/java/org/apache/bookkeeper/meta/zk/ZKMetadataDriverTestBase.java

@@ -57,6 +57,7 @@ public void setup(AbstractConfiguration<?> conf) throws Exception {
         when(mockZkBuilder.operationRetryPolicy(any(RetryPolicy.class)))
             .thenReturn(mockZkBuilder);
         when(mockZkBuilder.requestRateLimit(anyDouble())).thenReturn(mockZkBuilder);
+        when(mockZkBuilder.watchers(any())).thenReturn(mockZkBuilder);
         when(mockZkBuilder.statsLogger(any(StatsLogger.class))).thenReturn(mockZkBuilder);
 
         this.mockZkc = mock(ZooKeeperClient.class);

bookkeeper-server/src/test/java/org/apache/bookkeeper/sasl/MD5DigestBookKeeperTest.java

@@ -118,7 +118,6 @@ public void testSingleMessageAuth() throws Exception {
         bookieConf.setBookieAuthProviderFactoryClass(
             SASLBookieAuthProviderFactory.class.getName());
         bookieConf.setProperty(JAAS_CLIENT_ALLOWED_IDS, ".*hd.*");
-
         ClientConfiguration clientConf = newClientConfiguration();
         clientConf.setClientAuthProviderFactoryClass(
             SASLClientProviderFactory.class.getName());