HBASE-24032 [RSGroup] Assign created tables to respective rsgroup automatically instead of manual operations (#1319)

Signed-off-by: binlijin <binlijin@gmail.com>

hbase-server/src/main/java/org/apache/hadoop/hbase/rsgroup/DisabledRSGroupInfoManager.java

@@ -115,4 +115,10 @@ class DisabledRSGroupInfoManager implements RSGroupInfoManager {
   public void setRSGroup(Set<TableName> tables, String groupName) throws IOException {
     throw new DoNotRetryIOException("RSGroup is disabled");
   }
+
+  @Override
+  public String determineRSGroupInfoForTable(TableName tableName) {
+    return RSGroupInfo.DEFAULT_GROUP;
+  }
+
 }

hbase-server/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupInfoManager.java

@@ -98,4 +98,12 @@ public interface RSGroupInfoManager {
    * Set group for tables.
    */
   void setRSGroup(Set<TableName> tables, String groupName) throws IOException;
+
+  /**
+   * Determine {@code RSGroupInfo} for the given table.
+   * @param tableName table name
+   * @return rsgroup name
+   */
+  String determineRSGroupInfoForTable(TableName tableName);
+
 }

hbase-server/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupInfoManagerImpl.java

@@ -36,6 +36,7 @@ import java.util.TreeSet;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.StringUtils;
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.Coprocessor;
 import org.apache.hadoop.hbase.DoNotRetryIOException;
 import org.apache.hadoop.hbase.HConstants;
@@ -80,6 +81,7 @@ import org.apache.hadoop.hbase.util.Threads;
 import org.apache.hadoop.hbase.zookeeper.ZKUtil;
 import org.apache.hadoop.hbase.zookeeper.ZKWatcher;
 import org.apache.hadoop.hbase.zookeeper.ZNodePaths;
+import org.apache.hadoop.util.Shell;
 import org.apache.yetus.audience.InterfaceAudience;
 import org.apache.zookeeper.KeeperException;
 import org.slf4j.Logger;
@@ -200,11 +202,52 @@ final class RSGroupInfoManagerImpl implements RSGroupInfoManager {
   // contains list of groups that were last flushed to persistent store
   private Set<String> prevRSGroups = new HashSet<>();
 
+  // Package visibility for testing
+  static class RSGroupMappingScript {
+    static final String RS_GROUP_MAPPING_SCRIPT = "hbase.rsgroup.table.mapping.script";
+    static final String RS_GROUP_MAPPING_SCRIPT_TIMEOUT =
+      "hbase.rsgroup.table.mapping.script.timeout";
+    private Shell.ShellCommandExecutor rsgroupMappingScript;
+
+    RSGroupMappingScript(Configuration conf) {
+      String script = conf.get(RS_GROUP_MAPPING_SCRIPT);
+      if (script == null || script.isEmpty()) {
+        return;
+      }
+
+      rsgroupMappingScript = new Shell.ShellCommandExecutor(
+        new String[] { script, "", "" }, null, null,
+        conf.getLong(RS_GROUP_MAPPING_SCRIPT_TIMEOUT, 5000) // 5 seconds
+      );
+    }
+
+    String getRSGroup(String namespace, String tablename) {
+      if (rsgroupMappingScript == null) {
+        return RSGroupInfo.DEFAULT_GROUP;
+      }
+      String[] exec = rsgroupMappingScript.getExecString();
+      exec[1] = namespace;
+      exec[2] = tablename;
+      try {
+        rsgroupMappingScript.execute();
+      } catch (IOException e) {
+        // This exception may happen, like process doesn't have permission to run this script.
+        LOG.error("{}, placing {} back to default rsgroup",
+          e.getMessage(),
+          TableName.valueOf(namespace, tablename));
+        return RSGroupInfo.DEFAULT_GROUP;
+      }
+      return rsgroupMappingScript.getOutput().trim();
+    }
+  }
+  private RSGroupMappingScript script;
+
   private RSGroupInfoManagerImpl(MasterServices masterServices) {
     this.masterServices = masterServices;
     this.watcher = masterServices.getZooKeeper();
     this.conn = masterServices.getAsyncClusterConnection();
     this.rsGroupStartupWorker = new RSGroupStartupWorker();
+    this.script = new RSGroupMappingScript(masterServices.getConfiguration());
   }
 
   private synchronized void updateDefaultServers() {
@@ -1173,4 +1216,9 @@ final class RSGroupInfoManagerImpl implements RSGroupInfoManager {
     }
   }
 
+  @Override
+  public String determineRSGroupInfoForTable(TableName tableName) {
+    return script.getRSGroup(tableName.getNamespaceAsString(), tableName.getQualifierAsString());
+  }
+
 }

hbase-server/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupUtil.java

@@ -71,6 +71,7 @@ public final class RSGroupUtil {
     if (td == null) {
       return Optional.empty();
     }
+    // RSGroup information determined by client.
     Optional<String> optGroupNameOfTable = td.getRegionServerGroup();
     if (optGroupNameOfTable.isPresent()) {
       RSGroupInfo group = manager.getRSGroup(optGroupNameOfTable.get());
@@ -84,6 +85,13 @@ public final class RSGroupUtil {
     if (groupFromOldRSGroupInfo != null) {
       return Optional.of(groupFromOldRSGroupInfo);
     }
+    // RSGroup information determined by administrator.
+    String groupDeterminedByAdmin = manager.determineRSGroupInfoForTable(tableName);
+    RSGroupInfo groupInfo = manager.getRSGroup(groupDeterminedByAdmin);
+    if (groupInfo != null) {
+      return Optional.of(groupInfo);
+    }
+    // Finally, we will try to fall back to namespace as rsgroup if exists
     ClusterSchema clusterSchema = master.getClusterSchema();
     if (clusterSchema == null) {
       if (TableName.isMetaTableName(tableName)) {

hbase-server/src/test/java/org/apache/hadoop/hbase/rsgroup/TestRSGroupMappingScript.java

@@ -0,0 +1,129 @@
+/**
+ * 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.hadoop.hbase.rsgroup;
+
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.FileReader;
+import java.io.IOException;
+import java.io.PrintWriter;
+import org.apache.hadoop.hbase.HBaseClassTestRule;
+import org.apache.hadoop.hbase.HBaseTestingUtility;
+import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.rsgroup.RSGroupInfoManagerImpl.RSGroupMappingScript;
+import org.apache.hadoop.hbase.testclassification.RSGroupTests;
+import org.apache.hadoop.hbase.testclassification.SmallTests;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@Category({ RSGroupTests.class, SmallTests.class })
+public class TestRSGroupMappingScript {
+
+  @ClassRule
+  public static final HBaseClassTestRule CLASS_RULE =
+    HBaseClassTestRule.forClass(TestRSGroupMappingScript.class);
+  private static final Logger LOG = LoggerFactory.getLogger(TestRSGroupMappingScript.class);
+
+  private static final HBaseTestingUtility UTIL = new HBaseTestingUtility();
+  private File script;
+
+  @BeforeClass
+  public static void setupScript() throws Exception {
+    String currentDir = new File("").getAbsolutePath();
+    UTIL.getConfiguration().set(
+      RSGroupInfoManagerImpl.RSGroupMappingScript.RS_GROUP_MAPPING_SCRIPT,
+      currentDir + "/rsgroup_table_mapping.sh"
+    );
+  }
+
+  @Before
+  public void setup() throws Exception {
+    script = new File(UTIL.getConfiguration().get(RSGroupMappingScript.RS_GROUP_MAPPING_SCRIPT));
+    if (!script.createNewFile()) {
+      throw new IOException("Can't create script");
+    }
+
+    PrintWriter pw = new PrintWriter(new FileOutputStream(script));
+    try {
+      pw.println("#!/bin/bash");
+      pw.println("namespace=$1");
+      pw.println("tablename=$2");
+      pw.println("if [[ $namespace == test ]]; then");
+      pw.println("  echo test");
+      pw.println("elif [[ $tablename == *foo* ]]; then");
+      pw.println("  echo other");
+      pw.println("else");
+      pw.println("  echo default");
+      pw.println("fi");
+      pw.flush();
+    } finally {
+      pw.close();
+    }
+    boolean executable = script.setExecutable(true);
+    LOG.info("Created " + script  + ", executable=" + executable);
+    verifyScriptContent(script);
+  }
+
+  private void verifyScriptContent(File file) throws Exception {
+    BufferedReader reader = new BufferedReader(new FileReader(file));
+    String line;
+    while ((line = reader.readLine()) != null) {
+      LOG.info(line);
+    }
+  }
+
+  @Test
+  public void testScript() throws Exception {
+    RSGroupMappingScript script = new RSGroupMappingScript(UTIL.getConfiguration());
+    TableName testNamespace =
+      TableName.valueOf("test", "should_be_in_test");
+    String rsgroup = script.getRSGroup(
+      testNamespace.getNamespaceAsString(), testNamespace.getQualifierAsString()
+    );
+    Assert.assertEquals("test", rsgroup);
+
+    TableName otherName =
+      TableName.valueOf("whatever", "oh_foo_should_be_in_other");
+    rsgroup = script.getRSGroup(otherName.getNamespaceAsString(), otherName.getQualifierAsString());
+    Assert.assertEquals("other", rsgroup);
+
+    TableName defaultName =
+      TableName.valueOf("nono", "should_be_in_default");
+    rsgroup = script.getRSGroup(
+      defaultName.getNamespaceAsString(), defaultName.getQualifierAsString()
+    );
+    Assert.assertEquals("default", rsgroup);
+  }
+
+  @After
+  public void teardown() throws Exception {
+    if (script.exists()) {
+      script.delete();
+    }
+  }
+
+}