Allow DynamoDbStoreTransaction to span multiple stores

Author: Alexander Patrikalakis

src/main/java/com/amazon/janusgraph/diskstorage/dynamodb/AbstractDynamoDbStore.java

@@ -16,7 +16,6 @@
 
 import java.util.Collections;
 import java.util.List;
-import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 
@@ -48,7 +47,6 @@
 import com.google.common.util.concurrent.ExecutionError;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 
-import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
 /**
@@ -67,6 +65,10 @@ public abstract class AbstractDynamoDbStore implements AwsStore {
     @Getter
     private final String name;
     private final boolean forceConsistentRead;
+    /**
+     * The key column local lock cache maps key-column pairs to the DynamoDbStoreTransaction that first
+     * acquired a lock on those key-column pairs.
+     */
     private final Cache<Pair<StaticBuffer, StaticBuffer>, DynamoDbStoreTransaction> keyColumnLocalLocks;
 
     private enum ReportingRemovalListener implements RemovalListener<Pair<StaticBuffer, StaticBuffer>, DynamoDbStoreTransaction> {
@@ -152,28 +154,14 @@ public final void deleteStore() throws BackendException {
         client.getDelegate().ensureTableDeleted(getTableSchema().getTableName());
     }
 
-    @RequiredArgsConstructor
-    private class SetStoreIfTxMappingDoesntExist implements Callable<DynamoDbStoreTransaction> {
-        private final DynamoDbStoreTransaction tx;
-        private final Pair<StaticBuffer, StaticBuffer> keyColumn;
-        @Override
-        public DynamoDbStoreTransaction call() throws Exception {
-            tx.setStore(AbstractDynamoDbStore.this);
-            log.trace(String.format("acquiring lock on %s at " + System.nanoTime(), keyColumn.toString()));
-            // do not extend the expiry of an existing lock by the tx passed in this method
-            return tx;
-        }
-    }
-
     @Override
     public void acquireLock(final StaticBuffer key, final StaticBuffer column, final StaticBuffer expectedValue, final StoreTransaction txh) throws BackendException {
         final DynamoDbStoreTransaction tx = DynamoDbStoreTransaction.getTx(txh);
         final Pair<StaticBuffer, StaticBuffer> keyColumn = Pair.of(key, column);
 
         final DynamoDbStoreTransaction existing;
         try {
-            existing = keyColumnLocalLocks.get(keyColumn,
-                new SetStoreIfTxMappingDoesntExist(tx, keyColumn));
+            existing = keyColumnLocalLocks.get(keyColumn, () -> tx);
         } catch (ExecutionException | UncheckedExecutionException | ExecutionError e) {
             throw new TemporaryLockingException("Unable to acquire lock", e);
         }
@@ -182,17 +170,15 @@ public void acquireLock(final StaticBuffer key, final StaticBuffer column, final
         }
 
         // Titan's locking expects that only the first expectedValue for a given key/column should be used
-        if (!tx.contains(key, column)) {
-            tx.put(key, column, expectedValue);
-        }
+        tx.putKeyColumnOnlyIfItIsNotYetChangedInTx(this, key, column, expectedValue);
     }
 
     @Override
     public void close() throws BackendException {
         log.debug("Closing table:{}", tableName);
     }
 
-    protected String encodeKeyForLog(final StaticBuffer key) {
+    String encodeKeyForLog(final StaticBuffer key) {
         if (null == key) {
             return "";
         }

src/main/java/com/amazon/janusgraph/diskstorage/dynamodb/DynamoDbSingleRowStore.java

@@ -204,9 +204,8 @@ public Collection<MutateWorker> createMutationWorkers(final Map<StaticBuffer, KC
 
             // Using ExpectedAttributeValue map to handle large mutations in a single request
             // Large mutations would require multiple requests using expressions
-            final Map<String, ExpectedAttributeValue> expected = new SingleExpectedAttributeValueBuilder().key(hashKey)
-                                                                                                    .transaction(txh)
-                                                                                                    .build(mutation);
+            final Map<String, ExpectedAttributeValue> expected =
+                new SingleExpectedAttributeValueBuilder(this, txh, hashKey).build(mutation);
 
             final Map<String, AttributeValueUpdate> attributeValueUpdates =
                 new SingleUpdateBuilder().deletions(mutation.getDeletions())

src/main/java/com/amazon/janusgraph/diskstorage/dynamodb/DynamoDbStore.java

@@ -268,7 +268,7 @@ private Collection<MutateWorker> createWorkersForAdditions(final StaticBuffer ha
                     .rangeKey(rangeKey)
                     .build();
 
-                final Expression updateExpression = new MultiUpdateExpressionBuilder(txh).hashKey(hashKey)
+                final Expression updateExpression = new MultiUpdateExpressionBuilder(this, txh).hashKey(hashKey)
                     .rangeKey(rangeKey)
                     .value(addition.getValue())
                     .build();
@@ -290,7 +290,7 @@ private Collection<MutateWorker> createWorkersForDeletions(final StaticBuffer ha
                                                                       .rangeKey(rangeKey)
                                                                       .build();
 
-            final Expression updateExpression = new MultiUpdateExpressionBuilder(txh).hashKey(hashKey)
+            final Expression updateExpression = new MultiUpdateExpressionBuilder(this, txh).hashKey(hashKey)
                                                                                   .rangeKey(rangeKey)
                                                                                   .build();
 

src/main/java/com/amazon/janusgraph/diskstorage/dynamodb/DynamoDbStoreTransaction.java

@@ -14,6 +14,7 @@
  */
 package com.amazon.janusgraph.diskstorage.dynamodb;
 
+import java.util.HashMap;
 import java.util.Map;
 
 import org.apache.commons.lang.builder.EqualsBuilder;
@@ -39,7 +40,7 @@
  *
  */
 @Slf4j
-public class DynamoDbStoreTransaction extends AbstractStoreTransaction { //CHECKSTYLE SUPPRESS - for the DB in DynamoDB
+public class DynamoDbStoreTransaction extends AbstractStoreTransaction {
 
     public static DynamoDbStoreTransaction getTx(@NonNull final StoreTransaction txh) {
         Preconditions
@@ -51,8 +52,7 @@ public static DynamoDbStoreTransaction getTx(@NonNull final StoreTransaction txh
      * This is only used for toString for debugging purposes.
      */
     private final String id;
-    private final Map<StaticBuffer, Map<StaticBuffer, StaticBuffer>> expectedValues = Maps.newHashMap();
-    private AbstractDynamoDbStore store;
+    private final Map<AbstractDynamoDbStore, Map<StaticBuffer, Map<StaticBuffer, StaticBuffer>>> expectedValues = Maps.newHashMap();
 
     /**
      * Creates a DynamoDB Store transaction.
@@ -73,12 +73,9 @@ public void commit() throws BackendException {
     }
 
     private void releaseLocks() {
-        for (final Map.Entry<StaticBuffer, Map<StaticBuffer, StaticBuffer>> entry : expectedValues.entrySet()) {
-            final StaticBuffer key = entry.getKey();
-            for (final StaticBuffer column : entry.getValue().keySet()) {
-                store.releaseLock(key, column);
-            }
-        }
+        expectedValues.forEach((store, kcMap) -> kcMap.forEach((key, columnValueMap) -> {
+            columnValueMap.forEach((column, valueIgnored) -> store.releaseLock(key, column));
+        }));
     }
 
     /**
@@ -87,8 +84,10 @@ private void releaseLocks() {
      * @param column column to check for existence
      * @return true if both the key and column combination are in this transaction and false otherwise.
      */
-    public boolean contains(final StaticBuffer key, final StaticBuffer column) {
-        return expectedValues.containsKey(key) && expectedValues.get(key).containsKey(column);
+    public boolean contains(final AbstractDynamoDbStore store, final StaticBuffer key, final StaticBuffer column) {
+        return expectedValues.containsKey(store)
+            && expectedValues.get(store).containsKey(key)
+            && expectedValues.get(store).get(key).containsKey(column);
     }
 
     @Override
@@ -112,35 +111,29 @@ public int hashCode() {
 
     /**
      * Gets the expected value for a particular key and column, if any
+     * @param store the store to put the expected key column value
      * @param key the key to get the expected value for
      * @param column the column to get the expected value for
      * @return the expected value of the given key-column pair, if any.
      */
-    public StaticBuffer get(final StaticBuffer key, final StaticBuffer column) {
+    public StaticBuffer get(final AbstractDynamoDbStore store, final StaticBuffer key, final StaticBuffer column) {
         // This method assumes the caller has called contains(..) and received a positive response
-        return expectedValues.get(key)
-                             .get(column);
+        return expectedValues.get(store).get(key).get(column);
     }
 
     /**
      * Puts the expected value for a particular key and column
+     * @param store the store to put the expected key column value
      * @param key the key to put the expected value for
      * @param column the column to put the expected value for
      * @param expectedValue the expected value to put
      */
-    public void put(final StaticBuffer key, final StaticBuffer column, final StaticBuffer expectedValue) {
-        final Map<StaticBuffer, StaticBuffer> valueMap;
-        if (expectedValues.containsKey(key)) {
-            valueMap = expectedValues.get(key);
-        } else {
-            valueMap = Maps.newHashMap();
-            expectedValues.put(key, valueMap);
-        }
-
-        // Ignore any calls to put if we already have an expected value
-        if (!valueMap.containsKey(column)) {
-            valueMap.put(column, expectedValue);
-        }
+    public void putKeyColumnOnlyIfItIsNotYetChangedInTx(final AbstractDynamoDbStore store, final StaticBuffer key, final StaticBuffer column,
+        final StaticBuffer expectedValue) {
+        expectedValues.computeIfAbsent(store, s -> new HashMap<>());
+        expectedValues.get(store).computeIfAbsent(key, k -> new HashMap<>());
+        // Ignore any calls to putKeyColumnOnlyIfItIsNotYetChangedInTx if we already have an expected value
+        expectedValues.get(store).get(key).putIfAbsent(column, expectedValue);
     }
 
     @Override
@@ -155,8 +148,4 @@ public void rollback() throws BackendException {
     public String toString() {
         return new ToStringBuilder(this).append(id).append(expectedValues).toString();
     }
-
-    public void setStore(final AbstractDynamoDbStore abstractDynamoDbStore) {
-        this.store = abstractDynamoDbStore;
-    }
 }

src/main/java/com/amazon/janusgraph/diskstorage/dynamodb/builder/MultiUpdateExpressionBuilder.java

@@ -14,12 +14,12 @@
  */
 package com.amazon.janusgraph.diskstorage.dynamodb.builder;
 
-import java.util.Collections;
 import java.util.Map;
 
 import org.janusgraph.diskstorage.StaticBuffer;
 
 import com.amazon.janusgraph.diskstorage.dynamodb.Constants;
+import com.amazon.janusgraph.diskstorage.dynamodb.DynamoDbStore;
 import com.amazon.janusgraph.diskstorage.dynamodb.DynamoDbStoreTransaction;
 import com.amazon.janusgraph.diskstorage.dynamodb.Expression;
 import com.amazonaws.services.dynamodbv2.model.AttributeValue;
@@ -47,8 +47,8 @@ public class MultiUpdateExpressionBuilder extends AbstractBuilder {
     private static final String EXPECTED_VALUE_EXPR = String.format("%s = %s", Constants.JANUSGRAPH_VALUE, EXPECTED_VALUE_LABEL);
     private static final String SET_VALUE_EXPR = String.format("SET %s = %s", Constants.JANUSGRAPH_VALUE, VALUE_LABEL);
 
-    private static final Map<String, String> EMPTY_ARGUMENT_NAMES = Collections.emptyMap();
-
+    @NonNull
+    private final DynamoDbStore store;
     @NonNull
     private final DynamoDbStoreTransaction transaction;
     @Setter
@@ -75,8 +75,8 @@ public Expression build() {
 
         // Condition expression and attribute value
         String conditionExpression = null;
-        if (transaction.contains(hashKey, rangeKey)) {
-            final StaticBuffer expectedValue = transaction.get(hashKey, rangeKey);
+        if (transaction.contains(store, hashKey, rangeKey)) {
+            final StaticBuffer expectedValue = transaction.get(store, hashKey, rangeKey);
             if (expectedValue == null) {
                 conditionExpression = MISSING_VALUE_EXPR;
             } else {

src/main/java/com/amazon/janusgraph/diskstorage/dynamodb/builder/SingleExpectedAttributeValueBuilder.java

@@ -20,15 +20,15 @@
 import org.janusgraph.diskstorage.StaticBuffer;
 import org.janusgraph.diskstorage.keycolumnvalue.KCVMutation;
 
+import com.amazon.janusgraph.diskstorage.dynamodb.DynamoDbSingleRowStore;
 import com.amazon.janusgraph.diskstorage.dynamodb.DynamoDbStoreTransaction;
 import com.amazonaws.services.dynamodbv2.model.AttributeValue;
 import com.amazonaws.services.dynamodbv2.model.ComparisonOperator;
 import com.amazonaws.services.dynamodbv2.model.ExpectedAttributeValue;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Maps;
 
-import lombok.Setter;
-import lombok.experimental.Accessors;
+import lombok.RequiredArgsConstructor;
 
 
 /**
@@ -38,12 +38,12 @@
  * @author Michael Rodaitis
  * @author Alexander Patrikalakis
  */
-@Setter
-@Accessors(fluent = true, chain = true)
+@RequiredArgsConstructor
 public class SingleExpectedAttributeValueBuilder extends AbstractBuilder {
 
-    private DynamoDbStoreTransaction transaction = null;
-    private StaticBuffer key = null;
+    private final DynamoDbSingleRowStore store;
+    private final DynamoDbStoreTransaction transaction;
+    private final StaticBuffer key;
 
     public Map<String, ExpectedAttributeValue> build(final KCVMutation mutation) {
         Preconditions.checkState(transaction != null, "Transaction must not be null");
@@ -70,8 +70,8 @@ private void addExpectedValueIfPresent(final StaticBuffer column, final Map<Stri
             return;
         }
 
-        if (transaction.contains(key, column)) {
-            final StaticBuffer expectedValue = transaction.get(key, column);
+        if (transaction.contains(store, key, column)) {
+            final StaticBuffer expectedValue = transaction.get(store, key, column);
             final ExpectedAttributeValue expectedAttributeValue;
             if (expectedValue == null) {
                 expectedAttributeValue = new ExpectedAttributeValue().withExists(false);