More efficient token generation

Old code would create a new System.Random generator each time, which requires going to the system entropy store. New code caches an mwc-random gen at startup.
2014-10-20 10:34:38 +03:00 · 2014-10-20 10:34:38 +03:00 · 972e117005
commit 972e117005
parent 658f31e0d7
6 changed files with 55 additions and 20 deletions
--- a/yesod-core/Yesod/Core/Dispatch.hs
+++ b/yesod-core/Yesod/Core/Dispatch.hs
@ -68,6 +68,7 @@ import Control.Monad.Logger
 import Control.Monad (when)
 import qualified Paths_yesod_core
 import Data.Version (showVersion)
+import qualified System.Random.MWC as MWC

 -- | Convert the given argument into a WAI application, executable with any WAI
 -- handler. This function will provide no middlewares; if you want commonly
@ -76,10 +77,12 @@ toWaiAppPlain :: YesodDispatch site => site -> IO W.Application
 toWaiAppPlain site = do
    logger <- makeLogger site
    sb <- makeSessionBackend site
+    gen <- MWC.createSystemRandom
    return $ toWaiAppYre $ YesodRunnerEnv
            { yreLogger = logger
            , yreSite = site
            , yreSessionBackend = sb
+            , yreGen = gen
            }

 toWaiAppYre :: YesodDispatch site => YesodRunnerEnv site -> W.Application
@ -125,10 +128,12 @@ toWaiApp site = do
 toWaiAppLogger :: YesodDispatch site => Logger -> site -> IO W.Application
 toWaiAppLogger logger site = do
    sb <- makeSessionBackend site
+    gen <- MWC.createSystemRandom
    let yre = YesodRunnerEnv
                { yreLogger = logger
                , yreSite = site
                , yreSessionBackend = sb
+                , yreGen = gen
                }
    messageLoggerSource
        site
--- a/yesod-core/Yesod/Core/Internal/Request.hs
+++ b/yesod-core/Yesod/Core/Internal/Request.hs
@ -23,7 +23,6 @@ import Data.String (IsString)
 import Control.Arrow (second)
 import qualified Network.Wai.Parse as NWP
 import qualified Network.Wai as W
-import System.Random (RandomGen, randomRs)
 import Web.Cookie (parseCookiesText)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString.Char8 as S8
@ -33,7 +32,7 @@ import Data.Maybe (fromMaybe, catMaybes)
 import qualified Data.ByteString.Lazy as L
 import qualified Data.Set as Set
 import qualified Data.Text as T
-import Data.Text.Encoding (decodeUtf8With)
+import Data.Text.Encoding (decodeUtf8With, decodeUtf8)
 import Data.Text.Encoding.Error (lenientDecode)
 import Data.Conduit
 import Data.Conduit.List (sourceList)
@ -41,9 +40,16 @@ import Data.Conduit.Binary (sourceFile, sinkFile)
 import Data.Word (Word64)
 import Control.Monad.Trans.Resource (runResourceT, ResourceT)
 import Control.Exception (throwIO)
+import Control.Monad ((<=<), liftM)
 import Yesod.Core.Types
 import qualified Data.Map as Map
 import Data.IORef
+import qualified System.Random.MWC as MWC
+import Control.Monad.Primitive (PrimMonad, PrimState)
+import qualified Data.Vector.Storable as V
+import Data.Word (Word8)
+import Data.ByteString.Internal (ByteString (PS))
+import qualified Data.Word8 as Word8

 -- | Impose a limit on the size of the request body.
 limitRequestBody :: Word64 -> W.Request -> IO W.Request
@ -68,12 +74,11 @@ tooLargeResponse = W.responseLBS
    [("Content-Type", "text/plain")]
    "Request body too large to be processed."

-parseWaiRequest :: RandomGen g
-                => W.Request
+parseWaiRequest :: W.Request
                -> SessionMap
                -> Bool
                -> Maybe Word64 -- ^ max body size
-                -> (Either (IO YesodRequest) (g -> IO YesodRequest))
+                -> (Either (IO YesodRequest) (MWC.GenIO -> IO YesodRequest))
 parseWaiRequest env session useToken mmaxBodySize =
    -- In most cases, we won't need to generate any random values. Therefore,
    -- we split our results: if we need a random generator, return a Right
@ -81,7 +86,7 @@ parseWaiRequest env session useToken mmaxBodySize =
    -- acquisition.
    case etoken of
        Left token -> Left $ mkRequest token
-        Right mkToken -> Right $ mkRequest . mkToken
+        Right mkToken -> Right $ mkRequest <=< mkToken
  where
    mkRequest token' = do
        envLimited <- maybe return limitRequestBody mmaxBodySize env
@ -124,7 +129,7 @@ parseWaiRequest env session useToken mmaxBodySize =
                -- Already have a token, use it.
                Just bs -> Left $ Just $ decodeUtf8With lenientDecode bs
                -- Don't have a token, get a random generator and make a new one.
-                Nothing -> Right $ Just . pack . randomString 10
+                Nothing -> Right $ fmap Just . randomString 10
        | otherwise = Left Nothing

 textQueryString :: W.Request -> [(Text, Text)]
@ -153,13 +158,23 @@ addTwoLetters (toAdd, exist) (l:ls) =
 -- | Generate a random String of alphanumerical characters
 -- (a-z, A-Z, and 0-9) of the given length using the given
 -- random number generator.
