Snowchecked.hs

{-# LANGUAGE BlockArguments  #-}
{-# LANGUAGE RecordWildCards #-}
{-|
Description : Unique id generator derived from Twitter's Snowflake.
License     : Apache 2.0
Maintainer  : smokejumperit@gmail.com
Stability   : experimental

This generates unique (guaranteed) identifiers build from a timestamp,
counter, and node id.  Identifiers are convertible to values which are
monotonically increasing with respect to time.
-}
module Data.Snowchecked
( newSnowcheckedGen
, nextFlake
, SnowcheckedConfig(..)
, SnowcheckedGen
, Flake
, snowcheckedConfigBitCount
, uniqueFlakeCount
) where

import           Control.Concurrent               (threadDelay)
import           Control.Concurrent.MVar
import           Control.Monad.IO.Class           (MonadIO, liftIO)
import           Data.Snowchecked.Internal.Import
import           Data.Time.Clock.POSIX            (getPOSIXTime)
import GHC.Stack (HasCallStack)


currentTimestamp :: IO Word256
currentTimestamp = toMillisWord256 <$> getPOSIXTime
	where
		toMillisWord256 = round . (*1000)
{-# INLINE currentTimestamp #-}

currentTimestampBits :: Word8 -> IO Word256
currentTimestampBits n = (`cutBits` toInt n) <$> currentTimestamp
{-# INLINE currentTimestampBits #-}

-- | Create a new generator. Takes a configuration and node id.  The node id may be any
-- value that fits in a 'Word256', but it will be truncated to the number of bits specified
-- in the provided configuration.
newSnowcheckedGen :: (HasCallStack, MonadIO io) => SnowcheckedConfig -> Word256 -> io SnowcheckedGen
newSnowcheckedGen conf@SnowcheckedConfig{..} nodeId = liftIO $ do
	startTimeBits <- currentTimestampBits confTimeBits
	SnowcheckedGen <$> newMVar Flake
		{ flakeTime = startTimeBits
		, flakeCount = maxBound
		, flakeNodeId = cutBits nodeId (toInt confNodeBits)
		, flakeConfig = conf
		}
{-# INLINEABLE newSnowcheckedGen #-}
{-# SPECIALIZE newSnowcheckedGen :: SnowcheckedConfig -> Word256 -> IO SnowcheckedGen #-}

-- | Calculates the number of bits in each 'Flake' generated using a given configuration.
-- It returns a 'Word32' because there are 4 fields and the bitlength of each field fits
-- in a 'Word8', so the total bit count must fit within a 'Word32'.
snowcheckedConfigBitCount :: SnowcheckedConfig -> Word32
snowcheckedConfigBitCount SnowcheckedConfig{..} = foldr foldFunc 0
		[ confTimeBits
		, confCountBits
		, confNodeBits
		, confCheckBits
		]
	where
		foldFunc :: Word8 -> Word32 -> Word32
		foldFunc nxt memo = memo + toWord32 nxt
{-# INLINEABLE snowcheckedConfigBitCount #-}

-- | Generates the next id.
nextFlake :: (HasCallStack, MonadIO io) => SnowcheckedGen -> io Flake
nextFlake SnowcheckedGen{..} = liftIO $ modifyMVar genLastFlake mkNextFlake
	where
		-- TODO: Special case when confTimeBits is 0.
		-- TODO: Track the number of flakes generated and error out if we've exhausted them.
		mkNextFlake flake@Flake{..} =
			let SnowcheckedConfig{..} = flakeConfig in
			currentTimestampBits confTimeBits >>= \currentTimeBits ->
				if flakeTime < currentTimeBits then
					let newFlake = flake
						{ flakeTime = currentTimeBits
						, flakeCount = 0
						}
					in return (newFlake, newFlake)
				else if confCountBits == 0 then
					threadDelay 1000 >> mkNextFlake flake
				else
					let nextCount = cutBits (flakeCount + 1) (toInt confCountBits) in
					if flakeCount < nextCount then
						let newFlake = flake { flakeCount = nextCount }
						in return (newFlake, newFlake)
					else
						-- The count wrapped and we need to wait for the time to change.
						-- This assumes that the next millisecond will give us a new time.
						threadDelay 1000 >> mkNextFlake flake
{-# INLINEABLE nextFlake #-}
{-# SPECIALIZE nextFlake :: SnowcheckedGen -> IO Flake #-}


-- | Provides the count of total number of unique flakes possibly generated by a node using
-- this configuration.
uniqueFlakeCount :: SnowcheckedConfig -> Integer
uniqueFlakeCount SnowcheckedConfig{..}
	| confCountBits == 0 = 2^confTimeBits
	| confTimeBits  == 0 = 2^confCountBits
	| otherwise          = 2^confCountBits * 2^confTimeBits
{-# INLINE uniqueFlakeCount #-}