<commit_before>package com.thinkaurelius.titan.diskstorage.keycolumnvalue.cache;

import com.google.common.base.Preconditions;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.Weigher;
import com.thinkaurelius.titan.core.TitanException;
import com.thinkaurelius.titan.diskstorage.*;
import com.thinkaurelius.titan.diskstorage.keycolumnvalue.*;
import com.thinkaurelius.titan.diskstorage.util.CacheMetricsAction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import static com.thinkaurelius.titan.util.datastructures.ByteSize.*;

/**
 * @author <NAME> (<EMAIL>)
 */
public class ExpirationKCVSCache extends KCVSCache {

    private static final Logger log =
            LoggerFactory.getLogger(ExpirationKCVSCache.class);

    //Weight estimation
    private static final int STATICARRAYBUFFER_SIZE = STATICARRAYBUFFER_RAW_SIZE + 10; // 10 = last number is average length
    private static final int KEY_QUERY_SIZE = OBJECT_HEADER + 4 + 1 + 3 * (OBJECT_REFERENCE + STATICARRAYBUFFER_SIZE); // object_size + int + boolean + 3 static buffers

    private static final int INVALIDATE_KEY_FRACTION_PENALTY = 1000;
    private static final int PENALTY_THRESHOLD = 5;

    private volatile CountDownLatch penaltyCountdown;

    private final Cache<KeySliceQuery,EntryList> cache;
    private final ConcurrentHashMap<StaticBuffer,Long> expiredKeys;

    private final long cacheTimeMS;
    private final long expirationGracePeriodMS;
    private final CleanupThread cleanupThread;


    public ExpirationKCVSCache(final KeyColumnValueStore store, String metricsName, final long cacheTimeMS, final long expirationGracePeriodMS, final long maximumByteSize) {
        super(store, metricsName);
        Preconditions.checkArgument(cacheTimeMS > 0, "Cache expiration must be positive: %s", cacheTimeMS);
        Preconditions.checkArgument(System.currentTimeMillis()+1000l*3600*24*365*100+cacheTimeMS>0,"Cache expiration time too large, overflow may occur: %s",cacheTimeMS);
        this.cacheTimeMS = cacheTimeMS;
        int concurrencyLevel = Runtime.getRuntime().availableProcessors();
        Preconditions.checkArgument(expirationGracePeriodMS>=0,"Invalid expiration grace peiod: %s",expirationGracePeriodMS);
        this.expirationGracePeriodMS = expirationGracePeriodMS;
        CacheBuilder<KeySliceQuery,EntryList> cachebuilder = CacheBuilder.newBuilder()
                .maximumWeight(maximumByteSize)
                .concurrencyLevel(concurrencyLevel)
                .initialCapacity(1000)
                .expireAfterWrite(cacheTimeMS, TimeUnit.MILLISECONDS)
                .weigher(new Weigher<KeySliceQuery, EntryList>() {
                    @Override
                    public int weigh(KeySliceQuery keySliceQuery, EntryList entries) {
                        return GUAVA_CACHE_ENTRY_SIZE + KEY_QUERY_SIZE + entries.getByteSize();
                    }
                });

        cache = cachebuilder.build();
        expiredKeys = new ConcurrentHashMap<StaticBuffer, Long>(50,0.75f,concurrencyLevel);
        penaltyCountdown = new CountDownLatch(PENALTY_THRESHOLD);

        cleanupThread = new CleanupThread();
        cleanupThread.start();
    }

    @Override
    public EntryList getSlice(final KeySliceQuery query, final StoreTransaction txh) throws StorageException {
        incActionBy(1, CacheMetricsAction.RETRIEVAL,txh);
        if (isExpired(query)) {
            incActionBy(1, CacheMetricsAction.MISS,txh);
            return store.getSlice(query, getTx(txh));
        }

        try {
            return cache.get(query,new Callable<EntryList>() {
                @Override
                public EntryList call() throws Exception {
                    incActionBy(1, CacheMetricsAction.MISS,txh);
                    return store.getSlice(query, getTx(txh));
                }
            });
        } catch (Exception e) {
            if (e instanceof TitanException) throw (TitanException)e;
            else if (e.getCause() instanceof TitanException) throw (TitanException)e.getCause();
            else throw new TitanException(e);
        }
    }

