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