You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@uniffle.apache.org by ro...@apache.org on 2022/09/05 09:32:48 UTC

[incubator-uniffle] branch master updated: [IMPROVEMENT] Introduce the enumType in ConfigOptions (#199)

This is an automated email from the ASF dual-hosted git repository.

roryqi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git


The following commit(s) were added to refs/heads/master by this push:
     new 6b9f2d15 [IMPROVEMENT] Introduce the enumType in ConfigOptions (#199)
6b9f2d15 is described below

commit 6b9f2d156393abd52ec822498a0117856574481b
Author: Junfan Zhang <ju...@outlook.com>
AuthorDate: Mon Sep 5 17:32:44 2022 +0800

    [IMPROVEMENT] Introduce the enumType in ConfigOptions (#199)
    
    ### What changes were proposed in this pull request?
    Introduce the enumType in ConfigOptions
    
    ### Why are the changes needed?
    When I want to extend the zstd compression type in uniffle, I found we lack the enum class for configOptions.
    And I think it will be used for users to specify the concrete implementation in some pluggable class.
    
    ### Does this PR introduce _any_ user-facing change?
    No
    
    ### How was this patch tested?
    UTs.
---
 .../uniffle/common/config/ConfigOptions.java       |  9 +++++
 .../apache/uniffle/common/config/ConfigUtils.java  | 24 ++++++++++++++
 .../uniffle/common/config/ConfigOptionTest.java    | 38 ++++++++++++++++++++++
 .../coordinator/AssignmentStrategyFactory.java     |  8 ++---
 .../uniffle/coordinator/CoordinatorConf.java       |  9 +++--
 .../uniffle/coordinator/CoordinatorConfTest.java   |  2 +-
 .../apache/uniffle/test/CoordinatorGrpcTest.java   |  2 +-
 7 files changed, 83 insertions(+), 9 deletions(-)

diff --git a/common/src/main/java/org/apache/uniffle/common/config/ConfigOptions.java b/common/src/main/java/org/apache/uniffle/common/config/ConfigOptions.java
index 9193a74b..8f00204e 100644
--- a/common/src/main/java/org/apache/uniffle/common/config/ConfigOptions.java
+++ b/common/src/main/java/org/apache/uniffle/common/config/ConfigOptions.java
@@ -139,6 +139,15 @@ public class ConfigOptions {
     public TypedConfigOptionBuilder<String> stringType() {
       return new TypedConfigOptionBuilder<>(key, String.class);
     }
+
+    /**
+     * Defines that the value of the option should be of {@link Enum} type.
+     *
+     * @param enumClass Concrete type of the expected enum.
+     */
+    public <T extends Enum<T>> TypedConfigOptionBuilder<T> enumType(Class<T> enumClass) {
+      return new TypedConfigOptionBuilder<>(key, enumClass);
+    }
   }
 
   // ------------------------------------------------------------------------
diff --git a/common/src/main/java/org/apache/uniffle/common/config/ConfigUtils.java b/common/src/main/java/org/apache/uniffle/common/config/ConfigUtils.java
index 3de56927..c93334e0 100644
--- a/common/src/main/java/org/apache/uniffle/common/config/ConfigUtils.java
+++ b/common/src/main/java/org/apache/uniffle/common/config/ConfigUtils.java
@@ -18,7 +18,9 @@
 package org.apache.uniffle.common.config;
 
 import java.lang.reflect.Field;
+import java.util.Arrays;
 import java.util.List;
+import java.util.Locale;
 import java.util.function.Function;
 
 import com.google.common.collect.Lists;
@@ -63,10 +65,32 @@ public class ConfigUtils {
       return (T) convertToDouble(rawValue);
     } else if (String.class.equals(clazz)) {
       return (T) convertToString(rawValue);
+    } else if (clazz.isEnum()) {
+      return (T) convertToEnum(rawValue, (Class<? extends Enum<?>>) clazz);
     }
     throw new IllegalArgumentException("Unsupported type: " + clazz);
   }
 
+  public static <E extends Enum<?>> E convertToEnum(Object o, Class<E> clazz) {
+    if (o.getClass().equals(clazz)) {
+      return (E) o;
+    }
+
+    return Arrays.stream(clazz.getEnumConstants())
+        .filter(
+            e ->
+                e.toString()
+                    .toUpperCase(Locale.ROOT)
+                    .equals(o.toString().toUpperCase(Locale.ROOT)))
+        .findAny()
+        .orElseThrow(
+            () ->
+                new IllegalArgumentException(
+                    String.format(
+                        "Could not parse value for enum %s. Expected one of: [%s]",
+                        clazz, Arrays.toString(clazz.getEnumConstants()))));
+  }
+
   static String convertToString(Object o) {
     if (o.getClass() == String.class) {
       return (String) o;
diff --git a/common/src/test/java/org/apache/uniffle/common/config/ConfigOptionTest.java b/common/src/test/java/org/apache/uniffle/common/config/ConfigOptionTest.java
index 547d148a..9e982902 100644
--- a/common/src/test/java/org/apache/uniffle/common/config/ConfigOptionTest.java
+++ b/common/src/test/java/org/apache/uniffle/common/config/ConfigOptionTest.java
@@ -56,6 +56,44 @@ public class ConfigOptionTest {
     assertFalse(conf.get(booleanConfig));
   }
 
+  enum TestType {
+    TYPE_1,
+    TYPE_2,
+  }
+
+  @Test
+  public void testEnumType() {
+    final ConfigOption<TestType> enumConfigOption = ConfigOptions
+        .key("rss.enum")
+        .enumType(TestType.class)
+        .defaultValue(TestType.TYPE_1)
+        .withDescription("enum test");
+
+    RssBaseConf conf = new RssBaseConf();
+
+    // case1: default value
+    assertEquals(TestType.TYPE_1, conf.get(enumConfigOption));
+
+    // case2: return the user specified value
+    conf.set(enumConfigOption, TestType.TYPE_2);
+    assertEquals(TestType.TYPE_2, conf.get(enumConfigOption));
+
+    // case3: set enum val with string
+    conf = new RssBaseConf();
+    conf.setString("rss.enum", "TYPE_2");
+    assertEquals(TestType.TYPE_2, conf.get(enumConfigOption));
+
+    // case4: set the illegal enum val with string
+    conf = new RssBaseConf();
+    conf.setString("rss.enum", "TYPE_3");
+    try {
+      conf.get(enumConfigOption);
+      fail();
+    } catch (IllegalArgumentException e) {
+      // ignore
+    }
+  }
+
   @Test
   public void testListTypes() {
     // test the string type list.
diff --git a/coordinator/src/main/java/org/apache/uniffle/coordinator/AssignmentStrategyFactory.java b/coordinator/src/main/java/org/apache/uniffle/coordinator/AssignmentStrategyFactory.java
index 09eb589d..073689ca 100644
--- a/coordinator/src/main/java/org/apache/uniffle/coordinator/AssignmentStrategyFactory.java
+++ b/coordinator/src/main/java/org/apache/uniffle/coordinator/AssignmentStrategyFactory.java
@@ -28,17 +28,17 @@ public class AssignmentStrategyFactory {
   }
 
   public AssignmentStrategy getAssignmentStrategy() {
-    String strategy = conf.getString(CoordinatorConf.COORDINATOR_ASSIGNMENT_STRATEGY);
-    if (StrategyName.BASIC.name().equals(strategy)) {
+    StrategyName strategy = conf.get(CoordinatorConf.COORDINATOR_ASSIGNMENT_STRATEGY);
+    if (StrategyName.BASIC == strategy) {
       return new BasicAssignmentStrategy(clusterManager);
-    } else if (StrategyName.PARTITION_BALANCE.name().equals(strategy)) {
+    } else if (StrategyName.PARTITION_BALANCE == strategy) {
       return new PartitionBalanceAssignmentStrategy(clusterManager);
     } else {
       throw new UnsupportedOperationException("Unsupported assignment strategy.");
     }
   }
 
-  private enum StrategyName {
+  public enum StrategyName {
     BASIC,
     PARTITION_BALANCE
   }
diff --git a/coordinator/src/main/java/org/apache/uniffle/coordinator/CoordinatorConf.java b/coordinator/src/main/java/org/apache/uniffle/coordinator/CoordinatorConf.java
index 765f9702..34c65b98 100644
--- a/coordinator/src/main/java/org/apache/uniffle/coordinator/CoordinatorConf.java
+++ b/coordinator/src/main/java/org/apache/uniffle/coordinator/CoordinatorConf.java
@@ -26,6 +26,8 @@ import org.apache.uniffle.common.config.ConfigUtils;
 import org.apache.uniffle.common.config.RssBaseConf;
 import org.apache.uniffle.common.util.RssUtils;
 
+import static org.apache.uniffle.coordinator.AssignmentStrategyFactory.StrategyName.PARTITION_BALANCE;
+
 /**
  * Configuration for Coordinator Service and rss-cluster, including service port,
  * heartbeat interval and etc.
@@ -54,10 +56,11 @@ public class CoordinatorConf extends RssBaseConf {
       .defaultValue(30L)
       .withDescription("The periodic interval times of output alive nodes. The interval sec can be calculated by ("
           + COORDINATOR_HEARTBEAT_TIMEOUT.key() + "/3 * rss.coordinator.server.periodic.output.interval.times)");
-  public static final ConfigOption<String> COORDINATOR_ASSIGNMENT_STRATEGY = ConfigOptions
+  public static final ConfigOption<AssignmentStrategyFactory.StrategyName>
+      COORDINATOR_ASSIGNMENT_STRATEGY = ConfigOptions
       .key("rss.coordinator.assignment.strategy")
-      .stringType()
-      .defaultValue("PARTITION_BALANCE")
+      .enumType(AssignmentStrategyFactory.StrategyName.class)
+      .defaultValue(PARTITION_BALANCE)
       .withDescription("Strategy for assigning shuffle server to write partitions");
   public static final ConfigOption<Long> COORDINATOR_APP_EXPIRED = ConfigOptions
       .key("rss.coordinator.app.expired")
diff --git a/coordinator/src/test/java/org/apache/uniffle/coordinator/CoordinatorConfTest.java b/coordinator/src/test/java/org/apache/uniffle/coordinator/CoordinatorConfTest.java
index 71c0c9b6..3d06f79f 100644
--- a/coordinator/src/test/java/org/apache/uniffle/coordinator/CoordinatorConfTest.java
+++ b/coordinator/src/test/java/org/apache/uniffle/coordinator/CoordinatorConfTest.java
@@ -42,7 +42,7 @@ public class CoordinatorConfTest {
     assertEquals(123, conf.getLong(CoordinatorConf.COORDINATOR_HEARTBEAT_TIMEOUT));
 
     // test default conf
-    assertEquals("PARTITION_BALANCE", conf.getString(CoordinatorConf.COORDINATOR_ASSIGNMENT_STRATEGY));
+    assertEquals("PARTITION_BALANCE", conf.get(CoordinatorConf.COORDINATOR_ASSIGNMENT_STRATEGY).name());
     assertEquals(256, conf.getInteger(CoordinatorConf.JETTY_CORE_POOL_SIZE));
     assertEquals(60 * 1000, conf.getLong(CoordinatorConf.COORDINATOR_EXCLUDE_NODES_CHECK_INTERVAL));
 
diff --git a/integration-test/common/src/test/java/org/apache/uniffle/test/CoordinatorGrpcTest.java b/integration-test/common/src/test/java/org/apache/uniffle/test/CoordinatorGrpcTest.java
index aa4900b5..c09067ea 100644
--- a/integration-test/common/src/test/java/org/apache/uniffle/test/CoordinatorGrpcTest.java
+++ b/integration-test/common/src/test/java/org/apache/uniffle/test/CoordinatorGrpcTest.java
@@ -60,7 +60,7 @@ public class CoordinatorGrpcTest extends CoordinatorTestBase {
   public static void setupServers() throws Exception {
     CoordinatorConf coordinatorConf = getCoordinatorConf();
     coordinatorConf.set(RssBaseConf.RPC_METRICS_ENABLED, true);
-    coordinatorConf.set(CoordinatorConf.COORDINATOR_ASSIGNMENT_STRATEGY, "BASIC");
+    coordinatorConf.setString(CoordinatorConf.COORDINATOR_ASSIGNMENT_STRATEGY.key(), "BASIC");
     coordinatorConf.setLong("rss.coordinator.app.expired", 2000);
     coordinatorConf.setLong("rss.coordinator.server.heartbeat.timeout", 3000);
     createCoordinatorServer(coordinatorConf);