Merge branch 'master' into add-scanner-api-to-transaction-abstraction

Author: brfrn169

build.gradle

@@ -30,7 +30,7 @@ subprojects {
         awssdkVersion = '2.31.3'
         commonsDbcp2Version = '2.13.0'
         mysqlDriverVersion = '8.4.0'
-        postgresqlDriverVersion = '42.7.5'
+        postgresqlDriverVersion = '42.7.6'
         oracleDriverVersion = '23.8.0.25.04'
         sqlserverDriverVersion = '12.8.1.jre8'
         sqliteDriverVersion = '3.49.1.0'

data-loader/core/src/main/java/com/scalar/db/dataloader/core/dataimport/ImportEventListener.java

@@ -18,13 +18,6 @@ public interface ImportEventListener {
    */
   void onDataChunkStarted(ImportDataChunkStatus status);
 
-  /**
-   * Updates or adds new status information for a data chunk.
-   *
-   * @param status the updated status information for the data chunk
-   */
-  void addOrUpdateDataChunkStatus(ImportDataChunkStatus status);
-
   /**
    * Called when processing of a data chunk is completed.
    *

data-loader/core/src/main/java/com/scalar/db/dataloader/core/dataimport/ImportManager.java

@@ -17,7 +17,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 import lombok.AllArgsConstructor;
 import lombok.NonNull;
 
@@ -46,20 +45,15 @@ public class ImportManager implements ImportEventListener {
   private final ScalarDbMode scalarDbMode;
   private final DistributedStorage distributedStorage;
   private final DistributedTransactionManager distributedTransactionManager;
-  private final ConcurrentHashMap<Integer, ImportDataChunkStatus> importDataChunkStatusMap =
-      new ConcurrentHashMap<>();
 
   /**
    * Starts the import process using the configured parameters.
    *
    * <p>If the data chunk size in {@link ImportOptions} is set to 0, the entire file will be
    * processed as a single chunk. Otherwise, the file will be processed in chunks of the specified
    * size.
-   *
-   * @return a map of {@link ImportDataChunkStatus} objects containing the status of each processed
-   *     chunk
    */
-  public ConcurrentHashMap<Integer, ImportDataChunkStatus> startImport() {
+  public void startImport() {
     ImportProcessorParams params =
         ImportProcessorParams.builder()
             .scalarDbMode(scalarDbMode)
@@ -77,8 +71,7 @@ public ConcurrentHashMap<Integer, ImportDataChunkStatus> startImport() {
         importOptions.getDataChunkSize() == 0
             ? Integer.MAX_VALUE
             : importOptions.getDataChunkSize();
-    return processor.process(
-        dataChunkSize, importOptions.getTransactionBatchSize(), importFileReader);
+    processor.process(dataChunkSize, importOptions.getTransactionBatchSize(), importFileReader);
   }
 
   /**
@@ -108,15 +101,6 @@ public void onDataChunkStarted(ImportDataChunkStatus status) {
     }
   }
 
-  /**
-   * {@inheritDoc} Updates or adds the status of a data chunk in the status map. This method is
-   * thread-safe.
-   */
-  @Override
-  public void addOrUpdateDataChunkStatus(ImportDataChunkStatus status) {
-    importDataChunkStatusMap.put(status.getDataChunkId(), status);
-  }
-
   /** {@inheritDoc} Forwards the event to all registered listeners. */
   @Override
   public void onDataChunkCompleted(ImportDataChunkStatus status) {
@@ -152,18 +136,46 @@ public void onTaskComplete(ImportTaskResult taskResult) {
   /** {@inheritDoc} Forwards the event to all registered listeners. */
   @Override
   public void onAllDataChunksCompleted() {
+    Throwable firstException = null;
+
     for (ImportEventListener listener : listeners) {
-      listener.onAllDataChunksCompleted();
+      try {
+        listener.onAllDataChunksCompleted();
+      } catch (Throwable e) {
+        if (firstException == null) {
+          firstException = e;
+        } else {
+          firstException.addSuppressed(e);
+        }
+      }
+    }
+
+    try {
+      closeResources();
+    } catch (Throwable e) {
+      if (firstException != null) {
+        firstException.addSuppressed(e);
+      } else {
+        firstException = e;
+      }
+    }
+
+    if (firstException != null) {
+      throw new RuntimeException("Error during completion", firstException);
     }
   }
 
-  /**
-   * Returns the current map of import data chunk status objects.
-   *
-   * @return a map of {@link ImportDataChunkStatus} objects
-   */
-  public ConcurrentHashMap<Integer, ImportDataChunkStatus> getImportDataChunkStatus() {
-    return importDataChunkStatusMap;
+  /** Close resources properly once the process is completed */
+  public void closeResources() {
+    try {
+      if (distributedStorage != null) {
+        distributedStorage.close();
+      } else if (distributedTransactionManager != null) {
+        distributedTransactionManager.close();
+      }
+    } catch (Throwable e) {
+      throw new RuntimeException("Failed to close the resource", e);
+    }
   }
 
   /**

data-loader/core/src/main/java/com/scalar/db/dataloader/core/dataimport/dao/ScalarDbDao.java

@@ -22,19 +22,10 @@
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import javax.annotation.Nullable;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /** The generic DAO that is used to scan ScalarDB data */
 public class ScalarDbDao {
 
-  /* Class logger */
-  private static final Logger logger = LoggerFactory.getLogger(ScalarDbDao.class);
-  private static final String GET_COMPLETED_MSG = "GET completed for %s";
-  private static final String PUT_COMPLETED_MSG = "PUT completed for %s";
-  private static final String SCAN_START_MSG = "SCAN started...";
-  private static final String SCAN_END_MSG = "SCAN completed";
-
   /**
    * Retrieve record from ScalarDB instance in storage mode
    *
@@ -59,9 +50,7 @@ public Optional<Result> get(
 
     try {
       Get get = createGetWith(namespace, table, partitionKey, clusteringKey);
-      Optional<Result> result = storage.get(get);
-      logger.info(String.format(GET_COMPLETED_MSG, loggingKey));
-      return result;
+      return storage.get(get);
     } catch (ExecutionException e) {
       throw new ScalarDbDaoException("error GET " + loggingKey, e);
     }
@@ -90,9 +79,7 @@ public Optional<Result> get(
     // Retrieving the key data for logging
     String loggingKey = keysToString(partitionKey, clusteringKey);
     try {
-      Optional<Result> result = transaction.get(get);
-      logger.info(String.format(GET_COMPLETED_MSG, loggingKey));
-      return result;
+      return transaction.get(get);
     } catch (CrudException e) {
       throw new ScalarDbDaoException("error GET " + loggingKey, e.getCause());
     }
@@ -125,7 +112,6 @@ public void put(
       throw new ScalarDbDaoException(
           CoreError.DATA_LOADER_ERROR_CRUD_EXCEPTION.buildMessage(e.getMessage()), e);
     }
-    logger.info(String.format(PUT_COMPLETED_MSG, keysToString(partitionKey, clusteringKey)));
   }
 
   /**
@@ -154,7 +140,6 @@ public void put(
       throw new ScalarDbDaoException(
           CoreError.DATA_LOADER_ERROR_CRUD_EXCEPTION.buildMessage(e.getMessage()), e);
     }
-    logger.info(String.format(PUT_COMPLETED_MSG, keysToString(partitionKey, clusteringKey)));
   }
 
   /**
@@ -186,11 +171,8 @@ public List<Result> scan(
 
     // scan data
     try {
-      logger.info(SCAN_START_MSG);
       try (Scanner scanner = storage.scan(scan)) {
-        List<Result> allResults = scanner.all();
-        logger.info(SCAN_END_MSG);
-        return allResults;
+        return scanner.all();
       }
     } catch (ExecutionException | IOException e) {
       throw new ScalarDbDaoException(
@@ -229,10 +211,7 @@ public List<Result> scan(
 
     // scan data
     try {
-      logger.info(SCAN_START_MSG);
-      List<Result> results = transaction.scan(scan);
-      logger.info(SCAN_END_MSG);
-      return results;
+      return transaction.scan(scan);
     } catch (CrudException | NoSuchElementException e) {
       // No such element Exception is thrown when the scan is done in transaction mode but
       // ScalarDB is running in storage mode

data-loader/core/src/main/java/com/scalar/db/dataloader/core/dataimport/log/SingleFileImportLogger.java

@@ -70,15 +70,6 @@ public void onTaskComplete(ImportTaskResult taskResult) {
     }
   }
 
-  /**
-   * Called to add or update the status of a data chunk. This implementation does nothing as the
-   * status is only logged when the data chunk is completed.
-   *
-   * @param status the status of the data chunk
-   */
-  @Override
-  public void addOrUpdateDataChunkStatus(ImportDataChunkStatus status) {}
-
   /**
    * Called when a data chunk is completed. Logs the summary of the data chunk to the summary log
    * file.

data-loader/core/src/main/java/com/scalar/db/dataloader/core/dataimport/log/SplitByDataChunkImportLogger.java

@@ -112,15 +112,6 @@ private void writeLog(ImportTargetResult target, LogFileType logFileType, int da
     writer.write(jsonNode);
   }
 
-  /**
-   * Called to add or update the status of a data chunk. This implementation does nothing as the
-   * status is only logged when the data chunk is completed.
-   *
-   * @param status the status of the data chunk
-   */
-  @Override
-  public void addOrUpdateDataChunkStatus(ImportDataChunkStatus status) {}
-
   /**
    * Called when a data chunk is completed. Logs the summary of the data chunk and closes the log
    * writers for that data chunk.

data-loader/core/src/main/java/com/scalar/db/dataloader/core/dataimport/processor/CsvImportProcessor.java

@@ -5,7 +5,6 @@
 import com.scalar.db.common.error.CoreError;
 import com.scalar.db.dataloader.core.DataLoaderObjectMapper;
 import com.scalar.db.dataloader.core.dataimport.datachunk.ImportDataChunk;
-import com.scalar.db.dataloader.core.dataimport.datachunk.ImportDataChunkStatus;
 import com.scalar.db.dataloader.core.dataimport.datachunk.ImportRow;
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -14,12 +13,6 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.LinkedBlockingQueue;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -39,7 +32,7 @@
  */
 public class CsvImportProcessor extends ImportProcessor {
   private static final DataLoaderObjectMapper OBJECT_MAPPER = new DataLoaderObjectMapper();
-  private static final AtomicInteger dataChunkIdCounter = new AtomicInteger(0);
+  private final AtomicInteger dataChunkIdCounter = new AtomicInteger(0);
 
   /**
    * Creates a new CsvImportProcessor with the specified parameters.
@@ -50,60 +43,6 @@ public CsvImportProcessor(ImportProcessorParams params) {
     super(params);
   }
 
-  /**
-   * Processes the source data from the given import file.
-   *
-   * <p>This method reads data from the provided {@link BufferedReader}, processes it in chunks, and
-   * batches transactions according to the specified sizes. The method returns a list of {@link
-   * ImportDataChunkStatus} objects, each representing the status of a processed data chunk.
-   *
-   * @param dataChunkSize the number of records to include in each data chunk
-   * @param transactionBatchSize the number of records to include in each transaction batch
-   * @param reader the {@link BufferedReader} used to read the source file
-   * @return a map of {@link ImportDataChunkStatus} objects indicating the processing status of each
-   *     data chunk
-   */
-  @Override
-  public ConcurrentHashMap<Integer, ImportDataChunkStatus> process(
-      int dataChunkSize, int transactionBatchSize, BufferedReader reader) {
-    ExecutorService dataChunkExecutor = Executors.newSingleThreadExecutor();
-    BlockingQueue<ImportDataChunk> dataChunkQueue =
-        new LinkedBlockingQueue<>(params.getImportOptions().getDataChunkQueueSize());
-
-    try {
-      CompletableFuture<Void> readerFuture =
-          CompletableFuture.runAsync(
-              () -> readDataChunks(reader, dataChunkSize, dataChunkQueue), dataChunkExecutor);
-
-      ConcurrentHashMap<Integer, ImportDataChunkStatus> result = new ConcurrentHashMap<>();
-
-      while (!(dataChunkQueue.isEmpty() && readerFuture.isDone())) {
-        ImportDataChunk dataChunk = dataChunkQueue.poll(100, TimeUnit.MILLISECONDS);
-        if (dataChunk != null) {
-          ImportDataChunkStatus status = processDataChunk(dataChunk, transactionBatchSize);
-          result.put(status.getDataChunkId(), status);
-        }
-      }
-      readerFuture.join();
-      return result;
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
-      throw new RuntimeException(
-          CoreError.DATA_LOADER_DATA_CHUNK_PROCESS_FAILED.buildMessage(e.getMessage()), e);
-    } finally {
-      dataChunkExecutor.shutdown();
-      try {
-        if (!dataChunkExecutor.awaitTermination(60, TimeUnit.SECONDS)) {
-          dataChunkExecutor.shutdownNow();
-        }
-      } catch (InterruptedException e) {
-        dataChunkExecutor.shutdownNow();
-        Thread.currentThread().interrupt();
-      }
-      notifyAllDataChunksCompleted();
-    }
-  }
-
   /**
    * Reads and processes CSV data in chunks from the provided reader.
    *
@@ -122,7 +61,8 @@ public ConcurrentHashMap<Integer, ImportDataChunkStatus> process(
    * @param dataChunkQueue the queue where data chunks are placed for processing
    * @throws RuntimeException if there are errors reading the file or if interrupted
    */
-  private void readDataChunks(
+  @Override
+  protected void readDataChunks(
       BufferedReader reader, int dataChunkSize, BlockingQueue<ImportDataChunk> dataChunkQueue) {
     try {
       String delimiter =