Design a Key-Value Store

Client
  |
  v
[Load Balancer]
  |
  v
[API Gateway / Router]
  |
  +---> [Consistent Hash Ring] (determines which nodes store data)
  |
  v
[Storage Node 1] [Storage Node 2] ... [Storage Node N]
  |                |                        |
  +---> [Local Storage Engine (LSM Tree)]   |
  |                |                        |
  +---> [Replication Manager] <-----------+
  |
  v
[Failure Detector] (gossip protocol for node health)

public class ConsistentHashRing {
    private final TreeMap<Long, Node> ring = new TreeMap<>();
    private final int virtualNodes = 150; // per physical node
    
    public void addNode(Node node) {
        for (int i = 0; i < virtualNodes; i++) {
            long hash = hash(node.getId() + ":" + i);
            ring.put(hash, node);
        }
    }
    
    public void removeNode(Node node) {
        for (int i = 0; i < virtualNodes; i++) {
            long hash = hash(node.getId() + ":" + i);
            ring.remove(hash);
        }
    }
    
    public List<Node> getResponsibleNodes(String key, int count) {
        if (ring.isEmpty()) return Collections.emptyList();
        
        long keyHash = hash(key);
        List<Node> result = new ArrayList<>();
        Set<Node> seen = new HashSet<>();
        
        // Find the first node clockwise from the key hash
        NavigableMap<Long, Node> tailMap = ring.tailMap(keyHash);
        Iterator<Node> iterator = tailMap.values().iterator();
        
        // If no nodes after key hash, wrap around to beginning
        if (!iterator.hasNext()) {
            iterator = ring.values().iterator();
        }
        
        while (iterator.hasNext() && result.size() < count) {
            Node node = iterator.next();
            if (seen.add(node)) { // only add unique physical nodes
                result.add(node);
            }
            
            // Wrap around if we reach the end
            if (!iterator.hasNext() && result.size() < count) {
                iterator = ring.values().iterator();
            }
        }
        
        return result;
    }
    
    private long hash(String input) {
        return Hashing.murmur3_128().hashString(input, StandardCharsets.UTF_8)
            .asLong();
    }
}

public class LSMTree {
    private final WriteAheadLog wal;
    private volatile MemTable activeMemTable;
    private final List<SSTable> sstables = new CopyOnWriteArrayList<>();
    private final BlockingQueue<MemTable> flushQueue = new LinkedBlockingQueue<>();
    
    public void put(String key, String value) throws IOException {
        // 1. Write to WAL first for durability
        wal.append(key, value);
        
        // 2. Write to in-memory table
        synchronized (this) {
            activeMemTable.put(key, value);
            
            // 3. Flush to disk if memtable is full
            if (activeMemTable.size() >= MAX_MEMTABLE_SIZE) {
                flushQueue.offer(activeMemTable);
                activeMemTable = new MemTable();
            }
        }
    }
    
    public String get(String key) {
        // 1. Check active memtable first
        String value = activeMemTable.get(key);
        if (value != null) return value;
        
        // 2. Check flushing memtables
        for (MemTable flushing : flushQueue) {
            value = flushing.get(key);
            if (value != null) return value;
        }
        
        // 3. Check SSTables from newest to oldest
        for (int i = sstables.size() - 1; i >= 0; i--) {
            value = sstables.get(i).get(key);
            if (value != null) return value;
        }
        
        return null; // key not found
    }
}

public class ReplicationCoordinator {
    private final ConsistentHashRing hashRing;
    private final int replicationFactor = 3;
    
    public boolean put(String key, String value, int W) {
        List<Node> replicas = hashRing.getResponsibleNodes(key, replicationFactor);
        List<Future<Boolean>> futures = new ArrayList<>();
        
        for (Node node : replicas) {
            futures.add(executor.submit(() -> node.localPut(key, value)));
        }
        
        int successCount = 0;
        for (Future<Boolean> future : futures) {
            try {
                if (future.get(WRITE_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
                    successCount++;
                }
            } catch (Exception e) {
                // Log failure but continue
            }
        }
        
        return successCount >= W; // Success if W nodes acknowledge
    }
    
    public String get(String key, int R) {
        List<Node> replicas = hashRing.getResponsibleNodes(key, replicationFactor);
        List<Future<VersionedValue>> futures = new ArrayList<>();
        
        for (Node node : replicas) {
            futures.add(executor.submit(() -> node.localGet(key)));
        }
        
        List<VersionedValue> values = new ArrayList<>();
        for (Future<VersionedValue> future : futures) {
            try {
                VersionedValue value = future.get(READ_TIMEOUT_MS, TimeUnit.MILLISECONDS);
                if (value != null) {
                    values.add(value);
                    if (values.size() >= R) break; // Stop when we have R responses
                }
            } catch (Exception e) {
                // Log failure but continue
            }
        }
        
        if (values.size() < R) {
            throw new InsufficientReplicasException();
        }
        
        // Return the value with the highest vector clock
        return resolveConflicts(values);
    }
}

public class VectorClock {
    private final Map<String, Long> clock = new HashMap<>();
    
