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

   1 /*
   2  * Copyright (c) 2014, 2015 Cisco Systems, Inc. and others.  All rights reserved.
   3  *
   4  * This program and the accompanying materials are made available under the
   5  * terms of the Eclipse Public License v1.0 which accompanies this distribution,
   6  * and is available at http://www.eclipse.org/legal/epl-v10.html
   7  */
   8
   9 package org.opendaylight.controller.cluster.datastore;
  10
  11 import static org.junit.Assert.assertEquals;
  12 import static org.junit.Assert.assertNotNull;
  13 import static org.junit.Assert.assertNull;
  14 import static org.junit.Assert.assertSame;
  15 import static org.junit.Assert.assertTrue;
  16 import static org.mockito.Mockito.mock;
  17 import static org.mockito.Mockito.never;
  18 import static org.mockito.Mockito.times;
  19 import static org.mockito.Mockito.verify;
  20 import akka.actor.ActorRef;
  21 import akka.actor.ActorSystem;
  22 import akka.actor.AddressFromURIString;
  23 import akka.actor.Props;
  24 import akka.actor.Status;
  25 import akka.actor.Status.Failure;
  26 import akka.actor.Status.Success;
  27 import akka.actor.Terminated;
  28 import akka.cluster.Cluster;
  29 import akka.cluster.ClusterEvent;
  30 import akka.dispatch.Dispatchers;
  31 import akka.japi.Creator;
  32 import akka.pattern.Patterns;
  33 import akka.persistence.RecoveryCompleted;
  34 import akka.serialization.Serialization;
  35 import akka.testkit.JavaTestKit;
  36 import akka.testkit.TestActorRef;
  37 import akka.util.Timeout;
  38 import com.google.common.base.Function;
  39 import com.google.common.base.Optional;
  40 import com.google.common.base.Stopwatch;
  41 import com.google.common.collect.ImmutableMap;
  42 import com.google.common.collect.ImmutableSet;
  43 import com.google.common.collect.Lists;
  44 import com.google.common.collect.Sets;
  45 import com.google.common.util.concurrent.Uninterruptibles;
  46 import com.typesafe.config.ConfigFactory;
  47 import java.net.URI;
  48 import java.util.AbstractMap;
  49 import java.util.ArrayList;
  50 import java.util.Arrays;
  51 import java.util.Collection;
  52 import java.util.Collections;
  53 import java.util.HashMap;
  54 import java.util.List;
  55 import java.util.Map;
  56 import java.util.Map.Entry;
  57 import java.util.Set;
  58 import java.util.concurrent.CountDownLatch;
  59 import java.util.concurrent.TimeUnit;
  60 import java.util.concurrent.TimeoutException;
  61 import org.apache.commons.lang3.SerializationUtils;
  62 import org.junit.After;
  63 import org.junit.Before;
  64 import org.junit.Test;
  65 import org.mockito.Mock;
  66 import org.mockito.Mockito;
  67 import org.mockito.MockitoAnnotations;
  68 import org.opendaylight.controller.cluster.datastore.config.Configuration;
  69 import org.opendaylight.controller.cluster.datastore.config.ConfigurationImpl;
  70 import org.opendaylight.controller.cluster.datastore.config.EmptyModuleShardConfigProvider;
  71 import org.opendaylight.controller.cluster.datastore.config.ModuleShardConfiguration;
  72 import org.opendaylight.controller.cluster.datastore.exceptions.AlreadyExistsException;
  73 import org.opendaylight.controller.cluster.datastore.exceptions.NoShardLeaderException;
  74 import org.opendaylight.controller.cluster.datastore.exceptions.NotInitializedException;
  75 import org.opendaylight.controller.cluster.datastore.exceptions.PrimaryNotFoundException;
  76 import org.opendaylight.controller.cluster.datastore.identifiers.ShardIdentifier;
  77 import org.opendaylight.controller.cluster.datastore.identifiers.ShardManagerIdentifier;
  78 import org.opendaylight.controller.cluster.datastore.messages.ActorInitialized;
  79 import org.opendaylight.controller.cluster.datastore.messages.AddShardReplica;
  80 import org.opendaylight.controller.cluster.datastore.messages.CreateShard;
  81 import org.opendaylight.controller.cluster.datastore.messages.DatastoreSnapshot;
  82 import org.opendaylight.controller.cluster.datastore.messages.DatastoreSnapshot.ShardSnapshot;
  83 import org.opendaylight.controller.cluster.datastore.messages.FindLocalShard;
  84 import org.opendaylight.controller.cluster.datastore.messages.FindPrimary;
  85 import org.opendaylight.controller.cluster.datastore.messages.LocalPrimaryShardFound;
  86 import org.opendaylight.controller.cluster.datastore.messages.LocalShardFound;
  87 import org.opendaylight.controller.cluster.datastore.messages.LocalShardNotFound;
  88 import org.opendaylight.controller.cluster.datastore.messages.PeerDown;
  89 import org.opendaylight.controller.cluster.datastore.messages.PeerUp;
  90 import org.opendaylight.controller.cluster.datastore.messages.PrimaryShardInfo;
  91 import org.opendaylight.controller.cluster.datastore.messages.RemotePrimaryShardFound;
  92 import org.opendaylight.controller.cluster.datastore.messages.RemoveShardReplica;
  93 import org.opendaylight.controller.cluster.datastore.messages.ShardLeaderStateChanged;
  94 import org.opendaylight.controller.cluster.datastore.messages.SwitchShardBehavior;
  95 import org.opendaylight.controller.cluster.datastore.messages.UpdateSchemaContext;
  96 import org.opendaylight.controller.cluster.datastore.utils.ForwardingActor;
  97 import org.opendaylight.controller.cluster.datastore.utils.MockClusterWrapper;
  98 import org.opendaylight.controller.cluster.datastore.utils.MockConfiguration;
  99 import org.opendaylight.controller.cluster.datastore.utils.PrimaryShardInfoFutureCache;
 100 import org.opendaylight.controller.cluster.notifications.LeaderStateChanged;
 101 import org.opendaylight.controller.cluster.notifications.RegisterRoleChangeListener;
 102 import org.opendaylight.controller.cluster.notifications.RoleChangeNotification;
 103 import org.opendaylight.controller.cluster.raft.RaftState;
 104 import org.opendaylight.controller.cluster.raft.TestActorFactory;
 105 import org.opendaylight.controller.cluster.raft.base.messages.FollowerInitialSyncUpStatus;
 106 import org.opendaylight.controller.cluster.raft.base.messages.SwitchBehavior;
 107 import org.opendaylight.controller.cluster.raft.client.messages.GetSnapshot;
 108 import org.opendaylight.controller.cluster.raft.messages.AddServer;
 109 import org.opendaylight.controller.cluster.raft.messages.AddServerReply;
 110 import org.opendaylight.controller.cluster.raft.messages.RemoveServer;
 111 import org.opendaylight.controller.cluster.raft.messages.RemoveServerReply;
 112 import org.opendaylight.controller.cluster.raft.messages.ServerChangeStatus;
 113 import org.opendaylight.controller.cluster.raft.messages.ServerRemoved;
 114 import org.opendaylight.controller.cluster.raft.policy.DisableElectionsRaftPolicy;
 115 import org.opendaylight.controller.cluster.raft.utils.InMemoryJournal;
 116 import org.opendaylight.controller.cluster.raft.utils.InMemorySnapshotStore;
 117 import org.opendaylight.controller.cluster.raft.utils.MessageCollectorActor;
 118 import org.opendaylight.controller.md.cluster.datastore.model.TestModel;
 119 import org.opendaylight.yangtools.yang.data.api.schema.tree.DataTree;
 120 import org.opendaylight.yangtools.yang.model.api.SchemaContext;
 121 import org.slf4j.Logger;
 122 import org.slf4j.LoggerFactory;
 123 import scala.concurrent.Await;
 124 import scala.concurrent.Future;
 125 import scala.concurrent.duration.FiniteDuration;
 126
 127 public class ShardManagerTest extends AbstractActorTest {
 128     private static final Logger LOG = LoggerFactory.getLogger(ShardManagerTest.class);
 129
 130     private static int ID_COUNTER = 1;
 131
 132     private final String shardMrgIDSuffix = "config" + ID_COUNTER++;
 133     private final String shardMgrID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 134
 135     @Mock
 136     private static CountDownLatch ready;
 137
 138     private static TestActorRef<MessageCollectorActor> mockShardActor;
 139
 140     private static String mockShardName;
 141
 142     private final DatastoreContext.Builder datastoreContextBuilder = DatastoreContext.newBuilder().
 143             dataStoreName(shardMrgIDSuffix).shardInitializationTimeout(600, TimeUnit.MILLISECONDS)
 144                    .shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(6);
 145
 146     private final Collection<ActorSystem> actorSystems = new ArrayList<>();
 147
 148     private final TestActorFactory actorFactory = new TestActorFactory(getSystem());
 149
 150     @Before
 151     public void setUp() {
 152         MockitoAnnotations.initMocks(this);
 153
 154         InMemoryJournal.clear();
 155         InMemorySnapshotStore.clear();
 156
 157         if(mockShardActor == null) {
 158             mockShardName = new ShardIdentifier(Shard.DEFAULT_NAME, "member-1", "config").toString();
 159             mockShardActor = TestActorRef.create(getSystem(), Props.create(MessageCollectorActor.class), mockShardName);
 160         }
 161
 162         mockShardActor.underlyingActor().clear();
 163     }
 164
 165     @After
 166     public void tearDown() {
 167         InMemoryJournal.clear();
 168         InMemorySnapshotStore.clear();
 169
 170         for(ActorSystem system: actorSystems) {
 171             JavaTestKit.shutdownActorSystem(system, null, Boolean.TRUE);
 172         }
 173
 174         actorFactory.close();
 175     }
 176
 177     private ActorSystem newActorSystem(String config) {
 178         ActorSystem system = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig(config));
 179         actorSystems.add(system);
 180         return system;
 181     }
 182
 183     private ActorRef newMockShardActor(ActorSystem system, String shardName, String memberName) {
 184         String name = new ShardIdentifier(shardName, memberName,"config").toString();
 185         if(system == getSystem()) {
 186             return actorFactory.createTestActor(Props.create(MessageCollectorActor.class), name);
 187         }
 188
 189         return TestActorRef.create(system, Props.create(MessageCollectorActor.class), name);
 190     }
 191
 192     private Props newShardMgrProps() {
 193         return newShardMgrProps(new MockConfiguration());
 194     }
 195
 196     private static DatastoreContextFactory newDatastoreContextFactory(DatastoreContext datastoreContext) {
 197         DatastoreContextFactory mockFactory = mock(DatastoreContextFactory.class);
 198         Mockito.doReturn(datastoreContext).when(mockFactory).getBaseDatastoreContext();
 199         Mockito.doReturn(datastoreContext).when(mockFactory).getShardDatastoreContext(Mockito.anyString());
 200         return mockFactory;
 201     }
 202
 203     private TestShardManager.Builder newTestShardMgrBuilder() {
 204         return TestShardManager.builder(datastoreContextBuilder);
 205     }
 206
 207     private TestShardManager.Builder newTestShardMgrBuilder(Configuration config) {
 208         return TestShardManager.builder(datastoreContextBuilder).configuration(config);
 209     }
 210
 211     private Props newShardMgrProps(Configuration config) {
 212         return newTestShardMgrBuilder(config).props();
 213     }
 214
 215     private TestShardManager.Builder newTestShardMgrBuilderWithMockShardActor() {
 216         return newTestShardMgrBuilderWithMockShardActor(mockShardActor);
 217     }
 218
 219     private TestShardManager.Builder newTestShardMgrBuilderWithMockShardActor(ActorRef shardActor) {
 220         return TestShardManager.builder(datastoreContextBuilder).shardActor(shardActor);
 221     }
 222
 223
 224     private Props newPropsShardMgrWithMockShardActor() {
 225         return newTestShardMgrBuilderWithMockShardActor().props();
 226     }
 227
 228     private Props newPropsShardMgrWithMockShardActor(ActorRef shardActor) {
 229         return newTestShardMgrBuilderWithMockShardActor(shardActor).props();
 230     }
 231
 232
 233     private TestShardManager newTestShardManager() {
 234         return newTestShardManager(newShardMgrProps());
 235     }
 236
 237     private TestShardManager newTestShardManager(Props props) {
 238         TestActorRef<TestShardManager> shardManagerActor = actorFactory.createTestActor(props);
 239         TestShardManager shardManager = shardManagerActor.underlyingActor();
 240         shardManager.waitForRecoveryComplete();
 241         return shardManager;
 242     }
 243
 244     private void waitForShardInitialized(ActorRef shardManager, String shardName, JavaTestKit kit) {
 245         AssertionError last = null;
 246         Stopwatch sw = Stopwatch.createStarted();
 247         while(sw.elapsed(TimeUnit.SECONDS) <= 5) {
 248             try {
 249                 shardManager.tell(new FindLocalShard(shardName, true), kit.getRef());
 250                 kit.expectMsgClass(LocalShardFound.class);
 251                 return;
 252             } catch(AssertionError e) {
 253                 last = e;
 254             }
 255
 256             Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS);
 257         }
 258
 259         throw last;
 260     }
 261
 262     private <T> T expectMsgClassOrFailure(Class<T> msgClass, JavaTestKit kit, String msg) {
 263         Object reply = kit.expectMsgAnyClassOf(JavaTestKit.duration("5 sec"), msgClass, Failure.class);
 264         if(reply instanceof Failure) {
 265             throw new AssertionError(msg + " failed", ((Failure)reply).cause());
 266         }
 267
 268         return (T)reply;
 269     }
 270
 271     @Test
 272     public void testPerShardDatastoreContext() throws Exception {
 273         LOG.info("testPerShardDatastoreContext starting");
 274         final DatastoreContextFactory mockFactory = newDatastoreContextFactory(
 275                 datastoreContextBuilder.shardElectionTimeoutFactor(5).build());
 276
 277         Mockito.doReturn(DatastoreContext.newBuilderFrom(datastoreContextBuilder.build()).
 278                 shardElectionTimeoutFactor(6).build()).when(mockFactory).getShardDatastoreContext("default");
 279
 280         Mockito.doReturn(DatastoreContext.newBuilderFrom(datastoreContextBuilder.build()).
 281                 shardElectionTimeoutFactor(7).build()).when(mockFactory).getShardDatastoreContext("topology");
 282
 283         final MockConfiguration mockConfig = new MockConfiguration() {
 284             @Override
 285             public Collection<String> getMemberShardNames(String memberName) {
 286                 return Arrays.asList("default", "topology");
 287             }
 288
 289             @Override
 290             public Collection<String> getMembersFromShardName(String shardName) {
 291                 return Arrays.asList("member-1");
 292             }
 293         };
 294
 295         final TestActorRef<MessageCollectorActor> defaultShardActor = actorFactory.createTestActor(
 296                 Props.create(MessageCollectorActor.class), actorFactory.generateActorId("default"));
 297         final TestActorRef<MessageCollectorActor> topologyShardActor = actorFactory.createTestActor(
 298                 Props.create(MessageCollectorActor.class), actorFactory.generateActorId("topology"));
 299
 300         final Map<String, Entry<ActorRef, DatastoreContext>> shardInfoMap = Collections.synchronizedMap(
 301                 new HashMap<String, Entry<ActorRef, DatastoreContext>>());
 302         shardInfoMap.put("default", new AbstractMap.SimpleEntry<ActorRef, DatastoreContext>(defaultShardActor, null));
 303         shardInfoMap.put("topology", new AbstractMap.SimpleEntry<ActorRef, DatastoreContext>(topologyShardActor, null));
 304
 305         final PrimaryShardInfoFutureCache primaryShardInfoCache = new PrimaryShardInfoFutureCache();
 306         final CountDownLatch newShardActorLatch = new CountDownLatch(2);
 307         class LocalShardManager extends ShardManager {
 308             public LocalShardManager(AbstractBuilder<?> builder) {
 309                 super(builder);
 310             }
 311
 312             @Override
 313             protected ActorRef newShardActor(SchemaContext schemaContext, ShardInformation info) {
 314                 Entry<ActorRef, DatastoreContext> entry = shardInfoMap.get(info.getShardName());
 315                 ActorRef ref = null;
 316                 if(entry != null) {
 317                     ref = entry.getKey();
 318                     entry.setValue(info.getDatastoreContext());
 319                 }
 320
 321                 newShardActorLatch.countDown();
 322                 return ref;
 323             }
 324         }
 325
 326         final Creator<ShardManager> creator = new Creator<ShardManager>() {
 327             private static final long serialVersionUID = 1L;
 328             @Override
 329             public ShardManager create() throws Exception {
 330                 return new LocalShardManager(new GenericBuilder<LocalShardManager>(LocalShardManager.class).
 331                         datastoreContextFactory(mockFactory).primaryShardInfoCache(primaryShardInfoCache).
 332                         configuration(mockConfig));
 333             }
 334         };
 335
 336         JavaTestKit kit = new JavaTestKit(getSystem());
 337
 338         final ActorRef shardManager = actorFactory.createActor(Props.create(
 339                 new DelegatingShardManagerCreator(creator)).withDispatcher(Dispatchers.DefaultDispatcherId()));
 340
 341         shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), kit.getRef());
 342
 343         assertEquals("Shard actors created", true, newShardActorLatch.await(5, TimeUnit.SECONDS));
 344         assertEquals("getShardElectionTimeoutFactor", 6, shardInfoMap.get("default").getValue().
 345                 getShardElectionTimeoutFactor());
 346         assertEquals("getShardElectionTimeoutFactor", 7, shardInfoMap.get("topology").getValue().
 347                 getShardElectionTimeoutFactor());
 348
 349         DatastoreContextFactory newMockFactory = newDatastoreContextFactory(
 350                 datastoreContextBuilder.shardElectionTimeoutFactor(5).build());
 351         Mockito.doReturn(DatastoreContext.newBuilderFrom(datastoreContextBuilder.build()).
 352                 shardElectionTimeoutFactor(66).build()).when(newMockFactory).getShardDatastoreContext("default");
 353
 354         Mockito.doReturn(DatastoreContext.newBuilderFrom(datastoreContextBuilder.build()).
 355                 shardElectionTimeoutFactor(77).build()).when(newMockFactory).getShardDatastoreContext("topology");
 356
 357         shardManager.tell(newMockFactory, kit.getRef());
 358
 359         DatastoreContext newContext = MessageCollectorActor.expectFirstMatching(defaultShardActor, DatastoreContext.class);
 360         assertEquals("getShardElectionTimeoutFactor", 66, newContext.getShardElectionTimeoutFactor());
 361
 362         newContext = MessageCollectorActor.expectFirstMatching(topologyShardActor, DatastoreContext.class);
 363         assertEquals("getShardElectionTimeoutFactor", 77, newContext.getShardElectionTimeoutFactor());
 364
 365         LOG.info("testPerShardDatastoreContext ending");
 366     }
 367
 368     @Test
 369     public void testOnReceiveFindPrimaryForNonExistentShard() throws Exception {
 370         new JavaTestKit(getSystem()) {{
 371             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 372
 373             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 374
 375             shardManager.tell(new FindPrimary("non-existent", false), getRef());
 376
 377             expectMsgClass(duration("5 seconds"), PrimaryNotFoundException.class);
 378         }};
 379     }
 380
 381     @Test
 382     public void testOnReceiveFindPrimaryForLocalLeaderShard() throws Exception {
 383         LOG.info("testOnReceiveFindPrimaryForLocalLeaderShard starting");
 384         new JavaTestKit(getSystem()) {{
 385             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 386
 387             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 388
 389             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 390             shardManager.tell(new ActorInitialized(), mockShardActor);
 391
 392             DataTree mockDataTree = mock(DataTree.class);
 393             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 394                     DataStoreVersions.CURRENT_VERSION), getRef());
 395
 396             MessageCollectorActor.expectFirstMatching(mockShardActor, RegisterRoleChangeListener.class);
 397             shardManager.tell((new RoleChangeNotification(memberId, RaftState.Candidate.name(),
 398                     RaftState.Leader.name())), mockShardActor);
 399
 400             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 401
 402             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 403             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 404                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 405             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree());
 406         }};
 407
 408         LOG.info("testOnReceiveFindPrimaryForLocalLeaderShard ending");
 409     }
 410
 411     @Test
 412     public void testOnReceiveFindPrimaryForNonLocalLeaderShardBeforeMemberUp() throws Exception {
 413         LOG.info("testOnReceiveFindPrimaryForNonLocalLeaderShardBeforeMemberUp starting");
 414         new JavaTestKit(getSystem()) {{
 415             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 416
 417             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 418             shardManager.tell(new ActorInitialized(), mockShardActor);
 419
 420             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 421             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 422             shardManager.tell(new RoleChangeNotification(memberId1,
 423                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 424             shardManager.tell(new LeaderStateChanged(memberId1, memberId2, DataStoreVersions.CURRENT_VERSION), mockShardActor);
 425
 426             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 427
 428             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 429         }};
 430
 431         LOG.info("testOnReceiveFindPrimaryForNonLocalLeaderShardBeforeMemberUp ending");
 432     }
 433
 434     @Test
 435     public void testOnReceiveFindPrimaryForNonLocalLeaderShard() throws Exception {
 436         LOG.info("testOnReceiveFindPrimaryForNonLocalLeaderShard starting");
 437         new JavaTestKit(getSystem()) {{
 438             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 439
 440             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 441             shardManager.tell(new ActorInitialized(), mockShardActor);
 442
 443             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 444             MockClusterWrapper.sendMemberUp(shardManager, "member-2", getRef().path().toString());
 445
 446             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 447             shardManager.tell(new RoleChangeNotification(memberId1,
 448                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 449             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
 450             shardManager.tell(new ShardLeaderStateChanged(memberId1, memberId2, Optional.<DataTree>absent(),
 451                     leaderVersion), mockShardActor);
 452
 453             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 454
 455             RemotePrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 456             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 457                     primaryFound.getPrimaryPath().contains("member-2-shard-default"));
 458             assertEquals("getPrimaryVersion", leaderVersion, primaryFound.getPrimaryVersion());
 459         }};
 460
 461         LOG.info("testOnReceiveFindPrimaryForNonLocalLeaderShard ending");
 462     }
 463
 464     @Test
 465     public void testOnReceiveFindPrimaryForUninitializedShard() throws Exception {
 466         new JavaTestKit(getSystem()) {{
 467             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 468
 469             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 470
 471             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 472         }};
 473     }
 474
 475     @Test
 476     public void testOnReceiveFindPrimaryForInitializedShardWithNoRole() throws Exception {
 477         new JavaTestKit(getSystem()) {{
 478             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 479
 480             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 481             shardManager.tell(new ActorInitialized(), mockShardActor);
 482
 483             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 484
 485             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 486         }};
 487     }
 488
 489     @Test
 490     public void testOnReceiveFindPrimaryForFollowerShardWithNoInitialLeaderId() throws Exception {
 491         LOG.info("testOnReceiveFindPrimaryForFollowerShardWithNoInitialLeaderId starting");
 492         new JavaTestKit(getSystem()) {{
 493             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 494
 495             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 496             shardManager.tell(new ActorInitialized(), mockShardActor);
 497
 498             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 499             shardManager.tell(new RoleChangeNotification(memberId,
 500                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 501
 502             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 503
 504             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 505
 506             DataTree mockDataTree = mock(DataTree.class);
 507             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 508                     DataStoreVersions.CURRENT_VERSION), mockShardActor);
 509
 510             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 511
 512             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 513             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 514                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 515             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree());
 516         }};
 517
 518         LOG.info("testOnReceiveFindPrimaryForFollowerShardWithNoInitialLeaderId starting");
 519     }
 520
 521     @Test
 522     public void testOnReceiveFindPrimaryWaitForShardLeader() throws Exception {
 523         LOG.info("testOnReceiveFindPrimaryWaitForShardLeader starting");
 524         new JavaTestKit(getSystem()) {{
 525             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 526
 527             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 528
 529             // We're passing waitUntilInitialized = true to FindPrimary so the response should be
 530             // delayed until we send ActorInitialized and RoleChangeNotification.
 531             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 532
 533             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 534
 535             shardManager.tell(new ActorInitialized(), mockShardActor);
 536
 537             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 538
 539             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 540             shardManager.tell(new RoleChangeNotification(memberId,
 541                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor);
 542
 543             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 544
 545             DataTree mockDataTree = mock(DataTree.class);
 546             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 547                     DataStoreVersions.CURRENT_VERSION), mockShardActor);
 548
 549             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 550             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 551                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 552             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 553
 554             expectNoMsg(FiniteDuration.create(200, TimeUnit.MILLISECONDS));
 555         }};
 556
 557         LOG.info("testOnReceiveFindPrimaryWaitForShardLeader ending");
 558     }
 559
 560     @Test
 561     public void testOnReceiveFindPrimaryWaitForReadyWithUninitializedShard() throws Exception {
 562         LOG.info("testOnReceiveFindPrimaryWaitForReadyWithUninitializedShard starting");
 563         new JavaTestKit(getSystem()) {{
 564             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 565
 566             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 567
 568             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 569
 570             expectMsgClass(duration("2 seconds"), NotInitializedException.class);
 571
 572             shardManager.tell(new ActorInitialized(), mockShardActor);
 573
 574             expectNoMsg(FiniteDuration.create(200, TimeUnit.MILLISECONDS));
 575         }};
 576
 577         LOG.info("testOnReceiveFindPrimaryWaitForReadyWithUninitializedShard ending");
 578     }
 579
 580     @Test
 581     public void testOnReceiveFindPrimaryWaitForReadyWithCandidateShard() throws Exception {
 582         LOG.info("testOnReceiveFindPrimaryWaitForReadyWithCandidateShard starting");
 583         new JavaTestKit(getSystem()) {{
 584             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 585
 586             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 587             shardManager.tell(new ActorInitialized(), mockShardActor);
 588             shardManager.tell(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
 589                     null, RaftState.Candidate.name()), mockShardActor);
 590
 591             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 592
 593             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 594         }};
 595
 596         LOG.info("testOnReceiveFindPrimaryWaitForReadyWithCandidateShard ending");
 597     }
 598
 599     @Test
 600     public void testOnReceiveFindPrimaryWaitForReadyWithIsolatedLeaderShard() throws Exception {
 601         LOG.info("testOnReceiveFindPrimaryWaitForReadyWithIsolatedLeaderShard starting");
 602         new JavaTestKit(getSystem()) {{
 603             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 604
 605             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 606             shardManager.tell(new ActorInitialized(), mockShardActor);
 607             shardManager.tell(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
 608                     null, RaftState.IsolatedLeader.name()), mockShardActor);
 609
 610             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 611
 612             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 613         }};
 614
 615         LOG.info("testOnReceiveFindPrimaryWaitForReadyWithIsolatedLeaderShard ending");
 616     }
 617
 618     @Test
 619     public void testOnReceiveFindPrimaryWaitForReadyWithNoRoleShard() throws Exception {
 620         LOG.info("testOnReceiveFindPrimaryWaitForReadyWithNoRoleShard starting");
 621         new JavaTestKit(getSystem()) {{
 622             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 623
 624             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 625             shardManager.tell(new ActorInitialized(), mockShardActor);
 626
 627             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 628
 629             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 630         }};
 631
 632         LOG.info("testOnReceiveFindPrimaryWaitForReadyWithNoRoleShard ending");
 633     }
 634
 635     @Test
 636     public void testOnReceiveFindPrimaryForRemoteShard() throws Exception {
 637         LOG.info("testOnReceiveFindPrimaryForRemoteShard starting");
 638         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 639
 640         // Create an ActorSystem ShardManager actor for member-1.
 641
 642         final ActorSystem system1 = newActorSystem("Member1");
 643         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 644
 645         final TestActorRef<TestShardManager> shardManager1 = TestActorRef.create(system1,
 646                 newTestShardMgrBuilderWithMockShardActor().cluster(
 647                         new ClusterWrapperImpl(system1)).props(), shardManagerID);
 648
 649         // Create an ActorSystem ShardManager actor for member-2.
 650
 651         final ActorSystem system2 = newActorSystem("Member2");
 652
 653         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 654
 655         final ActorRef mockShardActor2 = newMockShardActor(system2, "astronauts", "member-2");
 656
 657         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 658                 put("default", Arrays.asList("member-1", "member-2")).
 659                 put("astronauts", Arrays.asList("member-2")).build());
 660
 661         final TestActorRef<TestShardManager> shardManager2 = TestActorRef.create(system2,
 662                 newTestShardMgrBuilder(mockConfig2).shardActor(mockShardActor2).cluster(
 663                         new ClusterWrapperImpl(system2)).props(), shardManagerID);
 664
 665         new JavaTestKit(system1) {{
 666
 667             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 668             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 669
 670             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 671
 672             String memberId2 = "member-2-shard-astronauts-" + shardMrgIDSuffix;
 673             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
 674             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2,
 675                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardActor2);
 676             shardManager2.tell(new RoleChangeNotification(memberId2,
 677                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 678
 679             shardManager1.underlyingActor().waitForMemberUp();
 680
 681             shardManager1.tell(new FindPrimary("astronauts", false), getRef());
 682
 683             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 684             String path = found.getPrimaryPath();
 685             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-astronauts-config"));
 686             assertEquals("getPrimaryVersion", leaderVersion, found.getPrimaryVersion());
 687
 688             shardManager2.underlyingActor().verifyFindPrimary();
 689
 690             Cluster.get(system2).down(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 691
 692             shardManager1.underlyingActor().waitForMemberRemoved();
 693
 694             shardManager1.tell(new FindPrimary("astronauts", false), getRef());
 695
 696             expectMsgClass(duration("5 seconds"), PrimaryNotFoundException.class);
 697         }};
 698
 699         LOG.info("testOnReceiveFindPrimaryForRemoteShard ending");
 700     }
 701
 702     @Test
 703     public void testShardAvailabilityOnChangeOfMemberReachability() throws Exception {
 704         LOG.info("testShardAvailabilityOnChangeOfMemberReachability starting");
 705         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 706
 707         // Create an ActorSystem ShardManager actor for member-1.
 708
 709         final ActorSystem system1 = newActorSystem("Member1");
 710         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 711
 712         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 713
 714         final TestActorRef<TestShardManager> shardManager1 = TestActorRef.create(system1,
 715                 newTestShardMgrBuilder().shardActor(mockShardActor1).cluster(
 716                         new ClusterWrapperImpl(system1)).props(), shardManagerID);
 717
 718         // Create an ActorSystem ShardManager actor for member-2.
 719
 720         final ActorSystem system2 = newActorSystem("Member2");
 721
 722         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 723
 724         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 725
 726         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 727             put("default", Arrays.asList("member-1", "member-2")).build());
 728
 729         final TestActorRef<TestShardManager> shardManager2 = TestActorRef.create(system2,
 730                 newTestShardMgrBuilder(mockConfig2).shardActor(mockShardActor2).cluster(
 731                         new ClusterWrapperImpl(system2)).props(), shardManagerID);
 732
 733         new JavaTestKit(system1) {{
 734
 735             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 736             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 737             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 738             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 739
 740             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 741             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 742             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 743                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 744             shardManager1.tell(new RoleChangeNotification(memberId1,
 745                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 746             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class)),
 747                     DataStoreVersions.CURRENT_VERSION),
 748                 mockShardActor2);
 749             shardManager2.tell(new RoleChangeNotification(memberId2,
 750                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 751             shardManager1.underlyingActor().waitForMemberUp();
 752
 753             shardManager1.tell(new FindPrimary("default", true), getRef());
 754
 755             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 756             String path = found.getPrimaryPath();
 757             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 758
 759             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 760                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 761
 762             shardManager1.underlyingActor().waitForUnreachableMember();
 763
 764             PeerDown peerDown = MessageCollectorActor.expectFirstMatching(mockShardActor1, PeerDown.class);
 765             assertEquals("getMemberName", "member-2", peerDown.getMemberName());
 766             MessageCollectorActor.clearMessages(mockShardActor1);
 767
 768             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 769                     createMemberRemoved("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 770
 771             MessageCollectorActor.expectFirstMatching(mockShardActor1, PeerDown.class);
 772
 773             shardManager1.tell(new FindPrimary("default", true), getRef());
 774
 775             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 776
 777             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 778                 createReachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 779
 780             shardManager1.underlyingActor().waitForReachableMember();
 781
 782             PeerUp peerUp = MessageCollectorActor.expectFirstMatching(mockShardActor1, PeerUp.class);
 783             assertEquals("getMemberName", "member-2", peerUp.getMemberName());
 784             MessageCollectorActor.clearMessages(mockShardActor1);
 785
 786             shardManager1.tell(new FindPrimary("default", true), getRef());
 787
 788             RemotePrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 789             String path1 = found1.getPrimaryPath();
 790             assertTrue("Unexpected primary path " + path1, path1.contains("member-2-shard-default-config"));
 791
 792             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 793                     createMemberUp("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 794
 795             MessageCollectorActor.expectFirstMatching(mockShardActor1, PeerUp.class);
 796
 797         }};
 798
 799         LOG.info("testShardAvailabilityOnChangeOfMemberReachability ending");
 800     }
 801
 802     @Test
 803     public void testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange() throws Exception {
 804         LOG.info("testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange starting");
 805         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 806
 807         // Create an ActorSystem ShardManager actor for member-1.
 808
 809         final ActorSystem system1 = newActorSystem("Member1");
 810         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 811
 812         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 813
 814         final PrimaryShardInfoFutureCache primaryShardInfoCache = new PrimaryShardInfoFutureCache();
 815         final TestActorRef<TestShardManager> shardManager1 = TestActorRef.create(system1,
 816                 newTestShardMgrBuilder().shardActor(mockShardActor1).cluster(
 817                         new ClusterWrapperImpl(system1)).primaryShardInfoCache(primaryShardInfoCache).props(),
 818                 shardManagerID);
 819
 820         // Create an ActorSystem ShardManager actor for member-2.
 821
 822         final ActorSystem system2 = newActorSystem("Member2");
 823
 824         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 825
 826         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 827
 828         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 829             put("default", Arrays.asList("member-1", "member-2")).build());
 830
 831         final TestActorRef<TestShardManager> shardManager2 = TestActorRef.create(system2,
 832                 newTestShardMgrBuilder(mockConfig2).shardActor(mockShardActor2).cluster(
 833                         new ClusterWrapperImpl(system2)).props(), shardManagerID);
 834
 835         new JavaTestKit(system1) {{
 836             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 837             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 838             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 839             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 840
 841             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 842             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 843             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 844                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 845             shardManager1.tell(new RoleChangeNotification(memberId1,
 846                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 847             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class)),
 848                     DataStoreVersions.CURRENT_VERSION),
 849                 mockShardActor2);
 850             shardManager2.tell(new RoleChangeNotification(memberId2,
 851                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 852             shardManager1.underlyingActor().waitForMemberUp();
 853
 854             shardManager1.tell(new FindPrimary("default", true), getRef());
 855
 856             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 857             String path = found.getPrimaryPath();
 858             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 859
 860             primaryShardInfoCache.putSuccessful("default", new PrimaryShardInfo(system1.actorSelection(
 861                     mockShardActor1.path()), DataStoreVersions.CURRENT_VERSION, Optional.<DataTree>absent()));
 862
 863             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 864                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 865
 866             shardManager1.underlyingActor().waitForUnreachableMember();
 867
 868             shardManager1.tell(new FindPrimary("default", true), getRef());
 869
 870             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 871
 872             assertNull("Expected primaryShardInfoCache entry removed", primaryShardInfoCache.getIfPresent("default"));
 873
 874             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId1, Optional.of(mock(DataTree.class)),
 875                     DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 876             shardManager1.tell(new RoleChangeNotification(memberId1,
 877                 RaftState.Follower.name(), RaftState.Leader.name()), mockShardActor1);
 878
 879             shardManager1.tell(new FindPrimary("default", true), getRef());
 880
 881             LocalPrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 882             String path1 = found1.getPrimaryPath();
 883             assertTrue("Unexpected primary path " + path1, path1.contains("member-1-shard-default-config"));
 884
 885         }};
 886
 887         LOG.info("testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange ending");
 888     }
 889
 890
 891     @Test
 892     public void testOnReceiveFindLocalShardForNonExistentShard() throws Exception {
 893         new JavaTestKit(getSystem()) {{
 894             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 895
 896             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 897
 898             shardManager.tell(new FindLocalShard("non-existent", false), getRef());
 899
 900             LocalShardNotFound notFound = expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
 901
 902             assertEquals("getShardName", "non-existent", notFound.getShardName());
 903         }};
 904     }
 905
 906     @Test
 907     public void testOnReceiveFindLocalShardForExistentShard() throws Exception {
 908         new JavaTestKit(getSystem()) {{
 909             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 910
 911             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 912             shardManager.tell(new ActorInitialized(), mockShardActor);
 913
 914             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 915
 916             LocalShardFound found = expectMsgClass(duration("5 seconds"), LocalShardFound.class);
 917
 918             assertTrue("Found path contains " + found.getPath().path().toString(),
 919                     found.getPath().path().toString().contains("member-1-shard-default-config"));
 920         }};
 921     }
 922
 923     @Test
 924     public void testOnReceiveFindLocalShardForNotInitializedShard() throws Exception {
 925         new JavaTestKit(getSystem()) {{
 926             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 927
 928             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 929
 930             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 931         }};
 932     }
 933
 934     @Test
 935     public void testOnReceiveFindLocalShardWaitForShardInitialized() throws Exception {
 936         LOG.info("testOnReceiveFindLocalShardWaitForShardInitialized starting");
 937         new JavaTestKit(getSystem()) {{
 938             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 939
 940             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 941
 942             // We're passing waitUntilInitialized = true to FindLocalShard so the response should be
 943             // delayed until we send ActorInitialized.
 944             Future<Object> future = Patterns.ask(shardManager, new FindLocalShard(Shard.DEFAULT_NAME, true),
 945                     new Timeout(5, TimeUnit.SECONDS));
 946
 947             shardManager.tell(new ActorInitialized(), mockShardActor);
 948
 949             Object resp = Await.result(future, duration("5 seconds"));
 950             assertTrue("Expected: LocalShardFound, Actual: " + resp, resp instanceof LocalShardFound);
 951         }};
 952
 953         LOG.info("testOnReceiveFindLocalShardWaitForShardInitialized starting");
 954     }
 955
 956     @Test
 957     public void testOnRecoveryJournalIsCleaned() {
 958         String persistenceID = "shard-manager-" + shardMrgIDSuffix;
 959         InMemoryJournal.addEntry(persistenceID, 1L, new ShardManager.SchemaContextModules(
 960                 ImmutableSet.of("foo")));
 961         InMemoryJournal.addEntry(persistenceID, 2L, new ShardManager.SchemaContextModules(
 962                 ImmutableSet.of("bar")));
 963         InMemoryJournal.addDeleteMessagesCompleteLatch(persistenceID);
 964
 965         TestShardManager shardManager = newTestShardManager();
 966
 967         InMemoryJournal.waitForDeleteMessagesComplete(persistenceID);
 968
 969         // Journal entries up to the last one should've been deleted
 970         Map<Long, Object> journal = InMemoryJournal.get(persistenceID);
 971         synchronized (journal) {
 972             assertEquals("Journal size", 0, journal.size());
 973         }
 974     }
 975
 976     @Test
 977     public void testRoleChangeNotificationAndShardLeaderStateChangedReleaseReady() throws Exception {
 978         TestShardManager shardManager = newTestShardManager();
 979
 980         String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 981         shardManager.onReceiveCommand(new RoleChangeNotification(
 982                 memberId, RaftState.Candidate.name(), RaftState.Leader.name()));
 983
 984         verify(ready, never()).countDown();
 985
 986         shardManager.onReceiveCommand(new ShardLeaderStateChanged(memberId, memberId,
 987                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION));
 988
 989         verify(ready, times(1)).countDown();
 990     }
 991
 992     @Test
 993     public void testRoleChangeNotificationToFollowerWithShardLeaderStateChangedReleaseReady() throws Exception {
 994         new JavaTestKit(getSystem()) {{
 995             TestShardManager shardManager = newTestShardManager();
 996
 997             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 998             shardManager.onReceiveCommand(new RoleChangeNotification(
 999                     memberId, null, RaftState.Follower.name()));
1000
1001             verify(ready, never()).countDown();
1002
1003             shardManager.onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
1004
1005             shardManager.onReceiveCommand(new ShardLeaderStateChanged(memberId,
1006                     "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class)),
1007                     DataStoreVersions.CURRENT_VERSION));
1008
1009             verify(ready, times(1)).countDown();
1010         }};
1011     }
1012
1013     @Test
1014     public void testReadyCountDownForMemberUpAfterLeaderStateChanged() throws Exception {
1015         new JavaTestKit(getSystem()) {{
1016             TestShardManager shardManager = newTestShardManager();
1017
1018             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
1019             shardManager.onReceiveCommand(new RoleChangeNotification(memberId, null, RaftState.Follower.name()));
1020
1021             verify(ready, never()).countDown();
1022
1023             shardManager.onReceiveCommand(new ShardLeaderStateChanged(memberId,
1024                     "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class)),
1025                     DataStoreVersions.CURRENT_VERSION));
1026
1027             shardManager.onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
1028
1029             verify(ready, times(1)).countDown();
1030         }};
1031     }
1032
1033     @Test
1034     public void testRoleChangeNotificationDoNothingForUnknownShard() throws Exception {
1035         TestShardManager shardManager = newTestShardManager();
1036
1037         shardManager.onReceiveCommand(new RoleChangeNotification(
1038                 "unknown", RaftState.Candidate.name(), RaftState.Leader.name()));
1039
1040         verify(ready, never()).countDown();
1041     }
1042
1043     @Test
1044     public void testByDefaultSyncStatusIsFalse() throws Exception{
1045         TestShardManager shardManager = newTestShardManager();
1046
1047         assertEquals(false, shardManager.getMBean().getSyncStatus());
1048     }
1049
1050     @Test
1051     public void testWhenShardIsLeaderSyncStatusIsTrue() throws Exception{
1052         TestShardManager shardManager = newTestShardManager();
1053
1054         shardManager.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
1055                 RaftState.Follower.name(), RaftState.Leader.name()));
1056
1057         assertEquals(true, shardManager.getMBean().getSyncStatus());
1058     }
1059
1060     @Test
1061     public void testWhenShardIsCandidateSyncStatusIsFalse() throws Exception{
1062         TestShardManager shardManager = newTestShardManager();
1063
1064         String shardId = "member-1-shard-default-" + shardMrgIDSuffix;
1065         shardManager.onReceiveCommand(new RoleChangeNotification(shardId,
1066                 RaftState.Follower.name(), RaftState.Candidate.name()));
1067
1068         assertEquals(false, shardManager.getMBean().getSyncStatus());
1069
1070         // Send a FollowerInitialSyncStatus with status = true for the replica whose current state is candidate
1071         shardManager.onReceiveCommand(new FollowerInitialSyncUpStatus(
1072                 true, shardId));
1073
1074         assertEquals(false, shardManager.getMBean().getSyncStatus());
1075     }
1076
1077     @Test
1078     public void testWhenShardIsFollowerSyncStatusDependsOnFollowerInitialSyncStatus() throws Exception{
1079         TestShardManager shardManager = newTestShardManager();
1080
1081         String shardId = "member-1-shard-default-" + shardMrgIDSuffix;
1082         shardManager.onReceiveCommand(new RoleChangeNotification(shardId,
1083                 RaftState.Candidate.name(), RaftState.Follower.name()));
1084
1085         // Initially will be false
1086         assertEquals(false, shardManager.getMBean().getSyncStatus());
1087
1088         // Send status true will make sync status true
1089         shardManager.onReceiveCommand(new FollowerInitialSyncUpStatus(true, shardId));
1090
1091         assertEquals(true, shardManager.getMBean().getSyncStatus());
1092
1093         // Send status false will make sync status false
1094         shardManager.onReceiveCommand(new FollowerInitialSyncUpStatus(false, shardId));
1095
1096         assertEquals(false, shardManager.getMBean().getSyncStatus());
1097
1098     }
1099
1100     @Test
1101     public void testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards() throws Exception{
1102         LOG.info("testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards starting");
1103         TestShardManager shardManager = newTestShardManager(newShardMgrProps(new MockConfiguration() {
1104             @Override
1105             public List<String> getMemberShardNames(String memberName) {
1106                 return Arrays.asList("default", "astronauts");
1107             }
1108         }));
1109
1110         // Initially will be false
1111         assertEquals(false, shardManager.getMBean().getSyncStatus());
1112
1113         // Make default shard leader
1114         String defaultShardId = "member-1-shard-default-" + shardMrgIDSuffix;
1115         shardManager.onReceiveCommand(new RoleChangeNotification(defaultShardId,
1116                 RaftState.Follower.name(), RaftState.Leader.name()));
1117
1118         // default = Leader, astronauts is unknown so sync status remains false
1119         assertEquals(false, shardManager.getMBean().getSyncStatus());
1120
1121         // Make astronauts shard leader as well
1122         String astronautsShardId = "member-1-shard-astronauts-" + shardMrgIDSuffix;
1123         shardManager.onReceiveCommand(new RoleChangeNotification(astronautsShardId,
1124                 RaftState.Follower.name(), RaftState.Leader.name()));
1125
1126         // Now sync status should be true
1127         assertEquals(true, shardManager.getMBean().getSyncStatus());
1128
1129         // Make astronauts a Follower
1130         shardManager.onReceiveCommand(new RoleChangeNotification(astronautsShardId,
1131                 RaftState.Leader.name(), RaftState.Follower.name()));
1132
1133         // Sync status is not true
1134         assertEquals(false, shardManager.getMBean().getSyncStatus());
1135
1136         // Make the astronauts follower sync status true
1137         shardManager.onReceiveCommand(new FollowerInitialSyncUpStatus(true, astronautsShardId));
1138
1139         // Sync status is now true
1140         assertEquals(true, shardManager.getMBean().getSyncStatus());
1141
1142         LOG.info("testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards ending");
1143     }
1144
1145     @Test
1146     public void testOnReceiveSwitchShardBehavior() throws Exception {
1147         new JavaTestKit(getSystem()) {{
1148             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
1149
1150             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1151             shardManager.tell(new ActorInitialized(), mockShardActor);
1152
1153             shardManager.tell(new SwitchShardBehavior(mockShardName, "Leader", 1000), getRef());
1154
1155             SwitchBehavior switchBehavior = MessageCollectorActor.expectFirstMatching(mockShardActor, SwitchBehavior.class);
1156
1157             assertEquals(RaftState.Leader, switchBehavior.getNewState());
1158             assertEquals(1000, switchBehavior.getNewTerm());
1159         }};
1160     }
1161
1162     @Test
1163     public void testOnCreateShard() {
1164         LOG.info("testOnCreateShard starting");
1165         new JavaTestKit(getSystem()) {{
1166             datastoreContextBuilder.shardInitializationTimeout(1, TimeUnit.MINUTES).persistent(true);
1167
1168             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1169                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1170
1171             SchemaContext schemaContext = TestModel.createTestContext();
1172             shardManager.tell(new UpdateSchemaContext(schemaContext), ActorRef.noSender());
1173
1174             DatastoreContext datastoreContext = DatastoreContext.newBuilder().shardElectionTimeoutFactor(100).
1175                     persistent(false).build();
1176             Shard.Builder shardBuilder = Shard.builder();
1177
1178             ModuleShardConfiguration config = new ModuleShardConfiguration(URI.create("foo-ns"), "foo-module",
1179                     "foo", null, Arrays.asList("member-1", "member-5", "member-6"));
1180             shardManager.tell(new CreateShard(config, shardBuilder, datastoreContext), getRef());
1181
1182             expectMsgClass(duration("5 seconds"), Success.class);
1183
1184             shardManager.tell(new FindLocalShard("foo", true), getRef());
1185
1186             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1187
1188             assertEquals("isRecoveryApplicable", false, shardBuilder.getDatastoreContext().isPersistent());
1189             assertTrue("Epxected ShardPeerAddressResolver", shardBuilder.getDatastoreContext().getShardRaftConfig().
1190                     getPeerAddressResolver() instanceof ShardPeerAddressResolver);
1191             assertEquals("peerMembers", Sets.newHashSet(new ShardIdentifier("foo", "member-5", shardMrgIDSuffix).toString(),
1192                     new ShardIdentifier("foo", "member-6", shardMrgIDSuffix).toString()),
1193                     shardBuilder.getPeerAddresses().keySet());
1194             assertEquals("ShardIdentifier", new ShardIdentifier("foo", "member-1", shardMrgIDSuffix),
1195                     shardBuilder.getId());
1196             assertSame("schemaContext", schemaContext, shardBuilder.getSchemaContext());
1197
1198             // Send CreateShard with same name - should return Success with a message.
1199
1200             shardManager.tell(new CreateShard(config, shardBuilder, null), getRef());
1201
1202             Success success = expectMsgClass(duration("5 seconds"), Success.class);
1203             assertNotNull("Success status is null", success.status());
1204         }};
1205
1206         LOG.info("testOnCreateShard ending");
1207     }
1208
1209     @Test
1210     public void testOnCreateShardWithLocalMemberNotInShardConfig() {
1211         LOG.info("testOnCreateShardWithLocalMemberNotInShardConfig starting");
1212         new JavaTestKit(getSystem()) {{
1213             datastoreContextBuilder.shardInitializationTimeout(1, TimeUnit.MINUTES).persistent(true);
1214
1215             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1216                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1217
1218             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), ActorRef.noSender());
1219
1220             Shard.Builder shardBuilder = Shard.builder();
1221             ModuleShardConfiguration config = new ModuleShardConfiguration(URI.create("foo-ns"), "foo-module",
1222                     "foo", null, Arrays.asList("member-5", "member-6"));
1223
1224             shardManager.tell(new CreateShard(config, shardBuilder, null), getRef());
1225             expectMsgClass(duration("5 seconds"), Success.class);
1226
1227             shardManager.tell(new FindLocalShard("foo", true), getRef());
1228             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1229
1230             assertEquals("peerMembers size", 0, shardBuilder.getPeerAddresses().size());
1231             assertEquals("schemaContext", DisableElectionsRaftPolicy.class.getName(),
1232                     shardBuilder.getDatastoreContext().getShardRaftConfig().getCustomRaftPolicyImplementationClass());
1233         }};
1234
1235         LOG.info("testOnCreateShardWithLocalMemberNotInShardConfig ending");
1236     }
1237
1238     @Test
1239     public void testOnCreateShardWithNoInitialSchemaContext() {
1240         LOG.info("testOnCreateShardWithNoInitialSchemaContext starting");
1241         new JavaTestKit(getSystem()) {{
1242             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1243                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1244
1245             Shard.Builder shardBuilder = Shard.builder();
1246
1247             ModuleShardConfiguration config = new ModuleShardConfiguration(URI.create("foo-ns"), "foo-module",
1248                     "foo", null, Arrays.asList("member-1"));
1249             shardManager.tell(new CreateShard(config, shardBuilder, null), getRef());
1250
1251             expectMsgClass(duration("5 seconds"), Success.class);
1252
1253             SchemaContext schemaContext = TestModel.createTestContext();
1254             shardManager.tell(new UpdateSchemaContext(schemaContext), ActorRef.noSender());
1255
1256             shardManager.tell(new FindLocalShard("foo", true), getRef());
1257
1258             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1259
1260             assertSame("schemaContext", schemaContext, shardBuilder.getSchemaContext());
1261             assertNotNull("schemaContext is null", shardBuilder.getDatastoreContext());
1262         }};
1263
1264         LOG.info("testOnCreateShardWithNoInitialSchemaContext ending");
1265     }
1266
1267     @Test
1268     public void testGetSnapshot() throws Throwable {
1269         LOG.info("testGetSnapshot starting");
1270         JavaTestKit kit = new JavaTestKit(getSystem());
1271
1272         MockConfiguration mockConfig = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1273                    put("shard1", Arrays.asList("member-1")).
1274                    put("shard2", Arrays.asList("member-1")).
1275                    put("astronauts", Collections.<String>emptyList()).build());
1276
1277         TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(newShardMgrProps(mockConfig).
1278                 withDispatcher(Dispatchers.DefaultDispatcherId()));
1279
1280         shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
1281         Failure failure = kit.expectMsgClass(Failure.class);
1282         assertEquals("Failure cause type", IllegalStateException.class, failure.cause().getClass());
1283
1284         shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), ActorRef.noSender());
1285
1286         waitForShardInitialized(shardManager, "shard1", kit);
1287         waitForShardInitialized(shardManager, "shard2", kit);
1288
1289         shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
1290
1291         DatastoreSnapshot datastoreSnapshot = expectMsgClassOrFailure(DatastoreSnapshot.class, kit, "GetSnapshot");
1292
1293         assertEquals("getType", shardMrgIDSuffix, datastoreSnapshot.getType());
1294         assertNull("Expected null ShardManagerSnapshot", datastoreSnapshot.getShardManagerSnapshot());
1295
1296         Function<ShardSnapshot, String> shardNameTransformer = new Function<ShardSnapshot, String>() {
1297             @Override
1298             public String apply(ShardSnapshot s) {
1299                 return s.getName();
1300             }
1301         };
1302
1303         assertEquals("Shard names", Sets.newHashSet("shard1", "shard2"), Sets.newHashSet(
1304                 Lists.transform(datastoreSnapshot.getShardSnapshots(), shardNameTransformer)));
1305
1306         // Add a new replica
1307
1308         JavaTestKit mockShardLeaderKit = new JavaTestKit(getSystem());
1309
1310         TestShardManager shardManagerInstance = shardManager.underlyingActor();
1311         shardManagerInstance.setMessageInterceptor(newFindPrimaryInterceptor(mockShardLeaderKit.getRef()));
1312
1313         shardManager.tell(new AddShardReplica("astronauts"), kit.getRef());
1314         mockShardLeaderKit.expectMsgClass(AddServer.class);
1315         mockShardLeaderKit.reply(new AddServerReply(ServerChangeStatus.OK, ""));
1316         kit.expectMsgClass(Status.Success.class);
1317         waitForShardInitialized(shardManager, "astronauts", kit);
1318
1319         // Send another GetSnapshot and verify
1320
1321         shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
1322         datastoreSnapshot = expectMsgClassOrFailure(DatastoreSnapshot.class, kit, "GetSnapshot");
1323
1324         assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"), Sets.newHashSet(
1325                 Lists.transform(datastoreSnapshot.getShardSnapshots(), shardNameTransformer)));
1326
1327         byte[] snapshotBytes = datastoreSnapshot.getShardManagerSnapshot();
1328         assertNotNull("Expected ShardManagerSnapshot", snapshotBytes);
1329         ShardManagerSnapshot snapshot = SerializationUtils.deserialize(snapshotBytes);
1330         assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"),
1331                 Sets.newHashSet(snapshot.getShardList()));
1332
1333         LOG.info("testGetSnapshot ending");
1334     }
1335
1336     @Test
1337     public void testRestoreFromSnapshot() throws Throwable {
1338         LOG.info("testRestoreFromSnapshot starting");
1339
1340         JavaTestKit kit = new JavaTestKit(getSystem());
1341
1342         MockConfiguration mockConfig = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1343                    put("shard1", Collections.<String>emptyList()).
1344                    put("shard2", Collections.<String>emptyList()).
1345                    put("astronauts", Collections.<String>emptyList()).build());
1346
1347
1348         ShardManagerSnapshot snapshot = new ShardManagerSnapshot(Arrays.asList("shard1", "shard2", "astronauts"));
1349         DatastoreSnapshot restoreFromSnapshot = new DatastoreSnapshot(shardMrgIDSuffix,
1350                 SerializationUtils.serialize(snapshot), Collections.<ShardSnapshot>emptyList());
1351         TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(newTestShardMgrBuilder(mockConfig).
1352                 restoreFromSnapshot(restoreFromSnapshot).props().withDispatcher(Dispatchers.DefaultDispatcherId()));
1353
1354         shardManager.underlyingActor().waitForRecoveryComplete();
1355
1356         shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), ActorRef.noSender());
1357
1358         waitForShardInitialized(shardManager, "shard1", kit);
1359         waitForShardInitialized(shardManager, "shard2", kit);
1360         waitForShardInitialized(shardManager, "astronauts", kit);
1361
1362         shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
1363
1364         DatastoreSnapshot datastoreSnapshot = expectMsgClassOrFailure(DatastoreSnapshot.class, kit, "GetSnapshot");
1365
1366         assertEquals("getType", shardMrgIDSuffix, datastoreSnapshot.getType());
1367
1368         byte[] snapshotBytes = datastoreSnapshot.getShardManagerSnapshot();
1369         assertNotNull("Expected ShardManagerSnapshot", snapshotBytes);
1370         snapshot = SerializationUtils.deserialize(snapshotBytes);
1371         assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"),
1372                 Sets.newHashSet(snapshot.getShardList()));
1373
1374         LOG.info("testRestoreFromSnapshot ending");
1375     }
1376
1377     @Test
1378     public void testAddShardReplicaForNonExistentShardConfig() throws Exception {
1379         new JavaTestKit(getSystem()) {{
1380             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1381                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1382
1383             shardManager.tell(new AddShardReplica("model-inventory"), getRef());
1384             Status.Failure resp = expectMsgClass(duration("2 seconds"), Status.Failure.class);
1385
1386             assertEquals("Failure obtained", true,
1387                           (resp.cause() instanceof IllegalArgumentException));
1388         }};
1389     }
1390
1391     @Test
1392     public void testAddShardReplica() throws Exception {
1393         MockConfiguration mockConfig =
1394                 new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1395                    put("default", Arrays.asList("member-1", "member-2")).
1396                    put("astronauts", Arrays.asList("member-2")).build());
1397
1398         final String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
1399         datastoreContextBuilder.shardManagerPersistenceId(shardManagerID);
1400
1401         // Create an ActorSystem ShardManager actor for member-1.
1402         final ActorSystem system1 = newActorSystem("Member1");
1403         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1404         ActorRef mockDefaultShardActor = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
1405         final TestActorRef<TestShardManager> newReplicaShardManager = TestActorRef.create(system1,
1406                 newTestShardMgrBuilder(mockConfig).shardActor(mockDefaultShardActor).cluster(
1407                         new ClusterWrapperImpl(system1)).props(), shardManagerID);
1408
1409         // Create an ActorSystem ShardManager actor for member-2.
1410         final ActorSystem system2 = newActorSystem("Member2");
1411         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1412
1413         String name = new ShardIdentifier("astronauts", "member-2", "config").toString();
1414         final TestActorRef<MockRespondActor> mockShardLeaderActor =
1415                 TestActorRef.create(system2, Props.create(MockRespondActor.class), name);
1416         final TestActorRef<TestShardManager> leaderShardManager = TestActorRef.create(system2,
1417                 newTestShardMgrBuilder(mockConfig).shardActor(mockShardLeaderActor).cluster(
1418                         new ClusterWrapperImpl(system2)).props(), shardManagerID);
1419
1420         new JavaTestKit(system1) {{
1421
1422             newReplicaShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1423             leaderShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1424
1425             leaderShardManager.tell(new ActorInitialized(), mockShardLeaderActor);
1426
1427             String memberId2 = "member-2-shard-astronauts-" + shardMrgIDSuffix;
1428             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
1429             leaderShardManager.tell(new ShardLeaderStateChanged(memberId2, memberId2,
1430                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardLeaderActor);
1431             leaderShardManager.tell(new RoleChangeNotification(memberId2,
1432                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardLeaderActor);
1433
1434             newReplicaShardManager.underlyingActor().waitForMemberUp();
1435             leaderShardManager.underlyingActor().waitForMemberUp();
1436
1437             //Have a dummy snapshot to be overwritten by the new data persisted.
1438             String[] restoredShards = {"default", "people"};
1439             ShardManagerSnapshot snapshot = new ShardManagerSnapshot(Arrays.asList(restoredShards));
1440             InMemorySnapshotStore.addSnapshot(shardManagerID, snapshot);
1441             Uninterruptibles.sleepUninterruptibly(2, TimeUnit.MILLISECONDS);
1442
1443             InMemorySnapshotStore.addSnapshotSavedLatch(shardManagerID);
1444             InMemorySnapshotStore.addSnapshotDeletedLatch(shardManagerID);
1445
1446             //construct a mock response message
1447             AddServerReply response = new AddServerReply(ServerChangeStatus.OK, memberId2);
1448             mockShardLeaderActor.underlyingActor().updateResponse(response);
1449             newReplicaShardManager.tell(new AddShardReplica("astronauts"), getRef());
1450             AddServer addServerMsg = MessageCollectorActor.expectFirstMatching(mockShardLeaderActor,
1451                 AddServer.class);
1452             String addServerId = "member-1-shard-astronauts-" + shardMrgIDSuffix;
1453             assertEquals("AddServer serverId", addServerId, addServerMsg.getNewServerId());
1454             expectMsgClass(duration("5 seconds"), Status.Success.class);
1455
1456             InMemorySnapshotStore.waitForSavedSnapshot(shardManagerID, ShardManagerSnapshot.class);
1457             InMemorySnapshotStore.waitForDeletedSnapshot(shardManagerID);
1458             List<ShardManagerSnapshot> persistedSnapshots =
1459                 InMemorySnapshotStore.getSnapshots(shardManagerID, ShardManagerSnapshot.class);
1460             assertEquals("Number of snapshots persisted", 1, persistedSnapshots.size());
1461             ShardManagerSnapshot shardManagerSnapshot = persistedSnapshots.get(0);
1462             assertEquals("Persisted local shards", Sets.newHashSet("default", "astronauts"),
1463                     Sets.newHashSet(shardManagerSnapshot.getShardList()));
1464         }};
1465
1466         LOG.info("testAddShardReplica ending");
1467     }
1468
1469     @Test
1470     public void testAddShardReplicaWithPreExistingReplicaInRemoteShardLeader() throws Exception {
1471         LOG.info("testAddShardReplicaWithPreExistingReplicaInRemoteShardLeader starting");
1472         new JavaTestKit(getSystem()) {{
1473             TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(
1474                     newPropsShardMgrWithMockShardActor(), shardMgrID);
1475
1476             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1477             shardManager.tell(new ActorInitialized(), mockShardActor);
1478
1479             String leaderId = "leader-member-shard-default-" + shardMrgIDSuffix;
1480             AddServerReply addServerReply = new AddServerReply(ServerChangeStatus.ALREADY_EXISTS, null);
1481             ActorRef leaderShardActor = shardManager.underlyingActor().getContext().actorOf(
1482                     Props.create(MockRespondActor.class, addServerReply), leaderId);
1483
1484             MockClusterWrapper.sendMemberUp(shardManager, "leader-member", leaderShardActor.path().toString());
1485
1486             String newReplicaId = "member-1-shard-default-" + shardMrgIDSuffix;
1487             shardManager.tell(new RoleChangeNotification(newReplicaId,
1488                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
1489             shardManager.tell(new ShardLeaderStateChanged(newReplicaId, leaderId, Optional.<DataTree>absent(),
1490                     DataStoreVersions.CURRENT_VERSION), mockShardActor);
1491
1492             shardManager.tell(new AddShardReplica(Shard.DEFAULT_NAME), getRef());
1493
1494             MessageCollectorActor.expectFirstMatching(leaderShardActor, AddServer.class);
1495
1496             Failure resp = expectMsgClass(duration("5 seconds"), Failure.class);
1497             assertEquals("Failure cause", AlreadyExistsException.class, resp.cause().getClass());
1498
1499             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
1500             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1501
1502             // Send message again to verify previous in progress state is cleared
1503
1504             shardManager.tell(new AddShardReplica(Shard.DEFAULT_NAME), getRef());
1505             resp = expectMsgClass(duration("5 seconds"), Failure.class);
1506             assertEquals("Failure cause", AlreadyExistsException.class, resp.cause().getClass());
1507
1508             // Send message again with an AddServer timeout to verify the pre-existing shard actor isn't terminated.
1509
1510             shardManager.tell(newDatastoreContextFactory(datastoreContextBuilder.
1511                     shardLeaderElectionTimeout(100, TimeUnit.MILLISECONDS).build()), getRef());
1512             leaderShardActor.tell(MockRespondActor.CLEAR_RESPONSE, ActorRef.noSender());
1513             shardManager.tell(new AddShardReplica(Shard.DEFAULT_NAME), getRef());
1514             expectMsgClass(duration("5 seconds"), Failure.class);
1515
1516             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
1517             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1518         }};
1519
1520         LOG.info("testAddShardReplicaWithPreExistingReplicaInRemoteShardLeader ending");
1521     }
1522
1523     @Test
1524     public void testAddShardReplicaWithPreExistingLocalReplicaLeader() throws Exception {
1525         LOG.info("testAddShardReplicaWithPreExistingLocalReplicaLeader starting");
1526         new JavaTestKit(getSystem()) {{
1527             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
1528             ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
1529
1530             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1531             shardManager.tell(new ActorInitialized(), mockShardActor);
1532             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mock(DataTree.class)),
1533                     DataStoreVersions.CURRENT_VERSION), getRef());
1534             shardManager.tell((new RoleChangeNotification(memberId, RaftState.Candidate.name(),
1535                     RaftState.Leader.name())), mockShardActor);
1536
1537             shardManager.tell(new AddShardReplica(Shard.DEFAULT_NAME), getRef());
1538             Failure resp = expectMsgClass(duration("5 seconds"), Failure.class);
1539             assertEquals("Failure cause", AlreadyExistsException.class, resp.cause().getClass());
1540
1541             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
1542             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1543         }};
1544
1545         LOG.info("testAddShardReplicaWithPreExistingLocalReplicaLeader ending");
1546     }
1547
1548     @Test
1549     public void testAddShardReplicaWithAddServerReplyFailure() throws Exception {
1550         LOG.info("testAddShardReplicaWithAddServerReplyFailure starting");
1551         new JavaTestKit(getSystem()) {{
1552             JavaTestKit mockShardLeaderKit = new JavaTestKit(getSystem());
1553
1554             MockConfiguration mockConfig =
1555                     new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1556                        put("astronauts", Arrays.asList("member-2")).build());
1557
1558             ActorRef mockNewReplicaShardActor = newMockShardActor(getSystem(), "astronauts", "member-1");
1559             final TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(
1560                     newTestShardMgrBuilder(mockConfig).shardActor(mockNewReplicaShardActor).props(), shardMgrID);
1561             shardManager.underlyingActor().setMessageInterceptor(newFindPrimaryInterceptor(mockShardLeaderKit.getRef()));
1562
1563             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1564
1565             JavaTestKit terminateWatcher = new JavaTestKit(getSystem());
1566             terminateWatcher.watch(mockNewReplicaShardActor);
1567
1568             shardManager.tell(new AddShardReplica("astronauts"), getRef());
1569
1570             AddServer addServerMsg = mockShardLeaderKit.expectMsgClass(AddServer.class);
1571             assertEquals("AddServer serverId", "member-1-shard-astronauts-" + shardMrgIDSuffix,
1572                     addServerMsg.getNewServerId());
1573             mockShardLeaderKit.reply(new AddServerReply(ServerChangeStatus.TIMEOUT, null));
1574
1575             Failure failure = expectMsgClass(duration("5 seconds"), Failure.class);
1576             assertEquals("Failure cause", TimeoutException.class, failure.cause().getClass());
1577
1578             shardManager.tell(new FindLocalShard("astronauts", false), getRef());
1579             expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
1580
1581             terminateWatcher.expectTerminated(mockNewReplicaShardActor);
1582
1583             shardManager.tell(new AddShardReplica("astronauts"), getRef());
1584             mockShardLeaderKit.expectMsgClass(AddServer.class);
1585             mockShardLeaderKit.reply(new AddServerReply(ServerChangeStatus.NO_LEADER, null));
1586             failure = expectMsgClass(duration("5 seconds"), Failure.class);
1587             assertEquals("Failure cause", NoShardLeaderException.class, failure.cause().getClass());
1588         }};
1589
1590         LOG.info("testAddShardReplicaWithAddServerReplyFailure ending");
1591     }
1592
1593     @Test
1594     public void testAddShardReplicaWithAlreadyInProgress() throws Exception {
1595         testServerChangeWhenAlreadyInProgress("astronauts", new AddShardReplica("astronauts"),
1596                 AddServer.class, new AddShardReplica("astronauts"));
1597     }
1598
1599     @Test
1600     public void testAddShardReplicaWithFindPrimaryTimeout() throws Exception {
1601         LOG.info("testAddShardReplicaWithFindPrimaryTimeout starting");
1602         new JavaTestKit(getSystem()) {{
1603             MockConfiguration mockConfig = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1604                        put("astronauts", Arrays.asList("member-2")).build());
1605
1606             final ActorRef newReplicaShardManager = actorFactory.createActor(newTestShardMgrBuilder(mockConfig).
1607                     shardActor(mockShardActor).props(), shardMgrID);
1608
1609             newReplicaShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1610             MockClusterWrapper.sendMemberUp(newReplicaShardManager, "member-2", getRef().path().toString());
1611
1612             newReplicaShardManager.tell(new AddShardReplica("astronauts"), getRef());
1613             Status.Failure resp = expectMsgClass(duration("5 seconds"), Status.Failure.class);
1614             assertEquals("Failure obtained", true,
1615                           (resp.cause() instanceof RuntimeException));
1616         }};
1617
1618         LOG.info("testAddShardReplicaWithFindPrimaryTimeout ending");
1619     }
1620
1621     @Test
1622     public void testRemoveShardReplicaForNonExistentShard() throws Exception {
1623         new JavaTestKit(getSystem()) {{
1624             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1625                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1626
1627             shardManager.tell(new RemoveShardReplica("model-inventory", "member-1"), getRef());
1628             Status.Failure resp = expectMsgClass(duration("10 seconds"), Status.Failure.class);
1629             assertEquals("Failure obtained", true,
1630                          (resp.cause() instanceof PrimaryNotFoundException));
1631         }};
1632     }
1633
1634     @Test
1635     /**
1636      * Primary is Local
1637      */
1638     public void testRemoveShardReplicaLocal() throws Exception {
1639         new JavaTestKit(getSystem()) {{
1640             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
1641
1642             final TestActorRef<MockRespondActor> respondActor =
1643                     TestActorRef.create(getSystem(), Props.create(MockRespondActor.class), memberId);
1644
1645             ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor(respondActor));
1646
1647             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1648             shardManager.tell(new ActorInitialized(), respondActor);
1649             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mock(DataTree.class)),
1650                     DataStoreVersions.CURRENT_VERSION), getRef());
1651             shardManager.tell((new RoleChangeNotification(memberId, RaftState.Candidate.name(),
1652                     RaftState.Leader.name())), respondActor);
1653
1654             respondActor.underlyingActor().updateResponse(new RemoveServerReply(ServerChangeStatus.OK, null));
1655             shardManager.tell(new RemoveShardReplica(Shard.DEFAULT_NAME, "member-1"), getRef());
1656             final RemoveServer removeServer = MessageCollectorActor.expectFirstMatching(respondActor, RemoveServer.class);
1657             assertEquals(new ShardIdentifier("default", "member-1", shardMrgIDSuffix).toString(),
1658                     removeServer.getServerId());
1659             expectMsgClass(duration("5 seconds"), Success.class);
1660         }};
1661     }
1662
1663     @Test
1664     public void testRemoveShardReplicaRemote() throws Exception {
1665         MockConfiguration mockConfig =
1666                 new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1667                         put("default", Arrays.asList("member-1", "member-2")).
1668                         put("astronauts", Arrays.asList("member-1")).build());
1669
1670         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
1671
1672         // Create an ActorSystem ShardManager actor for member-1.
1673         final ActorSystem system1 = newActorSystem("Member1");
1674         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1675         ActorRef mockDefaultShardActor = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
1676
1677         final TestActorRef<TestShardManager> newReplicaShardManager = TestActorRef.create(system1,
1678                 newTestShardMgrBuilder().configuration(mockConfig).shardActor(mockDefaultShardActor).cluster(
1679                         new ClusterWrapperImpl(system1)).props(),
1680                 shardManagerID);
1681
1682         // Create an ActorSystem ShardManager actor for member-2.
1683         final ActorSystem system2 = newActorSystem("Member2");
1684         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1685
1686         String name = new ShardIdentifier("default", "member-2", shardMrgIDSuffix).toString();
1687         final TestActorRef<MockRespondActor> mockShardLeaderActor =
1688                 TestActorRef.create(system2, Props.create(MockRespondActor.class), name);
1689
1690         LOG.error("Mock Shard Leader Actor : {}", mockShardLeaderActor);
1691
1692         final TestActorRef<TestShardManager> leaderShardManager = TestActorRef.create(system2,
1693                 newTestShardMgrBuilder().configuration(mockConfig).shardActor(mockShardLeaderActor).cluster(
1694                         new ClusterWrapperImpl(system2)).props(),
1695                 shardManagerID);
1696
1697         // Because mockShardLeaderActor is created at the top level of the actor system it has an address like so,
1698         //    akka.tcp://cluster-test@127.0.0.1:2559/user/member-2-shard-default-config1
1699         // However when a shard manager has a local shard which is a follower and a leader that is remote it will
1700         // try to compute an address for the remote shard leader using the ShardPeerAddressResolver. This address will
1701         // look like so,
1702         //    akka.tcp://cluster-test@127.0.0.1:2559/user/shardmanager-config1/member-2-shard-default-config1
1703         // In this specific case if we did a FindPrimary for shard default from member-1 we would come up
1704         // with the address of an actor which does not exist, therefore any message sent to that actor would go to
1705         // dead letters.
1706         // To work around this problem we create a ForwardingActor with the right address and pass to it the
1707         // mockShardLeaderActor. The ForwardingActor simply forwards all messages to the mockShardLeaderActor and every
1708         // thing works as expected
1709         final ActorRef actorRef = leaderShardManager.underlyingActor().context()
1710                 .actorOf(Props.create(ForwardingActor.class, mockShardLeaderActor), "member-2-shard-default-" + shardMrgIDSuffix);
1711
1712         LOG.error("Forwarding actor : {}", actorRef);
1713
1714         new JavaTestKit(system1) {{
1715
1716             newReplicaShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1717             leaderShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1718
1719             leaderShardManager.tell(new ActorInitialized(), mockShardLeaderActor);
1720             newReplicaShardManager.tell(new ActorInitialized(), mockShardLeaderActor);
1721
1722             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
1723             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
1724             leaderShardManager.tell(new ShardLeaderStateChanged(memberId2, memberId2,
1725                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardLeaderActor);
1726             leaderShardManager.tell(new RoleChangeNotification(memberId2,
1727                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardLeaderActor);
1728
1729             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
1730             newReplicaShardManager.tell(new ShardLeaderStateChanged(memberId1, memberId2,
1731                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardActor);
1732             newReplicaShardManager.tell(new RoleChangeNotification(memberId1,
1733                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
1734
1735             newReplicaShardManager.underlyingActor().waitForMemberUp();
1736             leaderShardManager.underlyingActor().waitForMemberUp();
1737
1738             //construct a mock response message
1739             RemoveServerReply response = new RemoveServerReply(ServerChangeStatus.OK, memberId2);
1740             mockShardLeaderActor.underlyingActor().updateResponse(response);
1741             newReplicaShardManager.tell(new RemoveShardReplica("default", "member-1"), getRef());
1742             RemoveServer removeServer = MessageCollectorActor.expectFirstMatching(mockShardLeaderActor,
1743                     RemoveServer.class);
1744             String removeServerId = new ShardIdentifier("default", "member-1", shardMrgIDSuffix).toString();
1745             assertEquals("RemoveServer serverId", removeServerId, removeServer.getServerId());
1746             expectMsgClass(duration("5 seconds"), Status.Success.class);
1747         }};
1748
1749     }
1750
1751     @Test
1752     public void testRemoveShardReplicaWhenAnotherRemoveShardReplicaAlreadyInProgress() throws Exception {
1753         testServerChangeWhenAlreadyInProgress("astronauts", new RemoveShardReplica("astronauts", "member-2"),
1754                 RemoveServer.class, new RemoveShardReplica("astronauts", "member-3"));
1755     }
1756
1757     @Test
1758     public void testRemoveShardReplicaWhenAddShardReplicaAlreadyInProgress() throws Exception {
1759         testServerChangeWhenAlreadyInProgress("astronauts", new AddShardReplica("astronauts"),
1760                 AddServer.class, new RemoveShardReplica("astronauts", "member-2"));
1761     }
1762
1763
1764     public void testServerChangeWhenAlreadyInProgress(final String shardName, final Object firstServerChange,
1765                                                       final Class firstForwardedServerChangeClass,
1766                                                       final Object secondServerChange) throws Exception {
1767         new JavaTestKit(getSystem()) {{
1768             JavaTestKit mockShardLeaderKit = new JavaTestKit(getSystem());
1769             JavaTestKit secondRequestKit = new JavaTestKit(getSystem());
1770
1771             MockConfiguration mockConfig =
1772                     new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1773                             put(shardName, Arrays.asList("member-2")).build());
1774
1775             final TestActorRef<TestShardManager> shardManager = TestActorRef.create(getSystem(),
1776                     newTestShardMgrBuilder().configuration(mockConfig).shardActor(mockShardActor).cluster(
1777                             new MockClusterWrapper()).props(),
1778                     shardMgrID);
1779
1780             shardManager.underlyingActor().setMessageInterceptor(newFindPrimaryInterceptor(mockShardLeaderKit.getRef()));
1781
1782             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1783
1784             shardManager.tell(firstServerChange, getRef());
1785
1786             mockShardLeaderKit.expectMsgClass(firstForwardedServerChangeClass);
1787
1788             shardManager.tell(secondServerChange, secondRequestKit.getRef());
1789
1790             secondRequestKit.expectMsgClass(duration("5 seconds"), Failure.class);
1791         }};
1792     }
1793
1794     @Test
1795     public void testServerRemovedShardActorNotRunning() throws Exception {
1796         LOG.info("testServerRemovedShardActorNotRunning starting");
1797         new JavaTestKit(getSystem()) {{
1798             MockConfiguration mockConfig =
1799                     new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1800                             put("default", Arrays.asList("member-1", "member-2")).
1801                             put("astronauts", Arrays.asList("member-2")).
1802                             put("people", Arrays.asList("member-1", "member-2")).build());
1803
1804             TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(newShardMgrProps(mockConfig));
1805
1806             shardManager.underlyingActor().waitForRecoveryComplete();
1807             shardManager.tell(new FindLocalShard("people", false), getRef());
1808             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
1809
1810             shardManager.tell(new FindLocalShard("default", false), getRef());
1811             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
1812
1813             // Removed the default shard replica from member-1
1814             ShardIdentifier.Builder builder = new ShardIdentifier.Builder();
1815             ShardIdentifier shardId = builder.shardName("default").memberName("member-1").type(shardMrgIDSuffix).build();
1816             shardManager.tell(new ServerRemoved(shardId.toString()), getRef());
1817
1818             shardManager.underlyingActor().verifySnapshotPersisted(Sets.newHashSet("people"));
1819         }};
1820
1821         LOG.info("testServerRemovedShardActorNotRunning ending");
1822     }
1823
1824     @Test
1825     public void testServerRemovedShardActorRunning() throws Exception {
1826         LOG.info("testServerRemovedShardActorRunning starting");
1827         new JavaTestKit(getSystem()) {{
1828             MockConfiguration mockConfig =
1829                     new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1830                             put("default", Arrays.asList("member-1", "member-2")).
1831                             put("astronauts", Arrays.asList("member-2")).
1832                             put("people", Arrays.asList("member-1", "member-2")).build());
1833
1834             String shardId = ShardIdentifier.builder().shardName("default").memberName("member-1").
1835                     type(shardMrgIDSuffix).build().toString();
1836             TestActorRef<MessageCollectorActor> shard = actorFactory.createTestActor(
1837                     MessageCollectorActor.props(), shardId);
1838
1839             TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(
1840                     newTestShardMgrBuilder(mockConfig).addShardActor("default", shard).props());
1841
1842             watch(shard);
1843
1844             shardManager.underlyingActor().waitForRecoveryComplete();
1845
1846             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1847             shardManager.tell(new ActorInitialized(), shard);
1848
1849             waitForShardInitialized(shardManager, "people", this);
1850             waitForShardInitialized(shardManager, "default", this);
1851
1852             // Removed the default shard replica from member-1
1853             shardManager.tell(new ServerRemoved(shardId), getRef());
1854
1855             shardManager.underlyingActor().verifySnapshotPersisted(Sets.newHashSet("people"));
1856
1857             expectMsgClass(duration("5 seconds"), Terminated.class);
1858         }};
1859
1860         LOG.info("testServerRemovedShardActorRunning ending");
1861     }
1862
1863
1864     @Test
1865     public void testShardPersistenceWithRestoredData() throws Exception {
1866         LOG.info("testShardPersistenceWithRestoredData starting");
1867         new JavaTestKit(getSystem()) {{
1868             MockConfiguration mockConfig =
1869                 new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1870                    put("default", Arrays.asList("member-1", "member-2")).
1871                    put("astronauts", Arrays.asList("member-2")).
1872                    put("people", Arrays.asList("member-1", "member-2")).build());
1873             String[] restoredShards = {"default", "astronauts"};
1874             ShardManagerSnapshot snapshot = new ShardManagerSnapshot(Arrays.asList(restoredShards));
1875             InMemorySnapshotStore.addSnapshot("shard-manager-" + shardMrgIDSuffix, snapshot);
1876
1877             //create shardManager to come up with restored data
1878             TestActorRef<TestShardManager> newRestoredShardManager = actorFactory.createTestActor(
1879                     newShardMgrProps(mockConfig));
1880
1881             newRestoredShardManager.underlyingActor().waitForRecoveryComplete();
1882
1883             newRestoredShardManager.tell(new FindLocalShard("people", false), getRef());
1884             LocalShardNotFound notFound = expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
1885             assertEquals("for uninitialized shard", "people", notFound.getShardName());
1886
1887             //Verify a local shard is created for the restored shards,
1888             //although we expect a NotInitializedException for the shards as the actor initialization
1889             //message is not sent for them
1890             newRestoredShardManager.tell(new FindLocalShard("default", false), getRef());
1891             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
1892
1893             newRestoredShardManager.tell(new FindLocalShard("astronauts", false), getRef());
1894             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
1895         }};
1896
1897         LOG.info("testShardPersistenceWithRestoredData ending");
1898     }
1899
1900
1901     private static class TestShardManager extends ShardManager {
1902         private final CountDownLatch recoveryComplete = new CountDownLatch(1);
1903         private final CountDownLatch snapshotPersist = new CountDownLatch(1);
1904         private ShardManagerSnapshot snapshot;
1905         private final Map<String, ActorRef> shardActors;
1906         private final ActorRef shardActor;
1907         private CountDownLatch findPrimaryMessageReceived = new CountDownLatch(1);
1908         private CountDownLatch memberUpReceived = new CountDownLatch(1);
1909         private CountDownLatch memberRemovedReceived = new CountDownLatch(1);
1910         private CountDownLatch memberUnreachableReceived = new CountDownLatch(1);
1911         private CountDownLatch memberReachableReceived = new CountDownLatch(1);
1912         private volatile MessageInterceptor messageInterceptor;
1913
1914         private TestShardManager(Builder builder) {
1915             super(builder);
1916             shardActor = builder.shardActor;
1917             shardActors = builder.shardActors;
1918         }
1919
1920         @Override
1921         public void handleRecover(Object message) throws Exception {
1922             try {
1923                 super.handleRecover(message);
1924             } finally {
1925                 if(message instanceof RecoveryCompleted) {
1926                     recoveryComplete.countDown();
1927                 }
1928             }
1929         }
1930
1931         @Override
1932         public void handleCommand(Object message) throws Exception {
1933             try{
1934                 if(messageInterceptor != null && messageInterceptor.canIntercept(message)) {
1935                     getSender().tell(messageInterceptor.apply(message), getSelf());
1936                 } else {
1937                     super.handleCommand(message);
1938                 }
1939             } finally {
1940                 if(message instanceof FindPrimary) {
1941                     findPrimaryMessageReceived.countDown();
1942                 } else if(message instanceof ClusterEvent.MemberUp) {
1943                     String role = ((ClusterEvent.MemberUp)message).member().roles().head();
1944                     if(!getCluster().getCurrentMemberName().equals(role)) {
1945                         memberUpReceived.countDown();
1946                     }
1947                 } else if(message instanceof ClusterEvent.MemberRemoved) {
1948                     String role = ((ClusterEvent.MemberRemoved)message).member().roles().head();
1949                     if(!getCluster().getCurrentMemberName().equals(role)) {
1950                         memberRemovedReceived.countDown();
1951                     }
1952                 } else if(message instanceof ClusterEvent.UnreachableMember) {
1953                     String role = ((ClusterEvent.UnreachableMember)message).member().roles().head();
1954                     if(!getCluster().getCurrentMemberName().equals(role)) {
1955                         memberUnreachableReceived.countDown();
1956                     }
1957                 } else if(message instanceof ClusterEvent.ReachableMember) {
1958                     String role = ((ClusterEvent.ReachableMember)message).member().roles().head();
1959                     if(!getCluster().getCurrentMemberName().equals(role)) {
1960                         memberReachableReceived.countDown();
1961                     }
1962                 }
1963             }
1964         }
1965
1966         void setMessageInterceptor(MessageInterceptor messageInterceptor) {
1967             this.messageInterceptor = messageInterceptor;
1968         }
1969
1970         void waitForRecoveryComplete() {
1971             assertEquals("Recovery complete", true,
1972                     Uninterruptibles.awaitUninterruptibly(recoveryComplete, 5, TimeUnit.SECONDS));
1973         }
1974
1975         void waitForMemberUp() {
1976             assertEquals("MemberUp received", true,
1977                     Uninterruptibles.awaitUninterruptibly(memberUpReceived, 5, TimeUnit.SECONDS));
1978             memberUpReceived = new CountDownLatch(1);
1979         }
1980
1981         void waitForMemberRemoved() {
1982             assertEquals("MemberRemoved received", true,
1983                     Uninterruptibles.awaitUninterruptibly(memberRemovedReceived, 5, TimeUnit.SECONDS));
1984             memberRemovedReceived = new CountDownLatch(1);
1985         }
1986
1987         void waitForUnreachableMember() {
1988             assertEquals("UnreachableMember received", true,
1989                 Uninterruptibles.awaitUninterruptibly(memberUnreachableReceived, 5, TimeUnit.SECONDS
1990                 ));
1991             memberUnreachableReceived = new CountDownLatch(1);
1992         }
1993
1994         void waitForReachableMember() {
1995             assertEquals("ReachableMember received", true,
1996                 Uninterruptibles.awaitUninterruptibly(memberReachableReceived, 5, TimeUnit.SECONDS));
1997             memberReachableReceived = new CountDownLatch(1);
1998         }
1999
2000         void verifyFindPrimary() {
2001             assertEquals("FindPrimary received", true,
2002                     Uninterruptibles.awaitUninterruptibly(findPrimaryMessageReceived, 5, TimeUnit.SECONDS));
2003             findPrimaryMessageReceived = new CountDownLatch(1);
2004         }
2005
2006         public static Builder builder(DatastoreContext.Builder datastoreContextBuilder) {
2007             return new Builder(datastoreContextBuilder);
2008         }
2009
2010         private static class Builder extends AbstractGenericBuilder<Builder, TestShardManager> {
2011             private ActorRef shardActor;
2012             private final Map<String, ActorRef> shardActors = new HashMap<>();
2013
2014             Builder(DatastoreContext.Builder datastoreContextBuilder) {
2015                 super(TestShardManager.class);
2016                 datastoreContextFactory(newDatastoreContextFactory(datastoreContextBuilder.build()));
2017             }
2018
2019             Builder shardActor(ActorRef shardActor) {
2020                 this.shardActor = shardActor;
2021                 return this;
2022             }
2023
2024             Builder addShardActor(String shardName, ActorRef actorRef){
2025                 shardActors.put(shardName, actorRef);
2026                 return this;
2027             }
2028         }
2029
2030         @Override
2031         public void saveSnapshot(Object obj) {
2032             snapshot = (ShardManagerSnapshot) obj;
2033             snapshotPersist.countDown();
2034             super.saveSnapshot(obj);
2035         }
2036
2037         void verifySnapshotPersisted(Set<String> shardList) {
2038             assertEquals("saveSnapshot invoked", true,
2039                     Uninterruptibles.awaitUninterruptibly(snapshotPersist, 5, TimeUnit.SECONDS));
2040             assertEquals("Shard Persisted", shardList, Sets.newHashSet(snapshot.getShardList()));
2041         }
2042
2043         @Override
2044         protected ActorRef newShardActor(SchemaContext schemaContext, ShardInformation info) {
2045             if(shardActors.get(info.getShardName()) != null){
2046                 return shardActors.get(info.getShardName());
2047             }
2048
2049             if(shardActor != null) {
2050                 return shardActor;
2051             }
2052
2053             return super.newShardActor(schemaContext, info);
2054         }
2055     }
2056
2057     private static abstract class AbstractGenericBuilder<T extends AbstractGenericBuilder<T, ?>, C extends ShardManager>
2058                                                      extends ShardManager.AbstractBuilder<T> {
2059         private final Class<C> shardManagerClass;
2060
2061         AbstractGenericBuilder(Class<C> shardManagerClass) {
2062             this.shardManagerClass = shardManagerClass;
2063             cluster(new MockClusterWrapper()).configuration(new MockConfiguration()).
2064                     waitTillReadyCountdownLatch(ready).primaryShardInfoCache(new PrimaryShardInfoFutureCache());
2065         }
2066
2067         @Override
2068         public Props props() {
2069             verify();
2070             return Props.create(shardManagerClass, this);
2071         }
2072     }
2073
2074     private static class GenericBuilder<C extends ShardManager> extends AbstractGenericBuilder<GenericBuilder<C>, C> {
2075         GenericBuilder(Class<C> shardManagerClass) {
2076             super(shardManagerClass);
2077         }
2078     }
2079
2080     private static class DelegatingShardManagerCreator implements Creator<ShardManager> {
2081         private static final long serialVersionUID = 1L;
2082         private final Creator<ShardManager> delegate;
2083
2084         public DelegatingShardManagerCreator(Creator<ShardManager> delegate) {
2085             this.delegate = delegate;
2086         }
2087
2088         @Override
2089         public ShardManager create() throws Exception {
2090             return delegate.create();
2091         }
2092     }
2093
2094     interface MessageInterceptor extends Function<Object, Object> {
2095         boolean canIntercept(Object message);
2096     }
2097
2098     private MessageInterceptor newFindPrimaryInterceptor(final ActorRef primaryActor) {
2099         return new MessageInterceptor(){
2100             @Override
2101             public Object apply(Object message) {
2102                 return new RemotePrimaryShardFound(Serialization.serializedActorPath(primaryActor), (short) 1);
2103             }
2104
2105             @Override
2106             public boolean canIntercept(Object message) {
2107                 return message instanceof FindPrimary;
2108             }
2109         };
2110     }
2111
2112     private static class MockRespondActor extends MessageCollectorActor {
2113         static final String CLEAR_RESPONSE = "clear-response";
2114         static final org.slf4j.Logger LOG = LoggerFactory.getLogger(MockRespondActor.class);
2115
2116         private volatile Object responseMsg;
2117
2118         @SuppressWarnings("unused")
2119         public MockRespondActor() {
2120         }
2121
2122         @SuppressWarnings("unused")
2123         public MockRespondActor(Object responseMsg) {
2124             this.responseMsg = responseMsg;
2125         }
2126
2127         public void updateResponse(Object response) {
2128             responseMsg = response;
2129         }
2130
2131         @Override
2132         public void onReceive(Object message) throws Exception {
2133             if(!"get-all-messages".equals(message)) {
2134                 LOG.debug("Received message : {}", message);
2135             }
2136             super.onReceive(message);
2137             if (message instanceof AddServer && responseMsg != null) {
2138                 getSender().tell(responseMsg, getSelf());
2139             } else if(message instanceof RemoveServer && responseMsg != null){
2140                 getSender().tell(responseMsg, getSelf());
2141             } else if(message.equals(CLEAR_RESPONSE)) {
2142                 responseMsg = null;
2143             }
2144         }
2145     }
2146 }