599 lines
18 KiB
Java
599 lines
18 KiB
Java
package it.cavallium.dbengine.database.memory;
|
|
|
|
import io.netty5.buffer.Buffer;
|
|
import io.netty5.buffer.BufferAllocator;
|
|
import io.netty5.util.Resource;
|
|
import io.netty5.util.Send;
|
|
import it.cavallium.dbengine.client.BadBlock;
|
|
import it.cavallium.dbengine.database.LLDelta;
|
|
import it.cavallium.dbengine.database.LLDictionary;
|
|
import it.cavallium.dbengine.database.LLDictionaryResultType;
|
|
import it.cavallium.dbengine.database.LLEntry;
|
|
import it.cavallium.dbengine.database.LLRange;
|
|
import it.cavallium.dbengine.database.LLSnapshot;
|
|
import it.cavallium.dbengine.database.LLUtils;
|
|
import it.cavallium.dbengine.database.OptionalBuf;
|
|
import it.cavallium.dbengine.database.UpdateMode;
|
|
import it.cavallium.dbengine.database.disk.BinarySerializationFunction;
|
|
import it.cavallium.dbengine.database.serialization.KVSerializationFunction;
|
|
import it.cavallium.dbengine.database.serialization.SerializationException;
|
|
import it.cavallium.dbengine.database.serialization.SerializationFunction;
|
|
import it.unimi.dsi.fastutil.bytes.ByteList;
|
|
import java.io.IOException;
|
|
import java.util.List;
|
|
import java.util.Map;
|
|
import java.util.Objects;
|
|
import java.util.Optional;
|
|
import java.util.SortedMap;
|
|
import java.util.concurrent.ConcurrentHashMap;
|
|
import java.util.concurrent.ConcurrentNavigableMap;
|
|
import java.util.concurrent.ConcurrentSkipListMap;
|
|
import java.util.concurrent.atomic.AtomicReference;
|
|
import java.util.function.Supplier;
|
|
import org.jetbrains.annotations.Nullable;
|
|
import reactor.core.publisher.Flux;
|
|
import reactor.core.publisher.Mono;
|
|
|
|
public class LLMemoryDictionary implements LLDictionary {
|
|
|
|
private final String databaseName;
|
|
private final String columnName;
|
|
private final BufferAllocator allocator;
|
|
private final UpdateMode updateMode;
|
|
private final Getter<Long, ConcurrentSkipListMap<ByteList, ByteList>> snapshots;
|
|
private final ConcurrentSkipListMap<ByteList, ByteList> mainDb;
|
|
|
|
private interface Getter<T, U> {
|
|
U get(T argument);
|
|
}
|
|
|
|
public LLMemoryDictionary(BufferAllocator allocator,
|
|
String databaseName,
|
|
String columnName,
|
|
UpdateMode updateMode,
|
|
ConcurrentHashMap<Long, ConcurrentHashMap<String, ConcurrentSkipListMap<ByteList, ByteList>>> snapshots,
|
|
ConcurrentHashMap<String, ConcurrentSkipListMap<ByteList, ByteList>> mainDb) {
|
|
this.databaseName = databaseName;
|
|
this.columnName = columnName;
|
|
this.allocator = allocator;
|
|
this.updateMode = updateMode;
|
|
this.snapshots = (snapshotId) -> snapshots.get(snapshotId).get(columnName);
|
|
this.mainDb = mainDb.get(columnName);
|
|
}
|
|
|
|
@Override
|
|
public String getColumnName() {
|
|
return columnName;
|
|
}
|
|
|
|
@Override
|
|
public BufferAllocator getAllocator() {
|
|
return allocator;
|
|
}
|
|
|
|
private long resolveSnapshot(LLSnapshot snapshot) {
|
|
if (snapshot == null) {
|
|
return Long.MIN_VALUE + 1L;
|
|
} else if (snapshot.getSequenceNumber() == Long.MIN_VALUE + 1L) {
|
|
throw new IllegalStateException();
|
|
} else {
|
|
return snapshot.getSequenceNumber();
|
|
}
|
|
}
|
|
|
|
private Mono<Buffer> transformResult(Mono<ByteList> result, LLDictionaryResultType resultType) {
|
|
if (resultType == LLDictionaryResultType.PREVIOUS_VALUE) {
|
|
// Don't retain the result because it has been removed from the skip list
|
|
return result.map(this::kkB);
|
|
} else if (resultType == LLDictionaryResultType.PREVIOUS_VALUE_EXISTENCE) {
|
|
return result
|
|
.map(prev -> true)
|
|
.defaultIfEmpty(false)
|
|
.map((Boolean bool) -> LLUtils.booleanToResponseByteBuffer(allocator, bool));
|
|
} else {
|
|
return result.then(Mono.empty());
|
|
}
|
|
}
|
|
|
|
private ByteList k(Send<Buffer> buf) {
|
|
try (var b = buf.receive()) {
|
|
return new BinaryLexicographicList(LLUtils.toArray(b));
|
|
}
|
|
}
|
|
|
|
private ByteList kShr(Buffer buf) {
|
|
return new BinaryLexicographicList(LLUtils.toArray(buf));
|
|
}
|
|
|
|
private ByteList kOwn(Buffer buf) {
|
|
try (buf) {
|
|
return new BinaryLexicographicList(LLUtils.toArray(buf));
|
|
}
|
|
}
|
|
|
|
private Send<Buffer> kk(ByteList bytesList) {
|
|
try (var buffer = getAllocator().allocate(bytesList.size())) {
|
|
buffer.writeBytes(bytesList.toByteArray());
|
|
return buffer.send();
|
|
}
|
|
}
|
|
|
|
private Buffer kkB(ByteList bytesList) {
|
|
var buffer = getAllocator().allocate(bytesList.size());
|
|
try {
|
|
buffer.writeBytes(bytesList.toByteArray());
|
|
return buffer;
|
|
} catch (Throwable t) {
|
|
buffer.close();
|
|
throw t;
|
|
}
|
|
}
|
|
|
|
private BLRange r(Supplier<LLRange> send) {
|
|
try(var range = send.get()) {
|
|
if (range.isAll()) {
|
|
return new BLRange(null, null, null);
|
|
} else if (range.isSingle()) {
|
|
return new BLRange(null, null, k(range.getSingle()));
|
|
} else if (range.hasMin() && range.hasMax()) {
|
|
return new BLRange(k(range.getMin()), k(range.getMax()), null);
|
|
} else if (range.hasMin()) {
|
|
return new BLRange(k(range.getMin()), null, null);
|
|
} else {
|
|
return new BLRange(k(range.getMax()), null, null);
|
|
}
|
|
}
|
|
}
|
|
|
|
private ConcurrentNavigableMap<ByteList, ByteList> mapSlice(LLSnapshot snapshot, LLRange range) {
|
|
if (range.isAll()) {
|
|
return snapshots.get(resolveSnapshot(snapshot));
|
|
} else if (range.isSingle()) {
|
|
var key = k(range.getSingle());
|
|
var value = snapshots
|
|
.get(resolveSnapshot(snapshot))
|
|
.get(key);
|
|
if (value != null) {
|
|
return new ConcurrentSkipListMap<>(Map.of(key, value));
|
|
} else {
|
|
return new ConcurrentSkipListMap<>(Map.of());
|
|
}
|
|
} else if (range.hasMin() && range.hasMax()) {
|
|
var min = k(range.getMin());
|
|
var max = k(range.getMax());
|
|
if (min.compareTo(max) > 0) {
|
|
return new ConcurrentSkipListMap<>(Map.of());
|
|
}
|
|
return snapshots
|
|
.get(resolveSnapshot(snapshot))
|
|
.subMap(min, true, max, false);
|
|
} else if (range.hasMin()) {
|
|
return snapshots
|
|
.get(resolveSnapshot(snapshot))
|
|
.tailMap(k(range.getMin()), true);
|
|
} else {
|
|
return snapshots
|
|
.get(resolveSnapshot(snapshot))
|
|
.headMap(k(range.getMax()), false);
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public Mono<Buffer> get(@Nullable LLSnapshot snapshot, Mono<Buffer> keyMono) {
|
|
return Mono.usingWhen(keyMono,
|
|
key -> Mono
|
|
.fromCallable(() -> snapshots.get(resolveSnapshot(snapshot)).get(kShr(key)))
|
|
.map(this::kkB)
|
|
.onErrorMap(cause -> new IOException("Failed to read", cause)),
|
|
LLUtils::finalizeResource
|
|
);
|
|
}
|
|
|
|
@Override
|
|
public Mono<Buffer> put(Mono<Buffer> keyMono, Mono<Buffer> valueMono, LLDictionaryResultType resultType) {
|
|
var kMono = keyMono.map(this::kOwn);
|
|
var vMono = valueMono.map(this::kOwn);
|
|
return Mono
|
|
.zip(kMono, vMono)
|
|
.mapNotNull(tuple -> mainDb.put(tuple.getT1(), tuple.getT2()))
|
|
.transform(result -> this.transformResult(result, resultType))
|
|
.onErrorMap(cause -> new IOException("Failed to read", cause));
|
|
}
|
|
|
|
@Override
|
|
public UpdateMode getUpdateMode() {
|
|
return updateMode;
|
|
}
|
|
|
|
@Override
|
|
public Mono<LLDelta> updateAndGetDelta(Mono<Buffer> keyMono, BinarySerializationFunction updater) {
|
|
return Mono.usingWhen(keyMono,
|
|
key -> Mono.fromCallable(() -> {
|
|
if (updateMode == UpdateMode.DISALLOW) {
|
|
throw new UnsupportedOperationException("update() is disallowed");
|
|
}
|
|
AtomicReference<ByteList> oldRef = new AtomicReference<>(null);
|
|
var newValue = mainDb.compute(kShr(key), (_unused, old) -> {
|
|
if (old != null) {
|
|
oldRef.set(old);
|
|
}
|
|
Buffer v;
|
|
var oldToSend = old != null ? kkB(old) : null;
|
|
try {
|
|
assert oldToSend == null || oldToSend.isAccessible();
|
|
v = updater.apply(oldToSend);
|
|
assert v == null || v.isAccessible();
|
|
} catch (SerializationException e) {
|
|
throw new IllegalStateException(e);
|
|
} finally {
|
|
if (oldToSend != null && oldToSend.isAccessible()) {
|
|
oldToSend.close();
|
|
}
|
|
}
|
|
if (v != null) {
|
|
return kOwn(v);
|
|
} else {
|
|
return null;
|
|
}
|
|
});
|
|
var oldVal = oldRef.get();
|
|
return LLDelta.of(oldVal != null ? kkB(oldRef.get()) : null, newValue != null ? kkB(newValue) : null);
|
|
}),
|
|
LLUtils::finalizeResource
|
|
);
|
|
}
|
|
|
|
@Override
|
|
public Mono<Void> clear() {
|
|
return Mono.fromRunnable(mainDb::clear);
|
|
}
|
|
|
|
@Override
|
|
public Mono<Buffer> remove(Mono<Buffer> keyMono, LLDictionaryResultType resultType) {
|
|
return Mono.usingWhen(keyMono,
|
|
key -> Mono
|
|
.fromCallable(() -> mainDb.remove(kShr(key)))
|
|
// Don't retain the result because it has been removed from the skip list
|
|
.mapNotNull(bytesList -> switch (resultType) {
|
|
case VOID -> null;
|
|
case PREVIOUS_VALUE_EXISTENCE -> LLUtils.booleanToResponseByteBuffer(allocator, true);
|
|
case PREVIOUS_VALUE -> kkB(bytesList);
|
|
})
|
|
.switchIfEmpty(Mono.defer(() -> {
|
|
if (resultType == LLDictionaryResultType.PREVIOUS_VALUE_EXISTENCE) {
|
|
return Mono.fromCallable(() -> LLUtils.booleanToResponseByteBuffer(allocator, false));
|
|
} else {
|
|
return Mono.empty();
|
|
}
|
|
}))
|
|
.onErrorMap(cause -> new IOException("Failed to read", cause)),
|
|
LLUtils::finalizeResource
|
|
);
|
|
}
|
|
|
|
@Override
|
|
public Flux<OptionalBuf> getMulti(@Nullable LLSnapshot snapshot, Flux<Buffer> keys) {
|
|
return keys.map(key -> {
|
|
try (key) {
|
|
ByteList v = snapshots.get(resolveSnapshot(snapshot)).get(k(key.copy().send()));
|
|
if (v != null) {
|
|
return OptionalBuf.of(kkB(v));
|
|
} else {
|
|
return OptionalBuf.empty();
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
@Override
|
|
public Mono<Void> putMulti(Flux<LLEntry> entries) {
|
|
return entries.doOnNext(entry -> {
|
|
try (entry) {
|
|
mainDb.put(k(entry.getKeyUnsafe().send()), k(entry.getValueUnsafe().send()));
|
|
}
|
|
}).then();
|
|
}
|
|
|
|
@Override
|
|
public <K> Flux<Boolean> updateMulti(Flux<K> keys,
|
|
Flux<Buffer> serializedKeys,
|
|
KVSerializationFunction<K, @Nullable Buffer, @Nullable Buffer> updateFunction) {
|
|
return Flux.error(new UnsupportedOperationException("Not implemented"));
|
|
}
|
|
|
|
@Override
|
|
public Flux<LLEntry> getRange(@Nullable LLSnapshot snapshot,
|
|
Mono<LLRange> rangeMono,
|
|
boolean reverse,
|
|
boolean smallRange) {
|
|
return Flux.usingWhen(rangeMono, range -> {
|
|
if (range.isSingle()) {
|
|
var singleToReceive = range.getSingle();
|
|
return Mono.fromCallable(() -> {
|
|
try (var single = singleToReceive.receive()) {
|
|
var element = snapshots.get(resolveSnapshot(snapshot)).get(k(single.copy().send()));
|
|
if (element != null) {
|
|
return LLEntry.of(single, kkB(element));
|
|
} else {
|
|
return null;
|
|
}
|
|
}
|
|
}).flux();
|
|
} else {
|
|
return Mono
|
|
.fromCallable(() -> mapSlice(snapshot, range))
|
|
.flatMapIterable(map -> {
|
|
if (reverse) {
|
|
return map.descendingMap().entrySet();
|
|
} else {
|
|
return map.entrySet();
|
|
}
|
|
})
|
|
.map(entry -> LLEntry.of(kkB(entry.getKey()), kkB(entry.getValue())));
|
|
}
|
|
}, LLUtils::finalizeResource);
|
|
}
|
|
|
|
@Override
|
|
public Flux<List<LLEntry>> getRangeGrouped(@Nullable LLSnapshot snapshot,
|
|
Mono<LLRange> rangeMono,
|
|
int prefixLength, boolean smallRange) {
|
|
return Flux.usingWhen(rangeMono, range -> {
|
|
try (range) {
|
|
if (range.isSingle()) {
|
|
var singleToReceive = range.getSingle();
|
|
return Mono.fromCallable(() -> {
|
|
try (var single = singleToReceive.receive()) {
|
|
var element = snapshots.get(resolveSnapshot(snapshot)).get(k(single.copy().send()));
|
|
if (element != null) {
|
|
return List.of(LLEntry.of(single, kkB(element)));
|
|
} else {
|
|
return List.<LLEntry>of();
|
|
}
|
|
}
|
|
}).flux();
|
|
} else {
|
|
return Mono
|
|
.fromCallable(() -> mapSlice(snapshot, range))
|
|
.flatMapIterable(SortedMap::entrySet)
|
|
.groupBy(k -> k.getKey().subList(0, prefixLength))
|
|
.flatMap(groupedFlux -> groupedFlux
|
|
.map(entry -> LLEntry.of(kkB(entry.getKey()), kkB(entry.getValue())))
|
|
.collectList()
|
|
);
|
|
}
|
|
}
|
|
}, LLUtils::finalizeResource);
|
|
}
|
|
|
|
@Override
|
|
public Flux<Buffer> getRangeKeys(@Nullable LLSnapshot snapshot,
|
|
Mono<LLRange> rangeMono,
|
|
boolean reverse,
|
|
boolean smallRange) {
|
|
return Flux.usingWhen(rangeMono,
|
|
range -> {
|
|
if (range.isSingle()) {
|
|
var singleToReceive = range.getSingle();
|
|
return Mono.fromCallable(() -> {
|
|
var single = singleToReceive.receive();
|
|
try {
|
|
var contains = snapshots.get(resolveSnapshot(snapshot)).containsKey(k(single.copy().send()));
|
|
return contains ? single : null;
|
|
} catch (Throwable ex) {
|
|
single.close();
|
|
throw ex;
|
|
}
|
|
}).flux();
|
|
} else {
|
|
return Mono
|
|
.fromCallable(() -> mapSlice(snapshot, range))
|
|
.<ByteList>flatMapIterable(map -> {
|
|
if (reverse) {
|
|
return map.descendingMap().keySet();
|
|
} else {
|
|
return map.keySet();
|
|
}
|
|
})
|
|
.map(this::kkB);
|
|
}
|
|
},
|
|
LLUtils::finalizeResource
|
|
);
|
|
}
|
|
|
|
@Override
|
|
public Flux<List<Buffer>> getRangeKeysGrouped(@Nullable LLSnapshot snapshot,
|
|
Mono<LLRange> rangeMono,
|
|
int prefixLength, boolean smallRange) {
|
|
return Flux.usingWhen(rangeMono, range -> {
|
|
try (range) {
|
|
if (range.isSingle()) {
|
|
var singleToReceive = range.getSingle();
|
|
return Mono.fromCallable(() -> {
|
|
var single = singleToReceive.receive();
|
|
try {
|
|
var containsElement = snapshots.get(resolveSnapshot(snapshot)).containsKey(k(single.copy().send()));
|
|
if (containsElement) {
|
|
return List.of(single);
|
|
} else {
|
|
return List.<Buffer>of();
|
|
}
|
|
} catch (Throwable ex) {
|
|
single.close();
|
|
throw ex;
|
|
}
|
|
}).flux();
|
|
} else {
|
|
return Mono
|
|
.fromCallable(() -> mapSlice(snapshot, range))
|
|
.flatMapIterable(SortedMap::entrySet)
|
|
.groupBy(k -> k.getKey().subList(0, prefixLength))
|
|
.flatMap(groupedFlux -> groupedFlux
|
|
.map(entry -> kkB(entry.getKey()))
|
|
.collectList()
|
|
);
|
|
}
|
|
}
|
|
}, LLUtils::finalizeResource);
|
|
}
|
|
|
|
@SuppressWarnings("RedundantCast")
|
|
@Override
|
|
public Flux<Buffer> getRangeKeyPrefixes(@Nullable LLSnapshot snapshot,
|
|
Mono<LLRange> rangeMono,
|
|
int prefixLength, boolean smallRange) {
|
|
return Flux.usingWhen(rangeMono, range -> {
|
|
try (range) {
|
|
if (range.isSingle()) {
|
|
var singleToReceive = range.getSingle();
|
|
return Mono.fromCallable(() -> {
|
|
try (var single = singleToReceive.receive()) {
|
|
var k = k(single.copy().send());
|
|
var containsElement = snapshots.get(resolveSnapshot(snapshot)).containsKey(k);
|
|
if (containsElement) {
|
|
return kkB(k.subList(0, prefixLength));
|
|
} else {
|
|
return null;
|
|
}
|
|
}
|
|
}).flux();
|
|
} else {
|
|
return Mono
|
|
.fromCallable(() -> mapSlice(snapshot, range))
|
|
.flatMapIterable(SortedMap::entrySet)
|
|
.map(k -> (ByteList) k.getKey().subList(0, prefixLength))
|
|
.distinctUntilChanged()
|
|
.map(this::kkB);
|
|
}
|
|
}
|
|
}, LLUtils::finalizeResource);
|
|
}
|
|
|
|
@Override
|
|
public Flux<BadBlock> badBlocks(Mono<LLRange> rangeMono) {
|
|
return Flux.empty();
|
|
}
|
|
|
|
@Override
|
|
public Mono<Void> setRange(Mono<LLRange> rangeMono, Flux<LLEntry> entries, boolean smallRange) {
|
|
return Mono.usingWhen(rangeMono, range -> {
|
|
Mono<Void> clearMono;
|
|
if (range.isSingle()) {
|
|
var singleToReceive = range.getSingle();
|
|
clearMono = Mono.fromRunnable(() -> {
|
|
try (var single = singleToReceive.receive()) {
|
|
var k = k(single.copy().send());
|
|
mainDb.remove(k);
|
|
}
|
|
});
|
|
} else {
|
|
clearMono = Mono.fromRunnable(() -> mapSlice(null, range).clear());
|
|
}
|
|
|
|
var r = r(range::copy);
|
|
|
|
return clearMono
|
|
.thenMany(entries)
|
|
.doOnNext(entry -> {
|
|
try (entry) {
|
|
if (!isInsideRange(r, kShr(entry.getKeyUnsafe()))) {
|
|
throw new IndexOutOfBoundsException("Trying to set a key outside the range!");
|
|
}
|
|
mainDb.put(kShr(entry.getKeyUnsafe()), kShr(entry.getValueUnsafe()));
|
|
}
|
|
})
|
|
.then();
|
|
}, LLUtils::finalizeResource);
|
|
}
|
|
|
|
private boolean isInsideRange(BLRange range, ByteList key) {
|
|
if (range.isAll()) {
|
|
return true;
|
|
} else if (range.isSingle()) {
|
|
var single = range.getSingle();
|
|
return Objects.equals(single, key);
|
|
} else if (range.hasMin() && range.hasMax()) {
|
|
var min = range.getMin();
|
|
var max = range.getMax();
|
|
return min.compareTo(key) <= 0 && max.compareTo(key) > 0;
|
|
} else if (range.hasMin()) {
|
|
var min = range.getMin();
|
|
return min.compareTo(key) <= 0;
|
|
} else {
|
|
var max = range.getMax();
|
|
return max.compareTo(key) > 0;
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public Mono<Boolean> isRangeEmpty(@Nullable LLSnapshot snapshot, Mono<LLRange> rangeMono, boolean fillCache) {
|
|
return getRangeKeys(snapshot, rangeMono, false, false)
|
|
.doOnNext(LLUtils::finalizeResourceNow)
|
|
.count()
|
|
.map(count -> count == 0);
|
|
}
|
|
|
|
@Override
|
|
public Mono<Long> sizeRange(@Nullable LLSnapshot snapshot, Mono<LLRange> rangeMono, boolean fast) {
|
|
return Mono.usingWhen(rangeMono,
|
|
range -> Mono.fromCallable(() -> (long) mapSlice(snapshot, range).size()),
|
|
LLUtils::finalizeResource
|
|
);
|
|
}
|
|
|
|
@Override
|
|
public Mono<LLEntry> getOne(@Nullable LLSnapshot snapshot, Mono<LLRange> rangeMono) {
|
|
return getRange(snapshot, rangeMono, false, false)
|
|
.take(1, true)
|
|
.singleOrEmpty();
|
|
}
|
|
|
|
@Override
|
|
public Mono<Buffer> getOneKey(@Nullable LLSnapshot snapshot, Mono<LLRange> rangeMono) {
|
|
return getRangeKeys(snapshot, rangeMono, false, false)
|
|
.take(1, true)
|
|
.singleOrEmpty();
|
|
}
|
|
|
|
@Override
|
|
public Mono<LLEntry> removeOne(Mono<LLRange> rangeMono) {
|
|
return Mono.usingWhen(rangeMono, range -> {
|
|
try (range) {
|
|
if (range.isSingle()) {
|
|
var singleToReceive = range.getSingle();
|
|
return Mono.fromCallable(() -> {
|
|
try (var single = singleToReceive.receive()) {
|
|
var element = mainDb.remove(k(single.copy().send()));
|
|
if (element != null) {
|
|
return LLEntry.of(single, kkB(element));
|
|
} else {
|
|
return null;
|
|
}
|
|
}
|
|
});
|
|
} else {
|
|
return Mono
|
|
.fromCallable(() -> mapSlice(null, range))
|
|
.mapNotNull(map -> {
|
|
var it = map.entrySet().iterator();
|
|
if (it.hasNext()) {
|
|
var next = it.next();
|
|
it.remove();
|
|
return next;
|
|
} else {
|
|
return null;
|
|
}
|
|
})
|
|
.map(entry -> LLEntry.of(kkB(entry.getKey()), kkB(entry.getValue())));
|
|
}
|
|
}
|
|
}, LLUtils::finalizeResource);
|
|
}
|
|
|
|
@Override
|
|
public String getDatabaseName() {
|
|
return databaseName;
|
|
}
|
|
}
|