Simplify data consistency check code (#22504)

Author: sandynz

kernel/data-pipeline/scenario/src/main/java/org/apache/shardingsphere/data/pipeline/core/check/consistency/SingleTableInventoryDataConsistencyChecker.java

@@ -104,12 +104,10 @@ private DataConsistencyCheckResult check(final DataConsistencyCalculateAlgorithm
         Collection<String> columnNames = tableMetaData.getColumnNames();
         Map<String, Object> tableCheckPositions = jobItemContext.getTableCheckPositions();
         DataConsistencyCalculateParameter sourceParam = buildParameter(
-                sourceDataSource, schemaName, sourceTableName, columnNames, sourceDatabaseType, targetDatabaseType, uniqueKey,
-                tableCheckPositions.get(sourceTableName));
+                sourceDataSource, schemaName, sourceTableName, columnNames, sourceDatabaseType, targetDatabaseType, uniqueKey, tableCheckPositions.get(sourceTableName));
         String targetTableName = targetTable.getTableName().getOriginal();
-        DataConsistencyCalculateParameter targetParam = buildParameter(
-                targetDataSource, targetTable.getSchemaName().getOriginal(), targetTableName, columnNames, targetDatabaseType, sourceDatabaseType, uniqueKey,
-                tableCheckPositions.get(targetTableName));
+        DataConsistencyCalculateParameter targetParam = buildParameter(targetDataSource, targetTable.getSchemaName().getOriginal(), targetTableName,
+                columnNames, targetDatabaseType, sourceDatabaseType, uniqueKey, tableCheckPositions.get(targetTableName));
         Iterator<DataConsistencyCalculatedResult> sourceCalculatedResults = calculateAlgorithm.calculate(sourceParam).iterator();
         Iterator<DataConsistencyCalculatedResult> targetCalculatedResults = calculateAlgorithm.calculate(targetParam).iterator();
         long sourceRecordsCount = 0;
@@ -131,10 +129,10 @@ private DataConsistencyCheckResult check(final DataConsistencyCalculateAlgorithm
                 break;
             }
             if (sourceCalculatedResult.getMaxUniqueKeyValue().isPresent()) {
-                jobItemContext.getTableCheckPositions().put(sourceTableName, sourceCalculatedResult.getMaxUniqueKeyValue().get());
+                tableCheckPositions.put(sourceTableName, sourceCalculatedResult.getMaxUniqueKeyValue().get());
             }
             if (targetCalculatedResult.getMaxUniqueKeyValue().isPresent()) {
-                jobItemContext.getTableCheckPositions().put(targetTableName, targetCalculatedResult.getMaxUniqueKeyValue().get());
+                tableCheckPositions.put(targetTableName, targetCalculatedResult.getMaxUniqueKeyValue().get());
             }
             jobItemContext.onProgressUpdated(new PipelineJobProgressUpdatedParameter(sourceCalculatedResult.getRecordsCount()));
         }

kernel/data-pipeline/scenario/src/main/java/org/apache/shardingsphere/data/pipeline/scenario/consistencycheck/ConsistencyCheckJobItemContext.java

@@ -60,7 +60,7 @@ public final class ConsistencyCheckJobItemContext implements PipelineJobItemCont
 
     private final long checkBeginTimeMillis;
 
-    private Long checkEndTimeMillis;
+    private volatile Long checkEndTimeMillis;
 
     private final Map<String, Object> tableCheckPositions = new ConcurrentHashMap<>();
 

kernel/data-pipeline/scenario/src/main/java/org/apache/shardingsphere/data/pipeline/scenario/consistencycheck/ConsistencyCheckTasksRunner.java

