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

   1 /*
   2  * Copyright (c) 2015 Brocade Communications 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 package org.opendaylight.controller.cluster.datastore.entityownership;
   9
  10 import static org.junit.Assert.assertEquals;
  11 import static org.mockito.AdditionalMatchers.or;
  12 import static org.mockito.Matchers.any;
  13 import static org.mockito.Mockito.mock;
  14 import static org.mockito.Mockito.never;
  15 import static org.mockito.Mockito.times;
  16 import static org.mockito.Mockito.reset;
  17 import static org.mockito.Mockito.timeout;
  18 import static org.mockito.Mockito.verify;
  19 import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.ENTITY_OWNERS_PATH;
  20 import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.candidatePath;
  21 import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.entityOwnersWithCandidate;
  22 import akka.actor.ActorRef;
  23 import akka.actor.PoisonPill;
  24 import akka.actor.Props;
  25 import akka.actor.Terminated;
  26 import akka.actor.UntypedActor;
  27 import akka.dispatch.Dispatchers;
  28 import akka.testkit.JavaTestKit;
  29 import akka.testkit.TestActorRef;
  30 import com.google.common.base.Function;
  31 import com.google.common.collect.ImmutableMap;
  32 import com.google.common.util.concurrent.Uninterruptibles;
  33 import java.util.ArrayList;
  34 import java.util.Collections;
  35 import java.util.List;
  36 import java.util.Map;
  37 import java.util.concurrent.CountDownLatch;
  38 import java.util.concurrent.TimeUnit;
  39 import java.util.concurrent.atomic.AtomicInteger;
  40 import org.junit.After;
  41 import org.junit.Test;
  42 import org.opendaylight.controller.cluster.datastore.AbstractShardTest;
  43 import org.opendaylight.controller.cluster.datastore.DataStoreVersions;
  44 import org.opendaylight.controller.cluster.datastore.DatastoreContext;
  45 import org.opendaylight.controller.cluster.datastore.DatastoreContext.Builder;
  46 import org.opendaylight.controller.cluster.datastore.ShardDataTree;
  47 import org.opendaylight.controller.cluster.datastore.ShardTestKit;
  48 import org.opendaylight.controller.cluster.datastore.entityownership.messages.RegisterCandidateLocal;
  49 import org.opendaylight.controller.cluster.datastore.entityownership.messages.RegisterListenerLocal;
  50 import org.opendaylight.controller.cluster.datastore.entityownership.messages.UnregisterCandidateLocal;
  51 import org.opendaylight.controller.cluster.datastore.entityownership.messages.UnregisterListenerLocal;
  52 import org.opendaylight.controller.cluster.datastore.entityownership.selectionstrategy.EntityOwnerSelectionStrategyConfig;
  53 import org.opendaylight.controller.cluster.datastore.entityownership.selectionstrategy.LastCandidateSelectionStrategy;
  54 import org.opendaylight.controller.cluster.datastore.identifiers.ShardIdentifier;
  55 import org.opendaylight.controller.cluster.datastore.messages.BatchedModifications;
  56 import org.opendaylight.controller.cluster.datastore.messages.CommitTransactionReply;
  57 import org.opendaylight.controller.cluster.datastore.messages.PeerAddressResolved;
  58 import org.opendaylight.controller.cluster.datastore.messages.PeerDown;
  59 import org.opendaylight.controller.cluster.datastore.messages.PeerUp;
  60 import org.opendaylight.controller.cluster.datastore.messages.SuccessReply;
  61 import org.opendaylight.controller.cluster.datastore.modification.MergeModification;
  62 import org.opendaylight.controller.cluster.datastore.modification.Modification;
  63 import org.opendaylight.controller.cluster.raft.ReplicatedLogEntry;
  64 import org.opendaylight.controller.cluster.raft.TestActorFactory;
  65 import org.opendaylight.controller.cluster.raft.base.messages.ElectionTimeout;
  66 import org.opendaylight.controller.cluster.raft.messages.AppendEntries;
  67 import org.opendaylight.controller.cluster.raft.messages.AppendEntriesReply;
  68 import org.opendaylight.controller.cluster.raft.messages.RequestVote;
  69 import org.opendaylight.controller.cluster.raft.messages.RequestVoteReply;
  70 import org.opendaylight.controller.md.cluster.datastore.model.SchemaContextHelper;
  71 import org.opendaylight.controller.md.sal.common.api.clustering.Entity;
  72 import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipChange;
  73 import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipListener;
  74 import org.opendaylight.yangtools.yang.common.QName;
  75 import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier;
  76 import org.opendaylight.yangtools.yang.data.api.schema.NormalizedNode;
  77 import org.opendaylight.yangtools.yang.model.api.SchemaContext;
  78
  79 /**
  80  * Unit tests for EntityOwnershipShard.
  81  *
  82  * @author Thomas Pantelis
  83  */
  84 public class EntityOwnershipShardTest extends AbstractEntityOwnershipTest {
  85     private static final String ENTITY_TYPE = "test type";
  86     private static final YangInstanceIdentifier ENTITY_ID1 =
  87             YangInstanceIdentifier.of(QName.create("test", "2015-08-14", "entity1"));
  88     private static final YangInstanceIdentifier ENTITY_ID2 =
  89             YangInstanceIdentifier.of(QName.create("test", "2015-08-14", "entity2"));
  90     private static final YangInstanceIdentifier ENTITY_ID3 =
  91             YangInstanceIdentifier.of(QName.create("test", "2015-08-14", "entity3"));
  92     private static final YangInstanceIdentifier ENTITY_ID4 =
  93             YangInstanceIdentifier.of(QName.create("test", "2015-08-14", "entity4"));
  94     private static final YangInstanceIdentifier ENTITY_ID5 =
  95             YangInstanceIdentifier.of(QName.create("test", "2015-08-14", "entity5"));
  96     private static final SchemaContext SCHEMA_CONTEXT = SchemaContextHelper.entityOwners();
  97     private static final AtomicInteger NEXT_SHARD_NUM = new AtomicInteger();
  98     private static final String LOCAL_MEMBER_NAME = "member-1";
  99
 100     private final Builder dataStoreContextBuilder = DatastoreContext.newBuilder();
 101     private final TestActorFactory actorFactory = new TestActorFactory(getSystem());
 102
 103     @After
 104     public void tearDown() {
 105         actorFactory.close();
 106     }
 107
 108     @Test
 109     public void testOnRegisterCandidateLocal() throws Exception {
 110         ShardTestKit kit = new ShardTestKit(getSystem());
 111
 112         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps());
 113
 114         kit.waitUntilLeader(shard);
 115
 116         YangInstanceIdentifier entityId = ENTITY_ID1;
 117         Entity entity = new Entity(ENTITY_TYPE, entityId);
 118
 119         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 120         kit.expectMsgClass(SuccessReply.class);
 121
 122         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, entityId, LOCAL_MEMBER_NAME);
 123         verifyOwner(shard, ENTITY_TYPE, entityId, LOCAL_MEMBER_NAME);
 124     }
 125
 126     @Test
 127     public void testOnRegisterCandidateLocalWithNoInitialLeader() throws Exception {
 128         ShardTestKit kit = new ShardTestKit(getSystem());
 129
 130         dataStoreContextBuilder.shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(2);
 131
 132         String peerId = newShardId("follower").toString();
 133         TestActorRef<MockFollower> peer = actorFactory.createTestActor(Props.create(MockFollower.class, peerId, false).
 134                 withDispatcher(Dispatchers.DefaultDispatcherId()), peerId);
 135
 136         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps(
 137                 ImmutableMap.<String, String>builder().put(peerId, peer.path().toString()).build()).
 138                 withDispatcher(Dispatchers.DefaultDispatcherId()));
 139
 140         YangInstanceIdentifier entityId = ENTITY_ID1;
 141         Entity entity = new Entity(ENTITY_TYPE, entityId);
 142
 143         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 144         kit.expectMsgClass(SuccessReply.class);
 145
 146         // Now grant the vote so the shard becomes the leader. This should retry the commit.
 147         peer.underlyingActor().grantVote = true;
 148
 149         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, entityId, LOCAL_MEMBER_NAME);
 150         verifyOwner(shard, ENTITY_TYPE, entityId, LOCAL_MEMBER_NAME);
 151     }
 152
 153     @Test
 154     public void testOnRegisterCandidateLocalWithNoInitialConsensus() throws Exception {
 155         ShardTestKit kit = new ShardTestKit(getSystem());
 156
 157         dataStoreContextBuilder.shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(2).
 158                 shardTransactionCommitTimeoutInSeconds(1);
 159
 160         String peerId = newShardId("follower").toString();
 161         TestActorRef<MockFollower> peer = actorFactory.createTestActor(Props.create(MockFollower.class, peerId).
 162                 withDispatcher(Dispatchers.DefaultDispatcherId()), peerId);
 163
 164         MockFollower follower = peer.underlyingActor();
 165
 166         // Drop AppendEntries so consensus isn't reached.
 167         follower.dropAppendEntries = true;
 168
 169         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps(
 170                 ImmutableMap.<String, String>builder().put(peerId, peer.path().toString()).build()).
 171                 withDispatcher(Dispatchers.DefaultDispatcherId()));
 172
 173         kit.waitUntilLeader(shard);
 174
 175         YangInstanceIdentifier entityId = ENTITY_ID1;
 176         Entity entity = new Entity(ENTITY_TYPE, entityId);
 177
 178         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 179         kit.expectMsgClass(SuccessReply.class);
 180
 181         // Wait enough time for the commit to timeout.
 182         Uninterruptibles.sleepUninterruptibly(1, TimeUnit.SECONDS);
 183
 184         // Resume AppendEntries - the follower should ack the commit which should then result in the candidate
 185         // write being applied to the state.
 186         follower.dropAppendEntries = false;
 187
 188         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, entityId, LOCAL_MEMBER_NAME);
 189         verifyOwner(shard, ENTITY_TYPE, entityId, LOCAL_MEMBER_NAME);
 190     }
 191
 192     @Test
 193     public void testOnRegisterCandidateLocalWithIsolatedLeader() throws Exception {
 194         ShardTestKit kit = new ShardTestKit(getSystem());
 195
 196         dataStoreContextBuilder.shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(2).
 197                 shardIsolatedLeaderCheckIntervalInMillis(50);
 198
 199         String peerId = newShardId("follower").toString();
 200         TestActorRef<MockFollower> peer = actorFactory.createTestActor(Props.create(MockFollower.class, peerId).
 201                 withDispatcher(Dispatchers.DefaultDispatcherId()), peerId);
 202
 203         MockFollower follower = peer.underlyingActor();
 204
 205         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps(
 206                 ImmutableMap.<String, String>builder().put(peerId, peer.path().toString()).build()).
 207                 withDispatcher(Dispatchers.DefaultDispatcherId()));
 208
 209         kit.waitUntilLeader(shard);
 210
 211         // Drop AppendEntries and wait enough time for the shard to switch to IsolatedLeader.
 212         follower.dropAppendEntries = true;
 213         Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
 214
 215         YangInstanceIdentifier entityId = ENTITY_ID1;
 216         Entity entity = new Entity(ENTITY_TYPE, entityId);
 217
 218         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 219         kit.expectMsgClass(SuccessReply.class);
 220
 221         // Resume AppendEntries - the candidate write should now be committed.
 222         follower.dropAppendEntries = false;
 223         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, entityId, LOCAL_MEMBER_NAME);
 224         verifyOwner(shard, ENTITY_TYPE, entityId, LOCAL_MEMBER_NAME);
 225     }
 226
 227     @Test
 228     public void testOnRegisterCandidateLocalWithRemoteLeader() throws Exception {
 229         ShardTestKit kit = new ShardTestKit(getSystem());
 230
 231         dataStoreContextBuilder.shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(2).
 232                 shardBatchedModificationCount(5);
 233
 234         String peerId = newShardId("leader").toString();
 235         TestActorRef<MockLeader> peer = actorFactory.createTestActor(Props.create(MockLeader.class).
 236                 withDispatcher(Dispatchers.DefaultDispatcherId()), peerId);
 237
 238         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(Props.create(
 239                 TestEntityOwnershipShard.class, newShardId(LOCAL_MEMBER_NAME),
 240                         ImmutableMap.<String, String>builder().put(peerId, peer.path().toString()).build(),
 241                         dataStoreContextBuilder.build()).withDispatcher(Dispatchers.DefaultDispatcherId()));
 242
 243         shard.tell(new AppendEntries(1L, peerId, -1L, -1L, Collections.<ReplicatedLogEntry>emptyList(), -1L, -1L,
 244                 DataStoreVersions.CURRENT_VERSION), peer);
 245
 246         shard.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID1)), kit.getRef());
 247         kit.expectMsgClass(SuccessReply.class);
 248
 249         MockLeader leader = peer.underlyingActor();
 250         assertEquals("Leader received BatchedModifications", true, Uninterruptibles.awaitUninterruptibly(
 251                 leader.modificationsReceived, 5, TimeUnit.SECONDS));
 252         verifyBatchedEntityCandidate(leader.getAndClearReceivedModifications(), ENTITY_TYPE, ENTITY_ID1,
 253                 LOCAL_MEMBER_NAME);
 254
 255         // Test with initial commit timeout and subsequent retry.
 256
 257         leader.modificationsReceived = new CountDownLatch(1);
 258         leader.sendReply = false;
 259
 260         shard.tell(dataStoreContextBuilder.shardTransactionCommitTimeoutInSeconds(1).build(), ActorRef.noSender());
 261
 262         shard.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID2)), kit.getRef());
 263         kit.expectMsgClass(SuccessReply.class);
 264
 265         assertEquals("Leader received BatchedModifications", true, Uninterruptibles.awaitUninterruptibly(
 266                 leader.modificationsReceived, 5, TimeUnit.SECONDS));
 267         verifyBatchedEntityCandidate(leader.getAndClearReceivedModifications(), ENTITY_TYPE, ENTITY_ID2,
 268                 LOCAL_MEMBER_NAME);
 269
 270         // Send a bunch of registration messages quickly and verify.
 271
 272         int max = 100;
 273         leader.delay = 4;
 274         leader.modificationsReceived = new CountDownLatch(max);
 275         List<YangInstanceIdentifier> entityIds = new ArrayList<>();
 276         for(int i = 1; i <= max; i++) {
 277             YangInstanceIdentifier id = YangInstanceIdentifier.of(QName.create("test", "2015-08-14", "test" + i));
 278             entityIds.add(id);
 279             shard.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, id)), kit.getRef());
 280         }
 281
 282         assertEquals("Leader received BatchedModifications", true, Uninterruptibles.awaitUninterruptibly(
 283                 leader.modificationsReceived, 10, TimeUnit.SECONDS));
 284
 285         // Sleep a little to ensure no additional BatchedModifications are received.
 286
 287         Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
 288
 289         List<Modification> receivedMods = leader.getAndClearReceivedModifications();
 290         for(int i = 0; i < max; i++) {
 291             verifyBatchedEntityCandidate(receivedMods.get(i), ENTITY_TYPE, entityIds.get(i), LOCAL_MEMBER_NAME);
 292         }
 293
 294         assertEquals("# modifications received", max, receivedMods.size());
 295     }
 296
 297     @Test
 298     public void testOnUnregisterCandidateLocal() throws Exception {
 299         ShardTestKit kit = new ShardTestKit(getSystem());
 300         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps());
 301         kit.waitUntilLeader(shard);
 302
 303         Entity entity = new Entity(ENTITY_TYPE, ENTITY_ID1);
 304
 305         // Register
 306
 307         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 308         kit.expectMsgClass(SuccessReply.class);
 309
 310         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 311         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 312
 313         // Unregister
 314
 315         shard.tell(new UnregisterCandidateLocal(entity), kit.getRef());
 316         kit.expectMsgClass(SuccessReply.class);
 317
 318         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, "");
 319
 320         // Register again
 321
 322         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 323         kit.expectMsgClass(SuccessReply.class);
 324
 325         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 326         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 327     }
 328
 329     @Test
 330     public void testOwnershipChanges() throws Exception {
 331         ShardTestKit kit = new ShardTestKit(getSystem());
 332         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps());
 333         kit.waitUntilLeader(shard);
 334
 335         Entity entity = new Entity(ENTITY_TYPE, ENTITY_ID1);
 336         ShardDataTree shardDataTree = shard.underlyingActor().getDataStore();
 337
 338         // Add a remote candidate
 339
 340         String remoteMemberName1 = "remoteMember1";
 341         writeNode(ENTITY_OWNERS_PATH, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, remoteMemberName1), shardDataTree);
 342
 343         // Register local
 344
 345         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 346         kit.expectMsgClass(SuccessReply.class);
 347
 348         // Verify the remote candidate becomes owner
 349
 350         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName1);
 351         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 352         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName1);
 353
 354         // Add another remote candidate and verify ownership doesn't change
 355
 356         String remoteMemberName2 = "remoteMember2";
 357         writeNode(ENTITY_OWNERS_PATH, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, remoteMemberName2), shardDataTree);
 358
 359         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName2);
 360         Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
 361         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName1);
 362
 363         // Remove the second remote candidate and verify ownership doesn't change
 364
 365         deleteNode(candidatePath(ENTITY_TYPE, ENTITY_ID1, remoteMemberName2), shardDataTree);
 366
 367         verifyEntityCandidateRemoved(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName2);
 368         Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
 369         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName1);
 370
 371         // Remove the first remote candidate and verify the local candidate becomes owner
 372
 373         deleteNode(candidatePath(ENTITY_TYPE, ENTITY_ID1, remoteMemberName1), shardDataTree);
 374
 375         verifyEntityCandidateRemoved(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName1);
 376         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 377
 378         // Add the second remote candidate back and verify ownership doesn't change
 379
 380         writeNode(ENTITY_OWNERS_PATH, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, remoteMemberName2), shardDataTree);
 381
 382         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName2);
 383         Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
 384         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 385
 386         // Unregister the local candidate and verify the second remote candidate becomes owner
 387
 388         shard.tell(new UnregisterCandidateLocal(entity), kit.getRef());
 389         kit.expectMsgClass(SuccessReply.class);
 390
 391         verifyEntityCandidateRemoved(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 392         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName2);
 393     }
 394
 395     @Test
 396     public void testOwnerChangesOnPeerAvailabilityChanges() throws Exception {
 397         ShardTestKit kit = new ShardTestKit(getSystem());
 398
 399         dataStoreContextBuilder.shardHeartbeatIntervalInMillis(500).shardElectionTimeoutFactor(10000);
 400
 401         String peerMemberName1 = "peerMember1";
 402         String peerMemberName2 = "peerMember2";
 403
 404         ShardIdentifier leaderId = newShardId(LOCAL_MEMBER_NAME);
 405         ShardIdentifier peerId1 = newShardId(peerMemberName1);
 406         ShardIdentifier peerId2 = newShardId(peerMemberName2);
 407
 408         TestActorRef<EntityOwnershipShard> peer1 = actorFactory.createTestActor(newShardProps(peerId1,
 409                 ImmutableMap.<String, String>builder().put(leaderId.toString(), ""). put(peerId2.toString(), "").build(),
 410                         peerMemberName1, EntityOwnerSelectionStrategyConfig.newBuilder().build()).withDispatcher(Dispatchers.DefaultDispatcherId()), peerId1.toString());
 411
 412         TestActorRef<EntityOwnershipShard> peer2 = actorFactory.createTestActor(newShardProps(peerId2,
 413                 ImmutableMap.<String, String>builder().put(leaderId.toString(), ""). put(peerId1.toString(), "").build(),
 414                         peerMemberName2, EntityOwnerSelectionStrategyConfig.newBuilder().build()). withDispatcher(Dispatchers.DefaultDispatcherId()), peerId2.toString());
 415
 416         TestActorRef<EntityOwnershipShard> leader = actorFactory.createTestActor(newShardProps(leaderId,
 417                 ImmutableMap.<String, String>builder().put(peerId1.toString(), peer1.path().toString()).
 418                         put(peerId2.toString(), peer2.path().toString()).build(), LOCAL_MEMBER_NAME, EntityOwnerSelectionStrategyConfig.newBuilder().build()).
 419                 withDispatcher(Dispatchers.DefaultDispatcherId()), leaderId.toString());
 420         leader.tell(new ElectionTimeout(), leader);
 421
 422         kit.waitUntilLeader(leader);
 423
 424         // Send PeerDown and PeerUp with no entities
 425
 426         leader.tell(new PeerDown(peerMemberName2, peerId2.toString()), ActorRef.noSender());
 427         leader.tell(new PeerUp(peerMemberName2, peerId2.toString()), ActorRef.noSender());
 428
 429         // Add candidates for entity1 with the local leader as the owner
 430
 431         leader.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID1)), kit.getRef());
 432         kit.expectMsgClass(SuccessReply.class);
 433         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 434
 435         peer2.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID1)), kit.getRef());
 436         kit.expectMsgClass(SuccessReply.class);
 437         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, peerMemberName2), kit);
 438         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID1, peerMemberName2);
 439
 440         peer1.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID1)), kit.getRef());
 441         kit.expectMsgClass(SuccessReply.class);
 442         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, peerMemberName1), kit);
 443         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID1, peerMemberName1);
 444         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 445
 446         // Add candidates for entity2 with peerMember2 as the owner
 447
 448         peer2.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID2)), kit.getRef());
 449         kit.expectMsgClass(SuccessReply.class);
 450         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID2, peerMemberName2), kit);
 451         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 452
 453         peer1.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID2)), kit.getRef());
 454         kit.expectMsgClass(SuccessReply.class);
 455         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID2, peerMemberName1), kit);
 456         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName1);
 457         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 458
 459         // Add candidates for entity3 with peerMember2 as the owner.
 460
 461         peer2.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID3)), kit.getRef());
 462         kit.expectMsgClass(SuccessReply.class);
 463         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID3, peerMemberName2), kit);
 464         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID3, peerMemberName2);
 465
 466         leader.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID3)), kit.getRef());
 467         kit.expectMsgClass(SuccessReply.class);
 468         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID3, LOCAL_MEMBER_NAME);
 469
 470         peer1.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID3)), kit.getRef());
 471         kit.expectMsgClass(SuccessReply.class);
 472         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID3, peerMemberName1), kit);
 473         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID3, peerMemberName1);
 474         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID3, peerMemberName2);
 475
 476         // Add only candidate peerMember2 for entity4.
 477
 478         peer2.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID4)), kit.getRef());
 479         kit.expectMsgClass(SuccessReply.class);
 480         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID4, peerMemberName2), kit);
 481         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID4, peerMemberName2);
 482         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID4, peerMemberName2);
 483
 484         // Add only candidate peerMember1 for entity5.
 485
 486         peer1.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID5)), kit.getRef());
 487         kit.expectMsgClass(SuccessReply.class);
 488         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID5, peerMemberName1), kit);
 489         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID5, peerMemberName1);
 490         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID5, peerMemberName1);
 491
 492         // Kill peerMember2 and send PeerDown - the entities (2, 3, 4) owned by peerMember2 should get a new
 493         // owner selected
 494
 495         kit.watch(peer2);
 496         peer2.tell(PoisonPill.getInstance(), ActorRef.noSender());
 497         kit.expectMsgClass(JavaTestKit.duration("5 seconds"), Terminated.class);
 498         kit.unwatch(peer2);
 499
 500         leader.tell(new PeerDown(peerMemberName2, peerId2.toString()), ActorRef.noSender());
 501         // Send PeerDown again - should be noop
 502         leader.tell(new PeerDown(peerMemberName2, peerId2.toString()), ActorRef.noSender());
 503         peer1.tell(new PeerDown(peerMemberName2, peerId2.toString()), ActorRef.noSender());
 504
 505         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID4, ""); // no other candidates so should clear
 506         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID3, LOCAL_MEMBER_NAME);
 507         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName1);
 508         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 509
 510         verifyNoEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID1, peerMemberName2);
 511         verifyNoEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 512         verifyNoEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID3, peerMemberName2);
 513         verifyNoEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID4, peerMemberName2);
 514
 515         // Reinstate peerMember2 - no owners should change
 516
 517         peer2 = actorFactory.createTestActor(newShardProps(peerId2,
 518                 ImmutableMap.<String, String>builder().put(leaderId.toString(), ""). put(peerId1.toString(), "").build(),
 519                         peerMemberName2, EntityOwnerSelectionStrategyConfig.newBuilder().build()). withDispatcher(Dispatchers.DefaultDispatcherId()), peerId2.toString());
 520         leader.tell(new PeerUp(peerMemberName2, peerId2.toString()), ActorRef.noSender());
 521         // Send PeerUp again - should be noop
 522         leader.tell(new PeerUp(peerMemberName2, peerId2.toString()), ActorRef.noSender());
 523         peer1.tell(new PeerUp(peerMemberName2, peerId2.toString()), ActorRef.noSender());
 524
 525         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID3, LOCAL_MEMBER_NAME);
 526         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName1);
 527         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 528         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID4, "");
 529
 530         peer2.tell(new PeerAddressResolved(peerId1.toString(), peer1.path().toString()), ActorRef.noSender());
 531         peer2.tell(new PeerAddressResolved(leaderId.toString(), leader.path().toString()), ActorRef.noSender());
 532         peer2.tell(new PeerUp(LOCAL_MEMBER_NAME, leaderId.toString()), ActorRef.noSender());
 533         peer2.tell(new PeerUp(peerMemberName1, peerId1.toString()), ActorRef.noSender());
 534
 535         // Add back candidate peerMember2 for entities 1, 2, & 3.
 536         peer2.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID1)), kit.getRef());
 537         kit.expectMsgClass(SuccessReply.class);
 538         peer2.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID2)), kit.getRef());
 539         kit.expectMsgClass(SuccessReply.class);
 540         peer2.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID3)), kit.getRef());
 541         kit.expectMsgClass(SuccessReply.class);
 542         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, peerMemberName2), kit);
 543         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID2, peerMemberName2), kit);
 544         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID3, peerMemberName2), kit);
 545         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID1, peerMemberName2);
 546         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 547         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID3, peerMemberName2);
 548         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 549         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName1);
 550         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID3, LOCAL_MEMBER_NAME);
 551
 552         // Kill peerMember1 and send PeerDown - entity 2 should get a new owner selected
 553
 554         peer1.tell(PoisonPill.getInstance(), ActorRef.noSender());
 555         leader.tell(new PeerDown(peerMemberName1, peerId1.toString()), ActorRef.noSender());
 556
 557         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 558
 559         // Verify the reinstated peerMember2 is fully synced.
 560
 561         verifyOwner(peer2, ENTITY_TYPE, ENTITY_ID4, "");
 562         verifyOwner(peer2, ENTITY_TYPE, ENTITY_ID3, LOCAL_MEMBER_NAME);
 563         verifyOwner(peer2, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 564         verifyOwner(peer2, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 565
 566         // Reinstate peerMember1 and verify no owner changes
 567
 568         peer1 = actorFactory.createTestActor(newShardProps(peerId1,
 569                 ImmutableMap.<String, String>builder().put(leaderId.toString(), ""). put(peerId2.toString(), "").build(),
 570                         peerMemberName1, EntityOwnerSelectionStrategyConfig.newBuilder().build()).withDispatcher(Dispatchers.DefaultDispatcherId()), peerId1.toString());
 571         leader.tell(new PeerUp(peerMemberName1, peerId1.toString()), ActorRef.noSender());
 572
 573         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID4, "");
 574         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID3, LOCAL_MEMBER_NAME);
 575         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 576         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 577
 578         // Verify the reinstated peerMember1 is fully synced.
 579
 580         verifyOwner(peer1, ENTITY_TYPE, ENTITY_ID4, "");
 581         verifyOwner(peer1, ENTITY_TYPE, ENTITY_ID3, LOCAL_MEMBER_NAME);
 582         verifyOwner(peer1, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 583         verifyOwner(peer1, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 584
 585         // Kill the local leader and elect peer2 the leader. This should cause a new owner to be selected for
 586         // the entities (1 and 3) previously owned by the local leader member.
 587
 588
 589         leader.tell(PoisonPill.getInstance(), ActorRef.noSender());
 590         peer2.tell(new PeerDown(LOCAL_MEMBER_NAME, leaderId.toString()), ActorRef.noSender());
 591         peer2.tell(new ElectionTimeout(), peer2);
 592
 593         kit.waitUntilLeader(peer2);
 594
 595         verifyOwner(peer2, ENTITY_TYPE, ENTITY_ID4, "");
 596         verifyOwner(peer2, ENTITY_TYPE, ENTITY_ID3, peerMemberName2);
 597         verifyOwner(peer2, ENTITY_TYPE, ENTITY_ID2, peerMemberName2);
 598         verifyOwner(peer2, ENTITY_TYPE, ENTITY_ID1, peerMemberName2);
 599     }
 600
 601     @Test
 602     public void testCandidateRemovedWhenCandidateNotRegisteredLocally() throws Exception {
 603         ShardTestKit kit = new ShardTestKit(getSystem());
 604
 605         dataStoreContextBuilder.shardHeartbeatIntervalInMillis(500).shardElectionTimeoutFactor(10000);
 606
 607         String peerMemberName1 = "peerMember1";
 608         String peerMemberName2 = "peerMember2";
 609
 610         ShardIdentifier leaderId = newShardId(LOCAL_MEMBER_NAME);
 611         ShardIdentifier peerId1 = newShardId(peerMemberName1);
 612         ShardIdentifier peerId2 = newShardId(peerMemberName2);
 613
 614         TestActorRef<EntityOwnershipShard> peer1 = actorFactory.createTestActor(newShardProps(peerId1,
 615                 ImmutableMap.<String, String>builder().put(leaderId.toString(), "").put(peerId2.toString(), "").build(),
 616                 peerMemberName1, EntityOwnerSelectionStrategyConfig.newBuilder().build()).withDispatcher(Dispatchers.DefaultDispatcherId()), peerId1.toString());
 617
 618         TestActorRef<EntityOwnershipShard> peer2 = actorFactory.createTestActor(newShardProps(peerId2,
 619                 ImmutableMap.<String, String>builder().put(leaderId.toString(), "").put(peerId1.toString(), "").build(),
 620                 peerMemberName2, EntityOwnerSelectionStrategyConfig.newBuilder().build()).withDispatcher(Dispatchers.DefaultDispatcherId()), peerId2.toString());
 621
 622         TestActorRef<EntityOwnershipShard> leader = actorFactory.createTestActor(newShardProps(leaderId,
 623                 ImmutableMap.<String, String>builder().put(peerId1.toString(), peer1.path().toString()).
 624                         put(peerId2.toString(), peer2.path().toString()).build(), LOCAL_MEMBER_NAME, EntityOwnerSelectionStrategyConfig.newBuilder().build()).
 625                 withDispatcher(Dispatchers.DefaultDispatcherId()), leaderId.toString());
 626         leader.tell(new ElectionTimeout(), leader);
 627
 628         kit.waitUntilLeader(leader);
 629
 630         peer2.tell(new PeerAddressResolved(peerId1.toString(), peer1.path().toString()), ActorRef.noSender());
 631         peer2.tell(new PeerAddressResolved(leaderId.toString(), leader.path().toString()), ActorRef.noSender());
 632         peer1.tell(new PeerAddressResolved(peerId2.toString(), peer2.path().toString()), ActorRef.noSender());
 633         peer1.tell(new PeerAddressResolved(leaderId.toString(), leader.path().toString()), ActorRef.noSender());
 634
 635         leader.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID1)), kit.getRef());
 636         kit.expectMsgClass(SuccessReply.class);
 637         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 638
 639         peer2.tell(new RegisterCandidateLocal(new Entity(ENTITY_TYPE, ENTITY_ID1)), kit.getRef());
 640         kit.expectMsgClass(SuccessReply.class);
 641         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, peerMemberName2), kit);
 642         verifyCommittedEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID1, peerMemberName2);
 643
 644         // Try adding  peerMemberName1 as candidate without registering locally
 645         commitModification(leader, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, peerMemberName1), kit);
 646
 647         verifyOwner(leader, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 648
 649         // confirm peerMemberName1 is not candidate as was not registered locally
 650         verifyNoEntityCandidate(leader, ENTITY_TYPE, ENTITY_ID1, peerMemberName1);
 651     }
 652
 653
 654     @Test
 655     public void testLocalCandidateRemovedWithCandidateRegistered() throws Exception {
 656         ShardTestKit kit = new ShardTestKit(getSystem());
 657
 658         dataStoreContextBuilder.shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(10000);
 659         ShardIdentifier leaderId = newShardId("leader");
 660         ShardIdentifier localId = newShardId(LOCAL_MEMBER_NAME);
 661
 662         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(Props.create(
 663                 TestEntityOwnershipShard.class, localId,
 664                 ImmutableMap.<String, String>builder().put(leaderId.toString(), "".toString()).build(),
 665                 dataStoreContextBuilder.build()).withDispatcher(Dispatchers.DefaultDispatcherId()));
 666
 667         TestActorRef<EntityOwnershipShard> leader = actorFactory.createTestActor(newShardProps(leaderId,
 668                 ImmutableMap.<String, String>builder().put(localId.toString(), shard.path().toString()).build(),
 669                     LOCAL_MEMBER_NAME, EntityOwnerSelectionStrategyConfig.newBuilder().build()).withDispatcher(Dispatchers.DefaultDispatcherId()), leaderId.toString());
 670         leader.tell(new ElectionTimeout(), leader);
 671
 672         kit.waitUntilLeader(leader);
 673
 674         shard.tell(new PeerAddressResolved(leaderId.toString(), leader.path().toString()), ActorRef.noSender());
 675
 676         Entity entity = new Entity(ENTITY_TYPE, ENTITY_ID1);
 677         EntityOwnershipListener listener = mock(EntityOwnershipListener.class);
 678
 679         shard.tell(new RegisterListenerLocal(listener, ENTITY_TYPE), kit.getRef());
 680         kit.expectMsgClass(SuccessReply.class);
 681
 682         // Register local candidate
 683
 684         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 685         kit.expectMsgClass(SuccessReply.class);
 686         verifyCommittedEntityCandidate(shard, entity.getType(), entity.getId(), LOCAL_MEMBER_NAME);
 687         verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity, false, true, true));
 688         reset(listener);
 689
 690         // Simulate a replicated commit from the leader to remove the local candidate that would occur after a
 691         // network partition is healed.
 692
 693         leader.tell(new PeerDown(LOCAL_MEMBER_NAME, localId.toString()), ActorRef.noSender());
 694
 695         verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity, true, false, false));
 696
 697         // Since the the shard has a local candidate registered, it should re-add its candidate to the entity.
 698
 699         verifyCommittedEntityCandidate(shard, entity.getType(), entity.getId(), LOCAL_MEMBER_NAME);
 700         verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity, false, true, true));
 701
 702         // Unregister the local candidate and verify it's removed and no re-added.
 703
 704         shard.tell(new UnregisterCandidateLocal(entity), kit.getRef());
 705         kit.expectMsgClass(SuccessReply.class);
 706
 707         verifyNoEntityCandidate(shard, entity.getType(), entity.getId(), LOCAL_MEMBER_NAME);
 708         Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
 709         verifyNoEntityCandidate(shard, entity.getType(), entity.getId(), LOCAL_MEMBER_NAME);
 710     }
 711
 712     @Test
 713     public void testListenerRegistration() throws Exception {
 714         ShardTestKit kit = new ShardTestKit(getSystem());
 715         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps());
 716         kit.waitUntilLeader(shard);
 717         ShardDataTree shardDataTree = shard.underlyingActor().getDataStore();
 718
 719         String otherEntityType = "otherEntityType";
 720         Entity entity1 = new Entity(ENTITY_TYPE, ENTITY_ID1);
 721         Entity entity2 = new Entity(ENTITY_TYPE, ENTITY_ID2);
 722         Entity entity3 = new Entity(ENTITY_TYPE, ENTITY_ID3);
 723         Entity entity4 = new Entity(otherEntityType, ENTITY_ID3);
 724         EntityOwnershipListener listener = mock(EntityOwnershipListener.class);
 725
 726         // Register listener
 727
 728         shard.tell(new RegisterListenerLocal(listener, ENTITY_TYPE), kit.getRef());
 729         kit.expectMsgClass(SuccessReply.class);
 730
 731         // Register a couple candidates for the desired entity type and verify listener is notified.
 732
 733         shard.tell(new RegisterCandidateLocal(entity1), kit.getRef());
 734         kit.expectMsgClass(SuccessReply.class);
 735
 736         verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity1, false, true, true));
 737
 738         shard.tell(new RegisterCandidateLocal(entity2), kit.getRef());
 739         kit.expectMsgClass(SuccessReply.class);
 740
 741         verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity2, false, true, true));
 742         reset(listener);
 743
 744         // Register another candidate for another entity type and verify listener is not notified.
 745
 746         shard.tell(new RegisterCandidateLocal(entity4), kit.getRef());
 747         kit.expectMsgClass(SuccessReply.class);
 748
 749         Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
 750         verify(listener, never()).ownershipChanged(ownershipChange(entity4));
 751
 752         // Register remote candidate for entity1
 753
 754         String remoteMemberName = "remoteMember";
 755         writeNode(ENTITY_OWNERS_PATH, entityOwnersWithCandidate(ENTITY_TYPE, entity1.getId(), remoteMemberName),
 756                 shardDataTree);
 757         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, entity1.getId(), remoteMemberName);
 758
 759         // Unregister the local candidate for entity1 and verify listener is notified
 760
 761         shard.tell(new UnregisterCandidateLocal(entity1), kit.getRef());
 762         kit.expectMsgClass(SuccessReply.class);
 763
 764         verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity1, true, false, true));
 765         reset(listener);
 766
 767         // Unregister the listener, add a candidate for entity3 and verify listener isn't notified
 768
 769         shard.tell(new UnregisterListenerLocal(listener, ENTITY_TYPE), kit.getRef());
 770         kit.expectMsgClass(SuccessReply.class);
 771
 772         shard.tell(new RegisterCandidateLocal(entity3), kit.getRef());
 773         kit.expectMsgClass(SuccessReply.class);
 774
 775         verifyOwner(shard, ENTITY_TYPE, entity3.getId(), LOCAL_MEMBER_NAME);
 776         Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
 777         verify(listener, never()).ownershipChanged(any(EntityOwnershipChange.class));
 778
 779         // Re-register the listener and verify it gets notified of currently owned entities
 780
 781         reset(listener);
 782
 783         shard.tell(new RegisterListenerLocal(listener, ENTITY_TYPE), kit.getRef());
 784         kit.expectMsgClass(SuccessReply.class);
 785
 786         verify(listener, timeout(5000).times(2)).ownershipChanged(or(ownershipChange(entity2, false, true, true),
 787                 ownershipChange(entity3, false, true, true)));
 788         Uninterruptibles.sleepUninterruptibly(300, TimeUnit.MILLISECONDS);
 789         verify(listener, never()).ownershipChanged(ownershipChange(entity4));
 790         verify(listener, times(1)).ownershipChanged(ownershipChange(entity1));
 791     }
 792
 793     private static void commitModification(TestActorRef<EntityOwnershipShard> shard, NormalizedNode<?, ?> node,
 794             JavaTestKit sender) {
 795         BatchedModifications modifications = newBatchedModifications();
 796         modifications.addModification(new MergeModification(ENTITY_OWNERS_PATH, node));
 797
 798         shard.tell(modifications, sender.getRef());
 799         sender.expectMsgClass(CommitTransactionReply.SERIALIZABLE_CLASS);
 800     }
 801
 802     private static BatchedModifications newBatchedModifications() {
 803         BatchedModifications modifications = new BatchedModifications("tnx", DataStoreVersions.CURRENT_VERSION, "");
 804         modifications.setDoCommitOnReady(true);
 805         modifications.setReady(true);
 806         modifications.setTotalMessagesSent(1);
 807         return modifications;
 808     }
 809
 810     private void verifyEntityCandidateRemoved(final TestActorRef<EntityOwnershipShard> shard, String entityType,
 811             YangInstanceIdentifier entityId, String candidateName) {
 812         verifyNodeRemoved(candidatePath(entityType, entityId, candidateName),
 813                 new Function<YangInstanceIdentifier, NormalizedNode<?,?>>() {
 814                     @Override
 815                     public NormalizedNode<?, ?> apply(YangInstanceIdentifier path) {
 816                         try {
 817                             return AbstractShardTest.readStore(shard, path);
 818                         } catch(Exception e) {
 819                             throw new AssertionError("Failed to read " + path, e);
 820                         }
 821                 }
 822         });
 823     }
 824
 825     private void verifyCommittedEntityCandidate(final TestActorRef<EntityOwnershipShard> shard, String entityType,
 826             YangInstanceIdentifier entityId, String candidateName) {
 827         verifyEntityCandidate(entityType, entityId, candidateName, new Function<YangInstanceIdentifier, NormalizedNode<?,?>>() {
 828             @Override
 829             public NormalizedNode<?, ?> apply(YangInstanceIdentifier path) {
 830                 try {
 831                     return AbstractShardTest.readStore(shard, path);
 832                 } catch(Exception e) {
 833                     throw new AssertionError("Failed to read " + path, e);
 834                 }
 835             }
 836         });
 837     }
 838
 839     private void verifyNoEntityCandidate(final TestActorRef<EntityOwnershipShard> shard, String entityType,
 840             YangInstanceIdentifier entityId, String candidateName) {
 841         verifyEntityCandidate(entityType, entityId, candidateName, new Function<YangInstanceIdentifier, NormalizedNode<?,?>>() {
 842             @Override
 843             public NormalizedNode<?, ?> apply(YangInstanceIdentifier path) {
 844                 try {
 845                     return AbstractShardTest.readStore(shard, path);
 846                 } catch(Exception e) {
 847                     throw new AssertionError("Failed to read " + path, e);
 848                 }
 849             }
 850         }, false);
 851     }
 852
 853     private void verifyBatchedEntityCandidate(List<Modification> mods, String entityType,
 854             YangInstanceIdentifier entityId, String candidateName) throws Exception {
 855         assertEquals("BatchedModifications size", 1, mods.size());
 856         verifyBatchedEntityCandidate(mods.get(0), entityType, entityId, candidateName);
 857     }
 858
 859     private void verifyBatchedEntityCandidate(Modification mod, String entityType,
 860             YangInstanceIdentifier entityId, String candidateName) throws Exception {
 861         assertEquals("Modification type", MergeModification.class, mod.getClass());
 862         verifyEntityCandidate(((MergeModification)mod).getData(), entityType,
 863                 entityId, candidateName, true);
 864     }
 865
 866     private static void verifyOwner(final TestActorRef<EntityOwnershipShard> shard, String entityType,
 867             YangInstanceIdentifier entityId, String localMemberName) {
 868         verifyOwner(localMemberName, entityType, entityId, new Function<YangInstanceIdentifier, NormalizedNode<?,?>>() {
 869             @Override
 870             public NormalizedNode<?, ?> apply(YangInstanceIdentifier path) {
 871                 try {
 872                     return AbstractShardTest.readStore(shard, path);
 873                 } catch(Exception e) {
 874                     return null;
 875                 }
 876             }
 877         });
 878     }
 879
 880     private Props newShardProps() {
 881         return newShardProps(Collections.<String,String>emptyMap());
 882     }
 883
 884     private Props newShardProps(EntityOwnerSelectionStrategyConfig strategyConfig) {
 885         return newShardProps(newShardId(LOCAL_MEMBER_NAME), Collections.<String,String>emptyMap(),
 886                 LOCAL_MEMBER_NAME, strategyConfig);
 887     }
 888
 889     private Props newShardProps(EntityOwnerSelectionStrategyConfig strategyConfig, Map<String, String> peers) {
 890         return newShardProps(newShardId(LOCAL_MEMBER_NAME), peers, LOCAL_MEMBER_NAME, strategyConfig);
 891     }
 892
 893
 894     private Props newShardProps(Map<String,String> peers) {
 895         return newShardProps(newShardId(LOCAL_MEMBER_NAME), peers, LOCAL_MEMBER_NAME, EntityOwnerSelectionStrategyConfig.newBuilder().build());
 896     }
 897
 898     private Props newShardProps(ShardIdentifier shardId, Map<String,String> peers, String memberName,
 899                                 EntityOwnerSelectionStrategyConfig config) {
 900         return EntityOwnershipShard.newBuilder().id(shardId).peerAddresses(peers).
 901                 datastoreContext(dataStoreContextBuilder.build()).schemaContext(SCHEMA_CONTEXT).
 902                 localMemberName(memberName).ownerSelectionStrategyConfig(config).props().withDispatcher(Dispatchers.DefaultDispatcherId());
 903     }
 904
 905     private static ShardIdentifier newShardId(String memberName) {
 906         return ShardIdentifier.builder().memberName(memberName).shardName("entity-ownership").
 907                 type("operational" + NEXT_SHARD_NUM.getAndIncrement()).build();
 908     }
 909
 910     public static class TestEntityOwnershipShard extends EntityOwnershipShard {
 911
 912         TestEntityOwnershipShard(ShardIdentifier name, Map<String, String> peerAddresses,
 913                 DatastoreContext datastoreContext) {
 914             super(newBuilder().id(name).peerAddresses(peerAddresses).datastoreContext(datastoreContext).
 915                     schemaContext(SCHEMA_CONTEXT).localMemberName(LOCAL_MEMBER_NAME));
 916         }
 917
 918         @Override
 919         public void onReceiveCommand(Object message) throws Exception {
 920             if(!(message instanceof ElectionTimeout)) {
 921                 super.onReceiveCommand(message);
 922             }
 923         }
 924
 925
 926     }
 927
 928     public static class MockFollower extends UntypedActor {
 929         volatile boolean grantVote;
 930         volatile boolean dropAppendEntries;
 931         private final String myId;
 932
 933         public MockFollower(String myId) {
 934             this(myId, true);
 935         }
 936
 937         public MockFollower(String myId, boolean grantVote) {
 938             this.myId = myId;
 939             this.grantVote = grantVote;
 940         }
 941
 942         @Override
 943         public void onReceive(Object message) {
 944             if(message instanceof RequestVote) {
 945                 if(grantVote) {
 946                     getSender().tell(new RequestVoteReply(((RequestVote)message).getTerm(), true), getSelf());
 947                 }
 948             } else if(message instanceof AppendEntries) {
 949                 if(!dropAppendEntries) {
 950                     AppendEntries req = (AppendEntries) message;
 951                     long lastIndex = req.getLeaderCommit();
 952                     if (req.getEntries().size() > 0) {
 953                         for(ReplicatedLogEntry entry : req.getEntries()) {
 954                             lastIndex = entry.getIndex();
 955                         }
 956                     }
 957
 958                     getSender().tell(new AppendEntriesReply(myId, req.getTerm(), true, lastIndex, req.getTerm(),
 959                             DataStoreVersions.CURRENT_VERSION), getSelf());
 960                 }
 961             }
 962         }
 963     }
 964
 965
 966     @Test
 967     public void testDelayedEntityOwnerSelectionWhenMaxPeerRequestsReceived() throws Exception {
 968         ShardTestKit kit = new ShardTestKit(getSystem());
 969         EntityOwnerSelectionStrategyConfig.Builder builder
 970                 = EntityOwnerSelectionStrategyConfig.newBuilder().addStrategy(ENTITY_TYPE, LastCandidateSelectionStrategy.class, 500);
 971
 972         String peerId = newShardId("follower").toString();
 973         TestActorRef<MockFollower> peer = actorFactory.createTestActor(Props.create(MockFollower.class, peerId, false).
 974                 withDispatcher(Dispatchers.DefaultDispatcherId()), peerId);
 975
 976         peer.underlyingActor().grantVote = true;
 977
 978         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps(builder.build(),
 979                 ImmutableMap.of(peerId.toString(), peer.path().toString())));
 980         kit.waitUntilLeader(shard);
 981
 982         Entity entity = new Entity(ENTITY_TYPE, ENTITY_ID1);
 983         ShardDataTree shardDataTree = shard.underlyingActor().getDataStore();
 984
 985         // Add a remote candidate
 986
 987         String remoteMemberName1 = "follower";
 988         writeNode(ENTITY_OWNERS_PATH, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, remoteMemberName1), shardDataTree);
 989
 990         // Register local
 991
 992         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
 993         kit.expectMsgClass(SuccessReply.class);
 994
 995         // Verify the local candidate becomes owner
 996
 997         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName1);
 998         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
 999         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
