Кислов Данил ИТМО ФИТиП HW5

src/main/java/ru/vk/itmo/kislovdanil/MemTable.java

@@ -0,0 +1,40 @@
+package ru.vk.itmo.kislovdanil;
+
+import ru.vk.itmo.Entry;
+
+import java.lang.foreign.MemorySegment;
+import java.util.Comparator;
+import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.atomic.AtomicLong;
+
+/* Basically, ConcurrentSkipList with counter of threads that putting value in it.
+    Necessary for preventing data loss while flushing.
+ */
+public class MemTable {
+    private final ConcurrentSkipListMap<MemorySegment, Entry<MemorySegment>> storage;
+    private final long threshold;
+    private final AtomicLong size = new AtomicLong(0);
+
+    private static long getEntrySize(Entry<MemorySegment> entry) {
+        long valueSize = entry.value() == null ? 0 : entry.value().byteSize();
+        return valueSize + entry.key().byteSize();
+    }
+
+    public MemTable(Comparator<MemorySegment> comparator, long threshold) {
+        this.storage = new ConcurrentSkipListMap<>(comparator);
+        this.threshold = threshold;
+    }
+
+    public boolean put(Entry<MemorySegment> entry) {
+        long entrySize = getEntrySize(entry);
+        if (size.addAndGet(entrySize) - entrySize > threshold) {
+            return false;
+        }
+        storage.put(entry.key(), entry);
+        return true;
+    }
+
+    public ConcurrentSkipListMap<MemorySegment, Entry<MemorySegment>> getStorage() {
+        return storage;
+    }
+}

src/main/java/ru/vk/itmo/kislovdanil/PersistentDao.java

@@ -4,51 +4,86 @@
 import ru.vk.itmo.Dao;
 import ru.vk.itmo.Entry;
 import ru.vk.itmo.kislovdanil.exceptions.DBException;
+import ru.vk.itmo.kislovdanil.exceptions.OverloadException;
 import ru.vk.itmo.kislovdanil.iterators.DatabaseIterator;
+import ru.vk.itmo.kislovdanil.iterators.MemTableIterator;
 import ru.vk.itmo.kislovdanil.iterators.MergeIterator;
+import ru.vk.itmo.kislovdanil.sstable.SSTable;
 
 import java.io.File;
 import java.io.IOException;