@@ -99,9 +99,13 @@ protected void runBlocking() {
             DataConsistencyCalculateAlgorithm calculateAlgorithm = jobAPI.buildDataConsistencyCalculateAlgorithm(
                     parentJobConfig, checkJobConfig.getAlgorithmTypeName(), checkJobConfig.getAlgorithmProps());
             setCalculateAlgorithm(calculateAlgorithm);
-            Map<String, DataConsistencyCheckResult> dataConsistencyCheckResult = jobAPI.dataConsistencyCheck(parentJobConfig, calculateAlgorithm, jobItemContext);
-            PipelineAPIFactory.getGovernanceRepositoryAPI().persistCheckJobResult(parentJobId, checkJobId, dataConsistencyCheckResult);
-            jobItemContext.setCheckEndTimeMillis(System.currentTimeMillis());
+            Map<String, DataConsistencyCheckResult> dataConsistencyCheckResult;
+            try {
+                dataConsistencyCheckResult = jobAPI.dataConsistencyCheck(parentJobConfig, calculateAlgorithm, jobItemContext);
+                PipelineAPIFactory.getGovernanceRepositoryAPI().persistCheckJobResult(parentJobId, checkJobId, dataConsistencyCheckResult);
+            } finally {
+                jobItemContext.setCheckEndTimeMillis(System.currentTimeMillis());
+            }
         }
 
         @Override

kernel/data-pipeline/scenario/src/main/java/org/apache/shardingsphere/data/pipeline/scenario/migration/MigrationDataConsistencyChecker.java

@@ -80,18 +80,12 @@ public Map<String, DataConsistencyCheckResult> check(final DataConsistencyCalcul
         try (
                 PipelineDataSourceWrapper sourceDataSource = PipelineDataSourceFactory.newInstance(jobConfig.getSource());
                 PipelineDataSourceWrapper targetDataSource = PipelineDataSourceFactory.newInstance(jobConfig.getTarget())) {
-            // TODO simplify code
-            MigrationJobAPI migrationJobAPI = MigrationJobAPIFactory.getInstance();
-            Map<Integer, InventoryIncrementalJobItemProgress> jobProgress = migrationJobAPI.getJobProgress(jobConfig);
-            long recordsCount = jobProgress.values().stream().filter(Objects::nonNull).mapToLong(InventoryIncrementalJobItemProgress::getProcessedRecordsCount).sum();
-            checkJobItemContext.setRecordsCount(recordsCount);
+            checkJobItemContext.setRecordsCount(getRecordsCount());
             checkJobItemContext.getTableNames().add(jobConfig.getSourceTableName());
             PipelineTableMetaDataLoader metaDataLoader = new StandardPipelineTableMetaDataLoader(sourceDataSource);
             SingleTableInventoryDataConsistencyChecker singleTableInventoryChecker = new SingleTableInventoryDataConsistencyChecker(jobConfig.getJobId(), sourceDataSource, targetDataSource,
                     sourceTable, targetTable, jobConfig.getUniqueKeyColumn(), metaDataLoader, readRateLimitAlgorithm, checkJobItemContext);
             result.put(sourceTable.getTableName().getOriginal(), singleTableInventoryChecker.check(calculateAlgorithm));
-            // TODO make sure checkEndTimeMillis will be set
-            checkJobItemContext.setCheckEndTimeMillis(System.currentTimeMillis());
         } catch (final SQLException ex) {
             throw new SQLWrapperException(ex);
         }
@@ -102,4 +96,9 @@ private void verifyPipelineDatabaseType(final DataConsistencyCalculateAlgorithm
         ShardingSpherePreconditions.checkState(calculateAlgorithm.getSupportedDatabaseTypes().contains(dataSourceConfig.getDatabaseType().getType()),
                 () -> new UnsupportedPipelineDatabaseTypeException(dataSourceConfig.getDatabaseType()));
     }
+    
+    private long getRecordsCount() {
+        Map<Integer, InventoryIncrementalJobItemProgress> jobProgress = MigrationJobAPIFactory.getInstance().getJobProgress(jobConfig);
+        return jobProgress.values().stream().filter(Objects::nonNull).mapToLong(InventoryIncrementalJobItemProgress::getProcessedRecordsCount).sum();
+    }
 }