ConsistentHashingV2.java

package org.rakeshopensource.consistenthashing;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

public class ConsistentHashingV2<K, V> {
    public static class Details<K, V> {
        K node, key;
        V value;

        public Details(K node, K key, V value) {
            this.node = node;
            this.key = key;
            this.value = value;
        }

        @Override
        public String toString() {
            return "Details{" +
                    "node='" + node + '\'' +
                    ", key='" + key + '\'' +
                    ", value='" + value + '\'' +
                    '}';
        }
    }

    public static class Node<K, V> {
        private final K identifier;
        private final Map<K, V> dataStore = new HashMap<>();

        public Node(K identifier) {
            this.identifier = identifier;
        }

        public K getIdentifier() {
            return identifier;
        }

        public Map<K, V> getDataStore() {
            return dataStore;
        }
    }

    private final NavigableMap<Integer, Node<K, V>> hashRing = new TreeMap<>();
    private final int numberOfReplicas;
    private final Function<K, Integer> hashFunction;

    public ConsistentHashingV2(int numberOfReplicas, Function<K, Integer> hashFunction) {
        this.numberOfReplicas = numberOfReplicas;
        this.hashFunction = hashFunction;
    }

    public int hashFunction(K key) {
        return this.hashFunction.apply(key);
    }


    public void addNode(Node<K, V> node) {
        int hash = hashFunction(node.getIdentifier());
        hashRing.put(hash, node);
        redistributeKeys(node);
    }

    public void removeNode(Node<K, V> node) {
        int hash = hashFunction(node.getIdentifier());
        hashRing.remove(hash);
        redistributeKeysAll();
    }

    public Details<K, V> get(K key) {
        Node<K, V> node = getCorrectNode(key);
        V value = node.getDataStore().get(key);
        if (value != null) {
            return new Details<>(node.getIdentifier(), key, value);
        }
        return null;
    }


    public void put(K key, V value) {
        Node<K, V> node = getCorrectNode(key);
        node.getDataStore().put(key, value);
        // Replicate data
        replicateData(key, value, node);
    }

    private Node<K, V> getCorrectNode(K key) {
        int hash = hashFunction(key);
        if (hashRing.containsKey(hash)) {
            return hashRing.get(hash);
        }
        Map.Entry<Integer, Node<K, V>> higherEntry = hashRing.higherEntry(hash);
        if (higherEntry == null) {
            return hashRing.firstEntry().getValue();
        }
        return higherEntry.getValue();
    }

    private void replicateData(K key, V value, Node<K, V> startNode) {
        int hash = hashFunction(startNode.getIdentifier());
        for (int i = 1; i < numberOfReplicas; i++) {
            hash = hashRing.higherKey(hash) == null ? hashRing.firstKey() : hashRing.higherKey(hash);
            hashRing.get(hash).getDataStore().put(key, value);
        }
    }


    private void redistributeKeysAll() {
        List<K> allKeys = new ArrayList<>();
        hashRing.values().forEach(node -> allKeys.addAll(node.getDataStore().keySet()));

        allKeys.forEach(key -> {
            V value = get(key).value;
            put(key, value);
        });
    }

    private void redistributeKeys(Node<K, V> newNode) {
        int newNodeHash = hashFunction(newNode.getIdentifier());
        hashRing.headMap(newNodeHash, true).values().forEach(node -> {
            List<K> keysToMove = new ArrayList<>(node.getDataStore().keySet());
            keysToMove.forEach(key -> {
                V value = node.getDataStore().get(key);
                int keyHash = hashFunction(key);
                if (keyHash > newNodeHash) {
                    node.getDataStore().remove(key);
                    newNode.getDataStore().put(key, value);
                }
            });
        });
    }

    public static void main(String[] args) {
        ConsistentHashingV2<String, String> consistentHashing = new ConsistentHashingV2<>(3, String::hashCode);

        Node<String, String> node1 = new Node<>("Node1");
        Node<String, String> node2 = new Node<>("Node2");
        Node<String, String> node3 = new Node<>("Node3");

        consistentHashing.addNode(node1);
        consistentHashing.addNode(node2);
        consistentHashing.addNode(node3);

        // Add data to nodes
        consistentHashing.put("Key1", "Value1");
        consistentHashing.put("Key2", "Value2");
        consistentHashing.put("Key3", "Value3");

        // Retrieve and display the values
        System.out.println("Data for Key1: " + consistentHashing.get("Key1"));
        System.out.println("Data for Key2: " + consistentHashing.get("Key2"));
        System.out.println("Data for Key3: " + consistentHashing.get("Key3"));

        // Add a new node and redistribute keys
        Node<String, String> node4 = new Node<>("Node4");
        consistentHashing.addNode(node4);

        System.out.println("After adding Node4:");
        System.out.println("Data for Key1: " + consistentHashing.get("Key1"));
        System.out.println("Data for Key2: " + consistentHashing.get("Key2"));
        System.out.println("Data for Key3: " + consistentHashing.get("Key3"));

        // Remove a node and redistribute keys
        consistentHashing.removeNode(node1);

        System.out.println("After removing Node1:");
        System.out.println("Data for Key1: " + consistentHashing.get("Key1"));
        System.out.println("Data for Key2: " + consistentHashing.get("Key2"));
        System.out.println("Data for Key3: " + consistentHashing.get("Key3"));
    }
}