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

Signed-off-by: Andrew Purtell <apurtell@apache.org>
(cherry picked from commit fcb2012)

Author: bharathv

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

@@ -32,6 +32,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.AtomicReference;
@@ -268,6 +269,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

@@ -193,6 +193,11 @@ default Table getTable(TableName tableName, ExecutorService pool) throws IOExcep
    */
   TableBuilder getTableBuilder(TableName tableName, ExecutorService pool);
 
+  /**
+   * @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/ConnectionImplementation.java

@@ -378,6 +378,16 @@ public Table build() {
     };
   }
 
+  @Override
+  public String getClusterId() {
+    try {
+      return registry.getClusterId().get();
+    } catch (Exception e) {
+      LOG.error("Error fetching cluster ID: ", e);
+    }
+    return null;
+  }
+
   @Override
   public BufferedMutator getBufferedMutator(BufferedMutatorParams params) {
     if (params.getTableName() == null) {

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

@@ -240,5 +240,10 @@ public TableBuilder getTableBuilder(TableName tableName, ExecutorService pool) {
     @Override
     public void clearRegionLocationCache() {
     }
+
+    @Override
+    public String getClusterId() {
+      return null;
+    }
   }
 }

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

@@ -290,5 +290,10 @@ public TableBuilder getTableBuilder(TableName tableName, ExecutorService pool) {
     @Override
     public void clearRegionLocationCache() {
     }
+
+    @Override
+    public String getClusterId() {
+      return null;
+    }
   }
 }

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

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

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

@@ -25,18 +25,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.
  */
@@ -206,7 +202,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);
     }
@@ -225,7 +221,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);
     }
@@ -244,7 +240,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);
@@ -257,19 +253,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-server/src/test/java/org/apache/hadoop/hbase/replication/regionserver/TestWALEntrySinkFilter.java

@@ -547,6 +547,11 @@ public RegionLocator getRegionLocator() throws IOException {
       };
     }
 
+    @Override
+    public String getClusterId() {
+      return null;
+    }
+
     @Override
     public void clearRegionLocationCache() {
     }

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

@@ -1276,6 +1276,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));