PlaytikaOSS · kptfh · Mar 13, 2024 · Mar 13, 2024
diff --git a/common/src/main/java/com/playtika/shepherd/common/AssignmentData.java b/common/src/main/java/com/playtika/shepherd/common/AssignmentData.java
@@ -0,0 +1,6 @@
+package com.playtika.shepherd.common;
+
+public record AssignmentData(
+        long populationVersion,
+        String memberId, int generation, boolean isLeader) {
+}
diff --git a/common/src/main/java/com/playtika/shepherd/common/PastureListener.java b/common/src/main/java/com/playtika/shepherd/common/PastureListener.java
@@ -12,7 +12,7 @@ public interface PastureListener<Breed> {
     /**
      * Invoked when new subpopulation assigned to this pasture
      */
-    void assigned(List<Breed> population, long version, int generation, boolean isLeader);
+    void assigned(List<Breed> population, AssignmentData populationAssignment);
 
     /**
      * Invoked on first phase of rebalance

diff --git a/kafka-example/src/test/java/com/playtika/shepherd/JoinPasture.java b/kafka-example/src/test/java/com/playtika/shepherd/JoinPasture.java
@@ -36,8 +36,8 @@ public static void main(String[] args) {
 
 
         Pasture<String> skyNet = kafkaPushFarm.addBreedingPasture("SkyNet", String.class,
-                (population, version, generation, isLeader) -> {
-                    logger.info("Assigned leader={} version={} [{}]", isLeader, version, population);
+                (population, assignmentData) -> {
+                    logger.info("Assigned assignmentData={} [{}]", assignmentData, population);
                 });
 
 

diff --git a/kafka-functional-tests/src/test/java/com/playtika/shepherd/KafkaPullFarmTest.java b/kafka-functional-tests/src/test/java/com/playtika/shepherd/KafkaPullFarmTest.java
@@ -1,5 +1,6 @@
 package com.playtika.shepherd;
 
+import com.playtika.shepherd.common.AssignmentData;
 import com.playtika.shepherd.common.pull.Herd;
 import com.playtika.shepherd.common.pull.Pasture;
 import com.playtika.shepherd.common.pull.Shepherd;
@@ -38,8 +39,8 @@ public void shouldBalanceStaticHerd() {
 
         AtomicReference<List<ByteBuffer>> cows1 = new AtomicReference<>(List.of());
 
-        Pasture pasture1 = kafkaRanch.addPasture(herd, (population, version, generation, isLeader) -> {
-            logPopulation(1, population, version, isLeader);
+        Pasture pasture1 = kafkaRanch.addPasture(herd, (population, assignmentData) -> {
+            logPopulation(1, population, assignmentData);
             cows1.set(population);
         });
 
@@ -52,8 +53,8 @@ public void shouldBalanceStaticHerd() {
 
         //setup another pasture
         AtomicReference<List<ByteBuffer>> cows2 = new AtomicReference<>(List.of());
-        Pasture pasture2 = kafkaRanch.addPasture(herd, (population, version, generation, isLeader) -> {
-            logPopulation(2, population, version, isLeader);
+        Pasture pasture2 = kafkaRanch.addPasture(herd, (population, assignmentData) -> {
+            logPopulation(2, population, assignmentData);
             cows2.set(population);
         });
         pasture2.start();
@@ -65,8 +66,8 @@ public void shouldBalanceStaticHerd() {
 
         //setup third pasture
         AtomicReference<List<ByteBuffer>> cows3 = new AtomicReference<>(List.of());
-        Pasture pasture3 = kafkaRanch.addPasture(herd, (population, version, generation, isLeader) -> {
-            logPopulation(3, population, version, isLeader);
+        Pasture pasture3 = kafkaRanch.addPasture(herd, (population, assignmentData) -> {
+            logPopulation(3, population, assignmentData);
             cows3.set(population);
         });
         pasture3.start();
@@ -113,20 +114,20 @@ private void shouldBalanceDynamicHerd(boolean versioned) {
 
         AtomicReference<List<ByteBuffer>> cows1 = new AtomicReference<>(List.of());
         AtomicLong version1 = new AtomicLong();
-        Pasture pasture1 = kafkaRanch.addPasture(herd, (population, version, generation, isLeader) -> {
-            logPopulation(1, population, version, isLeader);
+        Pasture pasture1 = kafkaRanch.addPasture(herd, (population, assignmentData) -> {
+            logPopulation(1, population, assignmentData);
             cows1.set(population);
-            version1.set(version);
+            version1.set(assignmentData.populationVersion());
         });
         populationGlobal.set(new Herd.Population<>(new ByteBuffer[]{cow1, cow2}, ver1));
         pasture1.start();
 
         AtomicReference<List<ByteBuffer>> cows2 = new AtomicReference<>(List.of());
         AtomicLong version2 = new AtomicLong();
-        Pasture pasture2 = kafkaRanch.addPasture(herd, (population, version, generation, isLeader) -> {
-            logPopulation(2, population, version, isLeader);
+        Pasture pasture2 = kafkaRanch.addPasture(herd, (population, assignmentData) -> {
+            logPopulation(2, population, assignmentData);
             cows2.set(population);
-            version2.set(version);
+            version2.set(assignmentData.populationVersion());
         });
         pasture2.start();
 
@@ -206,7 +207,7 @@ public void shouldBalanceBreedingStaticHerd() {
         AtomicReference<List<BlackSheep>> subHerd1 = new AtomicReference<>(List.of());
 
         Pasture pasture1 = kafkaRanch.addBreedingPasture(herd, BlackSheep.class,
-                (population, version, generation, isLeader) -> {
+                (population, assignmentData) -> {
             logger.info("Assigned sheep1 [{}]", population);
             subHerd1.set(population);
         });
@@ -219,7 +220,7 @@ public void shouldBalanceBreedingStaticHerd() {
         //setup another pasture
         AtomicReference<List<BlackSheep>> subHerd2 = new AtomicReference<>(List.of());
         Pasture pasture2 = kafkaRanch.addBreedingPasture(herd, BlackSheep.class,
-                (population, version, generation, isLeader) -> {
+                (population, assignmentData) -> {
             logger.info("Assigned sheep2 [{}]", population);
             subHerd2.set(population);
         });
@@ -232,7 +233,7 @@ public void shouldBalanceBreedingStaticHerd() {
 
         //setup third pasture
         AtomicReference<List<BlackSheep>> subHerd3 = new AtomicReference<>(List.of());
-        Pasture pasture3 = kafkaRanch.addBreedingPasture(herd, BlackSheep.class, (population, version, generation, isLeader) -> {
+        Pasture pasture3 = kafkaRanch.addBreedingPasture(herd, BlackSheep.class, (population, assignmentData) -> {
             logger.info("Assigned cows3 [{}]", population);
             subHerd3.set(population);
         });
@@ -253,8 +254,8 @@ public void shouldBalanceBreedingStaticHerd() {
         });
     }
 
-    private static void logPopulation(int pastureIndex, List<ByteBuffer> population, long version, boolean isLeader) {
-        logger.info("Assigned to pasture{} leader={} version={} [{}]", pastureIndex, isLeader, version, toBytes(population));
+    private static void logPopulation(int pastureIndex, List<ByteBuffer> population, AssignmentData assignmentData) {
+        logger.info("Assigned to pasture{} assignmentData={} [{}]", pastureIndex, assignmentData, toBytes(population));
     }
 
     public static class TestHerd<Breed> implements Herd<Breed> {

diff --git a/kafka-functional-tests/src/test/java/com/playtika/shepherd/KafkaPushFarmTest.java b/kafka-functional-tests/src/test/java/com/playtika/shepherd/KafkaPushFarmTest.java
@@ -1,5 +1,6 @@
 package com.playtika.shepherd;
 
+import com.playtika.shepherd.common.AssignmentData;
 import com.playtika.shepherd.common.push.Pasture;
 import com.playtika.shepherd.common.push.Shepherd;
 import org.junit.jupiter.api.Test;
@@ -37,8 +38,8 @@ public void shouldBalanceStaticHerd() {
         AtomicReference<List<ByteBuffer>> cows1 = new AtomicReference<>(List.of());
 
         String herdName = "push-static-herd";
-        Pasture<ByteBuffer> pasture1 = kafkaRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
-            logPopulation(1, population, version, isLeader);
+        Pasture<ByteBuffer> pasture1 = kafkaRanch.addPasture(herdName, (population, assignmentData) -> {
+            logPopulation(1, population, assignmentData);
             cows1.set(population);
         });
 
@@ -51,8 +52,8 @@ public void shouldBalanceStaticHerd() {
 
         //setup another pasture
         AtomicReference<List<ByteBuffer>> cows2 = new AtomicReference<>(List.of());
-        Pasture<ByteBuffer> pasture2 = kafkaRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
-            logPopulation(2, population, version, isLeader);
+        Pasture<ByteBuffer> pasture2 = kafkaRanch.addPasture(herdName, (population, assignmentData) -> {
+            logPopulation(2, population, assignmentData);
             cows2.set(population);
         });
         pasture2.getShepherd().setPopulation(cows, -1);
@@ -65,8 +66,8 @@ public void shouldBalanceStaticHerd() {
 
         //setup third pasture
         AtomicReference<List<ByteBuffer>> cows3 = new AtomicReference<>(List.of());
-        Pasture<ByteBuffer> pasture3 = kafkaRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
-            logPopulation(3, population, version, isLeader);
+        Pasture<ByteBuffer> pasture3 = kafkaRanch.addPasture(herdName, (population, assignmentData) -> {
+            logPopulation(3, population, assignmentData);
             cows3.set(population);
         });
         pasture3.getShepherd().setPopulation(cows, -1);
@@ -113,20 +114,20 @@ private void shouldBalanceDynamicHerd(boolean versioned) {
         AtomicReference<List<ByteBuffer>> cows1 = new AtomicReference<>(List.of());
         AtomicLong version1 = new AtomicLong();
         String herdName = versioned ? "push-dynamic-group-versioned" : "push-dynamic-group";
-        Pasture<ByteBuffer> pasture1 = kafkaRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
-            logPopulation(1, population, version, isLeader);
+        Pasture<ByteBuffer> pasture1 = kafkaRanch.addPasture(herdName, (population, assignmentData) -> {
+            logPopulation(1, population, assignmentData);
             cows1.set(population);
-            version1.set(version);
+            version1.set(assignmentData.populationVersion());
         });
         pasture1.getShepherd().setPopulation(cows, ver1);
         pasture1.start();
 
         AtomicReference<List<ByteBuffer>> cows2 = new AtomicReference<>(List.of());
         AtomicLong version2 = new AtomicLong();
-        Pasture<ByteBuffer> pasture2 = kafkaRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
-            logPopulation(2, population, version, isLeader);
+        Pasture<ByteBuffer> pasture2 = kafkaRanch.addPasture(herdName, (population, assignmentData) -> {
+            logPopulation(2, population, assignmentData);
             cows2.set(population);
-            version2.set(version);
+            version2.set(assignmentData.populationVersion());
         });
         pasture2.getShepherd().setPopulation(cows, ver1);
         pasture2.start();
@@ -212,16 +213,16 @@ public void shouldBalanceDynamicConcurrentSequenceHerd() {
 
         AtomicReference<List<ByteBuffer>> cows1 = new AtomicReference<>(List.of());
         String herdName = "push-random-group";
-        Pasture<ByteBuffer> pasture1 = kafkaRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
-            logPopulation(1, population, version, isLeader);
+        Pasture<ByteBuffer> pasture1 = kafkaRanch.addPasture(herdName, (population, assignmentData) -> {
+            logPopulation(1, population, assignmentData);
             cows1.set(population);
         });
         pasture1.getShepherd().setPopulation(new ByteBuffer[]{}, 0);
         pasture1.start();
 
         AtomicReference<List<ByteBuffer>> cows2 = new AtomicReference<>(List.of());
-        Pasture<ByteBuffer> pasture2 = kafkaRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
-            logPopulation(2, population, version, isLeader);
+        Pasture<ByteBuffer> pasture2 = kafkaRanch.addPasture(herdName, (population, assignmentData) -> {
+            logPopulation(2, population, assignmentData);
             cows2.set(population);
         });
         pasture2.getShepherd().setPopulation(new ByteBuffer[]{}, 0);
@@ -274,8 +275,7 @@ public void shouldBalanceBreedingStaticHerd() {
         AtomicReference<List<BlackSheep>> subHerd1 = new AtomicReference<>(List.of());
 
         String herdName = "push-static-breeding-herd";
-        Pasture<BlackSheep> pasture1 = kafkaRanch.addBreedingPasture(herdName, BlackSheep.class,
-                (population, version, generation, isLeader) -> {
+        Pasture<BlackSheep> pasture1 = kafkaRanch.addBreedingPasture(herdName, BlackSheep.class, (population, assignmentData) -> {
             logger.info("Assigned sheep1 [{}]", population);
             subHerd1.set(population);
         });
@@ -289,8 +289,7 @@ public void shouldBalanceBreedingStaticHerd() {
 
         //setup another pasture
         AtomicReference<List<BlackSheep>> subHerd2 = new AtomicReference<>(List.of());
-        Pasture<BlackSheep> pasture2 = kafkaRanch.addBreedingPasture(herdName, BlackSheep.class,
-                (population, version, generation, isLeader) -> {
+        Pasture<BlackSheep> pasture2 = kafkaRanch.addBreedingPasture(herdName, BlackSheep.class, (population, assignmentData) -> {
             logger.info("Assigned sheep2 [{}]", population);
             subHerd2.set(population);
         });
@@ -304,7 +303,7 @@ public void shouldBalanceBreedingStaticHerd() {
 
         //setup third pasture
         AtomicReference<List<BlackSheep>> subHerd3 = new AtomicReference<>(List.of());
-        Pasture<BlackSheep> pasture3 = kafkaRanch.addBreedingPasture(herdName, BlackSheep.class, (population, version, generation, isLeader) -> {
+        Pasture<BlackSheep> pasture3 = kafkaRanch.addBreedingPasture(herdName, BlackSheep.class, (population, assignmentData) -> {
             logger.info("Assigned cows3 [{}]", population);
             subHerd3.set(population);
         });
@@ -326,8 +325,8 @@ public void shouldBalanceBreedingStaticHerd() {
         });
     }
 
-    private static void logPopulation(int pastureIndex, List<ByteBuffer> population, long version, boolean isLeader) {
-        logger.info("Assigned to pasture{} leader={} version={} [{}]", pastureIndex, isLeader, version, toBytes(population));
+    private static void logPopulation(int pastureIndex, List<ByteBuffer> population, AssignmentData assignmentData) {
+        logger.info("Assigned to pasture{} assignmentData={} [{}]", pastureIndex, assignmentData, toBytes(population));
     }
 
     static class BlackSheep {

diff --git a/kafka-functional-tests/src/test/java/com/playtika/shepherd/KafkaToxiFarmTest.java b/kafka-functional-tests/src/test/java/com/playtika/shepherd/KafkaToxiFarmTest.java
@@ -123,7 +123,7 @@ public void shouldRestoreBalanceForStaticHerd() throws IOException {
         AtomicReference<List<ByteBuffer>> cows1 = new AtomicReference<>(List.of());
 
         String herdName = "static-herd";
-        Pasture<ByteBuffer> pasture1 = kafkaRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
+        Pasture<ByteBuffer> pasture1 = kafkaRanch.addPasture(herdName, (population, assignmentData) -> {
             logger.info("Assigned cows1 [{}]", toBytes(population));
             cows1.set(population);
         });
@@ -138,7 +138,7 @@ public void shouldRestoreBalanceForStaticHerd() throws IOException {
         //setup toxi pasture
         KafkaPushFarm kafkaToxiRanch = new KafkaPushFarm(getToxiBootstrapServers(), TEST_PROPERTIES);
         AtomicReference<List<ByteBuffer>> cows2 = new AtomicReference<>(List.of());
-        Pasture<ByteBuffer> pasture2 = kafkaToxiRanch.addPasture(herdName, (population, version, generation, isLeader) -> {
+        Pasture<ByteBuffer> pasture2 = kafkaToxiRanch.addPasture(herdName, (population, assignmentData) -> {
             logger.info("Assigned cows2 [{}]", toBytes(population));
             cows2.set(population);
         });

diff --git a/kafka-functional-tests/src/test/java/com/playtika/shepherd/inernal/ShepherdTest.java b/kafka-functional-tests/src/test/java/com/playtika/shepherd/inernal/ShepherdTest.java
@@ -1,6 +1,7 @@
 package com.playtika.shepherd.inernal;
 
 import com.playtika.shepherd.BasicKafkaTest;
+import com.playtika.shepherd.common.AssignmentData;
 import com.playtika.shepherd.common.PastureListener;
 import org.apache.kafka.common.message.JoinGroupResponseData;
 import org.junit.jupiter.api.Test;
@@ -50,7 +51,7 @@ public void reset() {
         LinkedBlockingQueue<ByteBuffer> cows1 = new LinkedBlockingQueue<>();
         PastureListener<ByteBuffer> rebalanceListener1 = new PastureListener<>() {
             @Override
-            public void assigned(List<ByteBuffer> population, long version, int generation, boolean isLeader) {
+            public void assigned(List<ByteBuffer> population, AssignmentData assignmentData) {
                 logger.info("Assigned cows1 [{}]", toBytes(population));
                 cows1.addAll(population);
             }
@@ -80,7 +81,7 @@ public void cleanup() {
         LinkedBlockingQueue<ByteBuffer> cows2 = new LinkedBlockingQueue<>();
         PastureListener<ByteBuffer> rebalanceListener2 = new PastureListener<>() {
             @Override
-            public void assigned(List<ByteBuffer> population, long version, int generation, boolean isLeader) {
+            public void assigned(List<ByteBuffer> population, AssignmentData assignmentData) {
                 logger.info("Assigned cows2 [{}]", toBytes(population));
                 cows2.addAll(population);
             }
@@ -142,7 +143,7 @@ public void reset() {
         LinkedBlockingQueue<ByteBuffer> cows1 = new LinkedBlockingQueue<>();
         PastureListener<ByteBuffer> rebalanceListener1 = new PastureListener<>() {
             @Override
-            public void assigned(List<ByteBuffer> population, long version, int generation, boolean isLeader) {
+            public void assigned(List<ByteBuffer> population, AssignmentData assignmentData) {
                 logger.info("Assigned cows1 [{}]", toBytes(population));
                 cows1.addAll(population);
             }
@@ -166,7 +167,7 @@ public void cleanup() {
         LinkedBlockingQueue<ByteBuffer> cows2 = new LinkedBlockingQueue<>();
         PastureListener<ByteBuffer> rebalanceListener2 = new PastureListener<>() {
             @Override
-            public void assigned(List<ByteBuffer> population, long version, int generation, boolean isLeader) {
+            public void assigned(List<ByteBuffer> population, AssignmentData assignmentData) {
                 logger.info("Assigned cows2 [{}]", toBytes(population));
                 cows2.addAll(population);
             }

diff --git a/kafka/src/main/java/com/playtika/shepherd/KafkaPullFarm.java b/kafka/src/main/java/com/playtika/shepherd/KafkaPullFarm.java
@@ -1,5 +1,6 @@
 package com.playtika.shepherd;
 
+import com.playtika.shepherd.common.AssignmentData;
 import com.playtika.shepherd.common.PastureListener;
 import com.playtika.shepherd.common.pull.Farm;
 import com.playtika.shepherd.common.pull.Herd;
@@ -106,8 +107,8 @@ public void rebalanceHerd() {
         }
 
         @Override
-        public void assigned(List<ByteBuffer> population, long version, int generation, boolean isLeader) {
-            pastureListener.assigned(serDe.deserialize(population), version, generation, isLeader);
+        public void assigned(List<ByteBuffer> population, AssignmentData assignmentData) {
+            pastureListener.assigned(serDe.deserialize(population), assignmentData);
         }
 
         @Override

diff --git a/kafka/src/main/java/com/playtika/shepherd/KafkaPushFarm.java b/kafka/src/main/java/com/playtika/shepherd/KafkaPushFarm.java
@@ -1,6 +1,7 @@
 package com.playtika.shepherd;
 
 import com.playtika.shepherd.common.PastureListener;
+import com.playtika.shepherd.common.AssignmentData;
 import com.playtika.shepherd.common.push.Farm;
 import com.playtika.shepherd.common.push.Pasture;
 import com.playtika.shepherd.common.push.Shepherd;
@@ -134,9 +135,9 @@ public synchronized void reset() {
         }
 
         @Override
-        public synchronized void assigned(List<ByteBuffer> population, long version, int generation, boolean isLeader) {
-            this.pastureListener.assigned(serDe.deserialize(population), version, generation, isLeader);
-            this.assignedVersion = version;
+        public synchronized void assigned(List<ByteBuffer> population, AssignmentData assignmentData) {
+            this.pastureListener.assigned(serDe.deserialize(population), assignmentData);
+            this.assignedVersion = assignmentData.populationVersion();
         }
 
         @Override