+import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.ValueLayout;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
-import java.util.concurrent.ConcurrentNavigableMap;
-import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 public class PersistentDao implements Dao<MemorySegment, Entry<MemorySegment>>, Iterable<Entry<MemorySegment>> {
 
     public static final MemorySegment DELETED_VALUE = null;
     private final Config config;
-    private final List<SSTable> tables = new ArrayList<>();
+    private volatile List<SSTable> tables = new ArrayList<>();
     private final Comparator<MemorySegment> comparator = new MemSegComparator();
-    private final ConcurrentNavigableMap<MemorySegment, Entry<MemorySegment>> storage =
-            new ConcurrentSkipListMap<>(comparator);
+    private volatile MemTable memTable;
+    // Temporary storage in case of main storage flushing (Read only)
+    private volatile MemTable additionalStorage;
+    // In case of additional table overload while main table is flushing
+    private final AtomicLong nextId = new AtomicLong();
+    private final ExecutorService commonExecutorService = Executors.newFixedThreadPool(2);
+    // To prevent parallel flushing
+    private volatile Future<?> compcatFuture;
+    // To make sure that flushing in close() will be started
+    private volatile Future<?> flushFuture;
+    // Have to take before any tables modification
+    private final Lock compactionLock = new ReentrantLock();
+    // Have to take read while upsert and write while flushing (to prevent data loss)
+    private final ReadWriteLock upsertLock = new ReentrantReadWriteLock();
+    private final Arena filesArena = Arena.ofShared();
 
-    private long lastTimestamp = System.currentTimeMillis();
+    private long getMaxTablesId(Iterable<SSTable> tableIterable) {
+        long curMaxId = -1;
+        for (SSTable table : tableIterable) {
+            curMaxId = Math.max(curMaxId, table.getTableId());
+        }
+        return curMaxId;
+    }
 
     public PersistentDao(Config config) throws IOException {
         this.config = config;
+        this.memTable = new MemTable(comparator, config.flushThresholdBytes());
         File basePathDirectory = new File(config.basePath().toString());
         String[] ssTablesIds = basePathDirectory.list();
         if (ssTablesIds == null) return;
         for (String tableID : ssTablesIds) {
-            // SSTable constructor with rewrite=false reads table data from disk if it exists
-            tables.add(new SSTable(config.basePath(), comparator, Long.parseLong(tableID),
-                    storage.values(), false));
+            // SSTable constructor without entries iterator reads table data from disk if it exists
+            tables.add(new SSTable(config.basePath(), comparator, Long.parseLong(tableID), filesArena));
         }
+        nextId.set(getMaxTablesId(tables) + 1);
         tables.sort(SSTable::compareTo);
     }
 
     @Override
     public Iterator<Entry<MemorySegment>> get(MemorySegment from, MemorySegment to) {
-        List<DatabaseIterator> iterators = new ArrayList<>(tables.size() + 1);
+        List<DatabaseIterator> iterators = new ArrayList<>(tables.size() + 2);
         for (SSTable table : tables) {
             iterators.add(table.getRange(from, to));
         }
-        iterators.add(new MemTableIterator(from, to));
+        iterators.add(new MemTableIterator(from, to, memTable, Long.MAX_VALUE));
+        if (additionalStorage != null) {
+            iterators.add(new MemTableIterator(from, to, additionalStorage, Long.MAX_VALUE - 1));
+        }
         return new MergeIterator(iterators, comparator);
     }
 
@@ -57,19 +92,23 @@ private static Entry<MemorySegment> wrapEntryIfDeleted(Entry<MemorySegment> entr
         return entry;
     }
 
-    private void updateId() {
-        lastTimestamp = Math.max(lastTimestamp + 1, System.currentTimeMillis());
+    private long getNextId() {
+        return nextId.getAndIncrement();
     }
 
     @Override
     public Entry<MemorySegment> get(MemorySegment key) {
-        Entry<MemorySegment> ans = storage.get(key);
+        Entry<MemorySegment> ans = memTable.getStorage().get(key);
         if (ans != null) return wrapEntryIfDeleted(ans);
+        if (additionalStorage != null) {
+            ans = additionalStorage.getStorage().get(key);
+            if (ans != null) return wrapEntryIfDeleted(ans);
+        }
         try {
-            for (SSTable table : tables) {
-                Entry<MemorySegment> data = table.find(key);
-                if (data != null) {
-                    return wrapEntryIfDeleted(data);
+            for (SSTable table : tables.reversed()) {
+                ans = table.find(key);
+                if (ans != null) {
+                    return wrapEntryIfDeleted(ans);
                 }
             }
         } catch (IOException e) {
@@ -80,39 +119,122 @@ public Entry<MemorySegment> get(MemorySegment key) {
 
     @Override
     public void upsert(Entry<MemorySegment> entry) {
-        storage.put(entry.key(), entry);
+        upsertLock.readLock().lock();
+        try {
+            if (memTable.put(entry)) {
+                return;
+            }
+        } finally {
+            upsertLock.readLock().unlock();
+        }
+        flush();
+        upsertLock.readLock().lock();
+        try {
+            if (!memTable.put(entry)) {
+                throw new OverloadException(entry);
+            }
+        } finally {
+            upsertLock.readLock().unlock();
+        }
     }
 
-    @Override
-    public void flush() throws IOException {
-        if (!storage.isEmpty()) {
-            updateId();
-            // SSTable constructor with rewrite=true writes MemTable data on disk deleting old data if it exists
+    private void makeFlush() throws IOException {
+        compactionLock.lock();
+        try {
+            if (additionalStorage == null) return;
+            // SSTable constructor with entries iterator writes MemTable data on disk deleting old data if it exists
             tables.add(new SSTable(config.basePath(), comparator,
-                    lastTimestamp, storage.values(), true));
-            storage.clear();
+                    getNextId(), additionalStorage.getStorage().values().iterator(), filesArena));
+            additionalStorage = null;
+        } finally {
+            compactionLock.unlock();
+        }
+    }
+
+    @Override
+    public void flush() {
+        upsertLock.writeLock().lock();
+        try {
+            if (additionalStorage != null || memTable.getStorage().isEmpty()) {
+                return;
+            }
+            additionalStorage = memTable;
+            memTable = new MemTable(comparator, config.flushThresholdBytes());
+            flushFuture = commonExecutorService.submit(
+                    () -> {
+                        try {
+                            makeFlush();
+                        } catch (IOException e) {
+                            throw new DBException(e);
+                        }
+                    });
+        } finally {
+            upsertLock.writeLock().unlock();
+        }
+    }
+
+    private void closeExecutorService(ExecutorService executorService) {
+        executorService.shutdown();
+        try {
+            if (!executorService.awaitTermination(30, TimeUnit.SECONDS)) {
+                executorService.shutdownNow();
+            }
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
         }
     }
 
     @Override
     public void close() throws IOException {
+        if (!filesArena.scope().isAlive()) {
+            return;
+        }
+        if (flushFuture != null) {
+            try {
+                flushFuture.get();
+            } catch (InterruptedException | ExecutionException e) {
+                Thread.currentThread().interrupt();
+            }
+        }
         flush();
+        closeExecutorService(commonExecutorService);
+        filesArena.close();
     }
 
-    @Override
-    public void compact() throws IOException {
-        if (!tables.isEmpty()) {
-            updateId();
-            SSTable compactedTable = new SSTable(config.basePath(), comparator, lastTimestamp,
-                    this, true);
-            storage.clear();
-            for (SSTable table : tables) {
+    private void makeCompaction() throws IOException {
+        compactionLock.lock();
+        try {
+            if (tables.size() <= 1) return;
+            long compactedTableId = getNextId();
+            SSTable compactedTable = new SSTable(config.basePath(), comparator, compactedTableId,
+                    new MergeIterator(tables, comparator), filesArena);
+            List<SSTable> oldTables = tables;
+            List<SSTable> newTables = new ArrayList<>();
+            newTables.add(compactedTable);
+            tables = newTables;
+            for (SSTable table : oldTables) {
                 table.deleteFromDisk();
             }
-            tables.add(compactedTable);
+        } finally {
+            compactionLock.unlock();
         }
     }
 
+    @Override
+    public void compact() {
+        if (compcatFuture != null && !compcatFuture.isDone()) {
+            compcatFuture.cancel(false);
+        }
+        compcatFuture = commonExecutorService.submit(
+                () -> {
+                    try {
+                        makeCompaction();
+                    } catch (IOException e) {
+                        throw new DBException(e);
+                    }
+                });
+    }
+
     @Override
     public Iterator<Entry<MemorySegment>> iterator() {
         return get(null, null);
@@ -131,35 +253,4 @@ public int compare(MemorySegment o1, MemorySegment o2) {
             return Byte.compare(o1.get(ValueLayout.JAVA_BYTE, mismatch), o2.get(ValueLayout.JAVA_BYTE, mismatch));
         }
     }
-
-    private class MemTableIterator implements DatabaseIterator {
-        Iterator<Entry<MemorySegment>> innerIter;
-
-        public MemTableIterator(MemorySegment from, MemorySegment to) {
-            if (from == null && to == null) {
-                innerIter = storage.values().iterator();
-            } else if (from != null && to == null) {
-                innerIter = storage.tailMap(from).values().iterator();
-            } else if (from == null) {
-                innerIter = storage.headMap(to).values().iterator();
-            } else {
-                innerIter = storage.subMap(from, to).values().iterator();
-            }
-        }
-
-        @Override
-        public long getPriority() {
-            return Long.MAX_VALUE;
-        }
-
-        @Override
-        public boolean hasNext() {
-            return innerIter.hasNext();
-        }
-
-        @Override
-        public Entry<MemorySegment> next() {
-            return innerIter.next();
-        }
-    }
 }