Fix/updating central registry

backend/src/main/java/com/bakdata/conquery/io/storage/DirectIdentifiableStore.java

@@ -46,4 +46,18 @@ protected void added(VALUE value) {
 			throw new RuntimeException("Failed to add "+value, e);
 		}
 	}
+
+	@Override
+	protected void updated(VALUE value) {
+		try {
+			if (value == null) {
+				return;
+			}
+
+			centralRegistry.update(value);
+			onAdd.accept(value);
+		} catch(Exception e) {
+			throw new RuntimeException("Failed to add "+value, e);
+		}
+	}
 }

backend/src/main/java/com/bakdata/conquery/io/storage/IdentifiableCachedStore.java

@@ -8,6 +8,8 @@
 import lombok.Setter;
 import lombok.experimental.Accessors;
 
+import java.util.function.Function;
+
 /**
  * Registers accessors of values instead of the value itself to the central registry.
  * Might be useful if the object are very large and should only be loaded on demand.
@@ -49,6 +51,18 @@ protected void added(VALUE value) {
 			throw new RuntimeException("Failed to add "+value, e);
 		}
 	}
+
+	@Override
+	protected void updated(VALUE value) {
+		try {
+			if(value != null) {
+				final IId<VALUE> key = extractKey(value);
+				centralRegistry.updateCacheable(key, this::get);
+			}
+		} catch(Exception e) {
+			throw new RuntimeException("Failed to add "+value, e);
+		}
+	}
 
 	@Override
 	public void loadData() {

backend/src/main/java/com/bakdata/conquery/io/storage/NamespaceStorage.java

@@ -71,7 +71,9 @@ public Dictionary getPrimaryDictionaryRaw() {
 
 
 	private void decorateIdMapping(SingletonStore<EntityIdMap> idMapping) {
-		idMapping.onAdd(mapping -> mapping.setStorage(this));
+		idMapping
+				.onAdd(mapping -> mapping.setStorage(this))
+				.onUpdate(mapping -> mapping.setStorage(this));
 	}
 
 

backend/src/main/java/com/bakdata/conquery/io/storage/NamespacedStorage.java

@@ -113,8 +113,10 @@ public void removeStorage() {
 	protected abstract boolean isRegisterImports();
 
 	private void decorateDatasetStore(SingletonStore<Dataset> store) {
-		store.onAdd(centralRegistry::register)
-			 .onRemove(centralRegistry::remove);
+		store
+				.onAdd(centralRegistry::register)
+				.onRemove(centralRegistry::remove)
+				.onUpdate(centralRegistry::update);
 	}
 
 	private void decorateSecondaryIdDescriptionStore(IdentifiableStore<SecondaryIdDescription> store) {

backend/src/main/java/com/bakdata/conquery/io/storage/xodus/stores/KeyIncludingStore.java

@@ -31,11 +31,8 @@ public void forEach(Consumer<VALUE> consumer) {
 	}
 
 	public void update(VALUE value) {
-		VALUE old = get(extractKey(value));
-		if(old != null)
-			removed(old);
+		updated(value);
 		store.update(extractKey(value), value);
-		added(value);
 	}
 
 	public void remove(KEY key) {
@@ -69,6 +66,8 @@ public String toString() {
 
 	protected abstract void added(VALUE value);
 
+	protected abstract void updated(VALUE value);
+
 	public void clear() {
 		store.clear();
 	}

backend/src/main/java/com/bakdata/conquery/io/storage/xodus/stores/SingletonStore.java

@@ -19,6 +19,9 @@ public class SingletonStore<VALUE> extends KeyIncludingStore<Boolean, VALUE> {
 
 	@NonNull
 	protected ThrowingConsumer<VALUE> onRemove = (v) -> {};
+
+	@NonNull
+	protected ThrowingConsumer<VALUE> onUpdate = (v) -> {};
 
 	public SingletonStore(Store<Boolean, VALUE> store, Injectable... injectables) {
 		super(store);
@@ -49,7 +52,8 @@ public void remove(Boolean key) {
 	public void remove() {
 		super.remove(Boolean.TRUE);
 	}
-
+
+	@Override
 	protected void removed(VALUE value) {
 		try {
 			if(value != null) {
@@ -60,6 +64,7 @@ protected void removed(VALUE value) {
 		}
 	}
 
+	@Override
 	protected void added(VALUE value) {
 		try {
 			if(value != null) {
@@ -69,4 +74,15 @@ protected void added(VALUE value) {
 			throw new RuntimeException("Failed to add "+value, e);
 		}
 	}
+
+	@Override
+	protected void updated(VALUE value) {
+		try {
+			if(value != null) {
+				onUpdate.accept(value);
+			}
+		} catch(Exception e) {
+			throw new RuntimeException("Failed to update "+value, e);
+		}
+	}
 }

backend/src/main/java/com/bakdata/conquery/models/config/XodusStoreFactory.java

@@ -263,6 +263,7 @@ public IdentifiableStore<Dictionary> createDictionaryStore(CentralRegistry centr
 
 		final Store<IId<Dictionary>, Dictionary> result;
 
+		// TODO this looks like dictionaries are double cached
 		if (useWeakDictionaryCaching) {
 			result = new WeakCachedStore<>(bigStore, getWeakCacheDuration());
 		}

backend/src/main/java/com/bakdata/conquery/models/identifiable/CentralRegistry.java

@@ -29,21 +29,8 @@ public synchronized CentralRegistry register(Identifiable<?> ident) {
 		return this;
 	}
 
-	public synchronized void registerCacheable(IId id, Function<IId, Identifiable> supplier) {
-		cacheables.put(id, supplier);
-	}
-
-	protected  <T extends Identifiable<?>> T get(IId<T> name) {
-		Object res = map.get(name);
-		if (res != null) {
-			return (T) res;
-		}
-
-		Function<IId, Identifiable> supplier = cacheables.get(name);
-		if (supplier == null) {
-			return null;
-		}
-		return (T) supplier.apply(name);
+	public synchronized Function<IId, Identifiable> registerCacheable(IId id, Function<IId, Identifiable> supplier) {
+		return cacheables.put(id, supplier);
 	}
 
 	public <T extends Identifiable<?>> T resolve(IId<T> name) {
@@ -56,16 +43,27 @@ public <T extends Identifiable<?>> T resolve(IId<T> name) {
 		return result;
 	}
 
+	public Identifiable update(Identifiable<?> ident){
+		return map.update(ident);
+	}
+
+	public synchronized Function<IId, Identifiable> updateCacheable(IId id, Function<IId, Identifiable> supplier) {
+		Function<IId, Identifiable> old = cacheables.put(id, supplier);
+		if (old == null) {
+			// The supplier might have been invoked already and the object gone into the IdMap
+			// So we invalidate it
+			map.remove(id);
+		}
+		return supplier;
+	}
+
 	public <T extends Identifiable<?>> Optional<T> getOptional(IId<T> name) {
 		return Optional.ofNullable(get(name));
 	}
 
-	public void remove(Identifiable<?> ident) {
+	public synchronized void remove(Identifiable<?> ident) {
 		IId<?> id = ident.getId();
-		synchronized (this) {
-			map.remove(id);
-			cacheables.remove(id);
-		}
+		map.remove(id);
 	}
 
 	@Override
@@ -89,4 +87,32 @@ public void clear() {
 		map.clear();
 		cacheables.clear();
 	}
+
+	/**
+	 * Needs to be protected in order to be overwritten by {@link InjectingCentralRegistry}
+	 */
+	protected  <T extends Identifiable<?>> T get(IId<T> name) {
+		Object res = map.get(name);
+		if (res != null) {
+			return (T) res;
+		}
+		synchronized (this) {
+			// Retry synchronized to make sure it has not been resolved from cacheables in the mean time
+			Object res2 = map.get(name);
+			if (res2 != null) {
+				return (T) res2;
+			}
+			Function<IId, Identifiable> supplier = cacheables.get(name);
+			if (supplier == null) {
+				return null;
+			}
+
+			// Transfer object to the IdMap
+			final T apply = (T) supplier.apply(name);
+			register(apply);
+			cacheables.remove(name);
+		}
+
+		return (T) map.get(name);
+	}
 }

backend/src/main/java/com/bakdata/conquery/models/identifiable/IdMap.java

@@ -79,8 +79,8 @@ public boolean add(V entry) {
 		return true;
 	}
 
-	public void update(V entry) {
-		map.put(entry.getId(), entry);
+	public V update(V entry) {
+		return map.put((ID)entry.getId(), entry);
 	}
 
 	public V remove(ID id) {