    @Override
    public Map<StaticBuffer,EntryList> getSlice(final List<StaticBuffer> keys, final SliceQuery query, final StoreTransaction txh) throws StorageException {
        Map<StaticBuffer,EntryList> results = new HashMap<StaticBuffer, EntryList>(keys.size());
        List<StaticBuffer> remainingKeys = new ArrayList<StaticBuffer>(keys.size());
        KeySliceQuery[] ksqs = new KeySliceQuery[keys.size()];
        incActionBy(keys.size(), CacheMetricsAction.RETRIEVAL,txh);
        //Find all cached queries
        for (int i=0;i<keys.size();i++) {
            StaticBuffer key = keys.get(i);
            ksqs[i] = new KeySliceQuery(key,query);
            EntryList result = null;
            if (!isExpired(ksqs[i])) result = cache.getIfPresent(ksqs[i]);
            else ksqs[i]=null;
            if (result!=null) results.put(key,result);
            else remainingKeys.add(key);
        }
        //Request remaining ones from backend
        incActionBy(remainingKeys.size(), CacheMetricsAction.MISS,txh);
        Map<StaticBuffer,EntryList> subresults = store.getSlice(remainingKeys, query, getTx(txh));
        for (int i=0;i<keys.size();i++) {
            StaticBuffer key = keys.get(i);
            EntryList subresult = subresults.get(key);
            if (subresult!=null) {
                results.put(key,subresult);
                if (ksqs[i]!=null) cache.put(ksqs[i],subresult);
            }
        }
        return results;
    }

    @Override
    public void clearCache() {
        cache.invalidateAll();
        expiredKeys.clear();
        penaltyCountdown = new CountDownLatch(PENALTY_THRESHOLD);
    }

    @Override
    public void invalidate(StaticBuffer key, List<CachableStaticBuffer> entries) {
        Preconditions.checkArgument(!hasValidateKeysOnly() || entries.isEmpty());
        expiredKeys.put(key,getExpirationTime());
        if (Math.random()<1.0/INVALIDATE_KEY_FRACTION_PENALTY) penaltyCountdown.countDown();
    }

