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             // Test FindPrimary wait succeeds after reachable member event.
 798
 799             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 800                     createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 801             shardManager1.underlyingActor().waitForUnreachableMember();
 802
 803             shardManager1.tell(new FindPrimary("default", true), getRef());
 804
 805             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 806                     createReachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 807
 808             RemotePrimaryShardFound found2 = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 809             String path2 = found2.getPrimaryPath();
 810             assertTrue("Unexpected primary path " + path2, path2.contains("member-2-shard-default-config"));
 811         }};
 812
 813         LOG.info("testShardAvailabilityOnChangeOfMemberReachability ending");
 814     }
 815
 816     @Test
 817     public void testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange() throws Exception {
 818         LOG.info("testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange starting");
 819         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 820
 821         // Create an ActorSystem ShardManager actor for member-1.
 822
 823         final ActorSystem system1 = newActorSystem("Member1");
 824         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 825
 826         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 827
 828         final PrimaryShardInfoFutureCache primaryShardInfoCache = new PrimaryShardInfoFutureCache();
 829         final TestActorRef<TestShardManager> shardManager1 = TestActorRef.create(system1,
 830                 newTestShardMgrBuilder().shardActor(mockShardActor1).cluster(
 831                         new ClusterWrapperImpl(system1)).primaryShardInfoCache(primaryShardInfoCache).props(),
 832                 shardManagerID);
 833
 834         // Create an ActorSystem ShardManager actor for member-2.
 835
 836         final ActorSystem system2 = newActorSystem("Member2");
 837
 838         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 839
 840         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 841
 842         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 843             put("default", Arrays.asList("member-1", "member-2")).build());
 844
 845         final TestActorRef<TestShardManager> shardManager2 = TestActorRef.create(system2,
 846                 newTestShardMgrBuilder(mockConfig2).shardActor(mockShardActor2).cluster(
 847                         new ClusterWrapperImpl(system2)).props(), shardManagerID);
 848
 849         new JavaTestKit(system1) {{
 850             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 851             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 852             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 853             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 854
 855             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 856             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 857             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 858                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 859             shardManager1.tell(new RoleChangeNotification(memberId1,
 860                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 861             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class)),
 862                     DataStoreVersions.CURRENT_VERSION),
 863                 mockShardActor2);
 864             shardManager2.tell(new RoleChangeNotification(memberId2,
 865                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 866             shardManager1.underlyingActor().waitForMemberUp();
 867
 868             shardManager1.tell(new FindPrimary("default", true), getRef());
 869
 870             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 871             String path = found.getPrimaryPath();
 872             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 873
 874             primaryShardInfoCache.putSuccessful("default", new PrimaryShardInfo(system1.actorSelection(
 875                     mockShardActor1.path()), DataStoreVersions.CURRENT_VERSION, Optional.<DataTree>absent()));
 876
 877             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 878                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 879
 880             shardManager1.underlyingActor().waitForUnreachableMember();
 881
 882             shardManager1.tell(new FindPrimary("default", true), getRef());
 883
 884             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 885
 886             assertNull("Expected primaryShardInfoCache entry removed", primaryShardInfoCache.getIfPresent("default"));
 887
 888             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId1, Optional.of(mock(DataTree.class)),
 889                     DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 890             shardManager1.tell(new RoleChangeNotification(memberId1,
 891                 RaftState.Follower.name(), RaftState.Leader.name()), mockShardActor1);
 892
 893             shardManager1.tell(new FindPrimary("default", true), getRef());
 894
 895             LocalPrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 896             String path1 = found1.getPrimaryPath();
 897             assertTrue("Unexpected primary path " + path1, path1.contains("member-1-shard-default-config"));
 898
 899         }};
 900
 901         LOG.info("testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange ending");
 902     }
 903
 904
 905     @Test
 906     public void testOnReceiveFindLocalShardForNonExistentShard() throws Exception {
 907         new JavaTestKit(getSystem()) {{
 908             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 909
 910             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 911
 912             shardManager.tell(new FindLocalShard("non-existent", false), getRef());
 913
 914             LocalShardNotFound notFound = expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
 915
 916             assertEquals("getShardName", "non-existent", notFound.getShardName());
 917         }};
 918     }
 919
 920     @Test
 921     public void testOnReceiveFindLocalShardForExistentShard() throws Exception {
 922         new JavaTestKit(getSystem()) {{
 923             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 924
 925             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 926             shardManager.tell(new ActorInitialized(), mockShardActor);
 927
 928             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 929
 930             LocalShardFound found = expectMsgClass(duration("5 seconds"), LocalShardFound.class);
 931
 932             assertTrue("Found path contains " + found.getPath().path().toString(),
 933                     found.getPath().path().toString().contains("member-1-shard-default-config"));
 934         }};
 935     }
 936
 937     @Test
 938     public void testOnReceiveFindLocalShardForNotInitializedShard() throws Exception {
 939         new JavaTestKit(getSystem()) {{
 940             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 941
 942             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 943
 944             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 945         }};
 946     }
 947
 948     @Test
 949     public void testOnReceiveFindLocalShardWaitForShardInitialized() throws Exception {
 950         LOG.info("testOnReceiveFindLocalShardWaitForShardInitialized starting");
 951         new JavaTestKit(getSystem()) {{
 952             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
 953
 954             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 955
 956             // We're passing waitUntilInitialized = true to FindLocalShard so the response should be
 957             // delayed until we send ActorInitialized.
 958             Future<Object> future = Patterns.ask(shardManager, new FindLocalShard(Shard.DEFAULT_NAME, true),
 959                     new Timeout(5, TimeUnit.SECONDS));
 960
 961             shardManager.tell(new ActorInitialized(), mockShardActor);
 962
 963             Object resp = Await.result(future, duration("5 seconds"));
 964             assertTrue("Expected: LocalShardFound, Actual: " + resp, resp instanceof LocalShardFound);
 965         }};
 966
 967         LOG.info("testOnReceiveFindLocalShardWaitForShardInitialized starting");
 968     }
 969
 970     @Test
 971     public void testOnRecoveryJournalIsCleaned() {
 972         String persistenceID = "shard-manager-" + shardMrgIDSuffix;
 973         InMemoryJournal.addEntry(persistenceID, 1L, new ShardManager.SchemaContextModules(
 974                 ImmutableSet.of("foo")));
 975         InMemoryJournal.addEntry(persistenceID, 2L, new ShardManager.SchemaContextModules(
 976                 ImmutableSet.of("bar")));
 977         InMemoryJournal.addDeleteMessagesCompleteLatch(persistenceID);
 978
 979         TestShardManager shardManager = newTestShardManager();
 980
 981         InMemoryJournal.waitForDeleteMessagesComplete(persistenceID);
 982
 983         // Journal entries up to the last one should've been deleted
 984         Map<Long, Object> journal = InMemoryJournal.get(persistenceID);
 985         synchronized (journal) {
 986             assertEquals("Journal size", 0, journal.size());
 987         }
 988     }
 989
 990     @Test
 991     public void testRoleChangeNotificationAndShardLeaderStateChangedReleaseReady() throws Exception {
 992         TestShardManager shardManager = newTestShardManager();
 993
 994         String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 995         shardManager.onReceiveCommand(new RoleChangeNotification(
 996                 memberId, RaftState.Candidate.name(), RaftState.Leader.name()));
 997
 998         verify(ready, never()).countDown();
 999
1000         shardManager.onReceiveCommand(new ShardLeaderStateChanged(memberId, memberId,
1001                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION));
1002
1003         verify(ready, times(1)).countDown();
1004     }
1005
1006     @Test
1007     public void testRoleChangeNotificationToFollowerWithShardLeaderStateChangedReleaseReady() throws Exception {
1008         new JavaTestKit(getSystem()) {{
1009             TestShardManager shardManager = newTestShardManager();
1010
1011             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
1012             shardManager.onReceiveCommand(new RoleChangeNotification(
1013                     memberId, null, RaftState.Follower.name()));
1014
1015             verify(ready, never()).countDown();
1016
1017             shardManager.onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
1018
1019             shardManager.onReceiveCommand(new ShardLeaderStateChanged(memberId,
1020                     "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class)),
1021                     DataStoreVersions.CURRENT_VERSION));
1022
1023             verify(ready, times(1)).countDown();
1024         }};
1025     }
1026
1027     @Test
1028     public void testReadyCountDownForMemberUpAfterLeaderStateChanged() throws Exception {
1029         new JavaTestKit(getSystem()) {{
1030             TestShardManager shardManager = newTestShardManager();
1031
1032             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
1033             shardManager.onReceiveCommand(new RoleChangeNotification(memberId, null, RaftState.Follower.name()));
1034
1035             verify(ready, never()).countDown();
1036
1037             shardManager.onReceiveCommand(new ShardLeaderStateChanged(memberId,
1038                     "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class)),
1039                     DataStoreVersions.CURRENT_VERSION));
1040
1041             shardManager.onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
1042
1043             verify(ready, times(1)).countDown();
1044         }};
1045     }
1046
1047     @Test
1048     public void testRoleChangeNotificationDoNothingForUnknownShard() throws Exception {
1049         TestShardManager shardManager = newTestShardManager();
1050
1051         shardManager.onReceiveCommand(new RoleChangeNotification(
1052                 "unknown", RaftState.Candidate.name(), RaftState.Leader.name()));
1053
1054         verify(ready, never()).countDown();
1055     }
1056
1057     @Test
1058     public void testByDefaultSyncStatusIsFalse() throws Exception{
1059         TestShardManager shardManager = newTestShardManager();
1060
1061         assertEquals(false, shardManager.getMBean().getSyncStatus());
1062     }
1063
1064     @Test
1065     public void testWhenShardIsLeaderSyncStatusIsTrue() throws Exception{
1066         TestShardManager shardManager = newTestShardManager();
1067
1068         shardManager.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
1069                 RaftState.Follower.name(), RaftState.Leader.name()));
1070
1071         assertEquals(true, shardManager.getMBean().getSyncStatus());
1072     }
1073
1074     @Test
1075     public void testWhenShardIsCandidateSyncStatusIsFalse() throws Exception{
1076         TestShardManager shardManager = newTestShardManager();
1077
1078         String shardId = "member-1-shard-default-" + shardMrgIDSuffix;
1079         shardManager.onReceiveCommand(new RoleChangeNotification(shardId,
1080                 RaftState.Follower.name(), RaftState.Candidate.name()));
1081
1082         assertEquals(false, shardManager.getMBean().getSyncStatus());
1083
1084         // Send a FollowerInitialSyncStatus with status = true for the replica whose current state is candidate
1085         shardManager.onReceiveCommand(new FollowerInitialSyncUpStatus(
1086                 true, shardId));
1087
1088         assertEquals(false, shardManager.getMBean().getSyncStatus());
1089     }
1090
1091     @Test
1092     public void testWhenShardIsFollowerSyncStatusDependsOnFollowerInitialSyncStatus() throws Exception{
1093         TestShardManager shardManager = newTestShardManager();
1094
1095         String shardId = "member-1-shard-default-" + shardMrgIDSuffix;
1096         shardManager.onReceiveCommand(new RoleChangeNotification(shardId,
1097                 RaftState.Candidate.name(), RaftState.Follower.name()));
1098
1099         // Initially will be false
1100         assertEquals(false, shardManager.getMBean().getSyncStatus());
1101
1102         // Send status true will make sync status true
1103         shardManager.onReceiveCommand(new FollowerInitialSyncUpStatus(true, shardId));
1104
1105         assertEquals(true, shardManager.getMBean().getSyncStatus());
1106
1107         // Send status false will make sync status false
1108         shardManager.onReceiveCommand(new FollowerInitialSyncUpStatus(false, shardId));
1109
1110         assertEquals(false, shardManager.getMBean().getSyncStatus());
1111
1112     }
1113
1114     @Test
1115     public void testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards() throws Exception{
1116         LOG.info("testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards starting");
1117         TestShardManager shardManager = newTestShardManager(newShardMgrProps(new MockConfiguration() {
1118             @Override
1119             public List<String> getMemberShardNames(String memberName) {
1120                 return Arrays.asList("default", "astronauts");
1121             }
1122         }));
1123
1124         // Initially will be false
1125         assertEquals(false, shardManager.getMBean().getSyncStatus());
1126
1127         // Make default shard leader
1128         String defaultShardId = "member-1-shard-default-" + shardMrgIDSuffix;
1129         shardManager.onReceiveCommand(new RoleChangeNotification(defaultShardId,
1130                 RaftState.Follower.name(), RaftState.Leader.name()));
1131
1132         // default = Leader, astronauts is unknown so sync status remains false
1133         assertEquals(false, shardManager.getMBean().getSyncStatus());
1134
1135         // Make astronauts shard leader as well
1136         String astronautsShardId = "member-1-shard-astronauts-" + shardMrgIDSuffix;
1137         shardManager.onReceiveCommand(new RoleChangeNotification(astronautsShardId,
1138                 RaftState.Follower.name(), RaftState.Leader.name()));
1139
1140         // Now sync status should be true
1141         assertEquals(true, shardManager.getMBean().getSyncStatus());
1142
1143         // Make astronauts a Follower
1144         shardManager.onReceiveCommand(new RoleChangeNotification(astronautsShardId,
1145                 RaftState.Leader.name(), RaftState.Follower.name()));
1146
1147         // Sync status is not true
1148         assertEquals(false, shardManager.getMBean().getSyncStatus());
1149
1150         // Make the astronauts follower sync status true
1151         shardManager.onReceiveCommand(new FollowerInitialSyncUpStatus(true, astronautsShardId));
1152
1153         // Sync status is now true
1154         assertEquals(true, shardManager.getMBean().getSyncStatus());
1155
1156         LOG.info("testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards ending");
1157     }
1158
1159     @Test
1160     public void testOnReceiveSwitchShardBehavior() throws Exception {
1161         new JavaTestKit(getSystem()) {{
1162             final ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
1163
1164             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1165             shardManager.tell(new ActorInitialized(), mockShardActor);
1166
1167             shardManager.tell(new SwitchShardBehavior(mockShardName, "Leader", 1000), getRef());
1168
1169             SwitchBehavior switchBehavior = MessageCollectorActor.expectFirstMatching(mockShardActor, SwitchBehavior.class);
1170
1171             assertEquals(RaftState.Leader, switchBehavior.getNewState());
1172             assertEquals(1000, switchBehavior.getNewTerm());
1173         }};
1174     }
1175
1176     @Test
1177     public void testOnCreateShard() {
1178         LOG.info("testOnCreateShard starting");
1179         new JavaTestKit(getSystem()) {{
1180             datastoreContextBuilder.shardInitializationTimeout(1, TimeUnit.MINUTES).persistent(true);
1181
1182             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1183                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1184
1185             SchemaContext schemaContext = TestModel.createTestContext();
1186             shardManager.tell(new UpdateSchemaContext(schemaContext), ActorRef.noSender());
1187
1188             DatastoreContext datastoreContext = DatastoreContext.newBuilder().shardElectionTimeoutFactor(100).
1189                     persistent(false).build();
1190             Shard.Builder shardBuilder = Shard.builder();
1191
1192             ModuleShardConfiguration config = new ModuleShardConfiguration(URI.create("foo-ns"), "foo-module",
1193                     "foo", null, Arrays.asList("member-1", "member-5", "member-6"));
1194             shardManager.tell(new CreateShard(config, shardBuilder, datastoreContext), getRef());
1195
1196             expectMsgClass(duration("5 seconds"), Success.class);
1197
1198             shardManager.tell(new FindLocalShard("foo", true), getRef());
1199
1200             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1201
1202             assertEquals("isRecoveryApplicable", false, shardBuilder.getDatastoreContext().isPersistent());
1203             assertTrue("Epxected ShardPeerAddressResolver", shardBuilder.getDatastoreContext().getShardRaftConfig().
1204                     getPeerAddressResolver() instanceof ShardPeerAddressResolver);
1205             assertEquals("peerMembers", Sets.newHashSet(new ShardIdentifier("foo", "member-5", shardMrgIDSuffix).toString(),
1206                     new ShardIdentifier("foo", "member-6", shardMrgIDSuffix).toString()),
1207                     shardBuilder.getPeerAddresses().keySet());
1208             assertEquals("ShardIdentifier", new ShardIdentifier("foo", "member-1", shardMrgIDSuffix),
1209                     shardBuilder.getId());
1210             assertSame("schemaContext", schemaContext, shardBuilder.getSchemaContext());
1211
1212             // Send CreateShard with same name - should return Success with a message.
1213
1214             shardManager.tell(new CreateShard(config, shardBuilder, null), getRef());
1215
1216             Success success = expectMsgClass(duration("5 seconds"), Success.class);
1217             assertNotNull("Success status is null", success.status());
1218         }};
1219
1220         LOG.info("testOnCreateShard ending");
1221     }
1222
1223     @Test
1224     public void testOnCreateShardWithLocalMemberNotInShardConfig() {
1225         LOG.info("testOnCreateShardWithLocalMemberNotInShardConfig starting");
1226         new JavaTestKit(getSystem()) {{
1227             datastoreContextBuilder.shardInitializationTimeout(1, TimeUnit.MINUTES).persistent(true);
1228
1229             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1230                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1231
1232             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), ActorRef.noSender());
1233
1234             Shard.Builder shardBuilder = Shard.builder();
1235             ModuleShardConfiguration config = new ModuleShardConfiguration(URI.create("foo-ns"), "foo-module",
1236                     "foo", null, Arrays.asList("member-5", "member-6"));
1237
1238             shardManager.tell(new CreateShard(config, shardBuilder, null), getRef());
1239             expectMsgClass(duration("5 seconds"), Success.class);
1240
1241             shardManager.tell(new FindLocalShard("foo", true), getRef());
1242             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1243
1244             assertEquals("peerMembers size", 0, shardBuilder.getPeerAddresses().size());
1245             assertEquals("schemaContext", DisableElectionsRaftPolicy.class.getName(),
1246                     shardBuilder.getDatastoreContext().getShardRaftConfig().getCustomRaftPolicyImplementationClass());
1247         }};
1248
1249         LOG.info("testOnCreateShardWithLocalMemberNotInShardConfig ending");
1250     }
1251
1252     @Test
1253     public void testOnCreateShardWithNoInitialSchemaContext() {
1254         LOG.info("testOnCreateShardWithNoInitialSchemaContext starting");
1255         new JavaTestKit(getSystem()) {{
1256             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1257                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1258
1259             Shard.Builder shardBuilder = Shard.builder();
1260
1261             ModuleShardConfiguration config = new ModuleShardConfiguration(URI.create("foo-ns"), "foo-module",
1262                     "foo", null, Arrays.asList("member-1"));
1263             shardManager.tell(new CreateShard(config, shardBuilder, null), getRef());
1264
1265             expectMsgClass(duration("5 seconds"), Success.class);
1266
1267             SchemaContext schemaContext = TestModel.createTestContext();
1268             shardManager.tell(new UpdateSchemaContext(schemaContext), ActorRef.noSender());
1269
1270             shardManager.tell(new FindLocalShard("foo", true), getRef());
1271
1272             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1273
1274             assertSame("schemaContext", schemaContext, shardBuilder.getSchemaContext());
1275             assertNotNull("schemaContext is null", shardBuilder.getDatastoreContext());
1276         }};
1277
1278         LOG.info("testOnCreateShardWithNoInitialSchemaContext ending");
1279     }
1280
1281     @Test
1282     public void testGetSnapshot() throws Throwable {
1283         LOG.info("testGetSnapshot starting");
1284         JavaTestKit kit = new JavaTestKit(getSystem());
1285
1286         MockConfiguration mockConfig = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1287                    put("shard1", Arrays.asList("member-1")).
1288                    put("shard2", Arrays.asList("member-1")).
1289                    put("astronauts", Collections.<String>emptyList()).build());
1290
1291         TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(newShardMgrProps(mockConfig).
1292                 withDispatcher(Dispatchers.DefaultDispatcherId()));
1293
1294         shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
1295         Failure failure = kit.expectMsgClass(Failure.class);
1296         assertEquals("Failure cause type", IllegalStateException.class, failure.cause().getClass());
1297
1298         shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), ActorRef.noSender());
1299
1300         waitForShardInitialized(shardManager, "shard1", kit);
1301         waitForShardInitialized(shardManager, "shard2", kit);
1302
1303         shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
1304
1305         DatastoreSnapshot datastoreSnapshot = expectMsgClassOrFailure(DatastoreSnapshot.class, kit, "GetSnapshot");
1306
1307         assertEquals("getType", shardMrgIDSuffix, datastoreSnapshot.getType());
1308         assertNull("Expected null ShardManagerSnapshot", datastoreSnapshot.getShardManagerSnapshot());
1309
1310         Function<ShardSnapshot, String> shardNameTransformer = new Function<ShardSnapshot, String>() {
1311             @Override
1312             public String apply(ShardSnapshot s) {
1313                 return s.getName();
1314             }
1315         };
1316
1317         assertEquals("Shard names", Sets.newHashSet("shard1", "shard2"), Sets.newHashSet(
1318                 Lists.transform(datastoreSnapshot.getShardSnapshots(), shardNameTransformer)));
1319
1320         // Add a new replica
1321
1322         JavaTestKit mockShardLeaderKit = new JavaTestKit(getSystem());
1323
1324         TestShardManager shardManagerInstance = shardManager.underlyingActor();
1325         shardManagerInstance.setMessageInterceptor(newFindPrimaryInterceptor(mockShardLeaderKit.getRef()));
1326
1327         shardManager.tell(new AddShardReplica("astronauts"), kit.getRef());
1328         mockShardLeaderKit.expectMsgClass(AddServer.class);
1329         mockShardLeaderKit.reply(new AddServerReply(ServerChangeStatus.OK, ""));
1330         kit.expectMsgClass(Status.Success.class);
1331         waitForShardInitialized(shardManager, "astronauts", kit);
1332
1333         // Send another GetSnapshot and verify
1334
1335         shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
1336         datastoreSnapshot = expectMsgClassOrFailure(DatastoreSnapshot.class, kit, "GetSnapshot");
1337
1338         assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"), Sets.newHashSet(
1339                 Lists.transform(datastoreSnapshot.getShardSnapshots(), shardNameTransformer)));
1340
1341         byte[] snapshotBytes = datastoreSnapshot.getShardManagerSnapshot();
1342         assertNotNull("Expected ShardManagerSnapshot", snapshotBytes);
1343         ShardManagerSnapshot snapshot = SerializationUtils.deserialize(snapshotBytes);
1344         assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"),
1345                 Sets.newHashSet(snapshot.getShardList()));
1346
1347         LOG.info("testGetSnapshot ending");
1348     }
1349
1350     @Test
1351     public void testRestoreFromSnapshot() throws Throwable {
1352         LOG.info("testRestoreFromSnapshot starting");
1353
1354         JavaTestKit kit = new JavaTestKit(getSystem());
1355
1356         MockConfiguration mockConfig = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1357                    put("shard1", Collections.<String>emptyList()).
1358                    put("shard2", Collections.<String>emptyList()).
1359                    put("astronauts", Collections.<String>emptyList()).build());
1360
1361
1362         ShardManagerSnapshot snapshot = new ShardManagerSnapshot(Arrays.asList("shard1", "shard2", "astronauts"));
1363         DatastoreSnapshot restoreFromSnapshot = new DatastoreSnapshot(shardMrgIDSuffix,
1364                 SerializationUtils.serialize(snapshot), Collections.<ShardSnapshot>emptyList());
1365         TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(newTestShardMgrBuilder(mockConfig).
1366                 restoreFromSnapshot(restoreFromSnapshot).props().withDispatcher(Dispatchers.DefaultDispatcherId()));
1367
1368         shardManager.underlyingActor().waitForRecoveryComplete();
1369
1370         shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), ActorRef.noSender());
1371
1372         waitForShardInitialized(shardManager, "shard1", kit);
1373         waitForShardInitialized(shardManager, "shard2", kit);
1374         waitForShardInitialized(shardManager, "astronauts", kit);
1375
1376         shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
1377
1378         DatastoreSnapshot datastoreSnapshot = expectMsgClassOrFailure(DatastoreSnapshot.class, kit, "GetSnapshot");
1379
1380         assertEquals("getType", shardMrgIDSuffix, datastoreSnapshot.getType());
1381
1382         byte[] snapshotBytes = datastoreSnapshot.getShardManagerSnapshot();
1383         assertNotNull("Expected ShardManagerSnapshot", snapshotBytes);
1384         snapshot = SerializationUtils.deserialize(snapshotBytes);
1385         assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"),
1386                 Sets.newHashSet(snapshot.getShardList()));
1387
1388         LOG.info("testRestoreFromSnapshot ending");
1389     }
1390
1391     @Test
1392     public void testAddShardReplicaForNonExistentShardConfig() throws Exception {
1393         new JavaTestKit(getSystem()) {{
1394             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1395                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1396
1397             shardManager.tell(new AddShardReplica("model-inventory"), getRef());
1398             Status.Failure resp = expectMsgClass(duration("2 seconds"), Status.Failure.class);
1399
1400             assertEquals("Failure obtained", true,
1401                           (resp.cause() instanceof IllegalArgumentException));
1402         }};
1403     }
1404
1405     @Test
1406     public void testAddShardReplica() throws Exception {
1407         MockConfiguration mockConfig =
1408                 new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1409                    put("default", Arrays.asList("member-1", "member-2")).
1410                    put("astronauts", Arrays.asList("member-2")).build());
1411
1412         final String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
1413         datastoreContextBuilder.shardManagerPersistenceId(shardManagerID);
1414
1415         // Create an ActorSystem ShardManager actor for member-1.
1416         final ActorSystem system1 = newActorSystem("Member1");
1417         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1418         ActorRef mockDefaultShardActor = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
1419         final TestActorRef<TestShardManager> newReplicaShardManager = TestActorRef.create(system1,
1420                 newTestShardMgrBuilder(mockConfig).shardActor(mockDefaultShardActor).cluster(
1421                         new ClusterWrapperImpl(system1)).props(), shardManagerID);
1422
1423         // Create an ActorSystem ShardManager actor for member-2.
1424         final ActorSystem system2 = newActorSystem("Member2");
1425         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1426
1427         String name = new ShardIdentifier("astronauts", "member-2", "config").toString();
1428         final TestActorRef<MockRespondActor> mockShardLeaderActor =
1429                 TestActorRef.create(system2, Props.create(MockRespondActor.class), name);
1430         final TestActorRef<TestShardManager> leaderShardManager = TestActorRef.create(system2,
1431                 newTestShardMgrBuilder(mockConfig).shardActor(mockShardLeaderActor).cluster(
1432                         new ClusterWrapperImpl(system2)).props(), shardManagerID);
1433
1434         new JavaTestKit(system1) {{
1435
1436             newReplicaShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1437             leaderShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1438
1439             leaderShardManager.tell(new ActorInitialized(), mockShardLeaderActor);
1440
1441             String memberId2 = "member-2-shard-astronauts-" + shardMrgIDSuffix;
1442             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
1443             leaderShardManager.tell(new ShardLeaderStateChanged(memberId2, memberId2,
1444                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardLeaderActor);
1445             leaderShardManager.tell(new RoleChangeNotification(memberId2,
1446                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardLeaderActor);
1447
1448             newReplicaShardManager.underlyingActor().waitForMemberUp();
1449             leaderShardManager.underlyingActor().waitForMemberUp();
1450
1451             //Have a dummy snapshot to be overwritten by the new data persisted.
1452             String[] restoredShards = {"default", "people"};
1453             ShardManagerSnapshot snapshot = new ShardManagerSnapshot(Arrays.asList(restoredShards));
1454             InMemorySnapshotStore.addSnapshot(shardManagerID, snapshot);
1455             Uninterruptibles.sleepUninterruptibly(2, TimeUnit.MILLISECONDS);
1456
1457             InMemorySnapshotStore.addSnapshotSavedLatch(shardManagerID);
1458             InMemorySnapshotStore.addSnapshotDeletedLatch(shardManagerID);
1459
1460             //construct a mock response message
1461             AddServerReply response = new AddServerReply(ServerChangeStatus.OK, memberId2);
1462             mockShardLeaderActor.underlyingActor().updateResponse(response);
1463             newReplicaShardManager.tell(new AddShardReplica("astronauts"), getRef());
1464             AddServer addServerMsg = MessageCollectorActor.expectFirstMatching(mockShardLeaderActor,
1465                 AddServer.class);
1466             String addServerId = "member-1-shard-astronauts-" + shardMrgIDSuffix;
1467             assertEquals("AddServer serverId", addServerId, addServerMsg.getNewServerId());
1468             expectMsgClass(duration("5 seconds"), Status.Success.class);
1469
1470             InMemorySnapshotStore.waitForSavedSnapshot(shardManagerID, ShardManagerSnapshot.class);
1471             InMemorySnapshotStore.waitForDeletedSnapshot(shardManagerID);
1472             List<ShardManagerSnapshot> persistedSnapshots =
1473                 InMemorySnapshotStore.getSnapshots(shardManagerID, ShardManagerSnapshot.class);
1474             assertEquals("Number of snapshots persisted", 1, persistedSnapshots.size());
1475             ShardManagerSnapshot shardManagerSnapshot = persistedSnapshots.get(0);
1476             assertEquals("Persisted local shards", Sets.newHashSet("default", "astronauts"),
1477                     Sets.newHashSet(shardManagerSnapshot.getShardList()));
1478         }};
1479
1480         LOG.info("testAddShardReplica ending");
1481     }
1482
1483     @Test
1484     public void testAddShardReplicaWithPreExistingReplicaInRemoteShardLeader() throws Exception {
1485         LOG.info("testAddShardReplicaWithPreExistingReplicaInRemoteShardLeader starting");
1486         new JavaTestKit(getSystem()) {{
1487             TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(
1488                     newPropsShardMgrWithMockShardActor(), shardMgrID);
1489
1490             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1491             shardManager.tell(new ActorInitialized(), mockShardActor);
1492
1493             String leaderId = "leader-member-shard-default-" + shardMrgIDSuffix;
1494             AddServerReply addServerReply = new AddServerReply(ServerChangeStatus.ALREADY_EXISTS, null);
1495             ActorRef leaderShardActor = shardManager.underlyingActor().getContext().actorOf(
1496                     Props.create(MockRespondActor.class, addServerReply), leaderId);
1497
1498             MockClusterWrapper.sendMemberUp(shardManager, "leader-member", leaderShardActor.path().toString());
1499
1500             String newReplicaId = "member-1-shard-default-" + shardMrgIDSuffix;
1501             shardManager.tell(new RoleChangeNotification(newReplicaId,
1502                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
1503             shardManager.tell(new ShardLeaderStateChanged(newReplicaId, leaderId, Optional.<DataTree>absent(),
1504                     DataStoreVersions.CURRENT_VERSION), mockShardActor);
1505
1506             shardManager.tell(new AddShardReplica(Shard.DEFAULT_NAME), getRef());
1507
1508             MessageCollectorActor.expectFirstMatching(leaderShardActor, AddServer.class);
1509
1510             Failure resp = expectMsgClass(duration("5 seconds"), Failure.class);
1511             assertEquals("Failure cause", AlreadyExistsException.class, resp.cause().getClass());
1512
1513             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
1514             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1515
1516             // Send message again to verify previous in progress state is cleared
1517
1518             shardManager.tell(new AddShardReplica(Shard.DEFAULT_NAME), getRef());
1519             resp = expectMsgClass(duration("5 seconds"), Failure.class);
1520             assertEquals("Failure cause", AlreadyExistsException.class, resp.cause().getClass());
1521
1522             // Send message again with an AddServer timeout to verify the pre-existing shard actor isn't terminated.
1523
1524             shardManager.tell(newDatastoreContextFactory(datastoreContextBuilder.
1525                     shardLeaderElectionTimeout(100, TimeUnit.MILLISECONDS).build()), getRef());
1526             leaderShardActor.tell(MockRespondActor.CLEAR_RESPONSE, ActorRef.noSender());
1527             shardManager.tell(new AddShardReplica(Shard.DEFAULT_NAME), getRef());
1528             expectMsgClass(duration("5 seconds"), Failure.class);
1529
1530             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
1531             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1532         }};
1533
1534         LOG.info("testAddShardReplicaWithPreExistingReplicaInRemoteShardLeader ending");
1535     }
1536
1537     @Test
1538     public void testAddShardReplicaWithPreExistingLocalReplicaLeader() throws Exception {
1539         LOG.info("testAddShardReplicaWithPreExistingLocalReplicaLeader starting");
1540         new JavaTestKit(getSystem()) {{
1541             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
1542             ActorRef shardManager = actorFactory.createActor(newPropsShardMgrWithMockShardActor());
1543
1544             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1545             shardManager.tell(new ActorInitialized(), mockShardActor);
1546             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mock(DataTree.class)),
1547                     DataStoreVersions.CURRENT_VERSION), getRef());
1548             shardManager.tell((new RoleChangeNotification(memberId, RaftState.Candidate.name(),
1549                     RaftState.Leader.name())), mockShardActor);
1550
1551             shardManager.tell(new AddShardReplica(Shard.DEFAULT_NAME), getRef());
1552             Failure resp = expectMsgClass(duration("5 seconds"), Failure.class);
1553             assertEquals("Failure cause", AlreadyExistsException.class, resp.cause().getClass());
1554
1555             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
1556             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1557         }};
1558
1559         LOG.info("testAddShardReplicaWithPreExistingLocalReplicaLeader ending");
1560     }
1561
1562     @Test
1563     public void testAddShardReplicaWithAddServerReplyFailure() throws Exception {
1564         LOG.info("testAddShardReplicaWithAddServerReplyFailure starting");
1565         new JavaTestKit(getSystem()) {{
1566             JavaTestKit mockShardLeaderKit = new JavaTestKit(getSystem());
1567
1568             MockConfiguration mockConfig =
1569                     new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1570                        put("astronauts", Arrays.asList("member-2")).build());
1571
1572             ActorRef mockNewReplicaShardActor = newMockShardActor(getSystem(), "astronauts", "member-1");
1573             final TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(
1574                     newTestShardMgrBuilder(mockConfig).shardActor(mockNewReplicaShardActor).props(), shardMgrID);
1575             shardManager.underlyingActor().setMessageInterceptor(newFindPrimaryInterceptor(mockShardLeaderKit.getRef()));
1576
1577             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1578
1579             JavaTestKit terminateWatcher = new JavaTestKit(getSystem());
1580             terminateWatcher.watch(mockNewReplicaShardActor);
1581
1582             shardManager.tell(new AddShardReplica("astronauts"), getRef());
1583
1584             AddServer addServerMsg = mockShardLeaderKit.expectMsgClass(AddServer.class);
1585             assertEquals("AddServer serverId", "member-1-shard-astronauts-" + shardMrgIDSuffix,
1586                     addServerMsg.getNewServerId());
1587             mockShardLeaderKit.reply(new AddServerReply(ServerChangeStatus.TIMEOUT, null));
1588
1589             Failure failure = expectMsgClass(duration("5 seconds"), Failure.class);
1590             assertEquals("Failure cause", TimeoutException.class, failure.cause().getClass());
1591
1592             shardManager.tell(new FindLocalShard("astronauts", false), getRef());
1593             expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
1594
1595             terminateWatcher.expectTerminated(mockNewReplicaShardActor);
1596
1597             shardManager.tell(new AddShardReplica("astronauts"), getRef());
1598             mockShardLeaderKit.expectMsgClass(AddServer.class);
1599             mockShardLeaderKit.reply(new AddServerReply(ServerChangeStatus.NO_LEADER, null));
1600             failure = expectMsgClass(duration("5 seconds"), Failure.class);
1601             assertEquals("Failure cause", NoShardLeaderException.class, failure.cause().getClass());
1602         }};
1603
1604         LOG.info("testAddShardReplicaWithAddServerReplyFailure ending");
1605     }
1606
1607     @Test
1608     public void testAddShardReplicaWithAlreadyInProgress() throws Exception {
1609         testServerChangeWhenAlreadyInProgress("astronauts", new AddShardReplica("astronauts"),
1610                 AddServer.class, new AddShardReplica("astronauts"));
1611     }
1612
1613     @Test
1614     public void testAddShardReplicaWithFindPrimaryTimeout() throws Exception {
1615         LOG.info("testAddShardReplicaWithFindPrimaryTimeout starting");
1616         new JavaTestKit(getSystem()) {{
1617             MockConfiguration mockConfig = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1618                        put("astronauts", Arrays.asList("member-2")).build());
1619
1620             final ActorRef newReplicaShardManager = actorFactory.createActor(newTestShardMgrBuilder(mockConfig).
1621                     shardActor(mockShardActor).props(), shardMgrID);
1622
1623             newReplicaShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1624             MockClusterWrapper.sendMemberUp(newReplicaShardManager, "member-2", getRef().path().toString());
1625
1626             newReplicaShardManager.tell(new AddShardReplica("astronauts"), getRef());
1627             Status.Failure resp = expectMsgClass(duration("5 seconds"), Status.Failure.class);
1628             assertEquals("Failure obtained", true,
1629                           (resp.cause() instanceof RuntimeException));
1630         }};
1631
1632         LOG.info("testAddShardReplicaWithFindPrimaryTimeout ending");
1633     }
1634
1635     @Test
1636     public void testRemoveShardReplicaForNonExistentShard() throws Exception {
1637         new JavaTestKit(getSystem()) {{
1638             ActorRef shardManager = actorFactory.createActor(newShardMgrProps(
1639                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1640
1641             shardManager.tell(new RemoveShardReplica("model-inventory", "member-1"), getRef());
1642             Status.Failure resp = expectMsgClass(duration("10 seconds"), Status.Failure.class);
1643             assertEquals("Failure obtained", true,
1644                          (resp.cause() instanceof PrimaryNotFoundException));
1645         }};
1646     }
1647
1648     @Test
1649     /**
1650      * Primary is Local
1651      */
1652     public void testRemoveShardReplicaLocal() throws Exception {
1653         new JavaTestKit(getSystem()) {{
1654             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
1655
1656             final TestActorRef<MockRespondActor> respondActor =
1657                     TestActorRef.create(getSystem(), Props.create(MockRespondActor.class), memberId);
1658
1659             ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor(respondActor));
1660
1661             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1662             shardManager.tell(new ActorInitialized(), respondActor);
1663             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mock(DataTree.class)),
1664                     DataStoreVersions.CURRENT_VERSION), getRef());
1665             shardManager.tell((new RoleChangeNotification(memberId, RaftState.Candidate.name(),
1666                     RaftState.Leader.name())), respondActor);
1667
1668             respondActor.underlyingActor().updateResponse(new RemoveServerReply(ServerChangeStatus.OK, null));
1669             shardManager.tell(new RemoveShardReplica(Shard.DEFAULT_NAME, "member-1"), getRef());
1670             final RemoveServer removeServer = MessageCollectorActor.expectFirstMatching(respondActor, RemoveServer.class);
1671             assertEquals(new ShardIdentifier("default", "member-1", shardMrgIDSuffix).toString(),
1672                     removeServer.getServerId());
1673             expectMsgClass(duration("5 seconds"), Success.class);
1674         }};
1675     }
1676
1677     @Test
1678     public void testRemoveShardReplicaRemote() throws Exception {
1679         MockConfiguration mockConfig =
1680                 new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1681                         put("default", Arrays.asList("member-1", "member-2")).
1682                         put("astronauts", Arrays.asList("member-1")).build());
1683
1684         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
1685
1686         // Create an ActorSystem ShardManager actor for member-1.
1687         final ActorSystem system1 = newActorSystem("Member1");
1688         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1689         ActorRef mockDefaultShardActor = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
1690
1691         final TestActorRef<TestShardManager> newReplicaShardManager = TestActorRef.create(system1,
1692                 newTestShardMgrBuilder().configuration(mockConfig).shardActor(mockDefaultShardActor).cluster(
1693                         new ClusterWrapperImpl(system1)).props(),
1694                 shardManagerID);
1695
1696         // Create an ActorSystem ShardManager actor for member-2.
1697         final ActorSystem system2 = newActorSystem("Member2");
1698         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1699
1700         String name = new ShardIdentifier("default", "member-2", shardMrgIDSuffix).toString();
1701         final TestActorRef<MockRespondActor> mockShardLeaderActor =
1702                 TestActorRef.create(system2, Props.create(MockRespondActor.class), name);
1703
1704         LOG.error("Mock Shard Leader Actor : {}", mockShardLeaderActor);
1705
1706         final TestActorRef<TestShardManager> leaderShardManager = TestActorRef.create(system2,
1707                 newTestShardMgrBuilder().configuration(mockConfig).shardActor(mockShardLeaderActor).cluster(
1708                         new ClusterWrapperImpl(system2)).props(),
1709                 shardManagerID);
1710
1711         // Because mockShardLeaderActor is created at the top level of the actor system it has an address like so,
1712         //    akka.tcp://cluster-test@127.0.0.1:2559/user/member-2-shard-default-config1
1713         // However when a shard manager has a local shard which is a follower and a leader that is remote it will
1714         // try to compute an address for the remote shard leader using the ShardPeerAddressResolver. This address will
1715         // look like so,
1716         //    akka.tcp://cluster-test@127.0.0.1:2559/user/shardmanager-config1/member-2-shard-default-config1
1717         // In this specific case if we did a FindPrimary for shard default from member-1 we would come up
1718         // with the address of an actor which does not exist, therefore any message sent to that actor would go to
1719         // dead letters.
1720         // To work around this problem we create a ForwardingActor with the right address and pass to it the
1721         // mockShardLeaderActor. The ForwardingActor simply forwards all messages to the mockShardLeaderActor and every
1722         // thing works as expected
1723         final ActorRef actorRef = leaderShardManager.underlyingActor().context()
1724                 .actorOf(Props.create(ForwardingActor.class, mockShardLeaderActor), "member-2-shard-default-" + shardMrgIDSuffix);
1725
1726         LOG.error("Forwarding actor : {}", actorRef);
1727
1728         new JavaTestKit(system1) {{
1729
1730             newReplicaShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1731             leaderShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1732
1733             leaderShardManager.tell(new ActorInitialized(), mockShardLeaderActor);
1734             newReplicaShardManager.tell(new ActorInitialized(), mockShardLeaderActor);
1735
1736             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
1737             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
1738             leaderShardManager.tell(new ShardLeaderStateChanged(memberId2, memberId2,
1739                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardLeaderActor);
1740             leaderShardManager.tell(new RoleChangeNotification(memberId2,
1741                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardLeaderActor);
1742
1743             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
1744             newReplicaShardManager.tell(new ShardLeaderStateChanged(memberId1, memberId2,
1745                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardActor);
1746             newReplicaShardManager.tell(new RoleChangeNotification(memberId1,
1747                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
1748
1749             newReplicaShardManager.underlyingActor().waitForMemberUp();
1750             leaderShardManager.underlyingActor().waitForMemberUp();
1751
1752             //construct a mock response message
1753             RemoveServerReply response = new RemoveServerReply(ServerChangeStatus.OK, memberId2);
1754             mockShardLeaderActor.underlyingActor().updateResponse(response);
1755             newReplicaShardManager.tell(new RemoveShardReplica("default", "member-1"), getRef());
1756             RemoveServer removeServer = MessageCollectorActor.expectFirstMatching(mockShardLeaderActor,
1757                     RemoveServer.class);
1758             String removeServerId = new ShardIdentifier("default", "member-1", shardMrgIDSuffix).toString();
1759             assertEquals("RemoveServer serverId", removeServerId, removeServer.getServerId());
1760             expectMsgClass(duration("5 seconds"), Status.Success.class);
1761         }};
1762
1763     }
1764
1765     @Test
1766     public void testRemoveShardReplicaWhenAnotherRemoveShardReplicaAlreadyInProgress() throws Exception {
1767         testServerChangeWhenAlreadyInProgress("astronauts", new RemoveShardReplica("astronauts", "member-2"),
1768                 RemoveServer.class, new RemoveShardReplica("astronauts", "member-3"));
1769     }
1770
1771     @Test
1772     public void testRemoveShardReplicaWhenAddShardReplicaAlreadyInProgress() throws Exception {
1773         testServerChangeWhenAlreadyInProgress("astronauts", new AddShardReplica("astronauts"),
1774                 AddServer.class, new RemoveShardReplica("astronauts", "member-2"));
1775     }
1776
1777
1778     public void testServerChangeWhenAlreadyInProgress(final String shardName, final Object firstServerChange,
1779                                                       final Class firstForwardedServerChangeClass,
1780                                                       final Object secondServerChange) throws Exception {
1781         new JavaTestKit(getSystem()) {{
1782             JavaTestKit mockShardLeaderKit = new JavaTestKit(getSystem());
1783             JavaTestKit secondRequestKit = new JavaTestKit(getSystem());
1784
1785             MockConfiguration mockConfig =
1786                     new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1787                             put(shardName, Arrays.asList("member-2")).build());
1788
1789             final TestActorRef<TestShardManager> shardManager = TestActorRef.create(getSystem(),
1790                     newTestShardMgrBuilder().configuration(mockConfig).shardActor(mockShardActor).cluster(
1791                             new MockClusterWrapper()).props(),
1792                     shardMgrID);
1793
1794             shardManager.underlyingActor().setMessageInterceptor(newFindPrimaryInterceptor(mockShardLeaderKit.getRef()));
1795
1796             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1797
1798             shardManager.tell(firstServerChange, getRef());
1799
1800             mockShardLeaderKit.expectMsgClass(firstForwardedServerChangeClass);
1801
1802             shardManager.tell(secondServerChange, secondRequestKit.getRef());
1803
1804             secondRequestKit.expectMsgClass(duration("5 seconds"), Failure.class);
1805         }};
1806     }
1807
1808     @Test
1809     public void testServerRemovedShardActorNotRunning() throws Exception {
1810         LOG.info("testServerRemovedShardActorNotRunning starting");
1811         new JavaTestKit(getSystem()) {{
1812             MockConfiguration mockConfig =
1813                     new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1814                             put("default", Arrays.asList("member-1", "member-2")).
1815                             put("astronauts", Arrays.asList("member-2")).
1816                             put("people", Arrays.asList("member-1", "member-2")).build());
1817
1818             TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(newShardMgrProps(mockConfig));
1819
1820             shardManager.underlyingActor().waitForRecoveryComplete();
1821             shardManager.tell(new FindLocalShard("people", false), getRef());
1822             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
1823
1824             shardManager.tell(new FindLocalShard("default", false), getRef());
1825             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
1826
1827             // Removed the default shard replica from member-1
1828             ShardIdentifier.Builder builder = new ShardIdentifier.Builder();
1829             ShardIdentifier shardId = builder.shardName("default").memberName("member-1").type(shardMrgIDSuffix).build();
1830             shardManager.tell(new ServerRemoved(shardId.toString()), getRef());
1831
1832             shardManager.underlyingActor().verifySnapshotPersisted(Sets.newHashSet("people"));
1833         }};
1834
1835         LOG.info("testServerRemovedShardActorNotRunning ending");
1836     }
1837
1838     @Test
1839     public void testServerRemovedShardActorRunning() throws Exception {
1840         LOG.info("testServerRemovedShardActorRunning starting");
1841         new JavaTestKit(getSystem()) {{
1842             MockConfiguration mockConfig =
1843                     new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1844                             put("default", Arrays.asList("member-1", "member-2")).
1845                             put("astronauts", Arrays.asList("member-2")).
1846                             put("people", Arrays.asList("member-1", "member-2")).build());
1847
1848             String shardId = ShardIdentifier.builder().shardName("default").memberName("member-1").
1849                     type(shardMrgIDSuffix).build().toString();
1850             TestActorRef<MessageCollectorActor> shard = actorFactory.createTestActor(
1851                     MessageCollectorActor.props(), shardId);
1852
1853             TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(
1854                     newTestShardMgrBuilder(mockConfig).addShardActor("default", shard).props());
1855
1856             watch(shard);
1857
1858             shardManager.underlyingActor().waitForRecoveryComplete();
1859
1860             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1861             shardManager.tell(new ActorInitialized(), shard);
1862
1863             waitForShardInitialized(shardManager, "people", this);
1864             waitForShardInitialized(shardManager, "default", this);
1865
1866             // Removed the default shard replica from member-1
1867             shardManager.tell(new ServerRemoved(shardId), getRef());
1868
1869             shardManager.underlyingActor().verifySnapshotPersisted(Sets.newHashSet("people"));
1870
1871             expectMsgClass(duration("5 seconds"), Terminated.class);
1872         }};
1873
1874         LOG.info("testServerRemovedShardActorRunning ending");
1875     }
1876
1877
1878     @Test
1879     public void testShardPersistenceWithRestoredData() throws Exception {
1880         LOG.info("testShardPersistenceWithRestoredData starting");
1881         new JavaTestKit(getSystem()) {{
1882             MockConfiguration mockConfig =
1883                 new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1884                    put("default", Arrays.asList("member-1", "member-2")).
1885                    put("astronauts", Arrays.asList("member-2")).
1886                    put("people", Arrays.asList("member-1", "member-2")).build());
1887             String[] restoredShards = {"default", "astronauts"};
1888             ShardManagerSnapshot snapshot = new ShardManagerSnapshot(Arrays.asList(restoredShards));
1889             InMemorySnapshotStore.addSnapshot("shard-manager-" + shardMrgIDSuffix, snapshot);
1890
1891             //create shardManager to come up with restored data
1892             TestActorRef<TestShardManager> newRestoredShardManager = actorFactory.createTestActor(
1893                     newShardMgrProps(mockConfig));
1894
1895             newRestoredShardManager.underlyingActor().waitForRecoveryComplete();
1896
1897             newRestoredShardManager.tell(new FindLocalShard("people", false), getRef());
1898             LocalShardNotFound notFound = expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
1899             assertEquals("for uninitialized shard", "people", notFound.getShardName());
1900
1901             //Verify a local shard is created for the restored shards,
1902             //although we expect a NotInitializedException for the shards as the actor initialization
1903             //message is not sent for them
1904             newRestoredShardManager.tell(new FindLocalShard("default", false), getRef());
1905             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
1906
1907             newRestoredShardManager.tell(new FindLocalShard("astronauts", false), getRef());
1908             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
1909         }};
1910
1911         LOG.info("testShardPersistenceWithRestoredData ending");
1912     }
1913
1914
1915     private static class TestShardManager extends ShardManager {
1916         private final CountDownLatch recoveryComplete = new CountDownLatch(1);
1917         private final CountDownLatch snapshotPersist = new CountDownLatch(1);
1918         private ShardManagerSnapshot snapshot;
1919         private final Map<String, ActorRef> shardActors;
1920         private final ActorRef shardActor;
1921         private CountDownLatch findPrimaryMessageReceived = new CountDownLatch(1);
1922         private CountDownLatch memberUpReceived = new CountDownLatch(1);
1923         private CountDownLatch memberRemovedReceived = new CountDownLatch(1);
1924         private CountDownLatch memberUnreachableReceived = new CountDownLatch(1);
1925         private CountDownLatch memberReachableReceived = new CountDownLatch(1);
1926         private volatile MessageInterceptor messageInterceptor;
1927
1928         private TestShardManager(Builder builder) {
1929             super(builder);
1930             shardActor = builder.shardActor;
1931             shardActors = builder.shardActors;
1932         }
1933
1934         @Override
1935         public void handleRecover(Object message) throws Exception {
1936             try {
1937                 super.handleRecover(message);
1938             } finally {
1939                 if(message instanceof RecoveryCompleted) {
1940                     recoveryComplete.countDown();
1941                 }
1942             }
1943         }
1944
1945         @Override
1946         public void handleCommand(Object message) throws Exception {
1947             try{
1948                 if(messageInterceptor != null && messageInterceptor.canIntercept(message)) {
1949                     getSender().tell(messageInterceptor.apply(message), getSelf());
1950                 } else {
1951                     super.handleCommand(message);
1952                 }
1953             } finally {
1954                 if(message instanceof FindPrimary) {
1955                     findPrimaryMessageReceived.countDown();
1956                 } else if(message instanceof ClusterEvent.MemberUp) {
1957                     String role = ((ClusterEvent.MemberUp)message).member().roles().head();
1958                     if(!getCluster().getCurrentMemberName().equals(role)) {
1959                         memberUpReceived.countDown();
1960                     }
1961                 } else if(message instanceof ClusterEvent.MemberRemoved) {
1962                     String role = ((ClusterEvent.MemberRemoved)message).member().roles().head();
1963                     if(!getCluster().getCurrentMemberName().equals(role)) {
1964                         memberRemovedReceived.countDown();
1965                     }
1966                 } else if(message instanceof ClusterEvent.UnreachableMember) {
1967                     String role = ((ClusterEvent.UnreachableMember)message).member().roles().head();
1968                     if(!getCluster().getCurrentMemberName().equals(role)) {
1969                         memberUnreachableReceived.countDown();
1970                     }
1971                 } else if(message instanceof ClusterEvent.ReachableMember) {
1972                     String role = ((ClusterEvent.ReachableMember)message).member().roles().head();
1973                     if(!getCluster().getCurrentMemberName().equals(role)) {
1974                         memberReachableReceived.countDown();
1975                     }
1976                 }
1977             }
1978         }
1979
1980         void setMessageInterceptor(MessageInterceptor messageInterceptor) {
1981             this.messageInterceptor = messageInterceptor;
1982         }
1983
1984         void waitForRecoveryComplete() {
1985             assertEquals("Recovery complete", true,
1986                     Uninterruptibles.awaitUninterruptibly(recoveryComplete, 5, TimeUnit.SECONDS));
1987         }
1988
1989         void waitForMemberUp() {
1990             assertEquals("MemberUp received", true,
1991                     Uninterruptibles.awaitUninterruptibly(memberUpReceived, 5, TimeUnit.SECONDS));
1992             memberUpReceived = new CountDownLatch(1);
1993         }
1994
1995         void waitForMemberRemoved() {
1996             assertEquals("MemberRemoved received", true,
1997                     Uninterruptibles.awaitUninterruptibly(memberRemovedReceived, 5, TimeUnit.SECONDS));
1998             memberRemovedReceived = new CountDownLatch(1);
1999         }
2000
2001         void waitForUnreachableMember() {
2002             assertEquals("UnreachableMember received", true,
2003                 Uninterruptibles.awaitUninterruptibly(memberUnreachableReceived, 5, TimeUnit.SECONDS
2004                 ));
2005             memberUnreachableReceived = new CountDownLatch(1);
2006         }
2007
2008         void waitForReachableMember() {
2009             assertEquals("ReachableMember received", true,
2010                 Uninterruptibles.awaitUninterruptibly(memberReachableReceived, 5, TimeUnit.SECONDS));
2011             memberReachableReceived = new CountDownLatch(1);
2012         }
2013
2014         void verifyFindPrimary() {
2015             assertEquals("FindPrimary received", true,
2016                     Uninterruptibles.awaitUninterruptibly(findPrimaryMessageReceived, 5, TimeUnit.SECONDS));
2017             findPrimaryMessageReceived = new CountDownLatch(1);
2018         }
2019
2020         public static Builder builder(DatastoreContext.Builder datastoreContextBuilder) {
2021             return new Builder(datastoreContextBuilder);
2022         }
2023
2024         private static class Builder extends AbstractGenericBuilder<Builder, TestShardManager> {
2025             private ActorRef shardActor;
2026             private final Map<String, ActorRef> shardActors = new HashMap<>();
2027
2028             Builder(DatastoreContext.Builder datastoreContextBuilder) {
2029                 super(TestShardManager.class);
2030                 datastoreContextFactory(newDatastoreContextFactory(datastoreContextBuilder.build()));
2031             }
2032
2033             Builder shardActor(ActorRef shardActor) {
2034                 this.shardActor = shardActor;
2035                 return this;
2036             }
2037
2038             Builder addShardActor(String shardName, ActorRef actorRef){
2039                 shardActors.put(shardName, actorRef);
2040                 return this;
2041             }
2042         }
2043
2044         @Override
2045         public void saveSnapshot(Object obj) {
2046             snapshot = (ShardManagerSnapshot) obj;
2047             snapshotPersist.countDown();
2048             super.saveSnapshot(obj);
2049         }
2050
2051         void verifySnapshotPersisted(Set<String> shardList) {
2052             assertEquals("saveSnapshot invoked", true,
2053                     Uninterruptibles.awaitUninterruptibly(snapshotPersist, 5, TimeUnit.SECONDS));
2054             assertEquals("Shard Persisted", shardList, Sets.newHashSet(snapshot.getShardList()));
2055         }
2056
2057         @Override
2058         protected ActorRef newShardActor(SchemaContext schemaContext, ShardInformation info) {
2059             if(shardActors.get(info.getShardName()) != null){
2060                 return shardActors.get(info.getShardName());
2061             }
2062
2063             if(shardActor != null) {
2064                 return shardActor;
2065             }
2066
2067             return super.newShardActor(schemaContext, info);
2068         }
2069     }
2070
2071     private static abstract class AbstractGenericBuilder<T extends AbstractGenericBuilder<T, ?>, C extends ShardManager>
2072                                                      extends ShardManager.AbstractBuilder<T> {
2073         private final Class<C> shardManagerClass;
2074
2075         AbstractGenericBuilder(Class<C> shardManagerClass) {
2076             this.shardManagerClass = shardManagerClass;
2077             cluster(new MockClusterWrapper()).configuration(new MockConfiguration()).
2078                     waitTillReadyCountdownLatch(ready).primaryShardInfoCache(new PrimaryShardInfoFutureCache());
2079         }
2080
2081         @Override
2082         public Props props() {
2083             verify();
2084             return Props.create(shardManagerClass, this);
2085         }
2086     }
2087
2088     private static class GenericBuilder<C extends ShardManager> extends AbstractGenericBuilder<GenericBuilder<C>, C> {
2089         GenericBuilder(Class<C> shardManagerClass) {
2090             super(shardManagerClass);
2091         }
2092     }
2093
2094     private static class DelegatingShardManagerCreator implements Creator<ShardManager> {
2095         private static final long serialVersionUID = 1L;
2096         private final Creator<ShardManager> delegate;
2097
2098         public DelegatingShardManagerCreator(Creator<ShardManager> delegate) {
2099             this.delegate = delegate;
2100         }
2101
2102         @Override
2103         public ShardManager create() throws Exception {
2104             return delegate.create();
2105         }
2106     }
2107
2108     interface MessageInterceptor extends Function<Object, Object> {
2109         boolean canIntercept(Object message);
2110     }
2111
2112     private MessageInterceptor newFindPrimaryInterceptor(final ActorRef primaryActor) {
2113         return new MessageInterceptor(){
2114             @Override
2115             public Object apply(Object message) {
2116                 return new RemotePrimaryShardFound(Serialization.serializedActorPath(primaryActor), (short) 1);
2117             }
2118
2119             @Override
2120             public boolean canIntercept(Object message) {
2121                 return message instanceof FindPrimary;
2122             }
2123         };
2124     }
2125
2126     private static class MockRespondActor extends MessageCollectorActor {
2127         static final String CLEAR_RESPONSE = "clear-response";
2128         static final org.slf4j.Logger LOG = LoggerFactory.getLogger(MockRespondActor.class);
2129
2130         private volatile Object responseMsg;
2131
2132         @SuppressWarnings("unused")
2133         public MockRespondActor() {
2134         }
2135
2136         @SuppressWarnings("unused")
2137         public MockRespondActor(Object responseMsg) {
2138             this.responseMsg = responseMsg;
2139         }
2140
2141         public void updateResponse(Object response) {
2142             responseMsg = response;
2143         }
2144
2145         @Override
2146         public void onReceive(Object message) throws Exception {
2147             if(!"get-all-messages".equals(message)) {
2148                 LOG.debug("Received message : {}", message);
2149             }
2150             super.onReceive(message);
2151             if (message instanceof AddServer && responseMsg != null) {
2152                 getSender().tell(responseMsg, getSelf());
2153             } else if(message instanceof RemoveServer && responseMsg != null){
2154                 getSender().tell(responseMsg, getSelf());
2155             } else if(message.equals(CLEAR_RESPONSE)) {
2156                 responseMsg = null;
2157             }
2158         }
2159     }
2160 }