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

) (apache#4598)

Signed-off-by: Andrew Purtell <[email protected]>
(cherry picked from commit d8b3f55)

Co-authored-by: Bharath Vissapragada <[email protected]>
(cherry picked from commit 2cd30e5)
Change-Id: I7cf85b4b611c71b0e87d0b23a1f85542e912e197

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;
@@ -298,6 +299,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

@@ -171,6 +171,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 each thread. This is a lightweight

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

@@ -2175,4 +2175,14 @@ private static <T> T get(CompletableFuture<T> future) throws IOException {
       throw new IOException(cause);
     }
   }
+
+  @Override
+  public String getClusterId() {
+    try {
+      return registry.getClusterId().get();
+    } catch (InterruptedException | ExecutionException e) {
+      LOG.error("Error fetching cluster ID: ", e);
+    }
+    return null;
+  }
 }

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

@@ -1689,6 +1689,10 @@ public void abort(String why, Throwable e) {
     public boolean isAborted() {
       return delegate.isAborted();
     }
-  }
 
+    @Override
+    public String getClusterId() {
+      return delegate.getClusterId();
+    }
+  }
 }

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

@@ -233,5 +233,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

@@ -314,5 +314,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

@@ -24,14 +24,11 @@
 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;
 
@@ -197,7 +194,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);
     }
@@ -215,7 +212,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);
     }
@@ -233,7 +230,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);
@@ -246,19 +243,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

@@ -203,6 +203,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

@@ -570,5 +570,10 @@ public RegionLocator getRegionLocator() throws IOException {
     @Override
     public void clearRegionLocationCache() {
     }
+
+    @Override
+    public String getClusterId() {
+      return null;
+    }
   }
 }

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

@@ -1217,6 +1217,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.setCanRetry(!(throwable instanceof DoNotRetryIOException));