Make Scalar DB transactions be able to abort eagerly

core/src/main/java/com/scalar/db/api/DistributedTransactionManager.java

@@ -122,6 +122,15 @@ DistributedTransaction start(String txId, Isolation isolation, SerializableStrat
    */
   TransactionState getState(String txId) throws TransactionException;
 
+  /**
+   * Aborts a given transaction.
+   *
+   * @param txId a transaction ID
+   * @return {@link TransactionState}
+   * @throws TransactionException if aborting the given transaction failed
+   */
+  TransactionState abort(String txId) throws TransactionException;
+
   /**
    * Closes connections to the cluster. The connections are shared among multiple services such as
    * StorageService and TransactionService, thus this should only be used when closing applications.

core/src/main/java/com/scalar/db/service/TransactionService.java

@@ -96,6 +96,11 @@ public TransactionState getState(String txId) throws TransactionException {
     return manager.getState(txId);
   }
 
+  @Override
+  public TransactionState abort(String txId) throws TransactionException {
+    return manager.abort(txId);
+  }
+
   @Override
   public void close() {
     manager.close();

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

@@ -79,7 +79,7 @@ public void commit(Snapshot snapshot) throws CommitException, UnknownTransaction
     }
   }
 
-  private TransactionState abort(String id) throws UnknownTransactionStatusException {
+  public TransactionState abort(String id) throws UnknownTransactionStatusException {
     try {
       abortState(id);
       return TransactionState.ABORTED;

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

@@ -11,6 +11,7 @@
 import com.scalar.db.api.TransactionState;
 import com.scalar.db.config.DatabaseConfig;
 import com.scalar.db.exception.transaction.CoordinatorException;
+import com.scalar.db.exception.transaction.UnknownTransactionStatusException;
 import com.scalar.db.transaction.consensuscommit.Coordinator.State;
 import java.util.Optional;
 import java.util.UUID;
@@ -160,6 +161,16 @@ public TransactionState getState(String txId) {
     return TransactionState.UNKNOWN;
   }
 
+  @Override
+  public TransactionState abort(String txId) {
+    checkArgument(!Strings.isNullOrEmpty(txId));
+    try {
+      return commit.abort(txId);
+    } catch (UnknownTransactionStatusException ignored) {
+      return TransactionState.UNKNOWN;
+    }
+  }
+
   @Override
   public void close() {
     storage.close();

core/src/main/java/com/scalar/db/transaction/jdbc/JdbcTransactionManager.java

@@ -137,7 +137,12 @@ public JdbcTransaction start(String txId, Isolation isolation, SerializableStrat
 
   @Override
   public TransactionState getState(String txId) {
-    return TransactionState.UNKNOWN; // always returns UNKNOWN
+    throw new UnsupportedOperationException("this method is not supported in JDBC transaction");
+  }
+
+  @Override
+  public TransactionState abort(String txId) {
+    throw new UnsupportedOperationException("this method is not supported in JDBC transaction");
   }
 
   @Override

core/src/main/java/com/scalar/db/transaction/rpc/GrpcTransactionManager.java

@@ -8,6 +8,8 @@
 import com.scalar.db.api.TransactionState;
 import com.scalar.db.config.DatabaseConfig;
 import com.scalar.db.exception.transaction.TransactionException;
+import com.scalar.db.rpc.AbortRequest;
+import com.scalar.db.rpc.AbortResponse;
 import com.scalar.db.rpc.DistributedStorageAdminGrpc;
 import com.scalar.db.rpc.DistributedTransactionGrpc;
 import com.scalar.db.rpc.GetTransactionStateRequest;
@@ -156,6 +158,16 @@ public TransactionState getState(String txId) throws TransactionException {
         });
   }
 
+  @Override
+  public TransactionState abort(String txId) throws TransactionException {
+    return execute(
+        () -> {
+          AbortResponse response =
+              blockingStub.abort(AbortRequest.newBuilder().setTransactionId(txId).build());
+          return ProtoUtil.toTransactionState(response.getState());
+        });
+  }
+
   private static <T> T execute(ThrowableSupplier<T, Throwable> throwableSupplier)
       throws TransactionException {
     try {

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

@@ -10,6 +10,7 @@
 import com.scalar.db.api.TransactionState;
 import com.scalar.db.config.DatabaseConfig;
 import com.scalar.db.exception.transaction.CoordinatorException;
+import com.scalar.db.exception.transaction.UnknownTransactionStatusException;
 import com.scalar.db.transaction.consensuscommit.Coordinator.State;
 import java.util.Optional;
 import org.junit.Before;
@@ -130,4 +131,45 @@ public void check_CoordinatorExceptionThrown_ReturnUnknown() throws CoordinatorE
     // Assert
     assertThat(actual).isEqualTo(TransactionState.UNKNOWN);
   }
+
+  @Test
+  public void abort_CommitHandlerReturnsAborted_ShouldReturnTheState()
+      throws UnknownTransactionStatusException {
+    // Arrange
+    TransactionState expected = TransactionState.ABORTED;
+    when(commit.abort(ANY_TX_ID)).thenReturn(expected);
+
+    // Act
+    TransactionState actual = manager.abort(ANY_TX_ID);
+
+    // Assert
+    assertThat(actual).isEqualTo(expected);
+  }
+
+  @Test
+  public void abort_CommitHandlerReturnsCommitted_ShouldReturnTheState()
+      throws UnknownTransactionStatusException {
+    // Arrange
+    TransactionState expected = TransactionState.COMMITTED;
+    when(commit.abort(ANY_TX_ID)).thenReturn(expected);
+
+    // Act
+    TransactionState actual = manager.abort(ANY_TX_ID);
+
+    // Assert
+    assertThat(actual).isEqualTo(expected);
+  }
+
+  @Test
+  public void abort_CommitHandlerThrowsUnknownTransactionStatusException_ShouldReturnUnknown()
+      throws UnknownTransactionStatusException {
+    // Arrange
+    when(commit.abort(ANY_TX_ID)).thenThrow(UnknownTransactionStatusException.class);
+
+    // Act
+    TransactionState actual = manager.abort(ANY_TX_ID);
+
+    // Assert
+    assertThat(actual).isEqualTo(TransactionState.UNKNOWN);
+  }
 }

core/src/test/java/com/scalar/db/transaction/rpc/GrpcTransactionManagerTest.java

@@ -7,6 +7,7 @@
 import static org.mockito.Mockito.when;
 
 import com.scalar.db.exception.transaction.TransactionException;
+import com.scalar.db.rpc.AbortResponse;
 import com.scalar.db.rpc.DistributedTransactionGrpc;
 import com.scalar.db.rpc.GetTransactionStateResponse;
 import com.scalar.db.rpc.TransactionState;
@@ -68,4 +69,37 @@ public void getState_StubThrowsInternalError_ShouldThrowTransactionException() {
     // Act Assert
     assertThatThrownBy(() -> manager.getState(ANY_ID)).isInstanceOf(TransactionException.class);
   }
+
+  @Test
+  public void abort_IsCalledWithoutAnyArguments_StubShouldBeCalledProperly()
+      throws TransactionException {
+    // Arrange
+    AbortResponse response = mock(AbortResponse.class);
+    when(response.getState()).thenReturn(TransactionState.TRANSACTION_STATE_ABORTED);
+    when(blockingStub.abort(any())).thenReturn(response);
+
+    // Act
+    manager.abort(ANY_ID);
+
+    // Assert
+    verify(blockingStub).abort(any());
+  }
+
+  @Test
+  public void abort_StubThrowsInvalidArgumentError_ShouldThrowIllegalArgumentException() {
+    // Arrange
+    when(blockingStub.abort(any())).thenThrow(Status.INVALID_ARGUMENT.asRuntimeException());
+
+    // Act Assert
+    assertThatThrownBy(() -> manager.abort(ANY_ID)).isInstanceOf(IllegalArgumentException.class);
+  }
+
+  @Test
+  public void abort_StubThrowsInternalError_ShouldThrowTransactionException() {
+    // Arrange
+    when(blockingStub.abort(any())).thenThrow(Status.INTERNAL.asRuntimeException());
+
+    // Act Assert
+    assertThatThrownBy(() -> manager.abort(ANY_ID)).isInstanceOf(TransactionException.class);
+  }
 }

rpc/src/main/grpc/com/scalar/db/rpc/DistributedTransactionGrpc.java

@@ -76,6 +76,37 @@ com.scalar.db.rpc.GetTransactionStateResponse> getGetStateMethod() {
     return getGetStateMethod;
   }
 
+  private static volatile io.grpc.MethodDescriptor<com.scalar.db.rpc.AbortRequest,
+      com.scalar.db.rpc.AbortResponse> getAbortMethod;
+
+  @io.grpc.stub.annotations.RpcMethod(
+      fullMethodName = SERVICE_NAME + '/' + "Abort",
+      requestType = com.scalar.db.rpc.AbortRequest.class,
+      responseType = com.scalar.db.rpc.AbortResponse.class,
+      methodType = io.grpc.MethodDescriptor.MethodType.UNARY)
+  public static io.grpc.MethodDescriptor<com.scalar.db.rpc.AbortRequest,
+      com.scalar.db.rpc.AbortResponse> getAbortMethod() {
+    io.grpc.MethodDescriptor<com.scalar.db.rpc.AbortRequest, com.scalar.db.rpc.AbortResponse> getAbortMethod;
+    if ((getAbortMethod = DistributedTransactionGrpc.getAbortMethod) == null) {
+      synchronized (DistributedTransactionGrpc.class) {
+        if ((getAbortMethod = DistributedTransactionGrpc.getAbortMethod) == null) {
+          DistributedTransactionGrpc.getAbortMethod = getAbortMethod =
+              io.grpc.MethodDescriptor.<com.scalar.db.rpc.AbortRequest, com.scalar.db.rpc.AbortResponse>newBuilder()
+              .setType(io.grpc.MethodDescriptor.MethodType.UNARY)
+              .setFullMethodName(generateFullMethodName(SERVICE_NAME, "Abort"))
+              .setSampledToLocalTracing(true)
+              .setRequestMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
+                  com.scalar.db.rpc.AbortRequest.getDefaultInstance()))
+              .setResponseMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
+                  com.scalar.db.rpc.AbortResponse.getDefaultInstance()))
+              .setSchemaDescriptor(new DistributedTransactionMethodDescriptorSupplier("Abort"))
+              .build();
+        }
+      }
+    }
+    return getAbortMethod;
+  }
+
   /**
    * Creates a new async stub that supports all call types for the service
    */
@@ -138,6 +169,13 @@ public void getState(com.scalar.db.rpc.GetTransactionStateRequest request,
       io.grpc.stub.ServerCalls.asyncUnimplementedUnaryCall(getGetStateMethod(), responseObserver);
     }
 
+    /**
+     */
+    public void abort(com.scalar.db.rpc.AbortRequest request,
+        io.grpc.stub.StreamObserver<com.scalar.db.rpc.AbortResponse> responseObserver) {
+      io.grpc.stub.ServerCalls.asyncUnimplementedUnaryCall(getAbortMethod(), responseObserver);
+    }
+
     @java.lang.Override public final io.grpc.ServerServiceDefinition bindService() {
       return io.grpc.ServerServiceDefinition.builder(getServiceDescriptor())
           .addMethod(
@@ -154,6 +192,13 @@ public void getState(com.scalar.db.rpc.GetTransactionStateRequest request,
                 com.scalar.db.rpc.GetTransactionStateRequest,
                 com.scalar.db.rpc.GetTransactionStateResponse>(
                   this, METHODID_GET_STATE)))
+          .addMethod(
+            getAbortMethod(),
+            io.grpc.stub.ServerCalls.asyncUnaryCall(
+              new MethodHandlers<
+                com.scalar.db.rpc.AbortRequest,
+                com.scalar.db.rpc.AbortResponse>(
+                  this, METHODID_ABORT)))
           .build();
     }
   }
