A (bad) version of snowflake in (bad) haskell

snowflake.hs

import Data.Binary
import Data.Bits
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Int
import Network hiding (accept)
import Network.Socket
import Network.Socket.ByteString (sendAll)
import System.Console.GetOpt
import System.Environment
import System.IO
import System.Time

data Opts = Opts
  { workerBits :: Int
  , dcBits :: Int
  , seqBits :: Int
  , workerId :: Int64
  , dcId :: Int64
  , port :: Integer
  } deriving Show

defaults = Opts 5 5 12 1 1 5002
options :: [OptDescr (Opts -> Opts)]
options =
  [ Option [] ["workerBits"] (ReqArg (\v opts -> opts { workerBits = read v :: Int }) "bits")        "bits to use for worker value"
  , Option [] ["dcBits"]     (ReqArg (\v opts -> opts { dcBits = read v :: Int }) "bits")            "bits to use for DC value"
  , Option [] ["seqBits"]    (ReqArg (\v opts -> opts { seqBits = read v :: Int }) "bits")           "bits to use for the sequence value"
  , Option [] ["worker"]     (ReqArg (\v opts -> opts { workerId = read v :: Int64 }) "id")          "Worker ID"
  , Option [] ["dc"]         (ReqArg (\v opts -> opts { dcId = read v :: Int64 }) "id")              "DataCenter ID"
  , Option [] ["port"]       (ReqArg (\v opts -> opts { port = read v :: Integer }) "port") "Port to listen on"
  ]

data SnowflakeId = SnowflakeId
  { lastTs :: Int64
  , tsShift :: Int
  , dcStamp :: Int64
  , workerStamp :: Int64
  , seqNum :: Int64
  , seqMask :: Int64
  }

main = do
    (Opts wb dcb sb wid dcid port, _) <- getArgs >>= getOpts
    withSocketsDo $ do
      sock <- listenOn $ PortNumber (fromInteger port)
      loop sock (newId wb dcb sb dcid wid)

getOpts :: [String] -> IO (Opts, [String])
getOpts argv =
  case getOpt Permute options argv of
    (o, n, []  ) -> return (foldl (flip id) defaults o, n)
    (_, _, errs) -> ioError (userError (concat errs ++ usageInfo header options))
  where header = "Usage: snowflake [OPTION...]"

loop sock id = do
    (conn, _) <- accept sock
    next <- nextId id
    lastId <- case next of
      Nothing ->
        return id
      Just next -> do
        sendAll conn (idToBytes next)
        return next
    sClose conn
    loop sock lastId

newId :: Int -> Int -> Int -> Int64 -> Int64 -> SnowflakeId
newId dcBits workerBits seqBits dc worker =
  let seqMask = xor (complement 0) (shiftL (complement 0) seqBits) :: Int64
      workerStamp = shiftL worker seqBits
      dcStamp = shiftL dc (workerBits + seqBits)
      tsShift = dcBits + workerBits + seqBits
  in SnowflakeId 0 tsShift dcStamp workerStamp 0 seqMask

nextId :: SnowflakeId -> IO (Maybe SnowflakeId)
nextId id = curTime >>= nextId' id

nextId' :: SnowflakeId -> Int64 -> IO (Maybe SnowflakeId)
nextId' id time
    | time < lastTime = return Nothing
    | time > lastTime = return $ Just id { lastTs = time, seqNum = 0 }
    | time == lastTime = (nextSeq id time) >>= return . Just
  where lastTime = lastTs id

nextSeq :: SnowflakeId -> Int64 -> IO SnowflakeId
nextSeq id time =
    if newSeq == 0
      then (tillNextMillis time) >>= \x -> return id { seqNum = 0, lastTs = x }
      else return id { seqNum = newSeq }
  where newSeq = (seqNum id) + 1 .&. (seqMask id)

idToBytes :: SnowflakeId -> B.ByteString
idToBytes (SnowflakeId ts tsShift dc worker seq _) = B.concat . BL.toChunks $ encode $
    (shiftL ts tsShift) .|. dc .|. worker .|. seq

curTime :: IO Int64
curTime = do
    (TOD s ps) <- getClockTime
    let sec = fromIntegral s :: Int64
    let psec = fromIntegral ps :: Int64
    return $ (sec * 1000) + (psec `div` 1000000000)

tillNextMillis :: Int64 -> IO Int64
tillNextMillis last = do
  next <- curTime
  if next <= last
    then tillNextMillis next
    else return next