<commit_before>package peergos.fuse;

import jnr.ffi.Pointer;
import jnr.ffi.types.off_t;
import jnr.ffi.types.size_t;
import peergos.user.UserContext;
import peergos.user.fs.Chunk;
import peergos.user.fs.FileProperties;
import ru.serce.jnrfuse.ErrorCodes;
import ru.serce.jnrfuse.struct.*;

import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.*;

public class CachingPeergosFS extends PeergosFS {

    private static final int DEFAULT_SYNC_SLEEP = 1000*30;
    private static final int DEFAULT_CACHE_SIZE = 1;

    private final ConcurrentMap<String, CacheEntryHolder> entryMap;
    private final int chunkCacheSize, syncSleep;
//    private final Thread syncRunner;

    public CachingPeergosFS(UserContext userContext) {
        this(userContext, DEFAULT_CACHE_SIZE, DEFAULT_SYNC_SLEEP);
    }

    public CachingPeergosFS(UserContext userContext, int chunkCacheSize, int syncSleep) {
        super(userContext);

        this.chunkCacheSize = chunkCacheSize;
        this.syncSleep = syncSleep;
        this.entryMap = new ConcurrentHashMap<>();
        //TODO add GC thread to optimistically reduce size
//        this.syncRunner = new Thread(new Syncher());
//        this.syncRunner.start();
    }

    @Override
    public int read(String s, Pointer pointer, @size_t long size, @off_t long offset, FuseFileInfo fuseFileInfo) {
        if (!containedInOneChunk(offset, offset + size))
            throw new IllegalStateException("write op. straddles boundary : offset " + offset + " with size " + size);

        long startPos = alignToChunkSize(offset);
        int chunkOffset  = intraChunkOffset(offset);
        int iSize = (int) size;

        CacheEntryHolder cacheEntryHolder = entryMap.computeIfAbsent(s, path -> new CacheEntryHolder(new CacheEntry(path, startPos)));
        return cacheEntryHolder.apply(c -> c != null && c.offset == startPos, () -> new CacheEntry(s, startPos), ce -> ce.read(pointer, chunkOffset, iSize));
    }

    @Override
    public int write(String s, Pointer pointer, @size_t long size, @off_t long offset, FuseFileInfo fuseFileInfo) {
        if  (! containedInOneChunk(offset, offset+size))
            throw new  IllegalStateException("write op. straddles boundary : offset "+ offset  +" with size "+ size);

        long startPos  = alignToChunkSize(offset);
        int  chunkOffset  = intraChunkOffset(offset);
        int iSize = (int) size;

        CacheEntryHolder cacheEntry = entryMap.computeIfAbsent(s, path -> new CacheEntryHolder(new CacheEntry(path, startPos)));
        return cacheEntry.apply(c -> c != null && c.offset == startPos, () -> new CacheEntry(s, startPos), ce -> ce.write(pointer, chunkOffset, iSize));
    }

    @Override
    public int lock(String s, FuseFileInfo fuseFileInfo, int i, Flock flock) {
        System.out.println("LOCK: "+s);
        CacheEntryHolder cacheEntryHolder = entryMap.get(s);
        if (cacheEntryHolder != null)
            cacheEntryHolder.syncAndClear();
        return 0;
    }

    @Override
    public int flush(String s, FuseFileInfo fuseFileInfo) {
        CacheEntryHolder cacheEntry = entryMap.get(s);
        if  (cacheEntry != null) {
            cacheEntry.syncAndClear();
        }
        return super.flush(s, fuseFileInfo);
    }

    @Override
    protected int annotateAttributes(String fullPath, PeergosStat peergosStat, FileStat fileStat) {
        CacheEntryHolder cacheEntry = entryMap.get(fullPath);
        Optional<PeergosStat> updatedStat = Optional.empty();
        if (cacheEntry != null) {
            updatedStat = cacheEntry.applyIfPresent(ce -> {
                if (ce == null)
                    return peergosStat;
                long maxSize = ce.offset + ce.maxDirtyPos;
                if (peergosStat.properties.size < maxSize) {
                    FileProperties updated = peergosStat.properties.withSize(maxSize);
                    return new PeergosStat(peergosStat.treeNode, updated);
                }
                return peergosStat;
            });
        }
        return super.annotateAttributes(fullPath, updatedStat.orElse(peergosStat), fileStat);
    }