@@ -187,6 +232,14 @@ public void getState(com.scalar.db.rpc.GetTransactionStateRequest request,
       io.grpc.stub.ClientCalls.asyncUnaryCall(
           getChannel().newCall(getGetStateMethod(), getCallOptions()), request, responseObserver);
     }
+
+    /**
+     */
+    public void abort(com.scalar.db.rpc.AbortRequest request,
+        io.grpc.stub.StreamObserver<com.scalar.db.rpc.AbortResponse> responseObserver) {
+      io.grpc.stub.ClientCalls.asyncUnaryCall(
+          getChannel().newCall(getAbortMethod(), getCallOptions()), request, responseObserver);
+    }
   }
 
   /**
@@ -209,6 +262,13 @@ public com.scalar.db.rpc.GetTransactionStateResponse getState(com.scalar.db.rpc.
       return io.grpc.stub.ClientCalls.blockingUnaryCall(
           getChannel(), getGetStateMethod(), getCallOptions(), request);
     }
+
+    /**
+     */
+    public com.scalar.db.rpc.AbortResponse abort(com.scalar.db.rpc.AbortRequest request) {
+      return io.grpc.stub.ClientCalls.blockingUnaryCall(
+          getChannel(), getAbortMethod(), getCallOptions(), request);
+    }
   }
 
   /**
@@ -232,10 +292,19 @@ public com.google.common.util.concurrent.ListenableFuture<com.scalar.db.rpc.GetT
       return io.grpc.stub.ClientCalls.futureUnaryCall(
           getChannel().newCall(getGetStateMethod(), getCallOptions()), request);
     }
+
+    /**
+     */
+    public com.google.common.util.concurrent.ListenableFuture<com.scalar.db.rpc.AbortResponse> abort(
+        com.scalar.db.rpc.AbortRequest request) {
+      return io.grpc.stub.ClientCalls.futureUnaryCall(
+          getChannel().newCall(getAbortMethod(), getCallOptions()), request);
+    }
   }
 
   private static final int METHODID_GET_STATE = 0;