    public VectorClock increment(String nodeId) {
        VectorClock result = new VectorClock();
        result.clock.putAll(this.clock);
        result.clock.put(nodeId, result.clock.getOrDefault(nodeId, 0L) + 1);
        return result;
    }
    
    public ClockComparison compareTo(VectorClock other) {
        boolean thisGreater = false;
        boolean otherGreater = false;
        
        Set<String> allNodes = new HashSet<>(this.clock.keySet());
        allNodes.addAll(other.clock.keySet());
        
        for (String node : allNodes) {
            long thisValue = this.clock.getOrDefault(node, 0L);
            long otherValue = other.clock.getOrDefault(node, 0L);
            
            if (thisValue > otherValue) thisGreater = true;
            if (otherValue > thisValue) otherGreater = true;
        }
        
        if (thisGreater && !otherGreater) return ClockComparison.AFTER;
        if (otherGreater && !thisGreater) return ClockComparison.BEFORE;
        if (!thisGreater && !otherGreater) return ClockComparison.EQUAL;
        return ClockComparison.CONCURRENT; // conflict!
    }
}

public enum ClockComparison {
    BEFORE, AFTER, EQUAL, CONCURRENT
}

public class GossipProtocol {
    private final Map<String, NodeInfo> nodeStates = new ConcurrentHashMap<>();
    private final String localNodeId;
    
    public void startGossip() {
        executor.scheduleAtFixedRate(() -> {
            // Pick random nodes to gossip with
            List<String> targets = selectRandomNodes(GOSSIP_FANOUT);
            
            for (String targetId : targets) {
                try {
                    GossipMessage message = createGossipMessage();
                    GossipMessage response = sendGossip(targetId, message);
                    mergeGossipResponse(response);
                } catch (Exception e) {
                    markNodeAsSuspected(targetId);
                }
            }
        }, 0, GOSSIP_INTERVAL_MS, TimeUnit.MILLISECONDS);
    }
    
    private GossipMessage createGossipMessage() {
        GossipMessage message = new GossipMessage();
        message.setSender(localNodeId);
        message.setTimestamp(System.currentTimeMillis());
        
        // Include state of all known nodes
        for (NodeInfo info : nodeStates.values()) {
            message.addNodeState(info.getNodeId(), info.getHeartbeat(), info.getStatus());
        }
        
        return message;
    }
    
    private void mergeGossipResponse(GossipMessage response) {
        for (NodeState state : response.getNodeStates()) {
            NodeInfo existing = nodeStates.get(state.getNodeId());
            
            if (existing == null || state.getHeartbeat() > existing.getHeartbeat()) {
                nodeStates.put(state.getNodeId(), 
                    new NodeInfo(state.getNodeId(), state.getHeartbeat(), state.getStatus()));
                
                // Trigger rebalancing if node status changed
                if (existing != null && existing.getStatus() != state.getStatus()) {
                    triggerRebalancing(state.getNodeId(), state.getStatus());
                }
            }
        }
    }
}

public class MerkleTree {
    private final int maxDepth;
    private final MerkleNode root;
    
    public static class MerkleNode {
        private final String hash;
        private final String startKey;
        private final String endKey;
        private final MerkleNode left;
        private final MerkleNode right;
        
        // Leaf node constructor (for actual key ranges)
        public MerkleNode(String startKey, String endKey, List<String> keys) {
            this.startKey = startKey;
            this.endKey = endKey;
            this.left = null;
            this.right = null;
            
            // Hash all keys in this range
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            keys.stream().sorted().forEach(key -> digest.update(key.getBytes()));
            this.hash = Base64.getEncoder().encodeToString(digest.digest());
        }
        
        // Internal node constructor
        public MerkleNode(MerkleNode left, MerkleNode right) {
            this.left = left;
            this.right = right;
            this.startKey = left.startKey;
            this.endKey = right.endKey;
            
            // Hash of left hash + right hash
            String combined = left.hash + right.hash;
            this.hash = sha256(combined);
        }
    }
    
    public List<String> findDifferences(MerkleTree otherTree) {
        List<String> differences = new ArrayList<>();
        findDifferencesRecursive(this.root, otherTree.root, differences);
        return differences;
    }
    
    private void findDifferencesRecursive(MerkleNode node1, MerkleNode node2, 
                                         List<String> differences) {
        // If hashes match, subtrees are identical
        if (Objects.equals(node1.hash, node2.hash)) {
            return;
        }
        
        // If leaf nodes differ, add range to differences
        if (node1.left == null && node1.right == null) {
            differences.add(node1.startKey + ":" + node1.endKey);
            return;
        }
        
        // Recurse on children
        if (node1.left != null && node2.left != null) {
            findDifferencesRecursive(node1.left, node2.left, differences);
        }
        if (node1.right != null && node2.right != null) {
            findDifferencesRecursive(node1.right, node2.right, differences);
        }
    }
}