HBASE-23330: Fix delegation token fetch with MasterRegistry (apache#1084

)

Signed-off-by: Andrew Purtell <apurtell@apache.org>

hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncConnectionImpl.java

@@ -33,6 +33,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -278,6 +279,15 @@ CompletableFuture<MasterService.Interface> getMasterStub() {
     }, stub -> true, "master stub");
   }
 
+  String getClusterId() {
+    try {
+      return registry.getClusterId().get();
+    } catch (InterruptedException | ExecutionException e) {
+      LOG.error("Error fetching cluster ID: ", e);
+    }
+    return null;
+  }
+
   void clearMasterStubCache(MasterService.Interface stub) {
     masterStub.compareAndSet(stub, null);
   }

hbase-client/src/main/java/org/apache/hadoop/hbase/client/Connection.java

@@ -204,6 +204,11 @@ default BufferedMutator getBufferedMutator(TableName tableName) throws IOExcepti
    */
   AsyncConnection toAsyncConnection();
 
+  /**
+   * @return the cluster ID unique to this HBase cluster.
+   */
+  String getClusterId();
+
   /**
    * Retrieve an Hbck implementation to fix an HBase cluster.
    * The returned Hbck is not guaranteed to be thread-safe. A new instance should be created by

hbase-client/src/main/java/org/apache/hadoop/hbase/client/ConnectionOverAsyncConnection.java

@@ -204,6 +204,11 @@ public AsyncConnection toAsyncConnection() {
     return conn;
   }
 
+  @Override
+  public String getClusterId() {
+    return conn.getClusterId();
+  }
+
   @Override
   public Hbck getHbck() throws IOException {
     return FutureUtils.get(conn.getHbck());

hbase-mapreduce/src/test/java/org/apache/hadoop/hbase/mapreduce/TestMultiTableInputFormatBase.java

@@ -246,5 +246,11 @@ public void clearRegionLocationCache() {
     public AsyncConnection toAsyncConnection() {
       return null;
     }
+
+    @Override
+    public String getClusterId() {
+      return null;
+    }
+
   }
 }

hbase-mapreduce/src/test/java/org/apache/hadoop/hbase/mapreduce/TestTableInputFormatBase.java

@@ -296,5 +296,10 @@ public void clearRegionLocationCache() {
     public AsyncConnection toAsyncConnection() {
       throw new UnsupportedOperationException();
     }
+
+    @Override
+    public String getClusterId() {
+      return null;
+    }
   }
 }

hbase-server/src/main/java/org/apache/hadoop/hbase/client/SharedConnection.java

@@ -105,4 +105,9 @@ public Hbck getHbck(ServerName masterServer) throws IOException {
   public AsyncConnection toAsyncConnection() {
     return new SharedAsyncConnection(conn.toAsyncConnection());
   }
+
+  @Override
+  public String getClusterId() {
+    return conn.getClusterId();
+  }
 }

hbase-server/src/main/java/org/apache/hadoop/hbase/security/token/TokenUtil.java

@@ -26,18 +26,14 @@
 import org.apache.hadoop.hbase.client.ConnectionFactory;
 import org.apache.hadoop.hbase.protobuf.generated.AuthenticationProtos;
 import org.apache.hadoop.hbase.security.User;
-import org.apache.hadoop.hbase.zookeeper.ZKClusterId;
-import org.apache.hadoop.hbase.zookeeper.ZKWatcher;
 import org.apache.hadoop.io.Text;
 import org.apache.hadoop.mapred.JobConf;
 import org.apache.hadoop.mapreduce.Job;
 import org.apache.hadoop.security.token.Token;
 import org.apache.yetus.audience.InterfaceAudience;
-import org.apache.zookeeper.KeeperException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-
 /**
  * Utility methods for obtaining authentication tokens.
  */