    private boolean containedInOneChunk(long start, long end) {
        return alignToChunkSize(start) == alignToChunkSize(end-1);
    }

    private long alignToChunkSize(long pos) {
        return Math.max(0, pos / Chunk.MAX_SIZE) * Chunk.MAX_SIZE;
    }
    private int intraChunkOffset(long  pos) {
        return (int) pos % Chunk.MAX_SIZE;
    }

    private class CacheEntryHolder {
        private CacheEntry entry;

        public CacheEntryHolder(CacheEntry entry) {
            this.entry = entry;
        }

        public synchronized <A> A apply(Predicate<CacheEntry> correctChunk, Supplier<CacheEntry> supplier, Function<CacheEntry, A> func) {
            if (correctChunk.test(entry)) {
                return func.apply(entry);
            }
            syncAndClear();
            setEntry(supplier.get());
            return func.apply(entry);
        }

        public synchronized <A> Optional<A> applyIfPresent(Function<CacheEntry, A> func) {
            if (entry != null) {
                return Optional.of(func.apply(entry));
            }
            return Optional.empty();
        }

        public synchronized void setEntry(CacheEntry entry) {
            this.entry = entry;
        }

        public synchronized void syncAndClear() {
            if (entry == null)
                return;
            entry.sync();
            entry = null;
        }
    }

    private class CacheEntry {
        private final String path;
        private final byte[] data;
        private final long offset;
        private int maxDirtyPos;


