From 8838fa8010c146af6aab014a41bc7e68318b4eb0 Mon Sep 17 00:00:00 2001
From: Ewen Cheslack-Postava <me@ewencp.org>
Date: Wed, 28 Oct 2015 12:42:03 -0700
Subject: [PATCH] MINOR: KAFKA-2371 follow-up, DistributedHerder should wakeup
 WorkerGroupMember after assignment to ensure work is started immediately

Author: Ewen Cheslack-Postava <me@ewencp.org>

Reviewers: Gwen Shapira

Closes #360 from ewencp/minor-kafka-2371-follow-up-wakeup-after-rebalance
---
 .../kafka/copycat/runtime/distributed/DistributedHerder.java  | 4 ++++
 .../copycat/runtime/distributed/DistributedHerderTest.java    | 2 ++
 2 files changed, 6 insertions(+)

diff --git a/copycat/runtime/src/main/java/org/apache/kafka/copycat/runtime/distributed/DistributedHerder.java b/copycat/runtime/src/main/java/org/apache/kafka/copycat/runtime/distributed/DistributedHerder.java
index 17bf7b7b2f2..46c7686ee72 100644
--- a/copycat/runtime/src/main/java/org/apache/kafka/copycat/runtime/distributed/DistributedHerder.java
+++ b/copycat/runtime/src/main/java/org/apache/kafka/copycat/runtime/distributed/DistributedHerder.java
@@ -614,6 +614,10 @@ public class DistributedHerder implements Herder, Runnable {
                 log.info("Joined group and got assignment: {}", assignment);
                 DistributedHerder.this.assignment = assignment;
                 rebalanceResolved = false;
+                // We *must* interrupt any poll() call since this could occur when the poll starts, and we might then
+                // sleep in the poll() for a long time. Forcing a wakeup ensures we'll get to process this event in the
+                // main thread.
+                member.wakeup();
             }
 
             @Override
diff --git a/copycat/runtime/src/test/java/org/apache/kafka/copycat/runtime/distributed/DistributedHerderTest.java b/copycat/runtime/src/test/java/org/apache/kafka/copycat/runtime/distributed/DistributedHerderTest.java
index 1213656ae33..c8b4874222f 100644
--- a/copycat/runtime/src/test/java/org/apache/kafka/copycat/runtime/distributed/DistributedHerderTest.java
+++ b/copycat/runtime/src/test/java/org/apache/kafka/copycat/runtime/distributed/DistributedHerderTest.java
@@ -371,6 +371,8 @@ public class DistributedHerderTest {
                 return null;
             }
         });
+        member.wakeup();
+        PowerMock.expectLastCall();
     }
 
     private void expectPostRebalanceCatchup(final ClusterConfigState readToEndSnapshot) {