HDFS-12009. Accept human-friendly units in dfsadmin -setBalancerBandwidth and -setQuota. Contributed by Andrew Wang.

hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/tools/DFSAdmin.java

@@ -195,7 +195,8 @@ public class DFSAdmin extends FsShell {
       super(conf);
       CommandFormat c = new CommandFormat(2, Integer.MAX_VALUE);
       List<String> parameters = c.parse(args, pos);
-      this.quota = Long.parseLong(parameters.remove(0));
+      this.quota =
+          StringUtils.TraditionalBinaryPrefix.string2long(parameters.remove(0));
       this.args = parameters.toArray(new String[parameters.size()]);
     }
     
@@ -938,7 +939,7 @@ public class DFSAdmin extends FsShell {
     int exitCode = -1;
 
     try {
-      bandwidth = Long.parseLong(argv[idx]);
+      bandwidth = StringUtils.TraditionalBinaryPrefix.string2long(argv[idx]);
     } catch (NumberFormatException nfe) {
       System.err.println("NumberFormatException: " + nfe.getMessage());
       System.err.println("Usage: hdfs dfsadmin"

hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestQuota.java

@@ -184,11 +184,20 @@ public class TestQuota {
     final short replication = 5;
     final long spaceQuota = fileLen * replication * 15 / 8;
 
-    // 1: create a directory test and set its quota to be 3
+    // 1: create a test directory
     final Path parent = new Path(dir, "test");
     assertTrue(dfs.mkdirs(parent));
-    String[] args = new String[]{"-setQuota", "3", parent.toString()};
+
+    // Try setting name quota with suffixes
+    String[] args;
+    args = new String[]{"-setQuota", "3K", parent.toString()};
+    runCommand(admin, args, false);
+    args = new String[]{"-setQuota", "3m", parent.toString()};
     runCommand(admin, args, false);
+    // Set the final name quota to 3
+    args = new String[]{"-setQuota", "3", parent.toString()};
+    runCommand(admin, args, false);
+
 
     //try setting space quota with a 'binary prefix'
     runCommand(admin, false, "-setSpaceQuota", "2t", parent.toString());

hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/tools/TestDFSAdmin.java

@@ -696,4 +696,40 @@ public class TestDFSAdmin {
         client.getDatanodeStorageReport(DatanodeReportType.DEAD).length);
     assertEquals(numCorruptBlocks, client.getCorruptBlocksCount());
   }
+
+  @Test
+  public void testSetBalancerBandwidth() throws Exception {
+    redirectStream();
+
+    final DFSAdmin dfsAdmin = new DFSAdmin(conf);
+    String outStr;
+
+    // Test basic case: 10000
+    assertEquals(0, ToolRunner.run(dfsAdmin,
+        new String[]{"-setBalancerBandwidth", "10000"}));
+    outStr = scanIntoString(out);
+    assertTrue("Did not set bandwidth!", outStr.contains("Balancer " +
+        "bandwidth is set to 10000"));
+
+    // Test parsing with units
+    resetStream();
+    assertEquals(0, ToolRunner.run(dfsAdmin,
+        new String[]{"-setBalancerBandwidth", "10m"}));
+    outStr = scanIntoString(out);
+    assertTrue("Did not set bandwidth!", outStr.contains("Balancer " +
+        "bandwidth is set to 10485760"));
+
+    resetStream();
+    assertEquals(0, ToolRunner.run(dfsAdmin,
+        new String[]{"-setBalancerBandwidth", "10k"}));
+    outStr = scanIntoString(out);
+    assertTrue("Did not set bandwidth!", outStr.contains("Balancer " +
+        "bandwidth is set to 10240"));
+
+    // Test negative numbers
+    assertEquals(-1, ToolRunner.run(dfsAdmin,
+        new String[]{"-setBalancerBandwidth", "-10000"}));
+    assertEquals(-1, ToolRunner.run(dfsAdmin,
+        new String[]{"-setBalancerBandwidth", "-10m"}));
+  }
 }