From: Tom Pantelis Date: Fri, 25 Sep 2015 02:05:24 +0000 (-0400) Subject: Add getOwnershipState method to EntityOwnershipService X-Git-Tag: release/beryllium~260 X-Git-Url: https://git.opendaylight.org/gerrit/gitweb?p=controller.git;a=commitdiff_plain;h=689370fdbfdd18d6bd25359b94b1a8f6d76aba1c Add getOwnershipState method to EntityOwnershipService Added a new method to gte the current ownership state for an entity. This was requested for OF clustering. The DistributedEntityOwnershipService obtains the EntityOwnershipShard's DataTree via a new message GetShardDataTree and reads the entity's owner leaf in order to build the resulting EntityOwnershipState. The DataTree is obtained once and cached. Change-Id: Ib4aa2f4e5370d8d5183908b836417936a51458f7 Signed-off-by: Tom Pantelis (cherry picked from commit dd6976c24f12c7cef7bed8fa6bc645dc699dda4f) --- diff --git a/opendaylight/md-sal/sal-common-api/src/main/java/org/opendaylight/controller/md/sal/common/api/clustering/EntityOwnershipService.java b/opendaylight/md-sal/sal-common-api/src/main/java/org/opendaylight/controller/md/sal/common/api/clustering/EntityOwnershipService.java index f200f5adc1..311d1b8595 100644 --- a/opendaylight/md-sal/sal-common-api/src/main/java/org/opendaylight/controller/md/sal/common/api/clustering/EntityOwnershipService.java +++ b/opendaylight/md-sal/sal-common-api/src/main/java/org/opendaylight/controller/md/sal/common/api/clustering/EntityOwnershipService.java @@ -8,6 +8,7 @@ package org.opendaylight.controller.md.sal.common.api.clustering; +import com.google.common.base.Optional; import javax.annotation.Nonnull; /** @@ -50,4 +51,11 @@ public interface EntityOwnershipService { */ EntityOwnershipListenerRegistration registerListener(@Nonnull String entityType, @Nonnull EntityOwnershipListener listener); + /** + * Gets the current ownership state information for an entity. + * + * @param forEntity the entity to query. + * @return an Optional EntityOwnershipState whose instance is present if the entity is found + */ + Optional getOwnershipState(@Nonnull Entity forEntity); } diff --git a/opendaylight/md-sal/sal-common-api/src/main/java/org/opendaylight/controller/md/sal/common/api/clustering/EntityOwnershipState.java b/opendaylight/md-sal/sal-common-api/src/main/java/org/opendaylight/controller/md/sal/common/api/clustering/EntityOwnershipState.java new file mode 100644 index 0000000000..e461921df7 --- /dev/null +++ b/opendaylight/md-sal/sal-common-api/src/main/java/org/opendaylight/controller/md/sal/common/api/clustering/EntityOwnershipState.java @@ -0,0 +1,45 @@ +/* + * Copyright (c) 2015 Brocade Communications Systems, Inc. and others. All rights reserved. + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License v1.0 which accompanies this distribution, + * and is available at http://www.eclipse.org/legal/epl-v10.html + */ +package org.opendaylight.controller.md.sal.common.api.clustering; + +/** + * A DTO that encapsulates ownership state for an entity. + * + * @author Thomas Pantelis + */ +public class EntityOwnershipState { + private final boolean isOwner; + private final boolean hasOwner; + + public EntityOwnershipState(boolean isOwner, boolean hasOwner) { + this.isOwner = isOwner; + this.hasOwner = hasOwner; + } + + /** + * Returns the current ownership status of the entity for this process instance. + * @return true if this process is the owner of the entity + */ + public boolean isOwner() { + return isOwner; + } + + /** + * Returns the current ownership status of the entity across all process instances. + * @return true if the entity has an owner which may or may not be this process. If false, then + * the entity has no candidates and thus no owner. + */ + public boolean hasOwner() { + return hasOwner; + } + + @Override + public String toString() { + return "EntityOwnershipState [isOwner=" + isOwner + ", hasOwner=" + hasOwner + "]"; + } +} diff --git a/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/Shard.java b/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/Shard.java index 5025369907..edf6f57760 100644 --- a/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/Shard.java +++ b/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/Shard.java @@ -41,6 +41,7 @@ import org.opendaylight.controller.cluster.datastore.messages.CommitTransactionR import org.opendaylight.controller.cluster.datastore.messages.CreateTransaction; import org.opendaylight.controller.cluster.datastore.messages.CreateTransactionReply; import org.opendaylight.controller.cluster.datastore.messages.ForwardedReadyTransaction; +import org.opendaylight.controller.cluster.datastore.messages.GetShardDataTree; import org.opendaylight.controller.cluster.datastore.messages.PeerAddressResolved; import org.opendaylight.controller.cluster.datastore.messages.ReadyLocalTransaction; import org.opendaylight.controller.cluster.datastore.messages.RegisterChangeListener; @@ -257,6 +258,8 @@ public class Shard extends RaftActor { context().parent().tell(message, self()); } else if(GET_SHARD_MBEAN_MESSAGE.equals(message)){ sender().tell(getShardMBean(), self()); + } else if(message instanceof GetShardDataTree){ + sender().tell(store.getDataTree(), self()); } else { super.onReceiveCommand(message); } diff --git a/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipService.java b/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipService.java index e272b80a67..09c4e50335 100644 --- a/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipService.java +++ b/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipService.java @@ -7,11 +7,16 @@ */ package org.opendaylight.controller.cluster.datastore.entityownership; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.ENTITY_OWNER_NODE_ID; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.entityPath; import akka.actor.ActorRef; import akka.dispatch.OnComplete; +import akka.pattern.Patterns; import akka.util.Timeout; import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Optional; import com.google.common.base.Preconditions; +import com.google.common.base.Strings; import java.util.Collection; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -24,6 +29,7 @@ import org.opendaylight.controller.cluster.datastore.entityownership.messages.Re import org.opendaylight.controller.cluster.datastore.entityownership.messages.UnregisterCandidateLocal; import org.opendaylight.controller.cluster.datastore.entityownership.messages.UnregisterListenerLocal; import org.opendaylight.controller.cluster.datastore.messages.CreateShard; +import org.opendaylight.controller.cluster.datastore.messages.GetShardDataTree; import org.opendaylight.controller.cluster.datastore.shardstrategy.ModuleShardStrategy; import org.opendaylight.controller.md.sal.common.api.clustering.CandidateAlreadyRegisteredException; import org.opendaylight.controller.md.sal.common.api.clustering.Entity; @@ -31,10 +37,18 @@ import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipC import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipListener; import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipListenerRegistration; import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipService; +import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipState; import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.controller.md.sal.clustering.entity.owners.rev150804.EntityOwners; +import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.PathArgument; +import org.opendaylight.yangtools.yang.data.api.schema.DataContainerChild; +import org.opendaylight.yangtools.yang.data.api.schema.MapEntryNode; +import org.opendaylight.yangtools.yang.data.api.schema.NormalizedNode; +import org.opendaylight.yangtools.yang.data.api.schema.tree.DataTree; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import scala.concurrent.Await; import scala.concurrent.Future; +import scala.concurrent.duration.Duration; /** * The distributed implementation of the EntityOwnershipService. @@ -49,6 +63,7 @@ public class DistributedEntityOwnershipService implements EntityOwnershipService private final DistributedDataStore datastore; private final ConcurrentMap registeredEntities = new ConcurrentHashMap<>(); private volatile ActorRef localEntityOwnershipShard; + private volatile DataTree localEntityOwnershipShardDataTree; public DistributedEntityOwnershipService(DistributedDataStore datastore) { this.datastore = datastore; @@ -149,6 +164,49 @@ public class DistributedEntityOwnershipService implements EntityOwnershipService return new DistributedEntityOwnershipListenerRegistration(listener, entityType, this); } + @Override + public Optional getOwnershipState(Entity forEntity) { + Preconditions.checkNotNull(forEntity, "forEntity cannot be null"); + + DataTree dataTree = getLocalEntityOwnershipShardDataTree(); + if(dataTree == null) { + return Optional.absent(); + } + + Optional> entityNode = dataTree.takeSnapshot().readNode( + entityPath(forEntity.getType(), forEntity.getId())); + if(!entityNode.isPresent()) { + return Optional.absent(); + } + + String localMemberName = datastore.getActorContext().getCurrentMemberName(); + Optional> ownerLeaf = ((MapEntryNode)entityNode.get()). + getChild(ENTITY_OWNER_NODE_ID); + String owner = ownerLeaf.isPresent() ? ownerLeaf.get().getValue().toString() : null; + boolean hasOwner = !Strings.isNullOrEmpty(owner); + boolean isOwner = hasOwner && localMemberName.equals(owner); + + return Optional.of(new EntityOwnershipState(isOwner, hasOwner)); + } + + private DataTree getLocalEntityOwnershipShardDataTree() { + if(localEntityOwnershipShardDataTree == null) { + try { + if(localEntityOwnershipShard == null) { + localEntityOwnershipShard = Await.result(datastore.getActorContext().findLocalShardAsync( + ENTITY_OWNERSHIP_SHARD_NAME), Duration.Inf()); + } + + localEntityOwnershipShardDataTree = (DataTree) Await.result(Patterns.ask(localEntityOwnershipShard, + GetShardDataTree.INSTANCE, MESSAGE_TIMEOUT), Duration.Inf()); + } catch (Exception e) { + LOG.error("Failed to find local {} shard", ENTITY_OWNERSHIP_SHARD_NAME, e); + } + } + + return localEntityOwnershipShardDataTree; + } + void unregisterListener(String entityType, EntityOwnershipListener listener) { LOG.debug("Unregistering listener {} for entity type {}", listener, entityType); diff --git a/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/entityownership/EntityOwnersModel.java b/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/entityownership/EntityOwnersModel.java index 1c97b4818b..404d9b9666 100644 --- a/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/entityownership/EntityOwnersModel.java +++ b/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/entityownership/EntityOwnersModel.java @@ -106,7 +106,7 @@ final class EntityOwnersModel { return ImmutableNodes.mapEntry(Candidate.QNAME, CANDIDATE_NAME_QNAME, candidateName); } - static NormalizedNode entityEntryWithOwner(YangInstanceIdentifier entityId, String owner) { + static MapEntryNode entityEntryWithOwner(YangInstanceIdentifier entityId, String owner) { return ImmutableNodes.mapEntryBuilder(ENTITY_QNAME, ENTITY_ID_QNAME, entityId).addChild( ImmutableNodes.leafNode(ENTITY_OWNER_QNAME, owner)).build(); } diff --git a/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/messages/GetShardDataTree.java b/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/messages/GetShardDataTree.java new file mode 100644 index 0000000000..0d63b6bc6d --- /dev/null +++ b/opendaylight/md-sal/sal-distributed-datastore/src/main/java/org/opendaylight/controller/cluster/datastore/messages/GetShardDataTree.java @@ -0,0 +1,20 @@ +/* + * Copyright (c) 2015 Brocade Communications Systems, Inc. and others. All rights reserved. + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License v1.0 which accompanies this distribution, + * and is available at http://www.eclipse.org/legal/epl-v10.html + */ +package org.opendaylight.controller.cluster.datastore.messages; + +/** + * Local message sent to a Shard to retrieve its data tree instance. + * + * @author Thomas Pantelis + */ +public class GetShardDataTree { + public static final GetShardDataTree INSTANCE = new GetShardDataTree(); + + private GetShardDataTree() { + } +} diff --git a/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipIntegrationTest.java b/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipIntegrationTest.java index 990769e017..108e0c5c6e 100644 --- a/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipIntegrationTest.java +++ b/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipIntegrationTest.java @@ -39,6 +39,7 @@ import org.opendaylight.controller.cluster.datastore.IntegrationTestKit; import org.opendaylight.controller.md.cluster.datastore.model.SchemaContextHelper; import org.opendaylight.controller.md.sal.common.api.clustering.Entity; import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipListener; +import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipState; import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.controller.md.sal.clustering.entity.owners.rev150804.entity.owners.entity.type.entity.Candidate; import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier; import org.opendaylight.yangtools.yang.data.api.schema.MapEntryNode; @@ -161,6 +162,9 @@ public class DistributedEntityOwnershipIntegrationTest { verify(follower1MockListener, timeout(5000)).ownershipChanged(ownershipChange(ENTITY1, false, false, true)); reset(leaderMockListener, follower1MockListener); + verifyGetOwnershipState(leaderEntityOwnershipService, ENTITY1, true, true); + verifyGetOwnershipState(follower1EntityOwnershipService, ENTITY1, false, true); + // Register leader candidate for entity1_2 (same id, different type) and verify it becomes owner leaderEntityOwnershipService.registerCandidate(ENTITY1_2); @@ -246,6 +250,14 @@ public class DistributedEntityOwnershipIntegrationTest { verify(follower1MockListener, timeout(5000)).ownershipChanged(ownershipChange(ENTITY2, false, false, false)); } + private void verifyGetOwnershipState(DistributedEntityOwnershipService service, Entity entity, + boolean isOwner, boolean hasOwner) { + Optional state = service.getOwnershipState(entity); + assertEquals("getOwnershipState present", true, state.isPresent()); + assertEquals("isOwner", isOwner, state.get().isOwner()); + assertEquals("hasOwner", hasOwner, state.get().hasOwner()); + } + private void verifyCandidates(DistributedDataStore dataStore, Entity entity, String... expCandidates) throws Exception { AssertionError lastError = null; Stopwatch sw = Stopwatch.createStarted(); diff --git a/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipServiceTest.java b/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipServiceTest.java index c1c16a98b9..5591772afc 100644 --- a/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipServiceTest.java +++ b/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/entityownership/DistributedEntityOwnershipServiceTest.java @@ -13,10 +13,18 @@ import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.ENTITY_ID_QNAME; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.ENTITY_OWNERS_PATH; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.ENTITY_QNAME; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.entityEntryWithOwner; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.entityOwnersWithEntityTypeEntry; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.entityPath; +import static org.opendaylight.controller.cluster.datastore.entityownership.EntityOwnersModel.entityTypeEntryWithEntityEntry; import akka.actor.ActorRef; import akka.actor.PoisonPill; import akka.actor.Props; import com.google.common.base.Function; +import com.google.common.base.Optional; import com.google.common.util.concurrent.Uninterruptibles; import java.util.Collections; import java.util.Map; @@ -28,6 +36,7 @@ import org.junit.Before; import org.junit.Test; import org.opendaylight.controller.cluster.datastore.DatastoreContext; import org.opendaylight.controller.cluster.datastore.DistributedDataStore; +import org.opendaylight.controller.cluster.datastore.ShardDataTree; import org.opendaylight.controller.cluster.datastore.config.Configuration; import org.opendaylight.controller.cluster.datastore.config.ConfigurationImpl; import org.opendaylight.controller.cluster.datastore.config.ModuleConfig; @@ -37,6 +46,7 @@ import org.opendaylight.controller.cluster.datastore.entityownership.messages.Re import org.opendaylight.controller.cluster.datastore.entityownership.messages.UnregisterCandidateLocal; import org.opendaylight.controller.cluster.datastore.entityownership.messages.UnregisterListenerLocal; import org.opendaylight.controller.cluster.datastore.identifiers.ShardIdentifier; +import org.opendaylight.controller.cluster.datastore.messages.GetShardDataTree; import org.opendaylight.controller.cluster.datastore.utils.MockClusterWrapper; import org.opendaylight.controller.md.cluster.datastore.model.SchemaContextHelper; import org.opendaylight.controller.md.sal.common.api.clustering.CandidateAlreadyRegisteredException; @@ -44,9 +54,12 @@ import org.opendaylight.controller.md.sal.common.api.clustering.Entity; import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipCandidateRegistration; import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipListener; import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipListenerRegistration; +import org.opendaylight.controller.md.sal.common.api.clustering.EntityOwnershipState; import org.opendaylight.yangtools.yang.common.QName; import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier; import org.opendaylight.yangtools.yang.data.api.schema.NormalizedNode; +import org.opendaylight.yangtools.yang.data.api.schema.tree.DataTree; +import org.opendaylight.yangtools.yang.data.impl.schema.ImmutableNodes; import org.opendaylight.yangtools.yang.model.api.SchemaContext; import scala.concurrent.Await; import scala.concurrent.Future; @@ -228,6 +241,53 @@ public class DistributedEntityOwnershipServiceTest extends AbstractEntityOwnersh service.close(); } + @Test + public void testGetOwnershipState() throws Exception { + final TestShardPropsCreator shardPropsCreator = new TestShardPropsCreator(); + DistributedEntityOwnershipService service = new DistributedEntityOwnershipService(dataStore) { + @Override + protected EntityOwnershipShardPropsCreator newShardPropsCreator() { + return shardPropsCreator; + } + }; + + service.start(); + + ShardDataTree shardDataTree = new ShardDataTree(SchemaContextHelper.entityOwners()); + shardPropsCreator.setDataTree(shardDataTree.getDataTree()); + + Entity entity1 = new Entity(ENTITY_TYPE, "one"); + writeNode(ENTITY_OWNERS_PATH, entityOwnersWithEntityTypeEntry(entityTypeEntryWithEntityEntry(entity1.getType(), + entityEntryWithOwner(entity1.getId(), "member-1"))), shardDataTree); + verifyGetOwnershipState(service, entity1, true, true); + + writeNode(entityPath(entity1.getType(), entity1.getId()), entityEntryWithOwner(entity1.getId(), "member-2"), + shardDataTree); + verifyGetOwnershipState(service, entity1, false, true); + + writeNode(entityPath(entity1.getType(), entity1.getId()), entityEntryWithOwner(entity1.getId(), ""), + shardDataTree); + verifyGetOwnershipState(service, entity1, false, false); + + Entity entity2 = new Entity(ENTITY_TYPE, "two"); + Optional state = service.getOwnershipState(entity2); + assertEquals("getOwnershipState present", false, state.isPresent()); + + writeNode(entityPath(entity2.getType(), entity2.getId()), ImmutableNodes.mapEntry(ENTITY_QNAME, + ENTITY_ID_QNAME, entity2.getId()), shardDataTree); + verifyGetOwnershipState(service, entity2, false, false); + + service.close(); + } + + private void verifyGetOwnershipState(DistributedEntityOwnershipService service, Entity entity, + boolean isOwner, boolean hasOwner) { + Optional state = service.getOwnershipState(entity); + assertEquals("getOwnershipState present", true, state.isPresent()); + assertEquals("isOwner", isOwner, state.get().isOwner()); + assertEquals("hasOwner", hasOwner, state.get().hasOwner()); + } + private void verifyEntityCandidate(ActorRef entityOwnershipShard, String entityType, YangInstanceIdentifier entityId, String candidateName) { verifyEntityCandidate(entityType, entityId, candidateName, @@ -261,12 +321,13 @@ public class DistributedEntityOwnershipServiceTest extends AbstractEntityOwnersh private final AtomicReference messageReceived = new AtomicReference<>(); private final AtomicReference receivedMessage = new AtomicReference<>(); private final AtomicReference> messageClass = new AtomicReference<>(); + private final AtomicReference dataTree = new AtomicReference<>(); @Override public Props newProps(ShardIdentifier shardId, Map peerAddresses, DatastoreContext datastoreContext, SchemaContext schemaContext) { return Props.create(TestEntityOwnershipShard.class, shardId, peerAddresses, datastoreContext, - schemaContext, "member-1", messageClass, messageReceived, receivedMessage); + schemaContext, "member-1", messageClass, messageReceived, receivedMessage, dataTree); } @SuppressWarnings("unchecked") @@ -282,27 +343,37 @@ public class DistributedEntityOwnershipServiceTest extends AbstractEntityOwnersh receivedMessage.set(null); messageClass.set(ofType); } + + void setDataTree(DataTree tree) { + this.dataTree.set(tree); + } } static class TestEntityOwnershipShard extends EntityOwnershipShard { private final AtomicReference messageReceived; private final AtomicReference receivedMessage; private final AtomicReference> messageClass; + private final AtomicReference dataTree; protected TestEntityOwnershipShard(ShardIdentifier name, Map peerAddresses, DatastoreContext datastoreContext, SchemaContext schemaContext, String localMemberName, AtomicReference> messageClass, AtomicReference messageReceived, - AtomicReference receivedMessage) { + AtomicReference receivedMessage, AtomicReference dataTree) { super(name, peerAddresses, datastoreContext, schemaContext, localMemberName); this.messageClass = messageClass; this.messageReceived = messageReceived; this.receivedMessage = receivedMessage; + this.dataTree = dataTree; } @Override public void onReceiveCommand(final Object message) throws Exception { try { - super.onReceiveCommand(message); + if(dataTree.get() != null && message instanceof GetShardDataTree) { + sender().tell(dataTree.get(), self()); + } else { + super.onReceiveCommand(message); + } } finally { Class expMsgClass = messageClass.get(); if(expMsgClass != null && expMsgClass.equals(message.getClass())) {