view grok_interview/consistent_hash.py @ 61:9df5587cf23b

[Color game] It can compile on windows now.
author June Park <me@mrjunejune.com>
date Sat, 20 Dec 2025 21:07:34 -0500
parents 68fa88ac73fe
children
line wrap: on
line source

# Write a function taht maps
# Key string to one of N servers
from bisect import bisect_left
from hashlib import sha256

class Server:

    def __init__(self):
        self.id = str(sha256())

class ConsistentHash:
    def __init__(self):
        self.server = []
        self.hash_key_to_server = {}
        self.hash_range = 10000

    def get_hash_pos(self, key: str) -> int:
        x = sha256(bytes(key, "utf-8"))
        return int.from_bytes(x.digest()[:8]) % self.hash_range

    def add_server(self, server: Server):
        self.server.append(server)
        for i in range(len(self.server) + 1):
            hash_key = self.get_hash_pos(server.id)
            self.hash_key_to_server[hash_key] = i

    def get_server(self, key):
        hash_key = self.get_hash_pos(key)
        return self.server[bisect_left(list(self.hash_key_to_server.keys()), hash_key)]

    def remove_server(self, server):
        hash_key = self.get_hash_pos(server.id)