BucketOperatorWrapper implements BucketOperator { @Override public void createRedundantBucket(final InternalDistributedMember targetMember, final int i, final Map<String, Long> colocatedRegionBytes, final Completion completion) { if (stats != null) { stats.startBucketCreate(regionCount); } final long start = System.nanoTime(); delegate.createRedundantBucket(targetMember, i, colocatedRegionBytes, new Completion() { @Override public void onSuccess() { long totalBytes = 0; long elapsed = System.nanoTime() - start; if (logger.isDebugEnabled()) { logger.debug("Rebalancing {} redundant bucket {} created on {}", leaderRegion, i, targetMember); } for (PartitionRebalanceDetailsImpl details : detailSet) { String regionPath = details.getRegionPath(); Long lrb = colocatedRegionBytes.get(regionPath); if (lrb != null) { long regionBytes = lrb.longValue(); details.incCreates(regionBytes, details.getRegion().equals(leaderRegion) ? elapsed : 0); totalBytes += regionBytes; } } if (stats != null) { stats.endBucketCreate(regionCount, true, totalBytes, elapsed); } completion.onSuccess(); } @Override public void onFailure() { long elapsed = System.nanoTime() - start; if (logger.isDebugEnabled()) { logger.info("Rebalancing {} redundant bucket {} failed creation on {}", leaderRegion, i, targetMember); } if (stats != null) { stats.endBucketCreate(regionCount, false, 0, elapsed); } completion.onFailure(); } }); } BucketOperatorWrapper(BucketOperator delegate,
      Set<PartitionRebalanceDetailsImpl> rebalanceDetails, ResourceManagerStats stats,
      PartitionedRegion leaderRegion); @Override boolean moveBucket(InternalDistributedMember sourceMember,
      InternalDistributedMember targetMember, int id, Map<String, Long> colocatedRegionBytes); @Override void createRedundantBucket(final InternalDistributedMember targetMember, final int i,
      final Map<String, Long> colocatedRegionBytes, final Completion completion); @Override boolean removeBucket(InternalDistributedMember targetMember, int i,
      Map<String, Long> colocatedRegionBytes); @Override boolean movePrimary(InternalDistributedMember source, InternalDistributedMember target,
      int bucketId); @Override void waitForOperations(); Set<PartitionRebalanceDetailsImpl> getDetailSet();  }
@Test public void bucketWrapperShouldDelegateCreateBucketToEnclosedOperator() { Completion completionSentToWrapper = mock(Completion.class); doNothing().when(delegate).createRedundantBucket(targetMember, bucketId, colocatedRegionBytes, completionSentToWrapper); wrapper.createRedundantBucket(targetMember, bucketId, colocatedRegionBytes, completionSentToWrapper); verify(delegate, times(1)).createRedundantBucket(eq(targetMember), eq(bucketId), eq(colocatedRegionBytes), any(Completion.class)); }
@Test public void bucketWrapperShouldRecordNumberOfBucketsCreatedIfCreateBucketSucceeds() { doAnswer(new Answer<Object>() { @Override public Object answer(InvocationOnMock invocation) { ((Completion) invocation.getArguments()[3]).onSuccess(); return null; } }).when(delegate).createRedundantBucket(eq(targetMember), eq(bucketId), eq(colocatedRegionBytes), any(Completion.class)); Completion completionSentToWrapper = mock(Completion.class); wrapper.createRedundantBucket(targetMember, bucketId, colocatedRegionBytes, completionSentToWrapper); for (PartitionRebalanceDetailsImpl details : rebalanceDetails) { if (details.getRegionPath().equalsIgnoreCase(PR_LEADER_REGION_NAME)) verify(details, times(1)).incCreates(eq(colocatedRegionBytes.get(PR_LEADER_REGION_NAME)), anyLong()); else if (details.getRegionPath().equals(PR_COLOCATED_REGION_NAME)) verify(details, times(1)).incTransfers(colocatedRegionBytes.get(PR_COLOCATED_REGION_NAME), 0); } }
@Test public void bucketWrapperShouldNotRecordNumberOfBucketsCreatedIfCreateBucketFails() { doAnswer(new Answer<Object>() { @Override public Object answer(InvocationOnMock invocation) { ((Completion) invocation.getArguments()[3]).onFailure(); return null; } }).when(delegate).createRedundantBucket(eq(targetMember), eq(bucketId), eq(colocatedRegionBytes), any(Completion.class)); Completion completionSentToWrapper = mock(Completion.class); wrapper.createRedundantBucket(targetMember, bucketId, colocatedRegionBytes, completionSentToWrapper); for (PartitionRebalanceDetailsImpl details : rebalanceDetails) { verify(details, times(0)).incTransfers(anyLong(), anyLong()); } }
@Test public void bucketWrapperShouldInvokeOnFailureWhenCreateBucketFails() { doAnswer(new Answer<Object>() { @Override public Object answer(InvocationOnMock invocation) { ((Completion) invocation.getArguments()[3]).onFailure(); return null; } }).when(delegate).createRedundantBucket(eq(targetMember), eq(bucketId), eq(colocatedRegionBytes), any(Completion.class)); Completion completionSentToWrapper = mock(Completion.class); wrapper.createRedundantBucket(targetMember, bucketId, colocatedRegionBytes, completionSentToWrapper); verify(completionSentToWrapper, times(1)).onFailure(); }
@Test public void bucketWrapperShouldInvokeOnSuccessWhenCreateBucketSucceeds() { doAnswer(new Answer<Object>() { @Override public Object answer(InvocationOnMock invocation) { ((Completion) invocation.getArguments()[3]).onSuccess(); return null; } }).when(delegate).createRedundantBucket(eq(targetMember), eq(bucketId), eq(colocatedRegionBytes), any(Completion.class)); Completion completionSentToWrapper = mock(Completion.class); wrapper.createRedundantBucket(targetMember, bucketId, colocatedRegionBytes, completionSentToWrapper); verify(completionSentToWrapper, times(1)).onSuccess(); }