CachedSpawnDataStore.java

/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.addthis.hydra.job.store;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ExecutionException;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.cache.Weigher;

import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class CachedSpawnDataStore implements SpawnDataStore {

    private static Pair<String, String> defaultKey(String path) {
        return ImmutablePair.of(path, null);
    }
    private static final Logger log = LoggerFactory.getLogger(CachedSpawnDataStore.class);

    private final SpawnDataStore dataStore;
    private final LoadingCache<Pair<String, String>, String> cache;


    public CachedSpawnDataStore(SpawnDataStore dataStore, long dataStoreCacheSize) {
        this.dataStore = dataStore;
        this.cache = CacheBuilder.newBuilder()
                .weigher(new Weigher<Pair<String, String>, String>() {
                    @Override
                    public int weigh(Pair<String, String> key, String value) {
                        int leftWeight = key.getLeft() != null ?
                                    key.getLeft().length() :
                                    0;

                        int rightWeight = key.getRight() != null ?
                                    key.getRight().length() :
                                    0;

                        // Multiply strlen by 2 (full width characters in java
                        return 2 * (value.length() + leftWeight + rightWeight);
                    }
                })
                .maximumWeight(dataStoreCacheSize)
                .build(new CacheLoader<Pair<String, String>, String>() {
                    @Override
                    public String load(Pair<String, String> key) throws Exception {
                        String path = key.getLeft();
                        String childId = key.getRight();

                        if (childId == null) {
                            return CachedSpawnDataStore.this.dataStore.get(path);
                        } else {
                            return CachedSpawnDataStore.this.dataStore.getChild(path, childId);
                        }
                    }
                });
    }

    @Override
    public String getDescription() {
        return dataStore.getDescription();
    }

    @Override
    public String get(String path) {
        try {
            return cache.get(defaultKey(path));
        } catch (ExecutionException | CacheLoader.InvalidCacheLoadException e) {
            log.error("failed to execute get from cache", e);
            return null;
        }
    }

    @Override
    public Map<String, String> get(String[] paths) {
        List<String> notCached = new ArrayList<>();
        Map<String, String> results = new TreeMap<>();

        // Load every path from the cache that we can. For all the paths that do not exist in the cache, continue to
        // do the batched query logic which should be implemented in the underlying DataStore's `get(String[])`
        for (String path : paths) {
            String result = cache.getIfPresent(defaultKey(path));
            if (result == null) {
                notCached.add(path);
            } else {
                results.put(path, result);
            }
        }

        String[] remainingPaths = new String[notCached.size()];
        remainingPaths = notCached.toArray(remainingPaths);

        Map<String, String> resultsFromDB = dataStore.get(remainingPaths);

        for (Map.Entry<String, String> entry : resultsFromDB.entrySet()) {
            String path = entry.getKey();
            cache.put(defaultKey(path), entry.getValue());
        }

        results.putAll(resultsFromDB);
        return results;
    }

    @Override
    public void put(String path, String value) throws Exception {
        dataStore.put(path, value);
        cache.put(defaultKey(path), value);
    }

    @Override
    public void putAsChild(String parent, String childId, String value) throws Exception {
        dataStore.putAsChild(parent, childId, value);
        cache.put(ImmutablePair.of(parent, childId), value);
    }

    @Override
    public String getChild(String parent, String childId) throws Exception {
        return cache.get(ImmutablePair.of(parent, childId));
    }

    @Override
    public void deleteChild(String parent, String childId) {
        dataStore.deleteChild(parent, childId);
        cache.invalidate(ImmutablePair.of(parent, childId));
    }

    @Override
    public void delete(String path) {
        dataStore.delete(path);
        cache.invalidate(defaultKey(path));
    }

    @Override
    public List<String> getChildrenNames(String path) {
        return dataStore.getChildrenNames(path);
    }

    @Override
    public Map<String, String> getAllChildren(String path) {
        return dataStore.getAllChildren(path);
    }

    @Override
    public void close() {
        dataStore.close();
    }
}