1000     }
1001
1002     @Test
1003     public void testDelayedEntityOwnerSelection() throws Exception {
1004         ShardTestKit kit = new ShardTestKit(getSystem());
1005         EntityOwnerSelectionStrategyConfig.Builder builder
1006                 = EntityOwnerSelectionStrategyConfig.newBuilder().addStrategy(ENTITY_TYPE, LastCandidateSelectionStrategy.class, 500);
1007
1008         String follower1Id = newShardId("follower1").toString();
1009         TestActorRef<MockFollower> follower1 = actorFactory.createTestActor(Props.create(MockFollower.class, follower1Id, false).
1010                 withDispatcher(Dispatchers.DefaultDispatcherId()), follower1Id);
1011
1012         follower1.underlyingActor().grantVote = true;
1013
1014         String follower2Id = newShardId("follower").toString();
1015         TestActorRef<MockFollower> follower2 = actorFactory.createTestActor(Props.create(MockFollower.class, follower2Id, false).
1016                 withDispatcher(Dispatchers.DefaultDispatcherId()), follower2Id);
1017
1018         follower2.underlyingActor().grantVote = true;
1019
1020
1021         TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps(builder.build(),
1022                 ImmutableMap.of(follower1Id.toString(), follower2.path().toString(), follower2Id.toString(), follower2.path().toString())));
1023         kit.waitUntilLeader(shard);
1024
1025         Entity entity = new Entity(ENTITY_TYPE, ENTITY_ID1);
1026         ShardDataTree shardDataTree = shard.underlyingActor().getDataStore();
1027
1028         // Add a remote candidate
1029
1030         String remoteMemberName1 = "follower";
1031         writeNode(ENTITY_OWNERS_PATH, entityOwnersWithCandidate(ENTITY_TYPE, ENTITY_ID1, remoteMemberName1), shardDataTree);
1032
1033         // Register local
1034
1035         shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
1036         kit.expectMsgClass(SuccessReply.class);
1037
1038         // Verify the local candidate becomes owner
1039
1040         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, remoteMemberName1);
1041         verifyCommittedEntityCandidate(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
1042         verifyOwner(shard, ENTITY_TYPE, ENTITY_ID1, LOCAL_MEMBER_NAME);
1043     }
1044
1045     public static class MockLeader extends UntypedActor {
1046         volatile CountDownLatch modificationsReceived = new CountDownLatch(1);
1047         List<Modification> receivedModifications = new ArrayList<>();
1048         volatile boolean sendReply = true;
1049         volatile long delay;
1050
1051         @Override
1052         public void onReceive(Object message) {
1053             if(message instanceof BatchedModifications) {
1054                 if(delay > 0) {
1055                     Uninterruptibles.sleepUninterruptibly(delay, TimeUnit.MILLISECONDS);
1056                 }
1057
1058                 if(sendReply) {
1059                     BatchedModifications mods = (BatchedModifications) message;
1060                     synchronized (receivedModifications) {
1061                         for(int i = 0; i < mods.getModifications().size(); i++) {
1062                             receivedModifications.add(mods.getModifications().get(i));
1063                             modificationsReceived.countDown();
1064                         }
1065                     }
1066
1067                     getSender().tell(CommitTransactionReply.INSTANCE.toSerializable(), getSelf());
1068                 } else {
1069                     sendReply = true;
1070                 }
1071             }
1072         }
1073
1074         List<Modification> getAndClearReceivedModifications() {
1075             synchronized (receivedModifications) {
1076                 List<Modification> ret = new ArrayList<>(receivedModifications);
1077                 receivedModifications.clear();
1078                 return ret;
1079             }
1080         }
1081     }
1082 }