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