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 testOnReceiveFindLocalShardForNonExistentShard() throws Exception {
 429         new JavaTestKit(getSystem()) {{
 430             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 431
 432             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 433
 434             shardManager.tell(new FindLocalShard("non-existent", false), getRef());
 435
 436             LocalShardNotFound notFound = expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
 437
 438             assertEquals("getShardName", "non-existent", notFound.getShardName());
 439         }};
 440     }
 441
 442     @Test
 443     public void testOnReceiveFindLocalShardForExistentShard() throws Exception {
 444         new JavaTestKit(getSystem()) {{
 445             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 446
 447             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 448             shardManager.tell(new ActorInitialized(), mockShardActor);
 449
 450             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 451
 452             LocalShardFound found = expectMsgClass(duration("5 seconds"), LocalShardFound.class);
 453
 454             assertTrue("Found path contains " + found.getPath().path().toString(),
 455                     found.getPath().path().toString().contains("member-1-shard-default-config"));
 456         }};
 457     }
 458
 459     @Test
 460     public void testOnReceiveFindLocalShardForNotInitializedShard() throws Exception {
 461         new JavaTestKit(getSystem()) {{
 462             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 463
 464             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 465
 466             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 467         }};
 468     }
 469
 470     @Test
 471     public void testOnReceiveFindLocalShardWaitForShardInitialized() throws Exception {
 472         new JavaTestKit(getSystem()) {{
 473             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 474
 475             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 476
 477             // We're passing waitUntilInitialized = true to FindLocalShard so the response should be
 478             // delayed until we send ActorInitialized.
 479             Future<Object> future = Patterns.ask(shardManager, new FindLocalShard(Shard.DEFAULT_NAME, true),
 480                     new Timeout(5, TimeUnit.SECONDS));
 481
 482             shardManager.tell(new ActorInitialized(), mockShardActor);
 483
 484             Object resp = Await.result(future, duration("5 seconds"));
 485             assertTrue("Expected: LocalShardFound, Actual: " + resp, resp instanceof LocalShardFound);
 486         }};
 487     }
 488
 489     @Test
 490     public void testOnRecoveryJournalIsCleaned() {
 491         InMemoryJournal.addEntry(shardMgrID, 1L, new ShardManager.SchemaContextModules(
 492                 ImmutableSet.of("foo")));
 493         InMemoryJournal.addEntry(shardMgrID, 2L, new ShardManager.SchemaContextModules(
 494                 ImmutableSet.of("bar")));
 495         InMemoryJournal.addDeleteMessagesCompleteLatch(shardMgrID);
 496
 497         new JavaTestKit(getSystem()) {{
 498             TestActorRef<TestShardManager> shardManager = TestActorRef.create(getSystem(),
 499                     Props.create(new TestShardManagerCreator(shardMrgIDSuffix)));
 500
 501             shardManager.underlyingActor().waitForRecoveryComplete();
 502             InMemoryJournal.waitForDeleteMessagesComplete(shardMgrID);
 503
 504             // Journal entries up to the last one should've been deleted
 505             Map<Long, Object> journal = InMemoryJournal.get(shardMgrID);
 506             synchronized (journal) {
 507                 assertEquals("Journal size", 1, journal.size());
 508                 assertEquals("Journal entry seq #", Long.valueOf(2), journal.keySet().iterator().next());
 509             }
 510         }};
 511     }
 512
 513     @Test
 514     public void testOnRecoveryPreviouslyKnownModulesAreDiscovered() throws Exception {
 515         final ImmutableSet<String> persistedModules = ImmutableSet.of("foo", "bar");
 516         InMemoryJournal.addEntry(shardMgrID, 1L, new ShardManager.SchemaContextModules(
 517                 persistedModules));
 518         new JavaTestKit(getSystem()) {{
 519             TestActorRef<TestShardManager> shardManager = TestActorRef.create(getSystem(),
 520                     Props.create(new TestShardManagerCreator(shardMrgIDSuffix)));
 521
 522             shardManager.underlyingActor().waitForRecoveryComplete();
 523
 524             Collection<String> knownModules = shardManager.underlyingActor().getKnownModules();
 525
 526             assertEquals("getKnownModules", persistedModules, Sets.newHashSet(knownModules));
 527         }};
 528     }
 529
 530     @Test
 531     public void testOnUpdateSchemaContextUpdateKnownModulesIfTheyContainASuperSetOfTheKnownModules()
 532             throws Exception {
 533         new JavaTestKit(getSystem()) {{
 534             final TestActorRef<ShardManager> shardManager =
 535                     TestActorRef.create(getSystem(), newShardMgrProps(true));
 536
 537             assertEquals("getKnownModules size", 0, shardManager.underlyingActor().getKnownModules().size());
 538
 539             ModuleIdentifier foo = mock(ModuleIdentifier.class);
 540             when(foo.getNamespace()).thenReturn(new URI("foo"));
 541
 542             Set<ModuleIdentifier> moduleIdentifierSet = new HashSet<>();
 543             moduleIdentifierSet.add(foo);
 544
 545             SchemaContext schemaContext = mock(SchemaContext.class);
 546             when(schemaContext.getAllModuleIdentifiers()).thenReturn(moduleIdentifierSet);
 547
 548             shardManager.underlyingActor().onReceiveCommand(new UpdateSchemaContext(schemaContext));
 549
 550             assertEquals("getKnownModules", Sets.newHashSet("foo"),
 551                     Sets.newHashSet(shardManager.underlyingActor().getKnownModules()));
 552
 553             ModuleIdentifier bar = mock(ModuleIdentifier.class);
 554             when(bar.getNamespace()).thenReturn(new URI("bar"));
 555
 556             moduleIdentifierSet.add(bar);
 557
 558             shardManager.underlyingActor().onReceiveCommand(new UpdateSchemaContext(schemaContext));
 559
 560             assertEquals("getKnownModules", Sets.newHashSet("foo", "bar"),
 561                     Sets.newHashSet(shardManager.underlyingActor().getKnownModules()));
 562         }};
 563     }
 564
 565     @Test
 566     public void testOnUpdateSchemaContextDoNotUpdateKnownModulesIfTheyDoNotContainASuperSetOfKnownModules()
 567             throws Exception {
 568         new JavaTestKit(getSystem()) {{
 569             final TestActorRef<ShardManager> shardManager =
 570                     TestActorRef.create(getSystem(), newShardMgrProps(true));
 571
 572             SchemaContext schemaContext = mock(SchemaContext.class);
 573             Set<ModuleIdentifier> moduleIdentifierSet = new HashSet<>();
 574
 575             ModuleIdentifier foo = mock(ModuleIdentifier.class);
 576             when(foo.getNamespace()).thenReturn(new URI("foo"));
 577
 578             moduleIdentifierSet.add(foo);
 579
 580             when(schemaContext.getAllModuleIdentifiers()).thenReturn(moduleIdentifierSet);
 581
 582             shardManager.underlyingActor().onReceiveCommand(new UpdateSchemaContext(schemaContext));
 583
 584             assertEquals("getKnownModules", Sets.newHashSet("foo"),
 585                     Sets.newHashSet(shardManager.underlyingActor().getKnownModules()));
 586
 587             //Create a completely different SchemaContext with only the bar module in it
 588             //schemaContext = mock(SchemaContext.class);
 589             moduleIdentifierSet.clear();
 590             ModuleIdentifier bar = mock(ModuleIdentifier.class);
 591             when(bar.getNamespace()).thenReturn(new URI("bar"));
 592
 593             moduleIdentifierSet.add(bar);
 594
 595             shardManager.underlyingActor().onReceiveCommand(new UpdateSchemaContext(schemaContext));
 596
 597             assertEquals("getKnownModules", Sets.newHashSet("foo"),
 598                     Sets.newHashSet(shardManager.underlyingActor().getKnownModules()));
 599
 600         }};
 601     }
 602
 603     @Test
 604     public void testRecoveryApplicable(){
 605         new JavaTestKit(getSystem()) {
 606             {
 607                 final Props persistentProps = newShardMgrProps(true);
 608                 final TestActorRef<ShardManager> persistentShardManager =
 609                         TestActorRef.create(getSystem(), persistentProps);
 610
 611                 DataPersistenceProvider dataPersistenceProvider1 = persistentShardManager.underlyingActor().getDataPersistenceProvider();
 612
 613                 assertTrue("Recovery Applicable", dataPersistenceProvider1.isRecoveryApplicable());
 614
 615                 final Props nonPersistentProps = newShardMgrProps(false);
 616                 final TestActorRef<ShardManager> nonPersistentShardManager =
 617                         TestActorRef.create(getSystem(), nonPersistentProps);
 618
 619                 DataPersistenceProvider dataPersistenceProvider2 = nonPersistentShardManager.underlyingActor().getDataPersistenceProvider();
 620
 621                 assertFalse("Recovery Not Applicable", dataPersistenceProvider2.isRecoveryApplicable());
 622
 623
 624             }};
 625
 626     }
 627
 628     @Test
 629     public void testOnUpdateSchemaContextUpdateKnownModulesCallsDataPersistenceProvider()
 630             throws Exception {
 631         final CountDownLatch persistLatch = new CountDownLatch(1);
 632         final Creator<ShardManager> creator = new Creator<ShardManager>() {
 633             private static final long serialVersionUID = 1L;
 634             @Override
 635             public ShardManager create() throws Exception {
 636                 return new ShardManager(new MockClusterWrapper(), new MockConfiguration(), DatastoreContext.newBuilder().build(),
 637                         ready, new PrimaryShardInfoFutureCache()) {
 638                     @Override
 639                     protected DataPersistenceProvider createDataPersistenceProvider(boolean persistent) {
 640                         DataPersistenceProviderMonitor dataPersistenceProviderMonitor
 641                                 = new DataPersistenceProviderMonitor();
 642                         dataPersistenceProviderMonitor.setPersistLatch(persistLatch);
 643                         return dataPersistenceProviderMonitor;
 644                     }
 645                 };
 646             }
 647         };
 648
 649         new JavaTestKit(getSystem()) {{
 650
 651             final TestActorRef<ShardManager> shardManager =
 652                     TestActorRef.create(getSystem(), Props.create(new DelegatingShardManagerCreator(creator)));
 653
 654             ModuleIdentifier foo = mock(ModuleIdentifier.class);
 655             when(foo.getNamespace()).thenReturn(new URI("foo"));
 656
 657             Set<ModuleIdentifier> moduleIdentifierSet = new HashSet<>();
 658             moduleIdentifierSet.add(foo);
 659
 660             SchemaContext schemaContext = mock(SchemaContext.class);
 661             when(schemaContext.getAllModuleIdentifiers()).thenReturn(moduleIdentifierSet);
 662
 663             shardManager.underlyingActor().onReceiveCommand(new UpdateSchemaContext(schemaContext));
 664
 665             assertEquals("Persisted", true,
 666                     Uninterruptibles.awaitUninterruptibly(persistLatch, 5, TimeUnit.SECONDS));
 667
 668         }};
 669     }
 670
 671     @Test
 672     public void testRoleChangeNotificationAndShardLeaderStateChangedReleaseReady() throws Exception {
 673         new JavaTestKit(getSystem()) {
 674             {
 675                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 676
 677                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 678                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 679                         memberId, RaftState.Candidate.name(), RaftState.Leader.name()));
 680
 681                 verify(ready, never()).countDown();
 682
 683                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId, memberId,
 684                         Optional.of(mock(DataTree.class))));
 685
 686                 verify(ready, times(1)).countDown();
 687
 688             }};
 689     }
 690
 691     @Test
 692     public void testRoleChangeNotificationToFollowerWithShardLeaderStateChangedReleaseReady() throws Exception {
 693         new JavaTestKit(getSystem()) {
 694             {
 695                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 696
 697                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 698                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 699                         memberId, null, RaftState.Follower.name()));
 700
 701                 verify(ready, never()).countDown();
 702
 703                 shardManager.underlyingActor().onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
 704
 705                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId,
 706                         "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class))));
 707
 708                 verify(ready, times(1)).countDown();
 709
 710             }};
 711     }
 712
 713     @Test
 714     public void testReadyCountDownForMemberUpAfterLeaderStateChanged() throws Exception {
 715         new JavaTestKit(getSystem()) {
 716             {
 717                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 718
 719                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 720                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 721                         memberId, null, RaftState.Follower.name()));
 722
 723                 verify(ready, never()).countDown();
 724
 725                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId,
 726                         "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class))));
 727
 728                 shardManager.underlyingActor().onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
 729
 730                 verify(ready, times(1)).countDown();
 731
 732             }};
 733     }
 734
 735     @Test
 736     public void testRoleChangeNotificationDoNothingForUnknownShard() throws Exception {
 737         new JavaTestKit(getSystem()) {
 738             {
 739                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 740
 741                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 742                         "unknown", RaftState.Candidate.name(), RaftState.Leader.name()));
 743
 744                 verify(ready, never()).countDown();
 745
 746             }};
 747     }
 748
 749
 750     @Test
 751     public void testByDefaultSyncStatusIsFalse() throws Exception{
 752         final Props persistentProps = newShardMgrProps(true);
 753         final TestActorRef<ShardManager> shardManager =
 754                 TestActorRef.create(getSystem(), persistentProps);
 755
 756         ShardManager shardManagerActor = shardManager.underlyingActor();
 757
 758         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 759     }
 760
 761     @Test
 762     public void testWhenShardIsLeaderSyncStatusIsTrue() throws Exception{
 763         final Props persistentProps = ShardManager.props(
 764                 new MockClusterWrapper(),
 765                 new MockConfiguration(),
 766                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 767         final TestActorRef<ShardManager> shardManager =
 768                 TestActorRef.create(getSystem(), persistentProps);
 769
 770         ShardManager shardManagerActor = shardManager.underlyingActor();
 771         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 772                 RaftState.Follower.name(), RaftState.Leader.name()));
 773
 774         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 775     }
 776
 777     @Test
 778     public void testWhenShardIsCandidateSyncStatusIsFalse() throws Exception{
 779         final Props persistentProps = newShardMgrProps(true);
 780         final TestActorRef<ShardManager> shardManager =
 781                 TestActorRef.create(getSystem(), persistentProps);
 782
 783         ShardManager shardManagerActor = shardManager.underlyingActor();
 784         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 785                 RaftState.Follower.name(), RaftState.Candidate.name()));
 786
 787         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 788
 789         // Send a FollowerInitialSyncStatus with status = true for the replica whose current state is candidate
 790         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-default-unknown"));
 791
 792         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 793     }
 794
 795     @Test
 796     public void testWhenShardIsFollowerSyncStatusDependsOnFollowerInitialSyncStatus() throws Exception{
 797         final Props persistentProps = ShardManager.props(
 798                 new MockClusterWrapper(),
 799                 new MockConfiguration(),
 800                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 801         final TestActorRef<ShardManager> shardManager =
 802                 TestActorRef.create(getSystem(), persistentProps);
 803
 804         ShardManager shardManagerActor = shardManager.underlyingActor();
 805         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 806                 RaftState.Candidate.name(), RaftState.Follower.name()));
 807
 808         // Initially will be false
 809         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 810
 811         // Send status true will make sync status true
 812         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-default-unknown"));
 813
 814         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 815
 816         // Send status false will make sync status false
 817         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(false, "member-1-shard-default-unknown"));
 818
 819         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 820
 821     }
 822
 823     @Test
 824     public void testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards() throws Exception{
 825         final Props persistentProps = ShardManager.props(
 826                 new MockClusterWrapper(),
 827                 new MockConfiguration() {
 828                     @Override
 829                     public List<String> getMemberShardNames(String memberName) {
 830                         return Arrays.asList("default", "astronauts");
 831                     }
 832                 },
 833                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 834         final TestActorRef<ShardManager> shardManager =
 835                 TestActorRef.create(getSystem(), persistentProps);
 836
 837         ShardManager shardManagerActor = shardManager.underlyingActor();
 838
 839         // Initially will be false
 840         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 841
 842         // Make default shard leader
 843         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 844                 RaftState.Follower.name(), RaftState.Leader.name()));
 845
 846         // default = Leader, astronauts is unknown so sync status remains false
 847         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 848
 849         // Make astronauts shard leader as well
 850         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-astronauts-unknown",
 851                 RaftState.Follower.name(), RaftState.Leader.name()));
 852
 853         // Now sync status should be true
 854         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 855
 856         // Make astronauts a Follower
 857         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-astronauts-unknown",
 858                 RaftState.Leader.name(), RaftState.Follower.name()));
 859
 860         // Sync status is not true
 861         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 862
 863         // Make the astronauts follower sync status true
 864         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-astronauts-unknown"));
 865
 866         // Sync status is now true
 867         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 868
 869     }
 870
 871     private static class TestShardManager extends ShardManager {
 872         private final CountDownLatch recoveryComplete = new CountDownLatch(1);
 873
 874         TestShardManager(String shardMrgIDSuffix) {
 875             super(new MockClusterWrapper(), new MockConfiguration(),
 876                     DatastoreContext.newBuilder().dataStoreType(shardMrgIDSuffix).build(), ready,
 877                     new PrimaryShardInfoFutureCache());
 878         }
 879
 880         @Override
 881         public void handleRecover(Object message) throws Exception {
 882             try {
 883                 super.handleRecover(message);
 884             } finally {
 885                 if(message instanceof RecoveryCompleted) {
 886                     recoveryComplete.countDown();
 887                 }
 888             }
 889         }
 890
 891         void waitForRecoveryComplete() {
 892             assertEquals("Recovery complete", true,
 893                     Uninterruptibles.awaitUninterruptibly(recoveryComplete, 5, TimeUnit.SECONDS));
 894         }
 895     }
 896
 897     @SuppressWarnings("serial")
 898     static class TestShardManagerCreator implements Creator<TestShardManager> {
 899         String shardMrgIDSuffix;
 900
 901         TestShardManagerCreator(String shardMrgIDSuffix) {
 902             this.shardMrgIDSuffix = shardMrgIDSuffix;
 903         }
 904
 905         @Override
 906         public TestShardManager create() throws Exception {
 907             return new TestShardManager(shardMrgIDSuffix);
 908         }
 909
 910     }
 911
 912     private static class DelegatingShardManagerCreator implements Creator<ShardManager> {
 913         private static final long serialVersionUID = 1L;
 914         private final Creator<ShardManager> delegate;
 915
 916         public DelegatingShardManagerCreator(Creator<ShardManager> delegate) {
 917             this.delegate = delegate;
 918         }
 919
 920         @Override
 921         public ShardManager create() throws Exception {
 922             return delegate.create();
 923         }
 924     }
 925
 926     private static class ForwardingShardManager extends ShardManager {
 927         private CountDownLatch findPrimaryMessageReceived = new CountDownLatch(1);
 928         private CountDownLatch memberUpReceived = new CountDownLatch(1);
 929         private CountDownLatch memberRemovedReceived = new CountDownLatch(1);
 930         private final ActorRef shardActor;
 931         private final String name;
 932
 933         protected ForwardingShardManager(ClusterWrapper cluster, Configuration configuration,
 934                 DatastoreContext datastoreContext, CountDownLatch waitTillReadyCountdownLatch, String name,
 935                 ActorRef shardActor, PrimaryShardInfoFutureCache primaryShardInfoCache) {
 936             super(cluster, configuration, datastoreContext, waitTillReadyCountdownLatch, primaryShardInfoCache);
 937             this.shardActor = shardActor;
 938             this.name = name;
 939         }
 940
 941         @Override
 942         public void handleCommand(Object message) throws Exception {
 943             try{
 944                 super.handleCommand(message);
 945             } finally {
 946                 if(message instanceof FindPrimary) {
 947                     findPrimaryMessageReceived.countDown();
 948                 } else if(message instanceof ClusterEvent.MemberUp) {
 949                     String role = ((ClusterEvent.MemberUp)message).member().roles().head();
 950                     if(!getCluster().getCurrentMemberName().equals(role)) {
 951                         memberUpReceived.countDown();
 952                     }
 953                 } else if(message instanceof ClusterEvent.MemberRemoved) {
 954                     String role = ((ClusterEvent.MemberRemoved)message).member().roles().head();
 955                     if(!getCluster().getCurrentMemberName().equals(role)) {
 956                         memberRemovedReceived.countDown();
 957                     }
 958                 }
 959             }
 960         }
 961
 962         @Override
 963         public String persistenceId() {
 964             return name;
 965         }
 966
 967         @Override
 968         protected ActorRef newShardActor(SchemaContext schemaContext, ShardInformation info) {
 969             return shardActor;
 970         }
 971
 972         void waitForMemberUp() {
 973             assertEquals("MemberUp received", true,
 974                     Uninterruptibles.awaitUninterruptibly(memberUpReceived, 5, TimeUnit.SECONDS));
 975             memberUpReceived = new CountDownLatch(1);
 976         }
 977
 978         void waitForMemberRemoved() {
 979             assertEquals("MemberRemoved received", true,
 980                     Uninterruptibles.awaitUninterruptibly(memberRemovedReceived, 5, TimeUnit.SECONDS));
 981             memberRemovedReceived = new CountDownLatch(1);
 982         }
 983
 984         void verifyFindPrimary() {
 985             assertEquals("FindPrimary received", true,
 986                     Uninterruptibles.awaitUninterruptibly(findPrimaryMessageReceived, 5, TimeUnit.SECONDS));
 987             findPrimaryMessageReceived = new CountDownLatch(1);
 988         }
 989     }
 990 }