Hw5 Данилин Андрей ФИТиП

src/main/java/ru/vk/itmo/danilinandrew/DiskStorage.java

@@ -0,0 +1,250 @@
+package ru.vk.itmo.danilinandrew;
+
+import ru.vk.itmo.Entry;
+
+import java.io.IOException;
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.ValueLayout;
+import java.nio.channels.FileChannel;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardCopyOption;
+import java.nio.file.StandardOpenOption;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.Iterator;
+import java.util.List;
+
+public class DiskStorage {
+
+    public static final String SSTABLE_PREFIX = "sstable_";
+    private final List<MemorySegment> segmentList;
+
+    public DiskStorage(List<MemorySegment> segmentList) {
+        this.segmentList = segmentList;
+    }
+
+    public Iterator<Entry<MemorySegment>> range(
+            List<Iterator<Entry<MemorySegment>>> iters,
+            MemorySegment from,
+            MemorySegment to) {
+        List<Iterator<Entry<MemorySegment>>> iterators = new ArrayList<>(segmentList.size() + 1);
+        for (MemorySegment memorySegment : segmentList) {
+            iterators.add(DiskStorageExtension.iterator(memorySegment, from, to));
+        }
+        iterators.addAll(iters);
+
+        return new MergeIterator<>(iterators, Comparator.comparing(Entry::key, StorageDao::compare)) {
+            @Override
+            protected boolean shouldSkip(Entry<MemorySegment> memorySegmentEntry) {
+                return memorySegmentEntry.value() == null;
+            }
+        };
+    }
+
+    public void saveNextSSTable(
+            Path storagePath,
+            Iterable<Entry<MemorySegment>> iterable,
+            Arena arena
+    ) throws IOException {
+        final Path indexTmp = storagePath.resolve("index.tmp");
+        final Path indexFile = storagePath.resolve("index.idx");
+
+        try {
+            Files.createFile(indexFile);
+        } catch (FileAlreadyExistsException ignored) {
+            // it is ok, actually it is normal state
+        }
+        List<String> existedFiles = Files.readAllLines(indexFile, StandardCharsets.UTF_8);
+
+        String newFileName = SSTABLE_PREFIX + existedFiles.size();
+
+        long dataSize = 0;
+        long count = 0;
+        for (Entry<MemorySegment> entry : iterable) {
+            dataSize += entry.key().byteSize();
+            MemorySegment value = entry.value();
+            if (value != null) {
+                dataSize += value.byteSize();
+            }
+            count++;
+        }
+        long indexSize = count * 2 * Long.BYTES;
+
+        try (
+                FileChannel fileChannel = FileChannel.open(
+                        storagePath.resolve(newFileName),
+                        StandardOpenOption.WRITE,
+                        StandardOpenOption.READ,
+                        StandardOpenOption.CREATE
+                );
+                Arena writeArena = Arena.ofConfined()
+        ) {
+            MemorySegment fileSegment = fileChannel.map(
+                    FileChannel.MapMode.READ_WRITE,
+                    0,
+                    indexSize + dataSize,
+                    writeArena
+            );
+
+            // index:
+            // |key0_Start|value0_Start|key1_Start|value1_Start|key2_Start|value2_Start|...
+            // key0_Start = data start = end of index
+            long dataOffset = indexSize;
+            int indexOffset = 0;
+            for (Entry<MemorySegment> entry : iterable) {
+                fileSegment.set(ValueLayout.JAVA_LONG_UNALIGNED, indexOffset, dataOffset);
+                dataOffset += entry.key().byteSize();
+                indexOffset += Long.BYTES;
+
+                MemorySegment value = entry.value();
+                if (value == null) {
+                    fileSegment.set(
+                            ValueLayout.JAVA_LONG_UNALIGNED,
+                            indexOffset,
+                            DiskStorageExtension.tombstone(dataOffset)
+                    );
+                } else {
+                    fileSegment.set(ValueLayout.JAVA_LONG_UNALIGNED, indexOffset, dataOffset);
+                    dataOffset += value.byteSize();
+                }
+                indexOffset += Long.BYTES;
+            }
+
+            // data:
+            // |key0|value0|key1|value1|...
+            dataOffset = indexSize;
+            for (Entry<MemorySegment> entry : iterable) {
+                MemorySegment key = entry.key();
+                MemorySegment.copy(key, 0, fileSegment, dataOffset, key.byteSize());
+                dataOffset += key.byteSize();
+
+                MemorySegment value = entry.value();
+                if (value != null) {
+                    MemorySegment.copy(value, 0, fileSegment, dataOffset, value.byteSize());
+                    dataOffset += value.byteSize();
+                }
+            }
+        }
+
+        List<String> list = new ArrayList<>(existedFiles.size() + 1);
+        list.addAll(existedFiles);
+        list.add(newFileName);
+        Files.write(
+                indexTmp,
+                list,
+                StandardOpenOption.WRITE,
+                StandardOpenOption.CREATE,
+                StandardOpenOption.TRUNCATE_EXISTING
+        );
+
+        Files.deleteIfExists(indexFile);
+
+        Files.move(indexTmp, indexFile, StandardCopyOption.ATOMIC_MOVE);
+
+        if (arena.scope().isAlive()) {
+            openNewSSTable(storagePath.resolve(newFileName), arena);
+        }
+    }
+
+    public void openNewSSTable(Path file, Arena arena) {
+        try (FileChannel fileChannel = FileChannel.open(file, StandardOpenOption.READ, StandardOpenOption.WRITE)) {
+            MemorySegment fileSegment = fileChannel.map(
+                    FileChannel.MapMode.READ_WRITE,
+                    0,
+                    Files.size(file),
+                    arena
+            );
+            segmentList.add(fileSegment);
+        } catch (IOException e) {
+            throw new IllegalStateException("Error open after flush", e);
+        }
+    }
+
+    public static void compact(Path storagePath, Iterable<Entry<MemorySegment>> iterable)
+            throws IOException {
+
+        String newFileName = "compaction.tmp";
+        Path compactionTmpFile = storagePath.resolve(newFileName);
+
+        long dataSize = 0;
+        long count = 0;
+        for (Entry<MemorySegment> entry : iterable) {
+            dataSize += entry.key().byteSize();
+            MemorySegment value = entry.value();
+            if (value != null) {
+                dataSize += value.byteSize();
+            }
+            count++;
+        }
+        long indexSize = count * 2 * Long.BYTES;
+
+        try (
+                FileChannel fileChannel = FileChannel.open(
+                        compactionTmpFile,
+                        StandardOpenOption.WRITE,
+                        StandardOpenOption.READ,
+                        StandardOpenOption.CREATE
+                );
+                Arena writeArena = Arena.ofConfined()
+        ) {
+            MemorySegment fileSegment = fileChannel.map(
+                    FileChannel.MapMode.READ_WRITE,
+                    0,
+                    indexSize + dataSize,
+                    writeArena
+            );
+
+            // index:
+            // |key0_Start|value0_Start|key1_Start|value1_Start|key2_Start|value2_Start|...
+            // key0_Start = data start = end of index
+            long dataOffset = indexSize;
+            int indexOffset = 0;
+            for (Entry<MemorySegment> entry : iterable) {
+                fileSegment.set(ValueLayout.JAVA_LONG_UNALIGNED, indexOffset, dataOffset);
+                dataOffset += entry.key().byteSize();
+                indexOffset += Long.BYTES;
+
+                MemorySegment value = entry.value();
+                if (value == null) {
+                    fileSegment.set(
+                            ValueLayout.JAVA_LONG_UNALIGNED,
+                            indexOffset,
+                            DiskStorageExtension.tombstone(dataOffset)
+                    );
+                } else {
+                    fileSegment.set(ValueLayout.JAVA_LONG_UNALIGNED, indexOffset, dataOffset);
+                    dataOffset += value.byteSize();
+                }
+                indexOffset += Long.BYTES;
+            }
+
+            // data:
+            // |key0|value0|key1|value1|...
+            dataOffset = indexSize;
+            for (Entry<MemorySegment> entry : iterable) {
+                MemorySegment key = entry.key();
+                MemorySegment.copy(key, 0, fileSegment, dataOffset, key.byteSize());
+                dataOffset += key.byteSize();
+
+                MemorySegment value = entry.value();
+                if (value != null) {
+                    MemorySegment.copy(value, 0, fileSegment, dataOffset, value.byteSize());
+                    dataOffset += value.byteSize();
+                }
+            }
+        }
+
+        Files.move(
+                compactionTmpFile,
+                storagePath.resolve("compaction"),
+                StandardCopyOption.ATOMIC_MOVE,
+                StandardCopyOption.REPLACE_EXISTING
+        );
+
+        DiskStorageExtension.finalizeCompaction(storagePath);
+    }
+}