-  private static final int METHODID_TRANSACTION = 1;
+  private static final int METHODID_ABORT = 1;
+  private static final int METHODID_TRANSACTION = 2;
 
   private static final class MethodHandlers<Req, Resp> implements
       io.grpc.stub.ServerCalls.UnaryMethod<Req, Resp>,
@@ -258,6 +327,10 @@ public void invoke(Req request, io.grpc.stub.StreamObserver<Resp> responseObserv
           serviceImpl.getState((com.scalar.db.rpc.GetTransactionStateRequest) request,
               (io.grpc.stub.StreamObserver<com.scalar.db.rpc.GetTransactionStateResponse>) responseObserver);
           break;
+        case METHODID_ABORT:
+          serviceImpl.abort((com.scalar.db.rpc.AbortRequest) request,
+              (io.grpc.stub.StreamObserver<com.scalar.db.rpc.AbortResponse>) responseObserver);
+          break;
         default:
           throw new AssertionError();
       }
@@ -324,6 +397,7 @@ public static io.grpc.ServiceDescriptor getServiceDescriptor() {
               .setSchemaDescriptor(new DistributedTransactionFileDescriptorSupplier())
               .addMethod(getTransactionMethod())
               .addMethod(getGetStateMethod())
+              .addMethod(getAbortMethod())
               .build();
         }
       }