consistent_hash.hs

module Database.Redis.ConsistentHash where

import Database.Redis.Protocol
import Database.Redis.Key
import Database.Redis.Serializable

import Data.List (sort,delete,(\\),findIndex,(!!))
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Digest.Pure.SHA
import Data.Maybe
import Data.ByteString.Lazy.Char8 (pack)

import Control.Applicative 

type Ring = Map Int Server

data HashRing = HashRing Int Ring [Server] [Int]

emptyRing :: Int -> HashRing
emptyRing i = HashRing i Map.empty [] []

create :: [Server] -> Int -> HashRing
create nodes replicas = foldl addServer (emptyRing replicas) nodes

addServer :: HashRing -> Server -> HashRing
addServer (HashRing rep ring nodes skeys) node = HashRing rep newRing newServers (sort newKeys) where
    (newRing,newKeys) = foldl (addEntry node) (ring,skeys) [0 .. rep]
    newServers = node : nodes

addEntry :: Server -> (Ring,[Int]) -> Int -> (Ring,[Int])
addEntry n (r,ks) i = (nr, key : ks) where
    key = hashServerId n i
    nr = Map.insert key n r

hashServerId :: Server -> Int -> Int
hashServerId node i = hashString h where
    h = "#{" ++ show node ++ "}:#{" ++ show i ++ "}"

hashString :: String -> Int
hashString s = (fromInteger (integerDigest (sha1 (pack s))) :: Int)

-- Removing nodes will not change the sort order so no sort necessary on newKeys
removeServer :: HashRing -> Server -> HashRing
removeServer (HashRing rep ring nodes skeys) node = HashRing rep newRing newServers newKeys where
    (newRing,newKeys) = foldl (removeEntry node) (ring,skeys) [0 .. rep]
    newServers = delete node nodes
                                                  
removeEntry :: Server -> (Ring,[Int]) -> Int -> (Ring,[Int]) 
removeEntry n (r,ks) i = (nr, ks \\ [key]) where
    key = hashServerId n i
    nr = Map.delete key r

getServer :: HashRing -> String -> Maybe Server
getServer hashRing k | isNothing x = Nothing
                     | otherwise = Just (fst (fromJust x))
getServer hashRing k = fst <$> getServerPos hashRing (hashString k)
                  
getServerPos :: HashRing -> Int -> Maybe (Server,Int)
getServerPos (HashRing _ ring _ skeys) key | Map.null ring = Nothing
                                         | otherwise = Just (x,y)
    where
      y = fromMaybe 0 (findIndex (>= key) skeys)
      x = ring Map.! (skeys !! y)