opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/ShardManagerTest.java

   1 package org.opendaylight.controller.cluster.datastore;
   2
   3 import static org.junit.Assert.assertEquals;
   4 import static org.junit.Assert.assertNull;
   5 import static org.junit.Assert.assertSame;
   6 import static org.junit.Assert.assertTrue;
   7 import static org.mockito.Mockito.mock;
   8 import static org.mockito.Mockito.never;
   9 import static org.mockito.Mockito.times;
  10 import static org.mockito.Mockito.verify;
  11 import akka.actor.ActorRef;
  12 import akka.actor.ActorSystem;
  13 import akka.actor.AddressFromURIString;
  14 import akka.actor.Props;
  15 import akka.cluster.Cluster;
  16 import akka.cluster.ClusterEvent;
  17 import akka.dispatch.Dispatchers;
  18 import akka.japi.Creator;
  19 import akka.pattern.Patterns;
  20 import akka.persistence.RecoveryCompleted;
  21 import akka.testkit.JavaTestKit;
  22 import akka.testkit.TestActorRef;
  23 import akka.util.Timeout;
  24 import com.google.common.base.Optional;
  25 import com.google.common.collect.ImmutableMap;
  26 import com.google.common.collect.ImmutableSet;
  27 import com.google.common.util.concurrent.Uninterruptibles;
  28 import com.typesafe.config.ConfigFactory;
  29 import java.util.Arrays;
  30 import java.util.List;
  31 import java.util.Map;
  32 import java.util.concurrent.CountDownLatch;
  33 import java.util.concurrent.TimeUnit;
  34 import org.junit.After;
  35 import org.junit.Before;
  36 import org.junit.Test;
  37 import org.mockito.Mock;
  38 import org.mockito.MockitoAnnotations;
  39 import org.opendaylight.controller.cluster.datastore.exceptions.NoShardLeaderException;
  40 import org.opendaylight.controller.cluster.datastore.exceptions.NotInitializedException;
  41 import org.opendaylight.controller.cluster.datastore.exceptions.PrimaryNotFoundException;
  42 import org.opendaylight.controller.cluster.datastore.identifiers.ShardIdentifier;
  43 import org.opendaylight.controller.cluster.datastore.identifiers.ShardManagerIdentifier;
  44 import org.opendaylight.controller.cluster.datastore.messages.ActorInitialized;
  45 import org.opendaylight.controller.cluster.datastore.messages.FindLocalShard;
  46 import org.opendaylight.controller.cluster.datastore.messages.FindPrimary;
  47 import org.opendaylight.controller.cluster.datastore.messages.LocalPrimaryShardFound;
  48 import org.opendaylight.controller.cluster.datastore.messages.LocalShardFound;
  49 import org.opendaylight.controller.cluster.datastore.messages.LocalShardNotFound;
  50 import org.opendaylight.controller.cluster.datastore.messages.PrimaryShardInfo;
  51 import org.opendaylight.controller.cluster.datastore.messages.RemotePrimaryShardFound;
  52 import org.opendaylight.controller.cluster.datastore.messages.ShardLeaderStateChanged;
  53 import org.opendaylight.controller.cluster.datastore.messages.UpdateSchemaContext;
  54 import org.opendaylight.controller.cluster.datastore.utils.MessageCollectorActor;
  55 import org.opendaylight.controller.cluster.datastore.utils.MockClusterWrapper;
  56 import org.opendaylight.controller.cluster.datastore.utils.MockConfiguration;
  57 import org.opendaylight.controller.cluster.datastore.utils.PrimaryShardInfoFutureCache;
  58 import org.opendaylight.controller.cluster.notifications.LeaderStateChanged;
  59 import org.opendaylight.controller.cluster.notifications.RegisterRoleChangeListener;
  60 import org.opendaylight.controller.cluster.notifications.RoleChangeNotification;
  61 import org.opendaylight.controller.cluster.raft.RaftState;
  62 import org.opendaylight.controller.cluster.raft.base.messages.FollowerInitialSyncUpStatus;
  63 import org.opendaylight.controller.cluster.raft.utils.InMemoryJournal;
  64 import org.opendaylight.controller.md.cluster.datastore.model.TestModel;
  65 import org.opendaylight.yangtools.yang.data.api.schema.tree.DataTree;
  66 import org.opendaylight.yangtools.yang.model.api.SchemaContext;
  67 import scala.concurrent.Await;
  68 import scala.concurrent.Future;
  69 import scala.concurrent.duration.FiniteDuration;
  70
  71 public class ShardManagerTest extends AbstractActorTest {
  72     private static int ID_COUNTER = 1;
  73
  74     private final String shardMrgIDSuffix = "config" + ID_COUNTER++;
  75     private final String shardMgrID = "shard-manager-" + shardMrgIDSuffix;
  76
  77     @Mock
  78     private static CountDownLatch ready;
  79
  80     private static TestActorRef<MessageCollectorActor> mockShardActor;
  81
  82     private final DatastoreContext.Builder datastoreContextBuilder = DatastoreContext.newBuilder().
  83             dataStoreType(shardMrgIDSuffix).shardInitializationTimeout(600, TimeUnit.MILLISECONDS);
  84
  85     private static ActorRef newMockShardActor(ActorSystem system, String shardName, String memberName) {
  86         String name = new ShardIdentifier(shardName, memberName,"config").toString();
  87         return TestActorRef.create(system, Props.create(MessageCollectorActor.class), name);
  88     }
  89
  90     private final PrimaryShardInfoFutureCache primaryShardInfoCache = new PrimaryShardInfoFutureCache();
  91
  92     @Before
  93     public void setUp() {
  94         MockitoAnnotations.initMocks(this);
  95
  96         InMemoryJournal.clear();
  97
  98         if(mockShardActor == null) {
  99             String name = new ShardIdentifier(Shard.DEFAULT_NAME, "member-1", "config").toString();
 100             mockShardActor = TestActorRef.create(getSystem(), Props.create(MessageCollectorActor.class), name);
 101         }
 102
 103         mockShardActor.underlyingActor().clear();
 104     }
 105
 106     @After
 107     public void tearDown() {
 108         InMemoryJournal.clear();
 109     }
 110
 111     private Props newShardMgrProps(boolean persistent) {
 112         return ShardManager.props(new MockClusterWrapper(), new MockConfiguration(),
 113                 datastoreContextBuilder.persistent(persistent).build(), ready, primaryShardInfoCache);
 114     }
 115
 116     private Props newPropsShardMgrWithMockShardActor() {
 117         return newPropsShardMgrWithMockShardActor("shardManager", mockShardActor, new MockClusterWrapper(),
 118                 new MockConfiguration());
 119     }
 120
 121     private Props newPropsShardMgrWithMockShardActor(final String name, final ActorRef shardActor,
 122             final ClusterWrapper clusterWrapper, final Configuration config) {
 123         Creator<ShardManager> creator = new Creator<ShardManager>() {
 124             private static final long serialVersionUID = 1L;
 125             @Override
 126             public ShardManager create() throws Exception {
 127                 return new ForwardingShardManager(clusterWrapper, config, datastoreContextBuilder.build(),
 128                         ready, name, shardActor, primaryShardInfoCache);
 129             }
 130         };
 131
 132         return Props.create(new DelegatingShardManagerCreator(creator)).withDispatcher(Dispatchers.DefaultDispatcherId());
 133     }
 134
 135     @Test
 136     public void testOnReceiveFindPrimaryForNonExistentShard() throws Exception {
 137         new JavaTestKit(getSystem()) {{
 138             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 139
 140             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 141
 142             shardManager.tell(new FindPrimary("non-existent", false), getRef());
 143
 144             expectMsgClass(duration("5 seconds"), PrimaryNotFoundException.class);
 145         }};
 146     }
 147
 148     @Test
 149     public void testOnReceiveFindPrimaryForLocalLeaderShard() throws Exception {
 150         new JavaTestKit(getSystem()) {{
 151             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 152
 153             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 154
 155             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 156             shardManager.tell(new ActorInitialized(), mockShardActor);
 157
 158             DataTree mockDataTree = mock(DataTree.class);
 159             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree)), getRef());
 160
 161             MessageCollectorActor.expectFirstMatching(mockShardActor, RegisterRoleChangeListener.class);
 162             shardManager.tell((new RoleChangeNotification(memberId, RaftState.Candidate.name(),
 163                     RaftState.Leader.name())), mockShardActor);
 164
 165             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 166
 167             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 168             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 169                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 170             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 171         }};
 172     }
 173
 174     @Test
 175     public void testOnReceiveFindPrimaryForNonLocalLeaderShardBeforeMemberUp() throws Exception {
 176         new JavaTestKit(getSystem()) {{
 177             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 178
 179             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 180             shardManager.tell(new ActorInitialized(), mockShardActor);
 181
 182             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 183             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 184             shardManager.tell(new RoleChangeNotification(memberId1,
 185                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 186             shardManager.tell(new LeaderStateChanged(memberId1, memberId2), mockShardActor);
 187
 188             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 189
 190             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 191         }};
 192     }
 193
 194     @Test
 195     public void testOnReceiveFindPrimaryForNonLocalLeaderShard() throws Exception {
 196         new JavaTestKit(getSystem()) {{
 197             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 198
 199             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 200             shardManager.tell(new ActorInitialized(), mockShardActor);
 201
 202             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 203             MockClusterWrapper.sendMemberUp(shardManager, "member-2", getRef().path().toString());
 204
 205             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 206             shardManager.tell(new RoleChangeNotification(memberId1,
 207                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 208             shardManager.tell(new ShardLeaderStateChanged(memberId1, memberId2, Optional.<DataTree>absent()), mockShardActor);
 209
 210             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 211
 212             RemotePrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 213             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 214                     primaryFound.getPrimaryPath().contains("member-2-shard-default"));
 215         }};
 216     }
 217
 218     @Test
 219     public void testOnReceiveFindPrimaryForUninitializedShard() throws Exception {
 220         new JavaTestKit(getSystem()) {{
 221             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 222
 223             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 224
 225             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 226         }};
 227     }
 228
 229     @Test
 230     public void testOnReceiveFindPrimaryForInitializedShardWithNoRole() throws Exception {
 231         new JavaTestKit(getSystem()) {{
 232             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 233
 234             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 235             shardManager.tell(new ActorInitialized(), mockShardActor);
 236
 237             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 238
 239             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 240         }};
 241     }
 242
 243     @Test
 244     public void testOnReceiveFindPrimaryForFollowerShardWithNoInitialLeaderId() throws Exception {
 245         new JavaTestKit(getSystem()) {{
 246             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 247
 248             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 249             shardManager.tell(new ActorInitialized(), mockShardActor);
 250
 251             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 252             shardManager.tell(new RoleChangeNotification(memberId,
 253                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 254
 255             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 256
 257             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 258
 259             DataTree mockDataTree = mock(DataTree.class);
 260             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree)), mockShardActor);
 261
 262             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 263
 264             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 265             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 266                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 267             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 268         }};
 269     }
 270
 271     @Test
 272     public void testOnReceiveFindPrimaryWaitForShardLeader() throws Exception {
 273         new JavaTestKit(getSystem()) {{
 274             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 275
 276             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 277
 278             // We're passing waitUntilInitialized = true to FindPrimary so the response should be
 279             // delayed until we send ActorInitialized and RoleChangeNotification.
 280             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 281
 282             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 283
 284             shardManager.tell(new ActorInitialized(), mockShardActor);
 285
 286             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 287
 288             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 289             shardManager.tell(new RoleChangeNotification(memberId,
 290                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor);
 291
 292             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 293
 294             DataTree mockDataTree = mock(DataTree.class);
 295             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree)), mockShardActor);
 296
 297             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 298             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 299                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 300             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 301
 302             expectNoMsg(FiniteDuration.create(200, TimeUnit.MILLISECONDS));
 303         }};
 304     }
 305
 306     @Test
 307     public void testOnReceiveFindPrimaryWaitForReadyWithUninitializedShard() throws Exception {
 308         new JavaTestKit(getSystem()) {{
 309             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 310
 311             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 312
 313             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 314
 315             expectMsgClass(duration("2 seconds"), NotInitializedException.class);
 316
 317             shardManager.tell(new ActorInitialized(), mockShardActor);
 318
 319             expectNoMsg(FiniteDuration.create(200, TimeUnit.MILLISECONDS));
 320         }};
 321     }
 322
 323     @Test
 324     public void testOnReceiveFindPrimaryWaitForReadyWithCandidateShard() throws Exception {
 325         new JavaTestKit(getSystem()) {{
 326             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 327
 328             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 329             shardManager.tell(new ActorInitialized(), mockShardActor);
 330             shardManager.tell(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
 331                     null, RaftState.Candidate.name()), mockShardActor);
 332
 333             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 334
 335             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 336         }};
 337     }
 338
 339     @Test
 340     public void testOnReceiveFindPrimaryWaitForReadyWithNoRoleShard() throws Exception {
 341         new JavaTestKit(getSystem()) {{
 342             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 343
 344             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 345             shardManager.tell(new ActorInitialized(), mockShardActor);
 346
 347             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 348
 349             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 350         }};
 351     }
 352
 353     @Test
 354     public void testOnReceiveFindPrimaryForRemoteShard() throws Exception {
 355         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 356
 357         // Create an ActorSystem ShardManager actor for member-1.
 358
 359         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 360         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 361
 362         ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 363
 364         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 365                 newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 366                         new MockConfiguration()), shardManagerID);
 367
 368         // Create an ActorSystem ShardManager actor for member-2.
 369
 370         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 371
 372         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 373
 374         final ActorRef mockShardActor2 = newMockShardActor(system2, "astronauts", "member-2");
 375
 376         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 377                 put("default", Arrays.asList("member-1", "member-2")).
 378                 put("astronauts", Arrays.asList("member-2")).build());
 379
 380         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 381                 newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 382                         mockConfig2), shardManagerID);
 383
 384         new JavaTestKit(system1) {{
 385
 386             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 387             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 388
 389             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 390
 391             String memberId2 = "member-2-shard-astronauts-" + shardMrgIDSuffix;
 392             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2,
 393                     Optional.of(mock(DataTree.class))), mockShardActor2);
 394             shardManager2.tell(new RoleChangeNotification(memberId2,
 395                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 396
 397             shardManager1.underlyingActor().waitForMemberUp();
 398
 399             shardManager1.tell(new FindPrimary("astronauts", false), getRef());
 400
 401             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 402             String path = found.getPrimaryPath();
 403             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-astronauts-config"));
 404
 405             shardManager2.underlyingActor().verifyFindPrimary();
 406
 407             Cluster.get(system2).down(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 408
 409             shardManager1.underlyingActor().waitForMemberRemoved();
 410
 411             shardManager1.tell(new FindPrimary("astronauts", false), getRef());
 412
 413             expectMsgClass(duration("5 seconds"), PrimaryNotFoundException.class);
 414         }};
 415
 416         JavaTestKit.shutdownActorSystem(system1);
 417         JavaTestKit.shutdownActorSystem(system2);
 418     }
 419
 420     @Test
 421     public void testShardAvailabilityOnChangeOfMemberReachability() throws Exception {
 422         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 423
 424         // Create an ActorSystem ShardManager actor for member-1.
 425
 426         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 427         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 428
 429         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 430
 431         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 432             newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 433                 new MockConfiguration()), shardManagerID);
 434
 435         // Create an ActorSystem ShardManager actor for member-2.
 436
 437         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 438
 439         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 440
 441         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 442
 443         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 444             put("default", Arrays.asList("member-1", "member-2")).build());
 445
 446         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 447             newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 448                 mockConfig2), shardManagerID);
 449
 450         new JavaTestKit(system1) {{
 451
 452             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 453             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 454             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 455             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 456
 457             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 458             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 459             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 460                 Optional.of(mock(DataTree.class))), mockShardActor1);
 461             shardManager1.tell(new RoleChangeNotification(memberId1,
 462                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 463             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class))),
 464                 mockShardActor2);
 465             shardManager2.tell(new RoleChangeNotification(memberId2,
 466                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 467             shardManager1.underlyingActor().waitForMemberUp();
 468
 469             shardManager1.tell(new FindPrimary("default", true), getRef());
 470
 471             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 472             String path = found.getPrimaryPath();
 473             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 474
 475             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 476                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 477
 478             shardManager1.underlyingActor().waitForUnreachableMember();
 479
 480             shardManager1.tell(new FindPrimary("default", true), getRef());
 481
 482             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 483
 484             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 485                 createReachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 486
 487             shardManager1.underlyingActor().waitForReachableMember();
 488
 489             shardManager1.tell(new FindPrimary("default", true), getRef());
 490
 491             RemotePrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 492             String path1 = found1.getPrimaryPath();
 493             assertTrue("Unexpected primary path " + path1, path1.contains("member-2-shard-default-config"));
 494
 495         }};
 496
 497         JavaTestKit.shutdownActorSystem(system1);
 498         JavaTestKit.shutdownActorSystem(system2);
 499     }
 500
 501     @Test
 502     public void testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange() throws Exception {
 503         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 504
 505         // Create an ActorSystem ShardManager actor for member-1.
 506
 507         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 508         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 509
 510         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 511
 512         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 513             newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 514                 new MockConfiguration()), shardManagerID);
 515
 516         // Create an ActorSystem ShardManager actor for member-2.
 517
 518         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 519
 520         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 521
 522         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 523
 524         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 525             put("default", Arrays.asList("member-1", "member-2")).build());
 526
 527         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 528             newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 529                 mockConfig2), shardManagerID);
 530
 531         new JavaTestKit(system1) {{
 532
 533             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 534             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 535             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 536             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 537
 538             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 539             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 540             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 541                 Optional.of(mock(DataTree.class))), mockShardActor1);
 542             shardManager1.tell(new RoleChangeNotification(memberId1,
 543                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 544             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class))),
 545                 mockShardActor2);
 546             shardManager2.tell(new RoleChangeNotification(memberId2,
 547                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 548             shardManager1.underlyingActor().waitForMemberUp();
 549
 550             shardManager1.tell(new FindPrimary("default", true), getRef());
 551
 552             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 553             String path = found.getPrimaryPath();
 554             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 555
 556             primaryShardInfoCache.putSuccessful("default", new PrimaryShardInfo(system1.actorSelection(
 557                     mockShardActor1.path()), Optional.<DataTree>absent()));
 558
 559             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 560                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 561
 562             shardManager1.underlyingActor().waitForUnreachableMember();
 563
 564             shardManager1.tell(new FindPrimary("default", true), getRef());
 565
 566             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 567
 568             assertNull("Expected primaryShardInfoCache entry removed", primaryShardInfoCache.getIfPresent("default"));
 569
 570             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId1, Optional.of(mock(DataTree.class))),
 571                 mockShardActor1);
 572             shardManager1.tell(new RoleChangeNotification(memberId1,
 573                 RaftState.Follower.name(), RaftState.Leader.name()), mockShardActor1);
 574
 575             shardManager1.tell(new FindPrimary("default", true), getRef());
 576
 577             LocalPrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 578             String path1 = found1.getPrimaryPath();
 579             assertTrue("Unexpected primary path " + path1, path1.contains("member-1-shard-default-config"));
 580
 581         }};
 582
 583         JavaTestKit.shutdownActorSystem(system1);
 584         JavaTestKit.shutdownActorSystem(system2);
 585     }
 586
 587
 588     @Test
 589     public void testOnReceiveFindLocalShardForNonExistentShard() throws Exception {
 590         new JavaTestKit(getSystem()) {{
 591             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 592
 593             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 594
 595             shardManager.tell(new FindLocalShard("non-existent", false), getRef());
 596
 597             LocalShardNotFound notFound = expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
 598
 599             assertEquals("getShardName", "non-existent", notFound.getShardName());
 600         }};
 601     }
 602
 603     @Test
 604     public void testOnReceiveFindLocalShardForExistentShard() throws Exception {
 605         new JavaTestKit(getSystem()) {{
 606             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 607
 608             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 609             shardManager.tell(new ActorInitialized(), mockShardActor);
 610
 611             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 612
 613             LocalShardFound found = expectMsgClass(duration("5 seconds"), LocalShardFound.class);
 614
 615             assertTrue("Found path contains " + found.getPath().path().toString(),
 616                     found.getPath().path().toString().contains("member-1-shard-default-config"));
 617         }};
 618     }
 619
 620     @Test
 621     public void testOnReceiveFindLocalShardForNotInitializedShard() throws Exception {
 622         new JavaTestKit(getSystem()) {{
 623             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 624
 625             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 626
 627             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 628         }};
 629     }
 630
 631     @Test
 632     public void testOnReceiveFindLocalShardWaitForShardInitialized() throws Exception {
 633         new JavaTestKit(getSystem()) {{
 634             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 635
 636             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 637
 638             // We're passing waitUntilInitialized = true to FindLocalShard so the response should be
 639             // delayed until we send ActorInitialized.
 640             Future<Object> future = Patterns.ask(shardManager, new FindLocalShard(Shard.DEFAULT_NAME, true),
 641                     new Timeout(5, TimeUnit.SECONDS));
 642
 643             shardManager.tell(new ActorInitialized(), mockShardActor);
 644
 645             Object resp = Await.result(future, duration("5 seconds"));
 646             assertTrue("Expected: LocalShardFound, Actual: " + resp, resp instanceof LocalShardFound);
 647         }};
 648     }
 649
 650     @Test
 651     public void testOnRecoveryJournalIsCleaned() {
 652         InMemoryJournal.addEntry(shardMgrID, 1L, new ShardManager.SchemaContextModules(
 653                 ImmutableSet.of("foo")));
 654         InMemoryJournal.addEntry(shardMgrID, 2L, new ShardManager.SchemaContextModules(
 655                 ImmutableSet.of("bar")));
 656         InMemoryJournal.addDeleteMessagesCompleteLatch(shardMgrID);
 657
 658         new JavaTestKit(getSystem()) {{
 659             TestActorRef<TestShardManager> shardManager = TestActorRef.create(getSystem(),
 660                     Props.create(new TestShardManagerCreator(shardMrgIDSuffix)));
 661
 662             shardManager.underlyingActor().waitForRecoveryComplete();
 663             InMemoryJournal.waitForDeleteMessagesComplete(shardMgrID);
 664
 665             // Journal entries up to the last one should've been deleted
 666             Map<Long, Object> journal = InMemoryJournal.get(shardMgrID);
 667             synchronized (journal) {
 668                 assertEquals("Journal size", 0, journal.size());
 669             }
 670         }};
 671     }
 672
 673     @Test
 674     public void testRoleChangeNotificationAndShardLeaderStateChangedReleaseReady() throws Exception {
 675         new JavaTestKit(getSystem()) {
 676             {
 677                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 678
 679                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 680                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 681                         memberId, RaftState.Candidate.name(), RaftState.Leader.name()));
 682
 683                 verify(ready, never()).countDown();
 684
 685                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId, memberId,
 686                         Optional.of(mock(DataTree.class))));
 687
 688                 verify(ready, times(1)).countDown();
 689
 690             }};
 691     }
 692
 693     @Test
 694     public void testRoleChangeNotificationToFollowerWithShardLeaderStateChangedReleaseReady() throws Exception {
 695         new JavaTestKit(getSystem()) {
 696             {
 697                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 698
 699                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 700                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 701                         memberId, null, RaftState.Follower.name()));
 702
 703                 verify(ready, never()).countDown();
 704
 705                 shardManager.underlyingActor().onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
 706
 707                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId,
 708                         "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class))));
 709
 710                 verify(ready, times(1)).countDown();
 711
 712             }};
 713     }
 714
 715     @Test
 716     public void testReadyCountDownForMemberUpAfterLeaderStateChanged() throws Exception {
 717         new JavaTestKit(getSystem()) {
 718             {
 719                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 720
 721                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 722                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 723                         memberId, null, RaftState.Follower.name()));
 724
 725                 verify(ready, never()).countDown();
 726
 727                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId,
 728                         "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class))));
 729
 730                 shardManager.underlyingActor().onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
 731
 732                 verify(ready, times(1)).countDown();
 733
 734             }};
 735     }
 736
 737     @Test
 738     public void testRoleChangeNotificationDoNothingForUnknownShard() throws Exception {
 739         new JavaTestKit(getSystem()) {
 740             {
 741                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 742
 743                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 744                         "unknown", RaftState.Candidate.name(), RaftState.Leader.name()));
 745
 746                 verify(ready, never()).countDown();
 747
 748             }};
 749     }
 750
 751
 752     @Test
 753     public void testByDefaultSyncStatusIsFalse() throws Exception{
 754         final Props persistentProps = newShardMgrProps(true);
 755         final TestActorRef<ShardManager> shardManager =
 756                 TestActorRef.create(getSystem(), persistentProps);
 757
 758         ShardManager shardManagerActor = shardManager.underlyingActor();
 759
 760         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 761     }
 762
 763     @Test
 764     public void testWhenShardIsLeaderSyncStatusIsTrue() throws Exception{
 765         final Props persistentProps = ShardManager.props(
 766                 new MockClusterWrapper(),
 767                 new MockConfiguration(),
 768                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 769         final TestActorRef<ShardManager> shardManager =
 770                 TestActorRef.create(getSystem(), persistentProps);
 771
 772         ShardManager shardManagerActor = shardManager.underlyingActor();
 773         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 774                 RaftState.Follower.name(), RaftState.Leader.name()));
 775
 776         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 777     }
 778
 779     @Test
 780     public void testWhenShardIsCandidateSyncStatusIsFalse() throws Exception{
 781         final Props persistentProps = newShardMgrProps(true);
 782         final TestActorRef<ShardManager> shardManager =
 783                 TestActorRef.create(getSystem(), persistentProps);
 784
 785         ShardManager shardManagerActor = shardManager.underlyingActor();
 786         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 787                 RaftState.Follower.name(), RaftState.Candidate.name()));
 788
 789         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 790
 791         // Send a FollowerInitialSyncStatus with status = true for the replica whose current state is candidate
 792         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-default-unknown"));
 793
 794         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 795     }
 796
 797     @Test
 798     public void testWhenShardIsFollowerSyncStatusDependsOnFollowerInitialSyncStatus() throws Exception{
 799         final Props persistentProps = ShardManager.props(
 800                 new MockClusterWrapper(),
 801                 new MockConfiguration(),
 802                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 803         final TestActorRef<ShardManager> shardManager =
 804                 TestActorRef.create(getSystem(), persistentProps);
 805
 806         ShardManager shardManagerActor = shardManager.underlyingActor();
 807         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 808                 RaftState.Candidate.name(), RaftState.Follower.name()));
 809
 810         // Initially will be false
 811         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 812
 813         // Send status true will make sync status true
 814         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-default-unknown"));
 815
 816         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 817
 818         // Send status false will make sync status false
 819         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(false, "member-1-shard-default-unknown"));
 820
 821         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 822
 823     }
 824
 825     @Test
 826     public void testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards() throws Exception{
 827         final Props persistentProps = ShardManager.props(
 828                 new MockClusterWrapper(),
 829                 new MockConfiguration() {
 830                     @Override
 831                     public List<String> getMemberShardNames(String memberName) {
 832                         return Arrays.asList("default", "astronauts");
 833                     }
 834                 },
 835                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 836         final TestActorRef<ShardManager> shardManager =
 837                 TestActorRef.create(getSystem(), persistentProps);
 838
 839         ShardManager shardManagerActor = shardManager.underlyingActor();
 840
 841         // Initially will be false
 842         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 843
 844         // Make default shard leader
 845         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 846                 RaftState.Follower.name(), RaftState.Leader.name()));
 847
 848         // default = Leader, astronauts is unknown so sync status remains false
 849         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 850
 851         // Make astronauts shard leader as well
 852         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-astronauts-unknown",
 853                 RaftState.Follower.name(), RaftState.Leader.name()));
 854
 855         // Now sync status should be true
 856         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 857
 858         // Make astronauts a Follower
 859         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-astronauts-unknown",
 860                 RaftState.Leader.name(), RaftState.Follower.name()));
 861
 862         // Sync status is not true
 863         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 864
 865         // Make the astronauts follower sync status true
 866         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-astronauts-unknown"));
 867
 868         // Sync status is now true
 869         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 870
 871     }
 872
 873     private static class TestShardManager extends ShardManager {
 874         private final CountDownLatch recoveryComplete = new CountDownLatch(1);
 875
 876         TestShardManager(String shardMrgIDSuffix) {
 877             super(new MockClusterWrapper(), new MockConfiguration(),
 878                     DatastoreContext.newBuilder().dataStoreType(shardMrgIDSuffix).build(), ready,
 879                     new PrimaryShardInfoFutureCache());
 880         }
 881
 882         @Override
 883         public void handleRecover(Object message) throws Exception {
 884             try {
 885                 super.handleRecover(message);
 886             } finally {
 887                 if(message instanceof RecoveryCompleted) {
 888                     recoveryComplete.countDown();
 889                 }
 890             }
 891         }
 892
 893         void waitForRecoveryComplete() {
 894             assertEquals("Recovery complete", true,
 895                     Uninterruptibles.awaitUninterruptibly(recoveryComplete, 5, TimeUnit.SECONDS));
 896         }
 897     }
 898
 899     @SuppressWarnings("serial")
 900     static class TestShardManagerCreator implements Creator<TestShardManager> {
 901         String shardMrgIDSuffix;
 902
 903         TestShardManagerCreator(String shardMrgIDSuffix) {
 904             this.shardMrgIDSuffix = shardMrgIDSuffix;
 905         }
 906
 907         @Override
 908         public TestShardManager create() throws Exception {
 909             return new TestShardManager(shardMrgIDSuffix);
 910         }
 911
 912     }
 913
 914     private static class DelegatingShardManagerCreator implements Creator<ShardManager> {
 915         private static final long serialVersionUID = 1L;
 916         private final Creator<ShardManager> delegate;
 917
 918         public DelegatingShardManagerCreator(Creator<ShardManager> delegate) {
 919             this.delegate = delegate;
 920         }
 921
 922         @Override
 923         public ShardManager create() throws Exception {
 924             return delegate.create();
 925         }
 926     }
 927
 928     private static class ForwardingShardManager extends ShardManager {
 929         private CountDownLatch findPrimaryMessageReceived = new CountDownLatch(1);
 930         private CountDownLatch memberUpReceived = new CountDownLatch(1);
 931         private CountDownLatch memberRemovedReceived = new CountDownLatch(1);
 932         private CountDownLatch memberUnreachableReceived = new CountDownLatch(1);
 933         private CountDownLatch memberReachableReceived = new CountDownLatch(1);
 934         private final ActorRef shardActor;
 935         private final String name;
 936
 937         protected ForwardingShardManager(ClusterWrapper cluster, Configuration configuration,
 938                 DatastoreContext datastoreContext, CountDownLatch waitTillReadyCountdownLatch, String name,
 939                 ActorRef shardActor, PrimaryShardInfoFutureCache primaryShardInfoCache) {
 940             super(cluster, configuration, datastoreContext, waitTillReadyCountdownLatch, primaryShardInfoCache);
 941             this.shardActor = shardActor;
 942             this.name = name;
 943         }
 944
 945         @Override
 946         public void handleCommand(Object message) throws Exception {
 947             try{
 948                 super.handleCommand(message);
 949             } finally {
 950                 if(message instanceof FindPrimary) {
 951                     findPrimaryMessageReceived.countDown();
 952                 } else if(message instanceof ClusterEvent.MemberUp) {
 953                     String role = ((ClusterEvent.MemberUp)message).member().roles().head();
 954                     if(!getCluster().getCurrentMemberName().equals(role)) {
 955                         memberUpReceived.countDown();
 956                     }
 957                 } else if(message instanceof ClusterEvent.MemberRemoved) {
 958                     String role = ((ClusterEvent.MemberRemoved)message).member().roles().head();
 959                     if(!getCluster().getCurrentMemberName().equals(role)) {
 960                         memberRemovedReceived.countDown();
 961                     }
 962                 } else if(message instanceof ClusterEvent.UnreachableMember) {
 963                     String role = ((ClusterEvent.UnreachableMember)message).member().roles().head();
 964                     if(!getCluster().getCurrentMemberName().equals(role)) {
 965                         memberUnreachableReceived.countDown();
 966                     }
 967                 } else if(message instanceof ClusterEvent.ReachableMember) {
 968                     String role = ((ClusterEvent.ReachableMember)message).member().roles().head();
 969                     if(!getCluster().getCurrentMemberName().equals(role)) {
 970                         memberReachableReceived.countDown();
 971                     }
 972                 }
 973             }
 974         }
 975
 976         @Override
 977         public String persistenceId() {
 978             return name;
 979         }
 980
 981         @Override
 982         protected ActorRef newShardActor(SchemaContext schemaContext, ShardInformation info) {
 983             return shardActor;
 984         }
 985
 986         void waitForMemberUp() {
 987             assertEquals("MemberUp received", true,
 988                     Uninterruptibles.awaitUninterruptibly(memberUpReceived, 5, TimeUnit.SECONDS));
 989             memberUpReceived = new CountDownLatch(1);
 990         }
 991
 992         void waitForMemberRemoved() {
 993             assertEquals("MemberRemoved received", true,
 994                     Uninterruptibles.awaitUninterruptibly(memberRemovedReceived, 5, TimeUnit.SECONDS));
 995             memberRemovedReceived = new CountDownLatch(1);
 996         }
 997
 998         void waitForUnreachableMember() {
 999             assertEquals("UnreachableMember received", true,
1000                 Uninterruptibles.awaitUninterruptibly(memberUnreachableReceived, 5, TimeUnit.SECONDS
1001                 ));
1002             memberUnreachableReceived = new CountDownLatch(1);
1003         }
1004
1005         void waitForReachableMember() {
1006             assertEquals("ReachableMember received", true,
1007                 Uninterruptibles.awaitUninterruptibly(memberReachableReceived, 5, TimeUnit.SECONDS));
1008             memberReachableReceived = new CountDownLatch(1);
1009         }
1010
1011         void verifyFindPrimary() {
1012             assertEquals("FindPrimary received", true,
1013                     Uninterruptibles.awaitUninterruptibly(findPrimaryMessageReceived, 5, TimeUnit.SECONDS));
1014             findPrimaryMessageReceived = new CountDownLatch(1);
1015         }
1016     }
1017 }