    @Override
    public void close() throws StorageException {<commit_msg>Minor optimization by checking whether a store multi query is actually needed.<commit_after>package com.thinkaurelius.titan.diskstorage.keycolumnvalue.cache;

import com.google.common.base.Preconditions;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.Weigher;
import com.thinkaurelius.titan.core.TitanException;
import com.thinkaurelius.titan.diskstorage.*;
import com.thinkaurelius.titan.diskstorage.keycolumnvalue.*;
import com.thinkaurelius.titan.diskstorage.util.CacheMetricsAction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import static com.thinkaurelius.titan.util.datastructures.ByteSize.*;

/**
 * @author <NAME> (<EMAIL>)
 */
public class ExpirationKCVSCache extends KCVSCache {

    private static final Logger log =
            LoggerFactory.getLogger(ExpirationKCVSCache.class);

    //Weight estimation
    private static final int STATICARRAYBUFFER_SIZE = STATICARRAYBUFFER_RAW_SIZE + 10; // 10 = last number is average length
    private static final int KEY_QUERY_SIZE = OBJECT_HEADER + 4 + 1 + 3 * (OBJECT_REFERENCE + STATICARRAYBUFFER_SIZE); // object_size + int + boolean + 3 static buffers

    private static final int INVALIDATE_KEY_FRACTION_PENALTY = 1000;
    private static final int PENALTY_THRESHOLD = 5;

    private volatile CountDownLatch penaltyCountdown;

    private final Cache<KeySliceQuery,EntryList> cache;
    private final ConcurrentHashMap<StaticBuffer,Long> expiredKeys;

    private final long cacheTimeMS;
    private final long expirationGracePeriodMS;
    private final CleanupThread cleanupThread;


    public ExpirationKCVSCache(final KeyColumnValueStore store, String metricsName, final long cacheTimeMS, final long expirationGracePeriodMS, final long maximumByteSize) {
        super(store, metricsName);
        Preconditions.checkArgument(cacheTimeMS > 0, "Cache expiration must be positive: %s", cacheTimeMS);
        Preconditions.checkArgument(System.currentTimeMillis()+1000l*3600*24*365*100+cacheTimeMS>0,"Cache expiration time too large, overflow may occur: %s",cacheTimeMS);
        this.cacheTimeMS = cacheTimeMS;
        int concurrencyLevel = Runtime.getRuntime().availableProcessors();
        Preconditions.checkArgument(expirationGracePeriodMS>=0,"Invalid expiration grace peiod: %s",expirationGracePeriodMS);
        this.expirationGracePeriodMS = expirationGracePeriodMS;
        CacheBuilder<KeySliceQuery,EntryList> cachebuilder = CacheBuilder.newBuilder()
                .maximumWeight(maximumByteSize)
                .concurrencyLevel(concurrencyLevel)
                .initialCapacity(1000)
                .expireAfterWrite(cacheTimeMS, TimeUnit.MILLISECONDS)
                .weigher(new Weigher<KeySliceQuery, EntryList>() {
                    @Override
                    public int weigh(KeySliceQuery keySliceQuery, EntryList entries) {
                        return GUAVA_CACHE_ENTRY_SIZE + KEY_QUERY_SIZE + entries.getByteSize();
                    }
                });

        cache = cachebuilder.build();
        expiredKeys = new ConcurrentHashMap<StaticBuffer, Long>(50,0.75f,concurrencyLevel);
        penaltyCountdown = new CountDownLatch(PENALTY_THRESHOLD);

        cleanupThread = new CleanupThread();
        cleanupThread.start();
    }

    @Override
    public EntryList getSlice(final KeySliceQuery query, final StoreTransaction txh) throws StorageException {
        incActionBy(1, CacheMetricsAction.RETRIEVAL,txh);
        if (isExpired(query)) {
            incActionBy(1, CacheMetricsAction.MISS,txh);
            return store.getSlice(query, getTx(txh));
        }

        try {
            return cache.get(query,new Callable<EntryList>() {
                @Override
                public EntryList call() throws Exception {
                    incActionBy(1, CacheMetricsAction.MISS,txh);
                    return store.getSlice(query, getTx(txh));
                }
            });
        } catch (Exception e) {
            if (e instanceof TitanException) throw (TitanException)e;
            else if (e.getCause() instanceof TitanException) throw (TitanException)e.getCause();
            else throw new TitanException(e);
        }
    }

    @Override
    public Map<StaticBuffer,EntryList> getSlice(final List<StaticBuffer> keys, final SliceQuery query, final StoreTransaction txh) throws StorageException {
        Map<StaticBuffer,EntryList> results = new HashMap<StaticBuffer, EntryList>(keys.size());
        List<StaticBuffer> remainingKeys = new ArrayList<StaticBuffer>(keys.size());
        KeySliceQuery[] ksqs = new KeySliceQuery[keys.size()];
        incActionBy(keys.size(), CacheMetricsAction.RETRIEVAL,txh);
        //Find all cached queries
        for (int i=0;i<keys.size();i++) {
            StaticBuffer key = keys.get(i);
            ksqs[i] = new KeySliceQuery(key,query);
            EntryList result = null;
            if (!isExpired(ksqs[i])) result = cache.getIfPresent(ksqs[i]);
            else ksqs[i]=null;
            if (result!=null) results.put(key,result);
            else remainingKeys.add(key);
        }
        //Request remaining ones from backend
        if (!remainingKeys.isEmpty()) {
            incActionBy(remainingKeys.size(), CacheMetricsAction.MISS,txh);
            Map<StaticBuffer,EntryList> subresults = store.getSlice(remainingKeys, query, getTx(txh));
            for (int i=0;i<keys.size();i++) {
                StaticBuffer key = keys.get(i);
                EntryList subresult = subresults.get(key);
                if (subresult!=null) {
                    results.put(key,subresult);
                    if (ksqs[i]!=null) cache.put(ksqs[i],subresult);
                }
            }
        }
        return results;
    }

    @Override
    public void clearCache() {
        cache.invalidateAll();
        expiredKeys.clear();
        penaltyCountdown = new CountDownLatch(PENALTY_THRESHOLD);
    }

    @Override
    public void invalidate(StaticBuffer key, List<CachableStaticBuffer> entries) {
        Preconditions.checkArgument(!hasValidateKeysOnly() || entries.isEmpty());
        expiredKeys.put(key,getExpirationTime());
        if (Math.random()<1.0/INVALIDATE_KEY_FRACTION_PENALTY) penaltyCountdown.countDown();
    }

    @Override
    public void close() throws StorageException {
{'fraction_non_alphanumeric': 0.09587369384486309, 'fraction_numerical': 0.007627183281214247, 'mean_word_length': 3.131064902331443, 'pattern_counts': {'":': 0, '<': 41, '<?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': '4159946', 'n_tokens_mistral': 3779, 'n_tokens_neox': 3628, 'n_words': 759}