Rebase and address comments

linkedin · Jun 18, 2021 · e6ebe87 · e6ebe87
1 parent 5a6c8c7
commit e6ebe87
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 21 deletions.
diff --git a/...c/main/java/com/linkedin/datastream/server/assignment/StickyMulticastStrategyFactory.java b/...c/main/java/com/linkedin/datastream/server/assignment/StickyMulticastStrategyFactory.java
@@ -20,7 +20,7 @@
  */
 public class StickyMulticastStrategyFactory implements AssignmentStrategyFactory {
   public static final String CFG_IMBALANCE_THRESHOLD = "imbalanceThreshold";
-  public static final Integer DEFAULT_IMBALANCE_THRESHOLD = 1;
+  public static final int DEFAULT_IMBALANCE_THRESHOLD = 1;
 
   @Override
   public AssignmentStrategy createStrategy(Properties assignmentStrategyProperties) {

diff --git a/...a/com/linkedin/datastream/server/assignment/TestLoadBasedPartitionAssignmentStrategy.java b/...a/com/linkedin/datastream/server/assignment/TestLoadBasedPartitionAssignmentStrategy.java
@@ -120,7 +120,7 @@ public void fallbackToBaseClassWhenThroughputFetchFailsTest() {
     int taskCapacityUtilizationPct = 90;
     int throughputInfoFetchTimeoutMs = 1000;
     int throughputInfoFetchRetryPeriodMs = 200;
-    ZkClient zkClient = null;
+    ZkClient zkClient = _zkClient;
 
     LoadBasedPartitionAssignmentStrategy strategy = Mockito.spy(new LoadBasedPartitionAssignmentStrategy(mockProvider,
         maxTasks, imbalanceThreshold, maxPartitionPerTask, enableElasticTaskAssignment, partitionsPerTask,
@@ -158,7 +158,7 @@ public void doesntFetchPartitionInfoOnIncrementalAssignmentTest() {
     int taskCapacityUtilizationPct = 90;
     int throughputInfoFetchTimeoutMs = 1000;
     int throughputInfoFetchRetryPeriodMs = 200;
-    ZkClient zkClient = null;
+    ZkClient zkClient = _zkClient;
 
     LoadBasedPartitionAssignmentStrategy strategy = new LoadBasedPartitionAssignmentStrategy(mockProvider,
         maxTasks, imbalanceThreshold, maxPartitionPerTask, enableElasticTaskAssignment, partitionsPerTask,
@@ -201,7 +201,7 @@ public void updatesNumTasksAndThrowsExceptionWhenNoSufficientTasksTest() {
     int taskCapacityUtilizationPct = 90;
     int throughputInfoFetchTimeoutMs = 1000;
     int throughputInfoFetchRetryPeriodMs = 200;
-    ZkClient zkClient = null;
+    ZkClient zkClient = _zkClient;
 
     LoadBasedPartitionAssignmentStrategy strategy = new LoadBasedPartitionAssignmentStrategy(mockProvider,
         maxTasks, imbalanceThreshold, maxPartitionPerTask, enableElasticTaskAssignment, partitionsPerTask,

diff --git a/...java/com/linkedin/datastream/server/assignment/TestPartitionAssignmentStrategyConfig.java b/...java/com/linkedin/datastream/server/assignment/TestPartitionAssignmentStrategyConfig.java
@@ -34,7 +34,6 @@ public class TestPartitionAssignmentStrategyConfig {
   private static final String CFG_CLUSTER_NAME_VALUE = "dummyCluster";
   private static final String INVALID_INTEGER_VALUE = "-1";
 
-
   @Test
   public void configValuesCorrectlyAssignedTest() {
     Properties props = new Properties();
@@ -75,7 +74,7 @@ public void configValuesRevertedToEmptyWhenInvalidTest() {
 
     PartitionAssignmentStrategyConfig config = new PartitionAssignmentStrategyConfig(props);
     Assert.assertEquals(config.getMaxTasks(), Optional.empty());
-    Assert.assertEquals(config.getImbalanceThreshold(), (int) DEFAULT_IMBALANCE_THRESHOLD);
+    Assert.assertEquals(config.getImbalanceThreshold(), DEFAULT_IMBALANCE_THRESHOLD);
     Assert.assertEquals(config.getMaxPartitions(), Integer.MAX_VALUE);
     Assert.assertEquals(config.getPartitionsPerTask(), DEFAULT_PARTITIONS_PER_TASK);
   }
@@ -84,7 +83,7 @@ public void configValuesRevertedToEmptyWhenInvalidTest() {
   public void configValuesSetToDefaultWhenNotProvidedTest() {
     PartitionAssignmentStrategyConfig config = new PartitionAssignmentStrategyConfig(new Properties());
     Assert.assertEquals(config.getMaxTasks(), Optional.empty());
-    Assert.assertEquals(config.getImbalanceThreshold(), (int) DEFAULT_IMBALANCE_THRESHOLD);
+    Assert.assertEquals(config.getImbalanceThreshold(), DEFAULT_IMBALANCE_THRESHOLD);
     Assert.assertEquals(config.getMaxPartitions(), Integer.MAX_VALUE);
     Assert.assertEquals(config.getPartitionsPerTask(), DEFAULT_PARTITIONS_PER_TASK);
   }

diff --git a/.../src/test/java/com/linkedin/datastream/server/assignment/TestStickyMulticastStrategy.java b/.../src/test/java/com/linkedin/datastream/server/assignment/TestStickyMulticastStrategy.java
@@ -19,7 +19,6 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
-import org.jetbrains.annotations.NotNull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.Assert;
@@ -288,7 +287,7 @@ public void testSameTaskIsNotAssignedToMoreThanOneInstance() {
     String[] instances = new String[]{"instance1", "instance2", "instance3"};
     int numDatastreams = 5;
     List<DatastreamGroup> datastreams = generateDatastreams("ds", numDatastreams);
-    StickyMulticastStrategy strategy = createStickyMulticastStrategyObject(Optional.empty());
+    StickyMulticastStrategy strategy = new StickyMulticastStrategy(Optional.empty());
     Map<String, Set<DatastreamTask>> assignment =
         strategy.assign(datastreams, Arrays.asList(instances), new HashMap<>());
     // Copying the assignment to simulate the scenario where two instances have the same task,
@@ -302,16 +301,11 @@ public void testSameTaskIsNotAssignedToMoreThanOneInstance() {
     Assert.assertEquals(newAssignmentTasks.size(), instances.length * numDatastreams);
   }
 
-  @NotNull
-  private StickyMulticastStrategy createStickyMulticastStrategyObject(Optional<Integer> maxTasks) {
-    return new StickyMulticastStrategy(maxTasks);
-  }
-
   @Test
   public void testRemoveDatastreamTasksWhenDatastreamIsDeleted() {
     List<String> instances = Arrays.asList("instance1", "instance2", "instance3");
     List<DatastreamGroup> datastreams = generateDatastreams("ds", 5);
-    StickyMulticastStrategy strategy = createStickyMulticastStrategyObject(Optional.empty());
+    StickyMulticastStrategy strategy = new StickyMulticastStrategy(Optional.empty());
     Map<String, Set<DatastreamTask>> assignment = strategy.assign(datastreams, instances, new HashMap<>());
 
     datastreams.remove(0);
@@ -328,7 +322,7 @@ public void testRemoveDatastreamTasksWhenDatastreamIsDeleted() {
     // test with strategy where dsTaskLimitPerInstance is greater than 1
     int maxTasksConfig = 12;
     datastreams = generateDatastreams("ds", 5);
-    strategy = createStickyMulticastStrategyObject(Optional.of(maxTasksConfig));
+    strategy = new StickyMulticastStrategy(Optional.of(maxTasksConfig));
     assignment = strategy.assign(datastreams, instances, new HashMap<>());
 
     datastreams.remove(0);
@@ -348,7 +342,7 @@ public void testRemoveDatastreamTasksWhenDatastreamIsDeleted() {
   public void testCreateNewTasksOnlyForNewDatastreamWhenDatastreamIsCreated() {
     List<String> instances = Arrays.asList("instance1", "instance2", "instance3");
     List<DatastreamGroup> datastreams = generateDatastreams("ds", 5);
-    StickyMulticastStrategy strategy = createStickyMulticastStrategyObject(Optional.empty());
+    StickyMulticastStrategy strategy = new StickyMulticastStrategy(Optional.empty());
     Map<String, Set<DatastreamTask>> assignment = strategy.assign(datastreams, instances, new HashMap<>());
 
     List<DatastreamGroup> newDatastreams = new ArrayList<>(datastreams);
@@ -369,7 +363,7 @@ public void testCreateNewTasksOnlyForNewDatastreamWhenDatastreamIsCreated() {
     // test with strategy where dsTaskLimitPerInstance is greater than 1
     int maxTasksConfig = 12;
     datastreams = generateDatastreams("ds", 5);
-    strategy = createStickyMulticastStrategyObject(Optional.of(maxTasksConfig));
+    strategy = new StickyMulticastStrategy(Optional.of(maxTasksConfig));
     assignment = strategy.assign(datastreams, instances, new HashMap<>());
 
     newDatastreams = new ArrayList<>(datastreams);
@@ -394,7 +388,7 @@ public void testCreateNewTasksOnlyForNewInstanceWhenInstanceIsAdded() {
     List<String> instances = Arrays.asList("instance1", "instance2", "instance3");
     String instance4 = "instance4";
     List<DatastreamGroup> datastreams = generateDatastreams("ds", 5);
-    StickyMulticastStrategy strategy = createStickyMulticastStrategyObject(Optional.empty());
+    StickyMulticastStrategy strategy = new StickyMulticastStrategy(Optional.empty());
     Map<String, Set<DatastreamTask>> assignment = strategy.assign(datastreams, instances, new HashMap<>());
     List<String> newInstances = new ArrayList<>(instances);
     newInstances.add(instance4);
@@ -417,7 +411,7 @@ public void testStickyFairDistributionWhenNewInstanceIsAdded() {
     String instance4 = "instance4";
     List<DatastreamGroup> datastreams = generateDatastreams("ds", 5);
     int maxTasksConfig = 12;
-    StickyMulticastStrategy strategy = createStickyMulticastStrategyObject(Optional.of(maxTasksConfig));
+    StickyMulticastStrategy strategy = new StickyMulticastStrategy(Optional.of(maxTasksConfig));
     Map<String, Set<DatastreamTask>> assignment = strategy.assign(datastreams, instances, new HashMap<>());
     List<String> newInstances = new ArrayList<>(instances);
     newInstances.add(instance4);
@@ -484,7 +478,7 @@ public void testTriggerRebalanceWhenDeletingDatastream() {
   public void testExtraTasksAreNotAssignedDuringReassignment() {
     String[] instances = new String[]{"instance1"};
     List<DatastreamGroup> datastreams = generateDatastreams("ds", 5);
-    StickyMulticastStrategy strategy = createStickyMulticastStrategyObject(Optional.of(4));
+    StickyMulticastStrategy strategy = new StickyMulticastStrategy(Optional.of(4));
     Map<String, Set<DatastreamTask>> assignment1 =
         strategy.assign(datastreams, Arrays.asList(instances), new HashMap<>());
     Map<String, Set<DatastreamTask>> assignment2 =