+ DOMStoreThreePhaseCommitCohort ready = transactionProxy.ready();
+
+ assertTrue(ready instanceof ThreePhaseCommitCohortProxy);
+
+ ThreePhaseCommitCohortProxy proxy = (ThreePhaseCommitCohortProxy) ready;
+
+ verifyCohortFutures(proxy, getSystem().actorSelection(actorRef.path()));
+ }
+
+ private static interface TransactionProxyOperation {
+ void run(TransactionProxy transactionProxy);
+ }
+
+ private void throttleOperation(TransactionProxyOperation operation) {
+ throttleOperation(operation, 1, true);
+ }
+
+ private void throttleOperation(TransactionProxyOperation operation, int outstandingOpsLimit, boolean shardFound){
+ ActorSystem actorSystem = getSystem();
+ ActorRef shardActorRef = actorSystem.actorOf(Props.create(DoNothingActor.class));
+
+ doReturn(outstandingOpsLimit).when(mockActorContext).getTransactionOutstandingOperationLimit();
+
+ doReturn(actorSystem.actorSelection(shardActorRef.path())).
+ when(mockActorContext).actorSelection(shardActorRef.path().toString());
+
+ if(shardFound) {
+ doReturn(Futures.successful(actorSystem.actorSelection(shardActorRef.path()))).
+ when(mockActorContext).findPrimaryShardAsync(eq(DefaultShardStrategy.DEFAULT_SHARD));
+ } else {
+ doReturn(Futures.failed(new Exception("not found")))
+ .when(mockActorContext).findPrimaryShardAsync(eq(DefaultShardStrategy.DEFAULT_SHARD));
+ }
+
+ String actorPath = "akka.tcp://system@127.0.0.1:2550/user/tx-actor";
+ CreateTransactionReply createTransactionReply = CreateTransactionReply.newBuilder().
+ setTransactionId("txn-1").setTransactionActorPath(actorPath).
+ setMessageVersion(DataStoreVersions.CURRENT_VERSION).build();
+
+ doReturn(Futures.successful(createTransactionReply)).when(mockActorContext).
+ executeOperationAsync(eq(actorSystem.actorSelection(shardActorRef.path())),
+ eqCreateTransaction(memberName, READ_WRITE));
+
+ doReturn(true).when(mockActorContext).isPathLocal(actorPath);
+
+ TransactionProxy transactionProxy = new TransactionProxy(mockActorContext, READ_WRITE);
+
+ long start = System.nanoTime();
+
+ operation.run(transactionProxy);
+
+ long end = System.nanoTime();
+
+ long expected = TimeUnit.SECONDS.toNanos(mockActorContext.getDatastoreContext().getOperationTimeoutInSeconds());
+ Assert.assertTrue(String.format("Expected elapsed time: %s. Actual: %s",
+ expected, (end-start)), (end - start) > expected);
+
+ }
+
+ private void completeOperation(TransactionProxyOperation operation){
+ completeOperation(operation, true);
+ }
+
+ private void completeOperation(TransactionProxyOperation operation, boolean shardFound){
+ ActorSystem actorSystem = getSystem();
+ ActorRef shardActorRef = actorSystem.actorOf(Props.create(DoNothingActor.class));
+
+ doReturn(1).when(mockActorContext).getTransactionOutstandingOperationLimit();
+
+ doReturn(actorSystem.actorSelection(shardActorRef.path())).
+ when(mockActorContext).actorSelection(shardActorRef.path().toString());
+
+ if(shardFound) {
+ doReturn(Futures.successful(actorSystem.actorSelection(shardActorRef.path()))).
+ when(mockActorContext).findPrimaryShardAsync(eq(DefaultShardStrategy.DEFAULT_SHARD));
+ } else {
+ doReturn(Futures.failed(new PrimaryNotFoundException("test")))
+ .when(mockActorContext).findPrimaryShardAsync(eq(DefaultShardStrategy.DEFAULT_SHARD));
+ }
+
+ ActorRef txActorRef = actorSystem.actorOf(Props.create(DoNothingActor.class));
+ String actorPath = txActorRef.path().toString();
+ CreateTransactionReply createTransactionReply = CreateTransactionReply.newBuilder().
+ setTransactionId("txn-1").setTransactionActorPath(actorPath).
+ setMessageVersion(DataStoreVersions.CURRENT_VERSION).build();
+
+ doReturn(actorSystem.actorSelection(actorPath)).when(mockActorContext).actorSelection(actorPath);
+
+ doReturn(Futures.successful(createTransactionReply)).when(mockActorContext).
+ executeOperationAsync(eq(actorSystem.actorSelection(shardActorRef.path())),
+ eqCreateTransaction(memberName, READ_WRITE));
+
+ doReturn(true).when(mockActorContext).isPathLocal(anyString());
+
+ TransactionProxy transactionProxy = new TransactionProxy(mockActorContext, READ_WRITE);
+
+ long start = System.nanoTime();
+
+ operation.run(transactionProxy);
+
+ long end = System.nanoTime();
+
+ long expected = TimeUnit.SECONDS.toNanos(mockActorContext.getDatastoreContext().getOperationTimeoutInSeconds());
+ Assert.assertTrue(String.format("Expected elapsed time: %s. Actual: %s",
+ expected, (end-start)), (end - start) <= expected);
+ }
+
+ public void testWriteThrottling(boolean shardFound){
+
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+ }
+ }, 1, shardFound);
+ }
+
+ @Test
+ public void testWriteThrottlingWhenShardFound(){
+ dataStoreContextBuilder.shardBatchedModificationCount(1);
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectIncompleteBatchedModifications();
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+ }
+ });
+ }
+
+ @Test
+ public void testWriteThrottlingWhenShardNotFound(){
+ // Confirm that there is no throttling when the Shard is not found
+ dataStoreContextBuilder.shardBatchedModificationCount(1);
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+ }
+ }, false);
+
+ }
+
+
+ @Test
+ public void testWriteCompletion(){
+ dataStoreContextBuilder.shardBatchedModificationCount(1);
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+ }
+ });
+ }
+
+ @Test
+ public void testMergeThrottlingWhenShardFound(){
+ dataStoreContextBuilder.shardBatchedModificationCount(1);
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToMerge = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectIncompleteBatchedModifications();
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+ }
+ });
+ }
+
+ @Test
+ public void testMergeThrottlingWhenShardNotFound(){
+ dataStoreContextBuilder.shardBatchedModificationCount(1);
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToMerge = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+ }
+ }, false);
+ }
+
+ @Test
+ public void testMergeCompletion(){
+ dataStoreContextBuilder.shardBatchedModificationCount(1);
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToMerge = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+ }
+ });
+
+ }
+
+ @Test
+ public void testDeleteThrottlingWhenShardFound(){
+
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ expectIncompleteBatchedModifications();
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+ }
+ });
+ }
+
+
+ @Test
+ public void testDeleteThrottlingWhenShardNotFound(){
+
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ expectBatchedModifications(2);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+ }
+ }, false);
+ }
+
+ @Test
+ public void testDeleteCompletion(){
+ dataStoreContextBuilder.shardBatchedModificationCount(1);
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ expectBatchedModifications(2);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+ }
+ });
+
+ }
+
+ @Test
+ public void testReadThrottlingWhenShardFound(){
+
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ doReturn(incompleteFuture()).when(mockActorContext).executeOperationAsync(
+ any(ActorSelection.class), eqReadData());
+
+ transactionProxy.read(TestModel.TEST_PATH);