@@ -219,7 +215,7 @@ public static void obtainTokenForJob(final Connection conn, final JobConf job, U
   public static void addTokenForJob(final Connection conn, final JobConf job, User user)
       throws IOException, InterruptedException {
 
-    Token<AuthenticationTokenIdentifier> token = getAuthToken(conn.getConfiguration(), user);
+    Token<AuthenticationTokenIdentifier> token = getAuthToken(conn, user);
     if (token == null) {
       token = ClientTokenUtil.obtainToken(conn, user);
     }
@@ -238,7 +234,7 @@ public static void addTokenForJob(final Connection conn, final JobConf job, User
    */
   public static void addTokenForJob(final Connection conn, User user, Job job)
       throws IOException, InterruptedException {
-    Token<AuthenticationTokenIdentifier> token = getAuthToken(conn.getConfiguration(), user);
+    Token<AuthenticationTokenIdentifier> token = getAuthToken(conn, user);
     if (token == null) {
       token = ClientTokenUtil.obtainToken(conn, user);
     }
@@ -257,7 +253,7 @@ public static void addTokenForJob(final Connection conn, User user, Job job)
    */
   public static boolean addTokenIfMissing(Connection conn, User user)
       throws IOException, InterruptedException {
-    Token<AuthenticationTokenIdentifier> token = getAuthToken(conn.getConfiguration(), user);
+    Token<AuthenticationTokenIdentifier> token = getAuthToken(conn, user);
     if (token == null) {
       token = ClientTokenUtil.obtainToken(conn, user);
       user.getUGI().addToken(token.getService(), token);
@@ -270,19 +266,12 @@ public static boolean addTokenIfMissing(Connection conn, User user)
    * Get the authentication token of the user for the cluster specified in the configuration
    * @return null if the user does not have the token, otherwise the auth token for the cluster.
    */
-  private static Token<AuthenticationTokenIdentifier> getAuthToken(Configuration conf, User user)
-      throws IOException, InterruptedException {
-    ZKWatcher zkw = new ZKWatcher(conf, "TokenUtil-getAuthToken", null);
-    try {
-      String clusterId = ZKClusterId.readClusterIdZNode(zkw);
-      if (clusterId == null) {
-        throw new IOException("Failed to get cluster ID");
-      }
-      return new AuthenticationTokenSelector().selectToken(new Text(clusterId), user.getTokens());
-    } catch (KeeperException e) {
-      throw new IOException(e);
-    } finally {
-      zkw.close();
+  private static Token<AuthenticationTokenIdentifier> getAuthToken(Connection conn, User user)
+      throws IOException {
+    final String clusterId = conn.getClusterId();
+    if (clusterId == null) {
+      throw new IOException("Failed to get cluster ID");
     }
+    return new AuthenticationTokenSelector().selectToken(new Text(clusterId), user.getTokens());
   }
 }

hbase-server/src/main/java/org/apache/hadoop/hbase/util/ConnectionCache.java

@@ -202,6 +202,19 @@ public boolean updateConnectionAccessTime() {
     return false;
   }
 
+  /**
+   * @return Cluster ID for the HBase cluster or null if there is an err making the connection.
+   */
+  public String getClusterId() {
+    try {
+      ConnectionInfo connInfo = getCurrentConnection();
+      return connInfo.connection.getClusterId();
+    } catch (IOException e) {
+      LOG.error("Error getting connection: ", e);
+    }
+    return null;
+  }
+
   class ConnectionInfo {
     final Connection connection;
     final String userName;

hbase-thrift/src/main/java/org/apache/hadoop/hbase/thrift/ThriftHBaseServiceHandler.java

@@ -1272,6 +1272,11 @@ public TThriftServerType getThriftServerType() {
     return TThriftServerType.ONE;
   }
 
+  @Override
+  public String getClusterId() throws TException {
+    return connectionCache.getClusterId();
+  }
+
   private static IOError getIOError(Throwable throwable) {
     IOError error = new IOErrorWithCause(throwable);
     error.setMessage(Throwables.getStackTraceAsString(throwable));