Merge branch 'master' into fix/data-loader/logger-rename

Author: inv-jishnu

core/src/main/java/com/scalar/db/transaction/consensuscommit/CrudHandler.java

@@ -28,6 +28,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import javax.annotation.concurrent.NotThreadSafe;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -76,17 +77,27 @@ public CrudHandler(
   public Optional<Result> get(Get originalGet) throws CrudException {
     List<String> originalProjections = new ArrayList<>(originalGet.getProjections());
     Get get = (Get) prepareStorageSelection(originalGet);
-    Snapshot.Key key = new Snapshot.Key(get);
-    readUnread(key, get);
 
     TableMetadata metadata = getTableMetadata(get);
+
+    Snapshot.Key key;
+    if (ScalarDbUtils.isSecondaryIndexSpecified(get, metadata)) {
+      // In case of a Get with index, we don't know the key until we read the record
+      key = null;
+    } else {
+      key = new Snapshot.Key(get);
+    }
+
+    readUnread(key, get);
+
     return snapshot
         .getResult(key, get)
         .map(r -> new FilteredResult(r, originalProjections, metadata, isIncludeMetadataEnabled));
   }
 
+  // Only for a Get with index, the argument `key` is null
   @VisibleForTesting
-  void readUnread(Snapshot.Key key, Get get) throws CrudException {
+  void readUnread(@Nullable Snapshot.Key key, Get get) throws CrudException {
     if (!snapshot.containsKeyInGetSet(get)) {
       read(key, get);
     }
@@ -95,7 +106,7 @@ void readUnread(Snapshot.Key key, Get get) throws CrudException {
   // Although this class is not thread-safe, this method is actually thread-safe, so we call it
   // concurrently in the implicit pre-read
   @VisibleForTesting
-  void read(Snapshot.Key key, Get get) throws CrudException {
+  void read(@Nullable Snapshot.Key key, Get get) throws CrudException {
     Optional<TransactionResult> result = getFromStorage(get);
     if (!result.isPresent() || result.get().isCommitted()) {
       if (result.isPresent() || get.getConjunctions().isEmpty()) {
@@ -104,7 +115,18 @@ void read(Snapshot.Key key, Get get) throws CrudException {
         // transaction read it first. However, we update it only if a get operation has no
         // conjunction or the result exists. This is because we don’t know whether the record
         // actually exists or not due to the conjunction.
-        snapshot.putIntoReadSet(key, result);
+        if (key != null) {
+          snapshot.putIntoReadSet(key, result);
+        } else {
+          // Only for a Get with index, the argument `key` is null
+
+          if (result.isPresent()) {
+            // Only when we can get the record with the Get with index, we can put it into the read
+            // set
+            key = new Snapshot.Key(get, result.get());
+            snapshot.putIntoReadSet(key, result);
+          }
+        }
       }
       snapshot.putIntoGetSet(get, result); // for re-read and validation
       return;

core/src/main/java/com/scalar/db/transaction/consensuscommit/Snapshot.java

@@ -720,6 +720,13 @@ public Key(Get get) {
       this((Operation) get);
     }
 
+    public Key(Get get, Result result) {
+      this.namespace = get.forNamespace().get();
+      this.table = get.forTable().get();
+      this.partitionKey = result.getPartitionKey().get();
+      this.clusteringKey = result.getClusteringKey();
+    }
+
     public Key(Put put) {
       this((Operation) put);
     }

core/src/test/java/com/scalar/db/transaction/consensuscommit/CrudHandlerTest.java

@@ -56,6 +56,7 @@ public class CrudHandlerTest {
   private static final String ANY_ID_2 = "id2";
   private static final String ANY_NAME_1 = "name1";
   private static final String ANY_NAME_2 = "name2";
+  private static final String ANY_NAME_3 = "name3";
   private static final String ANY_TEXT_1 = "text1";
   private static final String ANY_TEXT_2 = "text2";
   private static final String ANY_TEXT_3 = "text3";
@@ -66,8 +67,10 @@ public class CrudHandlerTest {
           TableMetadata.newBuilder()
               .addColumn(ANY_NAME_1, DataType.TEXT)
               .addColumn(ANY_NAME_2, DataType.TEXT)
+              .addColumn(ANY_NAME_3, DataType.TEXT)
               .addPartitionKey(ANY_NAME_1)
               .addClusteringKey(ANY_NAME_2)
+              .addSecondaryIndex(ANY_NAME_3)
               .build());
   private static final TransactionTableMetadata TRANSACTION_TABLE_METADATA =
       new TransactionTableMetadata(TABLE_METADATA);
@@ -928,6 +931,7 @@ public void readUnread_GetContainedInGetSet_ShouldCallAppropriateMethods()
 
     // Assert
     verify(storage, never()).get(any());
+    verify(snapshot, never()).putIntoReadSet(any(Snapshot.Key.class), any(Optional.class));
     verify(snapshot, never()).putIntoGetSet(any(Get.class), any(Optional.class));
   }
 
@@ -1014,6 +1018,7 @@ public void readUnread_GetContainedInGetSet_ShouldCallAppropriateMethods()
     // Assert
     verify(storage).get(any());
     verify(snapshot).putIntoReadSet(key, Optional.of(new TransactionResult(result)));
+    verify(snapshot).putIntoGetSet(getForKey, Optional.of(new TransactionResult(result)));
   }
 
   @Test
@@ -1050,6 +1055,88 @@ public void readUnread_GetContainedInGetSet_ShouldCallAppropriateMethods()
             });
   }
 
+  @Test
+  public void
+      readUnread_NullKeyAndGetWithIndexNotContainedInGetSet_EmptyResultReturnedByStorage_ShouldCallAppropriateMethods()
+          throws CrudException, ExecutionException {
+    // Arrange
+    Get getWithIndex =
+        Get.newBuilder()
+            .namespace(ANY_NAMESPACE_NAME)
+            .table(ANY_TABLE_NAME)
+            .indexKey(Key.ofText(ANY_NAME_3, ANY_TEXT_1))
+            .build();
+    when(snapshot.containsKeyInGetSet(getWithIndex)).thenReturn(false);
+    when(storage.get(any())).thenReturn(Optional.empty());
+
+    // Act
+    handler.readUnread(null, getWithIndex);
+
+    // Assert
+    verify(storage).get(any());
+    verify(snapshot, never()).putIntoReadSet(any(), any());
+    verify(snapshot).putIntoGetSet(getWithIndex, Optional.empty());
+  }
+
+  @Test
+  public void
+      readUnread_NullKeyAndGetWithIndexNotContainedInGetSet_CommittedRecordReturnedByStorage_ShouldCallAppropriateMethods()
+          throws CrudException, ExecutionException {
+    // Arrange
+    Result result = mock(Result.class);
+    when(result.getInt(Attribute.STATE)).thenReturn(TransactionState.COMMITTED.get());
+    when(result.getPartitionKey()).thenReturn(Optional.of(Key.ofText(ANY_NAME_1, ANY_TEXT_1)));
+    when(result.getClusteringKey()).thenReturn(Optional.of(Key.ofText(ANY_NAME_2, ANY_TEXT_2)));
+    when(storage.get(any())).thenReturn(Optional.of(result));
+
+    Get getWithIndex =
+        Get.newBuilder()
+            .namespace(ANY_NAMESPACE_NAME)
+            .table(ANY_TABLE_NAME)
+            .indexKey(Key.ofText(ANY_NAME_3, ANY_TEXT_1))
+            .build();
+    when(snapshot.containsKeyInGetSet(getWithIndex)).thenReturn(false);
+
+    // Act
+    handler.readUnread(null, getWithIndex);
+
+    // Assert
+    verify(storage).get(any());
+    verify(snapshot)
+        .putIntoReadSet(
+            new Snapshot.Key(getWithIndex, result), Optional.of(new TransactionResult(result)));
+    verify(snapshot).putIntoGetSet(getWithIndex, Optional.of(new TransactionResult(result)));
+  }
+
+  @Test
+  public void
+      readUnread_NullKeyAndGetWithIndexNotContainedInGetSet_UncommittedRecordReturnedByStorage_ShouldThrowUncommittedRecordException()
+          throws ExecutionException {
+    // Arrange
+    Result result = mock(Result.class);
+    when(result.getInt(Attribute.STATE)).thenReturn(TransactionState.PREPARED.get());
+    when(storage.get(any())).thenReturn(Optional.of(result));
+
+    Get getWithIndex =
+        Get.newBuilder()
+            .namespace(ANY_NAMESPACE_NAME)
+            .table(ANY_TABLE_NAME)
+            .indexKey(Key.ofText(ANY_NAME_3, ANY_TEXT_1))
+            .build();
+    when(snapshot.containsKeyInGetSet(getWithIndex)).thenReturn(false);
+
+    // Act Assert
+    assertThatThrownBy(() -> handler.readUnread(null, getWithIndex))
+        .isInstanceOf(UncommittedRecordException.class)
+        .satisfies(
+            e -> {
+              UncommittedRecordException exception = (UncommittedRecordException) e;
+              assertThat(exception.getSelection()).isEqualTo(getWithIndex);
+              assertThat(exception.getResults().size()).isEqualTo(1);
+              assertThat(exception.getResults().get(0)).isEqualTo(result);
+            });
+  }
+
   @Test
   public void readIfImplicitPreReadEnabled_ShouldCallAppropriateMethods() throws CrudException {
     // Arrange

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

@@ -153,8 +153,45 @@ 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);
+    }
+  }
+
+  /** 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/processor/CsvImportProcessor.java

@@ -39,7 +39,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.

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

@@ -43,7 +43,7 @@
 public class JsonImportProcessor 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);
 
   public JsonImportProcessor(ImportProcessorParams params) {
     super(params);

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

@@ -37,7 +37,7 @@
 public class JsonLinesImportProcessor 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 JsonLinesImportProcessor with the specified parameters.

data-loader/core/src/test/java/com/scalar/db/dataloader/core/dataimport/ImportManagerTest.java

@@ -0,0 +1,101 @@
+package com.scalar.db.dataloader.core.dataimport;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import com.scalar.db.api.DistributedStorage;
+import com.scalar.db.api.DistributedTransactionManager;
+import com.scalar.db.api.TableMetadata;
+import com.scalar.db.dataloader.core.ScalarDbMode;
+import com.scalar.db.dataloader.core.dataimport.processor.ImportProcessorFactory;
+import java.io.BufferedReader;
+import java.util.HashMap;
+import java.util.Map;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class ImportManagerTest {
+
+  private ImportManager importManager;
+  private ImportEventListener listener1;
+  private ImportEventListener listener2;
+  private DistributedStorage distributedStorage;
+  private DistributedTransactionManager distributedTransactionManager;
+
+  @BeforeEach
+  void setUp() {
+    Map<String, TableMetadata> tableMetadata = new HashMap<>();
+    BufferedReader reader = mock(BufferedReader.class);
+    ImportOptions options = mock(ImportOptions.class);
+    ImportProcessorFactory processorFactory = mock(ImportProcessorFactory.class);
+
+    listener1 = mock(ImportEventListener.class);
+    listener2 = mock(ImportEventListener.class);
+    distributedStorage = mock(DistributedStorage.class);
+    distributedTransactionManager = mock(DistributedTransactionManager.class);
+
+    importManager =
+        new ImportManager(
+            tableMetadata,
+            reader,
+            options,
+            processorFactory,
+            ScalarDbMode.STORAGE,
+            distributedStorage,
+            null); // Only one resource present
+    importManager.addListener(listener1);
+    importManager.addListener(listener2);
+  }
+
+  @Test
+  void onAllDataChunksCompleted_shouldNotifyListenersAndCloseStorage() {
+    importManager.onAllDataChunksCompleted();
+
+    verify(listener1).onAllDataChunksCompleted();
+    verify(listener2).onAllDataChunksCompleted();
+    verify(distributedStorage).close();
+  }
+
+  @Test
+  void onAllDataChunksCompleted_shouldAggregateListenerExceptionAndStillCloseResources() {
+    doThrow(new RuntimeException("Listener1 failed")).when(listener1).onAllDataChunksCompleted();
+
+    RuntimeException thrown =
+        assertThrows(RuntimeException.class, () -> importManager.onAllDataChunksCompleted());
+
+    assertTrue(thrown.getMessage().contains("Error during completion"));
+    assertEquals("Listener1 failed", thrown.getCause().getMessage());
+    verify(distributedStorage).close();
+  }
+
+  @Test
+  void closeResources_shouldCloseTransactionManagerIfStorageIsNull() {
+    ImportManager managerWithTx =
+        new ImportManager(
+            new HashMap<>(),
+            mock(BufferedReader.class),
+            mock(ImportOptions.class),
+            mock(ImportProcessorFactory.class),
+            ScalarDbMode.TRANSACTION,
+            null,
+            distributedTransactionManager);
+
+    managerWithTx.closeResources();
+    verify(distributedTransactionManager).close();
+  }
+
+  @Test
+  void closeResources_shouldThrowIfResourceCloseFails() {
+    doThrow(new RuntimeException("Close failed")).when(distributedStorage).close();
+
+    RuntimeException ex =
+        assertThrows(RuntimeException.class, () -> importManager.closeResources());
+
+    assertEquals("Failed to close the resource", ex.getMessage());
+    assertEquals("Close failed", ex.getCause().getMessage());
+  }
+}