DistributedCache using ConsistentHash

DistributedCache.java

package com.kamatama41.sandbox4j;

import java.io.ByteArrayInputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

/**
 * {@link ConsistentHash}を利用した分散キャッシュを提供します
 * 
 * @author kamatama41
 *
 */
public class DistributedCache {

    private ConsistentHash<Node> consistentHash;
    private Map<String, Node> nodes; 
    private final MD5Hash md5Hash = new MD5Hash();

    /**
     * {@link DistributedCache}を構築します
     * @param nodeKeys Nodeのキーになる文字列
     */
    public DistributedCache(Collection<String> nodeKeys) {
        nodes = new HashMap<>();
        consistentHash = new ConsistentHash<>(md5Hash, 10, nodes.values());
        for (String nodeKey : nodeKeys) {
            Node node = new Node(nodeKey);
            consistentHash.add(node);
            nodes.put(nodeKey, node);
        }
    }

    /**
     * Keyが格納されているNodeを取得します。
     * @param key
     * @return
     */
    public String getNode(Object key) {
        Node node = consistentHash.get(key);
        if(node.cache.containsKey(key)) {
            return node.nodeKey;
        }
        return null;
    }

    /**
     * Nodeを追加します
     * @param nodeKey
     */
    public void addNode(String nodeKey) {
        if(nodes.containsKey(nodeKey)) return;

        Node newNode = new Node(nodeKey);
        consistentHash.add(newNode);
        nodes.put(nodeKey, newNode);
        for (Node node : nodes.values()) {
            node.reconstruction(consistentHash);
        }
    }

    /**
     * Nodeを削除します
     * @param nodeKey
     */
    public void removeNode(String nodeKey) {
        Node removeNode = nodes.get(nodeKey);
        if(removeNode == null) return;

        consistentHash.remove(removeNode);
        nodes.remove(nodeKey);
        removeNode.reconstruction(consistentHash);
    }

    /**
     * キャッシュに値を追加します
     * @param key
     * @param value
     */
    public void put(Object key, Object value) {
        Node node = consistentHash.get(key);
        node.put(key, value);
    }

    /**
     * キャッシュから値を取得します
     * @param key
     * @return
     */
    public Object get(Object key) {
        Node node = consistentHash.get(key);
        return node.get(key);
    }

    /**
     * MD5ハッシュを利用した{@link HashFunction}実装です
     * @author kamatama41
     *
     */
    private static class MD5Hash implements HashFunction {
        private MessageDigest md;
        MD5Hash() {
            try {
                md = MessageDigest.getInstance("MD5");
            } catch (NoSuchAlgorithmException e) {
                throw new RuntimeException(e);
            }
        }
        @Override
        public int hash(Object o) {
            DataInputStream in = new DataInputStream(new ByteArrayInputStream(md.digest(o.toString().getBytes())));
            try {
                return in.readInt();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * {@link DistributedCache}における各Nodeの実装です
     * @author kamatama41
     *
     */
    private static final class Node {
        private String nodeKey;
        private Map<Object, Object> cache;
        
        Node(String nodeKey) {
            this.nodeKey = nodeKey;
            this.cache = new HashMap<>();
        }

        /**
         * Node上のValueを再構築します。
         */
        public void reconstruction(ConsistentHash<Node> consistentHash) {
            Map<Object, Object> newCache = new HashMap<>();
            for(Map.Entry<Object, Object> keyValue : cache.entrySet()) {
                Node newNode = consistentHash.get(keyValue.getKey());
                if(this.nodeKey.equals(newNode.nodeKey)) {
                    newCache.put(keyValue.getKey(), keyValue.getValue());
                } else {
                    newNode.put(keyValue.getKey(), keyValue.getValue());
                }
            }
            cache = newCache;
        }

        private void put(Object key, Object value) {
            cache.put(key, value);
        }
        private Object get(Object key) {
            return cache.get(key);
        }

        @Override
        public String toString() {
            return nodeKey;
        }
    }
}