package it.cavallium.dbengine.database.collections; import io.netty.buffer.ByteBuf; import io.netty.util.ReferenceCounted; import it.cavallium.dbengine.client.CompositeSnapshot; import it.cavallium.dbengine.database.Delta; import it.cavallium.dbengine.database.LLDictionary; import it.cavallium.dbengine.database.LLDictionaryResultType; import it.cavallium.dbengine.database.LLUtils; import it.cavallium.dbengine.database.UpdateMode; import it.cavallium.dbengine.database.UpdateReturnMode; import it.cavallium.dbengine.database.serialization.Serializer; import it.cavallium.dbengine.database.serialization.SerializerFixedBinaryLength; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.Map.Entry; import java.util.Objects; import java.util.function.Function; import org.jetbrains.annotations.Nullable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; /** * Optimized implementation of "DatabaseMapDictionary with SubStageGetterSingle" */ public class DatabaseMapDictionary extends DatabaseMapDictionaryDeep> { private final Serializer valueSerializer; protected DatabaseMapDictionary(LLDictionary dictionary, ByteBuf prefixKey, SerializerFixedBinaryLength keySuffixSerializer, Serializer valueSerializer) { // Do not retain or release or use the prefixKey here super(dictionary, prefixKey, keySuffixSerializer, new SubStageGetterSingle<>(valueSerializer), 0); prefixKey = null; this.valueSerializer = valueSerializer; } public static DatabaseMapDictionary simple(LLDictionary dictionary, SerializerFixedBinaryLength keySerializer, Serializer valueSerializer) { return new DatabaseMapDictionary<>(dictionary, dictionary.getAllocator().buffer(0), keySerializer, valueSerializer); } public static DatabaseMapDictionary tail(LLDictionary dictionary, ByteBuf prefixKey, SerializerFixedBinaryLength keySuffixSerializer, Serializer valueSerializer) { return new DatabaseMapDictionary<>(dictionary, prefixKey, keySuffixSerializer, valueSerializer); } private ByteBuf toKey(ByteBuf suffixKey) { try { assert suffixKeyConsistency(suffixKey.readableBytes()); return LLUtils.compositeBuffer(dictionary.getAllocator(), keyPrefix.retain(), suffixKey.retain()); } finally { suffixKey.release(); } } @Override public Mono> get(@Nullable CompositeSnapshot snapshot, boolean existsAlmostCertainly) { return Flux .defer(() -> dictionary.getRange(resolveSnapshot(snapshot), range.retain(), existsAlmostCertainly)) .collectMap( entry -> deserializeSuffix(stripPrefix(entry.getKey(), false)), entry -> deserialize(entry.getValue()), HashMap::new) .filter(map -> !map.isEmpty()) .doFirst(range::retain) .doAfterTerminate(range::release); } @Override public Mono> setAndGetPrevious(Map value) { return Mono .usingWhen( Mono.just(true), b -> get(null, false), b -> dictionary .setRange(range.retain(), Flux .fromIterable(Collections.unmodifiableMap(value).entrySet()) .map(entry -> Map .entry(this.toKey(serializeSuffix(entry.getKey())), serialize(entry.getValue())) ) ) ) .doFirst(range::retain) .doAfterTerminate(range::release); } @Override public Mono> clearAndGetPrevious() { return this .setAndGetPrevious(Map.of()); } @Override public Mono leavesCount(@Nullable CompositeSnapshot snapshot, boolean fast) { return Mono .defer(() -> dictionary.sizeRange(resolveSnapshot(snapshot), range.retain(), fast)) .doFirst(range::retain) .doAfterTerminate(range::release); } @Override public Mono isEmpty(@Nullable CompositeSnapshot snapshot) { return Mono .defer(() -> dictionary.isRangeEmpty(resolveSnapshot(snapshot), range.retain())) .doFirst(range::retain) .doAfterTerminate(range::release); } @Override public Mono> at(@Nullable CompositeSnapshot snapshot, T keySuffix) { return Mono .fromSupplier(() -> new DatabaseSingleMapped<>( new DatabaseSingle<>(dictionary, toKey(serializeSuffix(keySuffix)), Serializer.noop()) , valueSerializer) ); } @Override public Mono getValue(@Nullable CompositeSnapshot snapshot, T keySuffix, boolean existsAlmostCertainly) { return Mono .using( () -> toKey(serializeSuffix(keySuffix)), keyBuf -> dictionary .get(resolveSnapshot(snapshot), keyBuf.retain(), existsAlmostCertainly) .map(this::deserialize), ReferenceCounted::release ); } @Override public Mono putValue(T keySuffix, U value) { return Mono .using( () -> serializeSuffix(keySuffix), keySuffixBuf -> Mono .using( () -> toKey(keySuffixBuf.retain()), keyBuf -> Mono .using( () -> serialize(value), valueBuf -> dictionary .put(keyBuf.retain(), valueBuf.retain(), LLDictionaryResultType.VOID) .doOnNext(ReferenceCounted::release), ReferenceCounted::release ), ReferenceCounted::release ), ReferenceCounted::release ) .then(); } @Override public Mono getUpdateMode() { return dictionary.getUpdateMode(); } @Override public Mono updateValue(T keySuffix, UpdateReturnMode updateReturnMode, boolean existsAlmostCertainly, Function<@Nullable U, @Nullable U> updater) { return Mono .using( () -> toKey(serializeSuffix(keySuffix)), keyBuf -> dictionary .update(keyBuf.retain(), oldSerialized -> { try { var result = updater.apply(oldSerialized == null ? null : this.deserialize(oldSerialized.retain())); if (result == null) { return null; } else { return this.serialize(result); } } finally { if (oldSerialized != null) { oldSerialized.release(); } } }, updateReturnMode, existsAlmostCertainly) .map(this::deserialize), ReferenceCounted::release ); } @Override public Mono> updateValueAndGetDelta(T keySuffix, boolean existsAlmostCertainly, Function<@Nullable U, @Nullable U> updater) { return Mono .using( () -> toKey(serializeSuffix(keySuffix)), keyBuf -> dictionary .updateAndGetDelta(keyBuf.retain(), oldSerialized -> { try { var result = updater.apply(oldSerialized == null ? null : this.deserialize(oldSerialized.retain())); if (result == null) { return null; } else { return this.serialize(result); } } finally { if (oldSerialized != null) { oldSerialized.release(); } } }, existsAlmostCertainly) .transform(mono -> LLUtils.mapDelta(mono, this::deserialize)), ReferenceCounted::release ); } @Override public Mono putValueAndGetPrevious(T keySuffix, U value) { return Mono .using( () -> serializeSuffix(keySuffix), keySuffixBuf -> Mono .using( () -> toKey(keySuffixBuf.retain()), keyBuf -> Mono .using( () -> serialize(value), valueBuf -> dictionary .put(keyBuf.retain(), valueBuf.retain(), LLDictionaryResultType.PREVIOUS_VALUE) .map(this::deserialize), ReferenceCounted::release ), ReferenceCounted::release ), ReferenceCounted::release ); } @Override public Mono putValueAndGetChanged(T keySuffix, U value) { return Mono .using( () -> serializeSuffix(keySuffix), keySuffixBuf -> Mono .using( () -> toKey(keySuffixBuf.retain()), keyBuf -> Mono .using( () -> serialize(value), valueBuf -> dictionary .put(keyBuf.retain(), valueBuf.retain(), LLDictionaryResultType.PREVIOUS_VALUE) .map(this::deserialize) .map(oldValue -> !Objects.equals(oldValue, value)) .defaultIfEmpty(value != null), ReferenceCounted::release ), ReferenceCounted::release ), ReferenceCounted::release ); } @Override public Mono remove(T keySuffix) { return Mono .using( () -> toKey(serializeSuffix(keySuffix)), keyBuf -> dictionary .remove(keyBuf.retain(), LLDictionaryResultType.VOID) .doOnNext(ReferenceCounted::release) .then(), ReferenceCounted::release ); } @Override public Mono removeAndGetPrevious(T keySuffix) { return Mono .using( () -> toKey(serializeSuffix(keySuffix)), keyBuf -> dictionary .remove(keyBuf.retain(), LLDictionaryResultType.PREVIOUS_VALUE) .map(this::deserialize), ReferenceCounted::release ); } @Override public Mono removeAndGetStatus(T keySuffix) { return Mono .using( () -> toKey(serializeSuffix(keySuffix)), keyBuf -> dictionary .remove(keyBuf.retain(), LLDictionaryResultType.PREVIOUS_VALUE_EXISTENCE) .map(LLUtils::responseToBoolean), ReferenceCounted::release ); } @Override public Flux> getMulti(@Nullable CompositeSnapshot snapshot, Flux keys, boolean existsAlmostCertainly) { return Flux .defer(() -> dictionary .getMulti(resolveSnapshot(snapshot), keys.flatMap(keySuffix -> Mono.fromCallable(() -> { ByteBuf keySuffixBuf = serializeSuffix(keySuffix); try { return toKey(keySuffixBuf.retain()); } finally { keySuffixBuf.release(); } })), existsAlmostCertainly) ) .flatMap(entry -> Mono .fromCallable(() -> Map.entry(deserializeSuffix(stripPrefix(entry.getKey(), false)), deserialize(entry.getValue()))) ); } private Entry serializeEntry(T key, U value) { ByteBuf serializedKey = toKey(serializeSuffix(key)); try { ByteBuf serializedValue = serialize(value); try { return Map.entry(serializedKey.retain(), serializedValue.retain()); } finally { serializedValue.release(); } } finally { serializedKey.release(); } } @Override public Mono putMulti(Flux> entries) { var serializedEntries = entries .flatMap(entry -> Mono .fromCallable(() -> serializeEntry(entry.getKey(), entry.getValue())) .doOnDiscard(Entry.class, uncastedEntry -> { //noinspection unchecked var castedEntry = (Entry) uncastedEntry; castedEntry.getKey().release(); castedEntry.getValue().release(); }) ); return dictionary .putMulti(serializedEntries, false) .then(); } @Override public Flux>> getAllStages(@Nullable CompositeSnapshot snapshot) { return Flux .defer(() -> dictionary.getRangeKeys(resolveSnapshot(snapshot), range.retain())) .>>map(key -> { ByteBuf keySuffixWithExt = stripPrefix(key.retain(), false); try { try { return Map.entry(deserializeSuffix(keySuffixWithExt.retainedSlice()), new DatabaseSingleMapped<>(new DatabaseSingle<>(dictionary, toKey(keySuffixWithExt.retainedSlice()), Serializer.noop() ), valueSerializer) ); } finally { keySuffixWithExt.release(); } } finally { key.release(); } }) .doFirst(range::retain) .doAfterTerminate(range::release); } @Override public Flux> getAllValues(@Nullable CompositeSnapshot snapshot) { return Flux .defer(() -> dictionary.getRange(resolveSnapshot(snapshot), range.retain())) .map(serializedEntry -> Map.entry( deserializeSuffix(stripPrefix(serializedEntry.getKey(), false)), valueSerializer.deserialize(serializedEntry.getValue()) )) .doOnDiscard(Entry.class, entry -> { //noinspection unchecked var castedEntry = (Entry) entry; castedEntry.getKey().release(); castedEntry.getValue().release(); }) .doFirst(range::retain) .doAfterTerminate(range::release); } @Override public Flux> setAllValuesAndGetPrevious(Flux> entries) { return Flux .usingWhen( Mono.just(true), b -> getAllValues(null), b -> dictionary .setRange(range.retain(), entries.map(entry -> Map.entry(toKey(serializeSuffix(entry.getKey())), serialize(entry.getValue())) ) ) ) .doFirst(range::retain) .doAfterTerminate(range::release); } @Override public Mono clear() { return Mono .defer(() -> { if (range.isAll()) { return dictionary.clear(); } else if (range.isSingle()) { return dictionary .remove(range.getSingle().retain(), LLDictionaryResultType.VOID) .doOnNext(ReferenceCounted::release) .then(); } else { return dictionary.setRange(range.retain(), Flux.empty()); } }) .doFirst(range::retain) .doAfterTerminate(range::release); } /** * This method is just a shorter version than valueSerializer::deserialize */ private U deserialize(ByteBuf bytes) { return valueSerializer.deserialize(bytes); } /** * This method is just a shorter version than valueSerializer::serialize */ private ByteBuf serialize(U bytes) { return valueSerializer.serialize(bytes); } }