@SuppressWarnings("unchecked") @Test public void testRecovery() throws Exception { int partitionId = 0; long destinationMaxTransactionId = 15L; long destinationLowWaterMark = 10L; long sourceLowWaterMark = 32L; long sourceLocalLowWaterMark = 31L; long sessionId = 123L; SessionInfo sourceSessionInfo = new SessionInfo(sessionId, sourceLowWaterMark, sourceLocalLowWaterMark); SessionInfo destinationSessionInfo = new SessionInfo(sessionId, destinationLowWaterMark, sourceLocalLowWaterMark); ArrayList<Record> twentyRecordList = ClientUtil.makeRecords(10L, 30L); ArrayList<Record> oneRecordList = ClientUtil.makeRecords(30L, 31L); CompletableFuture<Object> sourceSessionInfoCompletableFuture = mock(CompletableFuture.class); CompletableFuture<Object> destinationSessionInfoCompletableFuture = mock(CompletableFuture.class); CompletableFuture<Object> maxTransactionIdCompletableFuture = mock(CompletableFuture.class); CompletableFuture<Object> lowWatermarkCompletableFuture = mock(CompletableFuture.class); CompletableFuture<Object> recordList20CompletableFuture = mock(CompletableFuture.class); CompletableFuture<Object> recordList1CompletableFuture = mock(CompletableFuture.class); CompletableFuture<Object> appendRecordsCompletableFuture = mock(CompletableFuture.class); CompletableFuture<Object> setLowWaterMarkCompletableFuture = mock(CompletableFuture.class); StorageAdminClient mockSourceStorageAdminClient = mock(StorageAdminClient.class); StorageAdminClient mockDestinationStorageAdminClient = mock(StorageAdminClient.class); StorageClient mockDestinationStorageClient = mock(StorageClient.class); InOrder inOrder = inOrder(mockDestinationStorageClient); when(sourceSessionInfoCompletableFuture.get()).thenReturn(sourceSessionInfo); when(destinationSessionInfoCompletableFuture.get()).thenReturn(destinationSessionInfo).thenReturn(sourceSessionInfo); when(maxTransactionIdCompletableFuture.get()).thenReturn(destinationMaxTransactionId); when(lowWatermarkCompletableFuture.get()).thenReturn(destinationLowWaterMark); when(mockDestinationStorageAdminClient.lastSessionInfo(partitionId)).thenReturn(destinationSessionInfoCompletableFuture); when(mockSourceStorageAdminClient.lastSessionInfo(partitionId)).thenReturn(sourceSessionInfoCompletableFuture); when(mockDestinationStorageClient.lastSessionInfo(sessionId, partitionId)).thenReturn(destinationSessionInfoCompletableFuture); when(mockDestinationStorageClient.getMaxTransactionId(sessionId, partitionId)).thenReturn(maxTransactionIdCompletableFuture); when(mockDestinationStorageClient.truncate(sessionId, partitionId, destinationLowWaterMark)).thenReturn(lowWatermarkCompletableFuture); when(mockDestinationStorageClient.appendRecords(anyLong(), anyInt(), any())).thenReturn(appendRecordsCompletableFuture); when(mockDestinationStorageClient.setLowWaterMark(anyLong(), anyInt(), anyLong())).thenReturn(setLowWaterMarkCompletableFuture); when(mockSourceStorageAdminClient.getRecordList(partitionId, 11L, 20 )) .thenReturn(recordList20CompletableFuture); when(recordList20CompletableFuture.get()).thenReturn(twentyRecordList); when(mockSourceStorageAdminClient.getRecordList(partitionId, 31L, 1)) .thenReturn(recordList1CompletableFuture); when(recordList1CompletableFuture.get()).thenReturn(oneRecordList); StorageRecoveryRunnable storageRecoveryRunnable = new StorageRecoveryRunnable(mockSourceStorageAdminClient, mockDestinationStorageAdminClient, mockDestinationStorageClient, partitionId); storageRecoveryRunnable.run(); verify(mockSourceStorageAdminClient, times(2)).getRecordList(eq(partitionId), anyLong(), anyInt()); inOrder.verify(mockDestinationStorageClient, times(1)).appendRecords(eq(sessionId), eq(partitionId), eq(twentyRecordList)); inOrder.verify(mockDestinationStorageClient, times(1)).setLowWaterMark(eq(sessionId), eq(partitionId), eq(30L)); inOrder.verify(mockDestinationStorageClient, times(1)).appendRecords(eq(sessionId), eq(partitionId), eq(oneRecordList)); inOrder.verify(mockDestinationStorageClient, times(1)).setLowWaterMark(eq(sessionId), eq(partitionId), eq(31L)); }
@SuppressWarnings("unchecked") @Override public void run() { ReplicaConnection destinationReplicaConnection = null; try { SessionInfo sessionInfo = (SessionInfo) destinationStorageAdminClient.lastSessionInfo(partitionId).get(); long destinationSessionId = sessionInfo.sessionId; long destinationLowWatermark = sessionInfo.lowWaterMark; destinationReplicaConnection = new ReplicaConnectionImpl(partitionId, destinationSessionId, destinationStorageClient); long currentHighWaterMark = destinationReplicaConnection.getMaxTransactionId(); logger.debug("destination high watermark = " + currentHighWaterMark); currentHighWaterMark = destinationReplicaConnection.truncate(destinationLowWatermark); logger.info("destination truncated to low watermark = " + destinationLowWatermark); SessionInfo sourceSessionInfo = (SessionInfo) sourceStorageAdminClient.lastSessionInfo(destinationReplicaConnection.getPartitionId()).get(); long targetLowWaterMark = sourceSessionInfo.localLowWaterMark; logger.info("source target low watermark = " + targetLowWaterMark); while (currentHighWaterMark < targetLowWaterMark) { int batchSize = (int) Math.min((targetLowWaterMark - currentHighWaterMark), (long) maxBatchSize); ArrayList<Record> records = (ArrayList<Record>) sourceStorageAdminClient .getRecordList(destinationReplicaConnection.getPartitionId(), currentHighWaterMark + 1, batchSize) .get(); destinationReplicaConnection.appendRecords(records); currentHighWaterMark += records.size(); destinationReplicaConnection.setLowWaterMark(currentHighWaterMark); logger.debug("destination high watermark = " + currentHighWaterMark); if (currentHighWaterMark % LOG_SIZE == 0) { float percentProgress = currentHighWaterMark / (float) targetLowWaterMark; logger.info(String.format("destination high watermark = %d (%.2f%%)", currentHighWaterMark, percentProgress)); } } } catch (StorageRpcException | InterruptedException | ExecutionException e) { throw new RuntimeException(e); } finally { if (destinationReplicaConnection != null) { destinationReplicaConnection.close(); } } }
StorageRecoveryRunnable implements Runnable { @SuppressWarnings("unchecked") @Override public void run() { ReplicaConnection destinationReplicaConnection = null; try { SessionInfo sessionInfo = (SessionInfo) destinationStorageAdminClient.lastSessionInfo(partitionId).get(); long destinationSessionId = sessionInfo.sessionId; long destinationLowWatermark = sessionInfo.lowWaterMark; destinationReplicaConnection = new ReplicaConnectionImpl(partitionId, destinationSessionId, destinationStorageClient); long currentHighWaterMark = destinationReplicaConnection.getMaxTransactionId(); logger.debug("destination high watermark = " + currentHighWaterMark); currentHighWaterMark = destinationReplicaConnection.truncate(destinationLowWatermark); logger.info("destination truncated to low watermark = " + destinationLowWatermark); SessionInfo sourceSessionInfo = (SessionInfo) sourceStorageAdminClient.lastSessionInfo(destinationReplicaConnection.getPartitionId()).get(); long targetLowWaterMark = sourceSessionInfo.localLowWaterMark; logger.info("source target low watermark = " + targetLowWaterMark); while (currentHighWaterMark < targetLowWaterMark) { int batchSize = (int) Math.min((targetLowWaterMark - currentHighWaterMark), (long) maxBatchSize); ArrayList<Record> records = (ArrayList<Record>) sourceStorageAdminClient .getRecordList(destinationReplicaConnection.getPartitionId(), currentHighWaterMark + 1, batchSize) .get(); destinationReplicaConnection.appendRecords(records); currentHighWaterMark += records.size(); destinationReplicaConnection.setLowWaterMark(currentHighWaterMark); logger.debug("destination high watermark = " + currentHighWaterMark); if (currentHighWaterMark % LOG_SIZE == 0) { float percentProgress = currentHighWaterMark / (float) targetLowWaterMark; logger.info(String.format("destination high watermark = %d (%.2f%%)", currentHighWaterMark, percentProgress)); } } } catch (StorageRpcException | InterruptedException | ExecutionException e) { throw new RuntimeException(e); } finally { if (destinationReplicaConnection != null) { destinationReplicaConnection.close(); } } } }
StorageRecoveryRunnable implements Runnable { @SuppressWarnings("unchecked") @Override public void run() { ReplicaConnection destinationReplicaConnection = null; try { SessionInfo sessionInfo = (SessionInfo) destinationStorageAdminClient.lastSessionInfo(partitionId).get(); long destinationSessionId = sessionInfo.sessionId; long destinationLowWatermark = sessionInfo.lowWaterMark; destinationReplicaConnection = new ReplicaConnectionImpl(partitionId, destinationSessionId, destinationStorageClient); long currentHighWaterMark = destinationReplicaConnection.getMaxTransactionId(); logger.debug("destination high watermark = " + currentHighWaterMark); currentHighWaterMark = destinationReplicaConnection.truncate(destinationLowWatermark); logger.info("destination truncated to low watermark = " + destinationLowWatermark); SessionInfo sourceSessionInfo = (SessionInfo) sourceStorageAdminClient.lastSessionInfo(destinationReplicaConnection.getPartitionId()).get(); long targetLowWaterMark = sourceSessionInfo.localLowWaterMark; logger.info("source target low watermark = " + targetLowWaterMark); while (currentHighWaterMark < targetLowWaterMark) { int batchSize = (int) Math.min((targetLowWaterMark - currentHighWaterMark), (long) maxBatchSize); ArrayList<Record> records = (ArrayList<Record>) sourceStorageAdminClient .getRecordList(destinationReplicaConnection.getPartitionId(), currentHighWaterMark + 1, batchSize) .get(); destinationReplicaConnection.appendRecords(records); currentHighWaterMark += records.size(); destinationReplicaConnection.setLowWaterMark(currentHighWaterMark); logger.debug("destination high watermark = " + currentHighWaterMark); if (currentHighWaterMark % LOG_SIZE == 0) { float percentProgress = currentHighWaterMark / (float) targetLowWaterMark; logger.info(String.format("destination high watermark = %d (%.2f%%)", currentHighWaterMark, percentProgress)); } } } catch (StorageRpcException | InterruptedException | ExecutionException e) { throw new RuntimeException(e); } finally { if (destinationReplicaConnection != null) { destinationReplicaConnection.close(); } } } StorageRecoveryRunnable(StorageAdminClient sourceStorageAdminClient, StorageAdminClient destinationStorageAdminClient, StorageClient destinationStorageClient, int partitionId);  StorageRecoveryRunnable(StorageAdminClient sourceStorageAdminClient, StorageAdminClient destinationStorageAdminClient, StorageClient destinationStorageClient, int partitionId, int maxBatchSize); }
StorageRecoveryRunnable implements Runnable { @SuppressWarnings("unchecked") @Override public void run() { ReplicaConnection destinationReplicaConnection = null; try { SessionInfo sessionInfo = (SessionInfo) destinationStorageAdminClient.lastSessionInfo(partitionId).get(); long destinationSessionId = sessionInfo.sessionId; long destinationLowWatermark = sessionInfo.lowWaterMark; destinationReplicaConnection = new ReplicaConnectionImpl(partitionId, destinationSessionId, destinationStorageClient); long currentHighWaterMark = destinationReplicaConnection.getMaxTransactionId(); logger.debug("destination high watermark = " + currentHighWaterMark); currentHighWaterMark = destinationReplicaConnection.truncate(destinationLowWatermark); logger.info("destination truncated to low watermark = " + destinationLowWatermark); SessionInfo sourceSessionInfo = (SessionInfo) sourceStorageAdminClient.lastSessionInfo(destinationReplicaConnection.getPartitionId()).get(); long targetLowWaterMark = sourceSessionInfo.localLowWaterMark; logger.info("source target low watermark = " + targetLowWaterMark); while (currentHighWaterMark < targetLowWaterMark) { int batchSize = (int) Math.min((targetLowWaterMark - currentHighWaterMark), (long) maxBatchSize); ArrayList<Record> records = (ArrayList<Record>) sourceStorageAdminClient .getRecordList(destinationReplicaConnection.getPartitionId(), currentHighWaterMark + 1, batchSize) .get(); destinationReplicaConnection.appendRecords(records); currentHighWaterMark += records.size(); destinationReplicaConnection.setLowWaterMark(currentHighWaterMark); logger.debug("destination high watermark = " + currentHighWaterMark); if (currentHighWaterMark % LOG_SIZE == 0) { float percentProgress = currentHighWaterMark / (float) targetLowWaterMark; logger.info(String.format("destination high watermark = %d (%.2f%%)", currentHighWaterMark, percentProgress)); } } } catch (StorageRpcException | InterruptedException | ExecutionException e) { throw new RuntimeException(e); } finally { if (destinationReplicaConnection != null) { destinationReplicaConnection.close(); } } } StorageRecoveryRunnable(StorageAdminClient sourceStorageAdminClient, StorageAdminClient destinationStorageAdminClient, StorageClient destinationStorageClient, int partitionId);  StorageRecoveryRunnable(StorageAdminClient sourceStorageAdminClient, StorageAdminClient destinationStorageAdminClient, StorageClient destinationStorageClient, int partitionId, int maxBatchSize); @SuppressWarnings("unchecked") @Override void run(); }
StorageRecoveryRunnable implements Runnable { @SuppressWarnings("unchecked") @Override public void run() { ReplicaConnection destinationReplicaConnection = null; try { SessionInfo sessionInfo = (SessionInfo) destinationStorageAdminClient.lastSessionInfo(partitionId).get(); long destinationSessionId = sessionInfo.sessionId; long destinationLowWatermark = sessionInfo.lowWaterMark; destinationReplicaConnection = new ReplicaConnectionImpl(partitionId, destinationSessionId, destinationStorageClient); long currentHighWaterMark = destinationReplicaConnection.getMaxTransactionId(); logger.debug("destination high watermark = " + currentHighWaterMark); currentHighWaterMark = destinationReplicaConnection.truncate(destinationLowWatermark); logger.info("destination truncated to low watermark = " + destinationLowWatermark); SessionInfo sourceSessionInfo = (SessionInfo) sourceStorageAdminClient.lastSessionInfo(destinationReplicaConnection.getPartitionId()).get(); long targetLowWaterMark = sourceSessionInfo.localLowWaterMark; logger.info("source target low watermark = " + targetLowWaterMark); while (currentHighWaterMark < targetLowWaterMark) { int batchSize = (int) Math.min((targetLowWaterMark - currentHighWaterMark), (long) maxBatchSize); ArrayList<Record> records = (ArrayList<Record>) sourceStorageAdminClient .getRecordList(destinationReplicaConnection.getPartitionId(), currentHighWaterMark + 1, batchSize) .get(); destinationReplicaConnection.appendRecords(records); currentHighWaterMark += records.size(); destinationReplicaConnection.setLowWaterMark(currentHighWaterMark); logger.debug("destination high watermark = " + currentHighWaterMark); if (currentHighWaterMark % LOG_SIZE == 0) { float percentProgress = currentHighWaterMark / (float) targetLowWaterMark; logger.info(String.format("destination high watermark = %d (%.2f%%)", currentHighWaterMark, percentProgress)); } } } catch (StorageRpcException | InterruptedException | ExecutionException e) { throw new RuntimeException(e); } finally { if (destinationReplicaConnection != null) { destinationReplicaConnection.close(); } } } StorageRecoveryRunnable(StorageAdminClient sourceStorageAdminClient, StorageAdminClient destinationStorageAdminClient, StorageClient destinationStorageClient, int partitionId);  StorageRecoveryRunnable(StorageAdminClient sourceStorageAdminClient, StorageAdminClient destinationStorageAdminClient, StorageClient destinationStorageClient, int partitionId, int maxBatchSize); @SuppressWarnings("unchecked") @Override void run();  }