        public CacheEntry(String path, long offset) {
            this.path = path;<commit_msg>Even nicer code, though still failing that test.<commit_after>package peergos.fuse;

import jnr.ffi.Pointer;
import jnr.ffi.types.off_t;
import jnr.ffi.types.size_t;
import peergos.user.UserContext;
import peergos.user.fs.Chunk;
import peergos.user.fs.FileProperties;
import ru.serce.jnrfuse.ErrorCodes;
import ru.serce.jnrfuse.struct.*;

import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.*;

public class CachingPeergosFS extends PeergosFS {

    private static final int DEFAULT_SYNC_SLEEP = 1000*30;
    private static final int DEFAULT_CACHE_SIZE = 1;

    private final ConcurrentMap<String, CacheEntryHolder> entryMap;
    private final int chunkCacheSize, syncSleep;
//    private final Thread syncRunner;

    public CachingPeergosFS(UserContext userContext) {
        this(userContext, DEFAULT_CACHE_SIZE, DEFAULT_SYNC_SLEEP);
    }

    public CachingPeergosFS(UserContext userContext, int chunkCacheSize, int syncSleep) {
        super(userContext);

        this.chunkCacheSize = chunkCacheSize;
        this.syncSleep = syncSleep;
        this.entryMap = new ConcurrentHashMap<>();
        //TODO add GC thread to optimistically reduce size
//        this.syncRunner = new Thread(new Syncher());
//        this.syncRunner.start();
    }

    @Override
    public int read(String s, Pointer pointer, @size_t long size, @off_t long offset, FuseFileInfo fuseFileInfo) {
        if (!containedInOneChunk(offset, offset + size))
            throw new IllegalStateException("write op. straddles boundary : offset " + offset + " with size " + size);

        long startPos = alignToChunkSize(offset);
        int chunkOffset  = intraChunkOffset(offset);
        int iSize = (int) size;

        CacheEntryHolder cacheEntryHolder = entryMap.computeIfAbsent(s, path -> new CacheEntryHolder(new CacheEntry(path, startPos)));
        return cacheEntryHolder.apply(c -> c != null && c.offset == startPos, () -> new CacheEntry(s, startPos), ce -> ce.read(pointer, chunkOffset, iSize));
    }

    @Override
    public int write(String s, Pointer pointer, @size_t long size, @off_t long offset, FuseFileInfo fuseFileInfo) {
        if  (! containedInOneChunk(offset, offset+size))
            throw new  IllegalStateException("write op. straddles boundary : offset "+ offset  +" with size "+ size);

        long startPos  = alignToChunkSize(offset);
        int  chunkOffset  = intraChunkOffset(offset);
        int iSize = (int) size;

        CacheEntryHolder cacheEntry = entryMap.computeIfAbsent(s, path -> new CacheEntryHolder(new CacheEntry(path, startPos)));
        return cacheEntry.apply(c -> c != null && c.offset == startPos, () -> new CacheEntry(s, startPos), ce -> ce.write(pointer, chunkOffset, iSize));
    }

    @Override
    public int lock(String s, FuseFileInfo fuseFileInfo, int i, Flock flock) {
        System.out.println("LOCK: "+s);
        CacheEntryHolder cacheEntryHolder = entryMap.get(s);
        if (cacheEntryHolder != null)
            cacheEntryHolder.syncAndClear();
        return 0;
    }

    @Override
    public int flush(String s, FuseFileInfo fuseFileInfo) {
        CacheEntryHolder cacheEntry = entryMap.get(s);
        if  (cacheEntry != null) {
            cacheEntry.syncAndClear();
        }
        return super.flush(s, fuseFileInfo);
    }

    @Override
    protected int annotateAttributes(String fullPath, PeergosStat peergosStat, FileStat fileStat) {
        CacheEntryHolder cacheEntry = entryMap.get(fullPath);
        Optional<PeergosStat> updatedStat = Optional.empty();
        if (cacheEntry != null) {
            updatedStat = cacheEntry.applyIfPresent(ce -> {
                if (ce == null)
                    return peergosStat;
                long maxSize = ce.offset + ce.maxDirtyPos;
                if (peergosStat.properties.size < maxSize) {
                    FileProperties updated = peergosStat.properties.withSize(maxSize);
                    return new PeergosStat(peergosStat.treeNode, updated);
                }
                return peergosStat;
            });
        }
        return super.annotateAttributes(fullPath, updatedStat.orElse(peergosStat), fileStat);
    }

    private boolean containedInOneChunk(long start, long end) {
        return alignToChunkSize(start) == alignToChunkSize(end-1);
    }

    private long alignToChunkSize(long pos) {
        return Math.max(0, pos / Chunk.MAX_SIZE) * Chunk.MAX_SIZE;
    }
    private int intraChunkOffset(long  pos) {
        return (int) pos % Chunk.MAX_SIZE;
    }

    private class CacheEntryHolder {
        private CacheEntry entry;

        public CacheEntryHolder(CacheEntry entry) {
            this.entry = entry;
        }

        public synchronized <A> A apply(Predicate<CacheEntry> correctChunk, Supplier<CacheEntry> supplier, Function<CacheEntry, A> func) {
            if (!correctChunk.test(entry)) {
                syncAndClear();
                setEntry(supplier.get());
            }
            return func.apply(entry);
        }

        public synchronized <A> Optional<A> applyIfPresent(Function<CacheEntry, A> func) {
            if (entry != null) {
                return Optional.of(func.apply(entry));
            }
            return Optional.empty();
        }

        public synchronized void setEntry(CacheEntry entry) {
            this.entry = entry;
        }

        public synchronized void syncAndClear() {
            if (entry == null)
                return;
            entry.sync();
            entry = null;
        }
    }

    private class CacheEntry {
        private final String path;
        private final byte[] data;
        private final long offset;
        private int maxDirtyPos;


        public CacheEntry(String path, long offset) {
            this.path = path;
{'fraction_non_alphanumeric': 0.09690915300546447, 'fraction_numerical': 0.0017076502732240437, 'mean_word_length': 2.564516129032258, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4877520', 'n_tokens_mistral': 3342, 'n_tokens_neox': 3224, 'n_words': 826}