Support multiple JpaDataStore (#2998)

* Add data store customizer

* Add jta transaction manager for testing

* Fix reverse the order that the transactions are committed

* Add support in annotation to list managed classes

* Add datastore customization

* Fix test due to lock contention

* Update pom

* Remove synchronized blocks

.gitignore

@@ -18,3 +18,4 @@ dependency-reduced-pom.xml
 *.factorypath
 *.vscode
 .DS_Store
+tmlog*.log

elide-core/src/main/java/com/yahoo/elide/core/datastore/inmemory/HashMapDataStore.java

@@ -14,7 +14,8 @@
 import com.yahoo.elide.core.type.ClassType;
 import com.yahoo.elide.core.type.Type;
 import com.yahoo.elide.core.utils.ClassScanner;
-import com.google.common.collect.Sets;
+import com.yahoo.elide.core.utils.ObjectCloner;
+import com.yahoo.elide.core.utils.ObjectCloners;
 
 import lombok.Getter;
 
@@ -27,6 +28,8 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 /**
  * Simple in-memory only database.
@@ -35,29 +38,42 @@ public class HashMapDataStore implements DataStore, DataStoreTestHarness {
     protected final Map<Type<?>, Map<String, Object>> dataStore = Collections.synchronizedMap(new HashMap<>());
     @Getter protected EntityDictionary dictionary;
     @Getter private final ConcurrentHashMap<Type<?>, AtomicLong> typeIds = new ConcurrentHashMap<>();
+    private final ReadWriteLock readWriteLock = new ReentrantReadWriteLock();
+    private final ObjectCloner objectCloner;
 
     public HashMapDataStore(ClassScanner scanner, Package beanPackage) {
-        this(scanner, Sets.newHashSet(beanPackage));
+        this(scanner, beanPackage, ObjectCloners::clone);
+    }
+
+    public HashMapDataStore(ClassScanner scanner, Package beanPackage, ObjectCloner objectCloner) {
+        this(scanner, Collections.singleton(beanPackage), objectCloner);
     }
 
     public HashMapDataStore(ClassScanner scanner, Set<Package> beanPackages) {
+        this(scanner, beanPackages, ObjectCloners::clone);
+    }
+
+    public HashMapDataStore(ClassScanner scanner, Set<Package> beanPackages, ObjectCloner objectCloner) {
+        this.objectCloner = objectCloner;
         for (Package beanPackage : beanPackages) {
-            scanner.getAllClasses(beanPackage.getName()).stream()
-                .map(ClassType::new)
-                .filter(modelType -> EntityDictionary.getFirstAnnotation(modelType,
-                        Arrays.asList(Include.class, Exclude.class)) instanceof Include)
-                .forEach(modelType -> dataStore.put(modelType,
-                        Collections.synchronizedMap(new LinkedHashMap<>())));
+            process(scanner.getAllClasses(beanPackage.getName()));
         }
     }
 
     public HashMapDataStore(Collection<Class<?>> beanClasses) {
-        beanClasses.stream()
-                .map(ClassType::new)
+        this(beanClasses, ObjectCloners::clone);
+    }
+
+    public HashMapDataStore(Collection<Class<?>> beanClasses, ObjectCloner objectCloner) {
+        this.objectCloner = objectCloner;
+        process(beanClasses);
+    }
+
+    protected void process(Collection<Class<?>> beanClasses) {
+        beanClasses.stream().map(ClassType::of)
                 .filter(modelType -> EntityDictionary.getFirstAnnotation(modelType,
                         Arrays.asList(Include.class, Exclude.class)) instanceof Include)
-                .forEach(modelType -> dataStore.put(modelType,
-                        Collections.synchronizedMap(new LinkedHashMap<>())));
+                .forEach(modelType -> dataStore.put(modelType, Collections.synchronizedMap(new LinkedHashMap<>())));
     }
 
     @Override
@@ -71,7 +87,14 @@ public void populateEntityDictionary(EntityDictionary dictionary) {
 
     @Override
     public DataStoreTransaction beginTransaction() {
-        return new HashMapStoreTransaction(dataStore, dictionary, typeIds);
+        return new HashMapStoreTransaction(this.readWriteLock, this.dataStore, this.dictionary,
+                this.typeIds, this.objectCloner, false);
+    }
+
+    @Override
+    public DataStoreTransaction beginReadTransaction() {
+        return new HashMapStoreTransaction(this.readWriteLock, this.dataStore, this.dictionary,
+                this.typeIds, this.objectCloner, true);
     }
 
     @Override

elide-core/src/main/java/com/yahoo/elide/core/datastore/inmemory/HashMapStoreTransaction.java

@@ -14,16 +14,20 @@
 import com.yahoo.elide.core.request.EntityProjection;
 import com.yahoo.elide.core.request.Relationship;
 import com.yahoo.elide.core.type.Type;
+import com.yahoo.elide.core.utils.ObjectCloner;
 import com.yahoo.elide.core.utils.coerce.converters.Serde;
 
 import jakarta.persistence.GeneratedValue;
 
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
 
 /**
  * HashMapDataStore transaction handler.
@@ -33,13 +37,27 @@ public class HashMapStoreTransaction implements DataStoreTransaction {
     private final List<Operation> operations;
     private final EntityDictionary dictionary;
     private final Map<Type<?>, AtomicLong> typeIds;
+    private final Lock lock;
+    private final boolean readOnly;
+    private final ObjectCloner objectCloner;
+    private boolean committed = false;
 
-    public HashMapStoreTransaction(Map<Type<?>, Map<String, Object>> dataStore,
-                                   EntityDictionary dictionary, Map<Type<?>, AtomicLong> typeIds) {
+    public HashMapStoreTransaction(ReadWriteLock readWriteLock, Map<Type<?>, Map<String, Object>> dataStore,
+            EntityDictionary dictionary, Map<Type<?>, AtomicLong> typeIds, ObjectCloner objectCloner,
+            boolean readOnly) {
+        this.readOnly = readOnly;
         this.dataStore = dataStore;
         this.dictionary = dictionary;
         this.operations = new ArrayList<>();
         this.typeIds = typeIds;
+        this.objectCloner = objectCloner;
+
+        if (readWriteLock != null) {
+            this.lock = readOnly ? readWriteLock.readLock() : readWriteLock.writeLock();
+            this.lock.lock();
+        } else {
+            this.lock = null;
+        }
     }
 
     @Override
@@ -74,24 +92,21 @@ public void delete(Object object, RequestScope requestScope) {
 
     @Override
     public void commit(RequestScope scope) {
-        synchronized (dataStore) {
-            operations.stream()
-                    .filter(op -> op.getInstance() != null)
-                    .forEach(op -> {
-                        Object instance = op.getInstance();
-                        String id = op.getId();
-                        Map<String, Object> data = dataStore.get(op.getType());
-                        if (op.getOpType() == Operation.OpType.DELETE) {
-                            data.remove(id);
-                        } else {
-                            if (op.getOpType() == Operation.OpType.CREATE && data.get(id) != null) {
-                                throw new TransactionException(new IllegalStateException("Duplicate key"));
-                            }
-                            data.put(id, instance);
-                        }
-                    });
-            operations.clear();
-        }
+        operations.stream().filter(op -> op.getInstance() != null).forEach(op -> {
+            Object instance = op.getInstance();
+            String id = op.getId();
+            Map<String, Object> data = dataStore.get(op.getType());
+            if (op.getOpType() == Operation.OpType.DELETE) {
+                data.remove(id);
+            } else {
+                if (op.getOpType() == Operation.OpType.CREATE && data.get(id) != null) {
+                    throw new TransactionException(new IllegalStateException("Duplicate key"));
+                }
+                data.put(id, instance);
+            }
+        });
+        operations.clear();
+        committed = true;
     }
 
     @Override
@@ -108,10 +123,7 @@ public void createObject(Object entity, RequestScope scope) {
         //GeneratedValue means the DB needs to assign the ID.
         if (dictionary.getAttributeOrRelationAnnotation(entityClass, GeneratedValue.class, idFieldName) != null) {
             // TODO: Id's are not necessarily numeric.
-            AtomicLong nextId;
-            synchronized (dataStore) {
-                nextId = getId(entityClass);
-            }
+            AtomicLong nextId = getId(entityClass);
             id = String.valueOf(nextId.getAndIncrement());
             setId(entity, id);
         } else {
@@ -138,32 +150,70 @@ public DataStoreIterable<Object> getToManyRelation(DataStoreTransaction relation
     @Override
     public DataStoreIterable<Object> loadObjects(EntityProjection projection,
                                                           RequestScope scope) {
-        synchronized (dataStore) {
-            Map<String, Object> data = dataStore.get(projection.getType());
-            return new DataStoreIterableBuilder<>(data.values()).allInMemory().build();
-        }
+        Map<String, Object> data = dataStore.get(projection.getType());
+        cacheForRollback(projection.getType(), data);
+        return new DataStoreIterableBuilder<>(data.values()).allInMemory().build();
     }
 
     @Override
     public Object loadObject(EntityProjection projection, Serializable id, RequestScope scope) {
 
         EntityDictionary dictionary = scope.getDictionary();
 
-        synchronized (dataStore) {
-            Map<String, Object> data = dataStore.get(projection.getType());
-            if (data == null) {
-                return null;
-            }
-            Serde serde = dictionary.getSerdeLookup().apply(id.getClass());
+        Map<String, Object> data = dataStore.get(projection.getType());
+        cacheForRollback(projection.getType(), data);
+        if (data == null) {
+            return null;
+        }
+        Serde serde = dictionary.getSerdeLookup().apply(id.getClass());
+
+        String idString = (serde == null) ? id.toString() : (String) serde.serialize(id);
+        return data.get(idString);
+    }
+
+    /**
+     * Contains a copy of the objects loaded from the hash map store as they may be
+     * updated like a persistent object. Since what is returned is a reference to
+     * the object in the underlying store when updated it immediately reflects in
+     * the store. As such a copy of the original objects need to made in order to
+     * rollback.
+     */
+    private Map<Type<?>, Map<String, Object>> rollbackCache = new HashMap<>();
 
-            String idString = (serde == null) ? id.toString() : (String) serde.serialize(id);
-            return data.get(idString);
+    protected void cacheForRollback(Type<?> type, Map<String, Object> data) {
+        if (!readOnly) {
+            this.rollbackCache.computeIfAbsent(type, key -> {
+                if (data != null) {
+                    Map<String, Object> copy = new HashMap<>();
+                    data.entrySet().stream().forEach(entry -> {
+                        Object value = this.objectCloner.clone(entry.getValue(), type);
+                        copy.put(entry.getKey(), value);
+                    });
+                    return copy;
+                }
+                return null;
+            });
         }
     }
 
     @Override
     public void close() throws IOException {
-        operations.clear();
+        try {
+            if (!committed && !readOnly) {
+                rollback();
+            }
+            operations.clear();
+        } finally {
+            if (this.lock != null) {
+                this.lock.unlock();
+            }
+        }
+    }
+
+    public void rollback() {
+        // Rollback data
+        dataStore.putAll(this.rollbackCache);
+        this.rollbackCache.clear();
     }
 
     private boolean containsObject(Object obj) {

elide-core/src/main/java/com/yahoo/elide/core/utils/ObjectCloner.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2023, the original author or authors.
+ * Licensed under the Apache License, Version 2.0
+ * See LICENSE file in project root for terms.
+ */
+package com.yahoo.elide.core.utils;
+
+import com.yahoo.elide.core.type.ClassType;
+import com.yahoo.elide.core.type.Type;
+
+/**
+ * Clones an object.
+ */
+@FunctionalInterface
+public interface ObjectCloner {
+    <T> T clone(T source, Type<?> cls);
+
+    default <T> T clone(T source) {
+        return clone(source, ClassType.of(source.getClass()));
+    }
+}