Change YarnKeys to enum

heron/schedulers/src/java/com/twitter/heron/scheduler/yarn/YarnContext.java

@@ -18,19 +18,17 @@
 import com.twitter.heron.spi.common.Context;
 
 public final class YarnContext extends Context {
-  private static final String PREFIX = "heron.scheduler.yarn.";
-
-  public static final String HERON_SCHEDULER_YARN_QUEUE = PREFIX + "queue";
-  public static final String YARN_SCHEDULER_DRIVER_MEMORY_MB = PREFIX + "driver.memory.mb";
 
   private YarnContext() {
   }
 
   public static String heronYarnQueue(Config cfg) {
-    return cfg.getStringValue(HERON_SCHEDULER_YARN_QUEUE, "default");
+    return cfg.getStringValue(YarnKey.HERON_SCHEDULER_YARN_QUEUE.value(),
+        YarnKey.HERON_SCHEDULER_YARN_QUEUE.getDefaultString());
   }
 
   public static int heronDriverMemoryMb(Config cfg) {
-    return cfg.getIntegerValue(YARN_SCHEDULER_DRIVER_MEMORY_MB, 2048);
+    return cfg.getIntegerValue(YarnKey.YARN_SCHEDULER_DRIVER_MEMORY_MB.value(),
+        YarnKey.YARN_SCHEDULER_DRIVER_MEMORY_MB.getDefaultInt());
   }
 }

heron/schedulers/src/java/com/twitter/heron/scheduler/yarn/YarnKey.java

@@ -0,0 +1,67 @@
+// Copyright 2017 Twitter. All rights reserved.
+//
+// Licensed 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 com.twitter.heron.scheduler.yarn;
+
+import com.twitter.heron.spi.common.Key;
+
+/**
+ * Keys specific to the YARN scheduler
+ */
+public enum YarnKey {
+  // yarn queue for submitting and launching the topology
+  HERON_SCHEDULER_YARN_QUEUE("heron.scheduler.yarn.queue", "default"),
+  // the amount of memory topology's driver (yarn application master) needs
+  YARN_SCHEDULER_DRIVER_MEMORY_MB("heron.scheduler.yarn.driver.memory.mb", 2048);
+
+  private final String value;
+  private final Key.Type type;
+  private final Object defaultValue;
+
+  YarnKey(String value, String defaultValue) {
+    this.value = value;
+    this.type = Key.Type.STRING;
+    this.defaultValue = defaultValue;
+  }
+
+  YarnKey(String value, Integer defaultValue) {
+    this.value = value;
+    this.type = Key.Type.INTEGER;
+    this.defaultValue = defaultValue;
+  }
+
+  public String value() {
+    return value;
+  }
+
+  public Object getDefault() {
+    return defaultValue;
+  }
+
+  public String getDefaultString() {
+    if (type != Key.Type.STRING) {
+      throw new IllegalAccessError(String.format(
+          "Config Key %s is type %s, getDefaultString() not supported", this.name(), this.type));
+    }
+    return (String) this.defaultValue;
+  }
+
+  public int getDefaultInt() {
+    if (type != Key.Type.INTEGER) {
+      throw new IllegalAccessError(String.format(
+          "Config Key %s is type %s, getDefaultInt() not supported", this.name(), this.type));
+    }
+    return (Integer) this.defaultValue;
+  }
+}

heron/schedulers/tests/java/com/twitter/heron/scheduler/yarn/YarnLauncherTest.java

@@ -67,8 +67,9 @@ public void getHMDriverConfConstructsReefConfig() throws Exception {
     setConfigs(inputConf, expected, Key.CLUSTER, "cluster", Cluster.class);
     setConfigs(inputConf, expected, Key.ROLE, "role", Role.class);
     setConfigs(inputConf, expected, Key.ENVIRON, "env", Environ.class);
-    setConfigs(inputConf, expected, YarnContext.HERON_SCHEDULER_YARN_QUEUE, "q", JobQueue.class);
-    setConfigs(inputConf, expected, YarnContext.YARN_SCHEDULER_DRIVER_MEMORY_MB,
+    setConfigs(inputConf, expected, YarnKey.HERON_SCHEDULER_YARN_QUEUE.value(),
+        "q", JobQueue.class);
+    setConfigs(inputConf, expected, YarnKey.YARN_SCHEDULER_DRIVER_MEMORY_MB.value(),
         "123", DriverMemory.class);
     setConfigs(inputConf, expected, Key.CORE_PACKAGE_URI,
         new File(".").getName(), HeronCorePackageName.class);