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.

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

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
+parseWaiRequest :: W.Request
-                => 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 :: PrimMonad m => Int -> MWC.Gen (PrimState m) -> m Text
-randomString len = take len . map toChar . randomRs (0, 61)
+randomString len gen =
+    liftM (decodeUtf8 . fromByteVector) $ V.replicateM len asciiChar
   where
-    toChar i
+    asciiChar = liftM toAscii $ MWC.uniformR (0, 61) gen
-        | i < 26 = toEnum $ i + fromEnum 'A'
+
-        | i < 52 = toEnum $ i + fromEnum 'a' - 26
+    toAscii i
-        | otherwise = toEnum $ i + fromEnum '0' - 52
+        | 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)

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 <-
+    let yreq =
-        case mkYesodReq of
+            case mkYesodReq of
-            Left yreq -> return yreq
+                Left yreq -> yreq
-            Right needGen -> liftIO $ needGen <$> newStdGen
+                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

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

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

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