-randomString :: RandomGen g => Int -> g -> String
-randomString len = take len . map toChar . randomRs (0, 61)
+randomString :: PrimMonad m => Int -> MWC.Gen (PrimState m) -> m Text
+randomString len gen =
+    liftM (decodeUtf8 . fromByteVector) $ V.replicateM len asciiChar
  where
-    toChar i
-        | i < 26 = toEnum $ i + fromEnum 'A'
-        | i < 52 = toEnum $ i + fromEnum 'a' - 26
-        | otherwise = toEnum $ i + fromEnum '0' - 52
+    asciiChar = liftM toAscii $ MWC.uniformR (0, 61) gen
+
+    toAscii i
+        | i < 26 = i + Word8._A
+        | i < 52 = i + Word8._a - 26
+        | otherwise = i + Word8._0 - 52
+
+fromByteVector :: V.Vector Word8 -> ByteString
+fromByteVector v =
+    PS fptr offset idx
+  where
+    (fptr, offset, idx) = V.unsafeToForeignPtr v
+{-# INLINE fromByteVector #-}

 mkFileInfoLBS :: Text -> Text -> L.ByteString -> FileInfo
 mkFileInfoLBS name ct lbs = FileInfo name ct (sourceList $ L.toChunks lbs) (\fp -> L.writeFile fp lbs)
--- a/yesod-core/Yesod/Core/Internal/Run.hs
+++ b/yesod-core/Yesod/Core/Internal/Run.hs
@ -256,10 +256,10 @@ yesodRunner handler' YesodRunnerEnv {..} route req sendResponse
    (session, saveSession) <- liftIO $ do
        maybe (return (Map.empty, dontSaveSession)) (\sb -> sbLoadSession sb req) yreSessionBackend
    let mkYesodReq = parseWaiRequest req session (isJust yreSessionBackend) mmaxLen
-    yreq <-
-        case mkYesodReq of
-            Left yreq -> return yreq
-            Right needGen -> liftIO $ needGen <$> newStdGen
+    let yreq =
+            case mkYesodReq of
+                Left yreq -> yreq
+                Right needGen -> needGen yreGen
    let ra = resolveApproot yreSite req
    let log' = messageLoggerSource yreSite yreLogger
        -- We set up two environments: the first one has a "safe" error handler
--- a/yesod-core/Yesod/Core/Types.hs
+++ b/yesod-core/Yesod/Core/Types.hs
@ -46,6 +46,7 @@ import           Network.Wai                        (FilePart,
 import qualified Network.Wai                        as W
 import qualified Network.Wai.Parse                  as NWP
 import           System.Log.FastLogger              (LogStr, LoggerSet, toLogStr, pushLogStr)
+import qualified System.Random.MWC                  as MWC
 import           Network.Wai.Logger                 (DateCacheGetter)
 import           Text.Blaze.Html                    (Html)
 import           Text.Hamlet                        (HtmlUrl)
@ -193,6 +194,7 @@ data YesodRunnerEnv site = YesodRunnerEnv
    { yreLogger         :: !Logger
    , yreSite           :: !site
    , yreSessionBackend :: !(Maybe SessionBackend)
+    , yreGen            :: !MWC.GenIO
    }

 data YesodSubRunnerEnv sub parent parentMonad = YesodSubRunnerEnv
--- a/yesod-core/test/YesodCoreTest/InternalRequest.hs
+++ b/yesod-core/test/YesodCoreTest/InternalRequest.hs
@ -13,6 +13,9 @@ import Data.Map (singleton)
 import Yesod.Core
 import Data.Word (Word64)
 import System.IO.Unsafe (unsafePerformIO)
+import qualified System.Random.MWC as MWC
+import Control.Monad.ST
+import Control.Monad (replicateM)

 randomStringSpecs :: Spec
 randomStringSpecs = describe "Yesod.Internal.Request.randomString" $ do
@ -22,14 +25,20 @@ randomStringSpecs = describe "Yesod.Internal.Request.randomString" $ do
 -- NOTE: this testcase may break on other systems/architectures if
 -- mkStdGen is not identical everywhere (is it?).
 looksRandom :: Bool
-looksRandom = randomString 20 (mkStdGen 0) == "VH9SkhtptqPs6GqtofVg"
+looksRandom = runST $ do
+    gen <- MWC.create
+    s <- randomString 20 gen
+    return $ s == "VH9SkhtptqPs6GqtofVg"

 noRepeat :: Int -> Int -> Bool
-noRepeat len n = length (nub $ map (randomString len . mkStdGen) [1..n]) == n
+noRepeat len n = runST $ do
+    gen <- MWC.create
+    ss <- replicateM n $ randomString len gen
+    return $ length (nub ss) == n


 -- For convenience instead of "(undefined :: StdGen)".
-g :: StdGen
+g :: MWC.GenIO
 g = error "test/YesodCoreTest/InternalRequest.g"

 parseWaiRequest' :: Request
--- a/yesod-core/yesod-core.cabal
+++ b/yesod-core/yesod-core.cabal
@ -65,6 +65,9 @@ library
                   , conduit-extra
                   , exceptions            >= 0.6
                   , deepseq
+                   , mwc-random
+                   , primitive
+                   , word8

    exposed-modules: Yesod.Core
                     Yesod.Core.Content
@ -154,6 +157,7 @@ test-suite tests
                  , shakespeare
                  , streaming-commons
                  , wai-extra
+                  , mwc-random
    ghc-options:     -Wall
    extensions: TemplateHaskell