BucketOperatorWrapper implements BucketOperator { @Override public boolean movePrimary(InternalDistributedMember source, InternalDistributedMember target, int bucketId) { boolean result = false; long elapsed = 0; if (stats != null) { stats.startPrimaryTransfer(regionCount); } try { long start = System.nanoTime(); result = delegate.movePrimary(source, target, bucketId); elapsed = System.nanoTime() - start; if (result) { if (logger.isDebugEnabled()) { logger.debug("Rebalancing {} primary bucket {} moved from {} to {}", leaderRegion, bucketId, source, target); } for (PartitionRebalanceDetailsImpl details : detailSet) { details.incPrimaryTransfers(details.getRegion().equals(leaderRegion) ? elapsed : 0); } } else { if (logger.isDebugEnabled()) { logger.debug("Rebalancing {} primary bucket {} failed to move from {} to {}", leaderRegion, bucketId, source, target); } } } finally { if (stats != null) { stats.endPrimaryTransfer(regionCount, result, elapsed); } } return result; } 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 bucketWrapperShouldDelegateMovePrimaryToEnclosedOperator() { wrapper.movePrimary(sourceMember, targetMember, bucketId); verify(delegate, times(1)).movePrimary(sourceMember, targetMember, bucketId); verify(stats, times(1)).startPrimaryTransfer(anyInt()); verify(stats, times(1)).endPrimaryTransfer(anyInt(), anyBoolean(), anyLong()); }
@Test public void bucketWrapperShouldRecordPrimaryTransfersPerRegionAfterMovePrimaryIsSuccessful() { doReturn(true).when(delegate).movePrimary(sourceMember, targetMember, bucketId); wrapper.movePrimary(sourceMember, targetMember, bucketId); for (PartitionRebalanceDetailsImpl details : rebalanceDetails) { if (details.getRegionPath().equalsIgnoreCase(PR_LEADER_REGION_NAME)) verify(details, times(1)).incPrimaryTransfers(anyLong()); else if (details.getRegionPath().equals(PR_COLOCATED_REGION_NAME)) verify(details, times(1)).incPrimaryTransfers(0); } verify(stats, times(1)).startPrimaryTransfer(anyInt()); verify(stats, times(1)).endPrimaryTransfer(anyInt(), anyBoolean(), anyLong()); }
@Test public void bucketWrapperShouldNotRecordPrimaryTransfersPerRegionAfterMovePrimaryFails() { doReturn(false).when(delegate).movePrimary(sourceMember, targetMember, bucketId); wrapper.movePrimary(sourceMember, targetMember, bucketId); for (PartitionRebalanceDetailsImpl details : rebalanceDetails) { verify(details, times(0)).incTransfers(anyLong(), anyLong()); } verify(stats, times(1)).startPrimaryTransfer(anyInt()); verify(stats, times(1)).endPrimaryTransfer(anyInt(), anyBoolean(), anyLong()); }