perf(admin-tokens): worse but faster selection of active users

This commit is contained in:
Gregor Kleen 2021-03-17 09:45:13 +01:00
parent 8df6143ced
commit f09f851e2b

View File

@ -8,7 +8,6 @@ import Handler.Utils
import qualified Data.HashSet as HashSet import qualified Data.HashSet as HashSet
import qualified Data.HashMap.Strict as HashMap import qualified Data.HashMap.Strict as HashMap
import qualified Data.Set as Set import qualified Data.Set as Set
import qualified Data.Map as Map
import qualified Data.Aeson as Aeson import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Encode.Pretty as Aeson import qualified Data.Aeson.Encode.Pretty as Aeson
@ -19,15 +18,12 @@ import qualified Data.Text as Text
import qualified Database.Esqueleto as E hiding (random_) import qualified Database.Esqueleto as E hiding (random_)
import qualified Database.Esqueleto.PostgreSQL as E import qualified Database.Esqueleto.PostgreSQL as E
import qualified Database.Esqueleto.Utils as E
import qualified Control.Monad.Random.Class as Random
import Control.Monad.Random.Strict (evalRand, Rand)
import qualified Data.Conduit.List as C (unfoldM)
import qualified Data.Conduit.Combinators as C import qualified Data.Conduit.Combinators as C
import qualified Crypto.Random as Crypto import Data.List (genericTake)
import System.Random.Shuffle (shuffleM)
data BTFImpersonate data BTFImpersonate
@ -136,22 +132,13 @@ postAdminTokensR = do
E.limit btfiCount E.limit btfiCount
return $ user E.^. UserId return $ user E.^. UserId
| otherwise -> do | otherwise -> do
users'' <- E.select . E.from $ \user -> users <- fmap (fmap E.unValue) . E.select . E.from $ \user -> do
return ( user E.^. UserId E.orderBy [ E.asc . E.isNothing $ user E.^. UserLastAuthentication
, E.maybe E.nothing (E.just . E.diffTimes (E.val now)) $ user E.^. UserLastAuthentication , E.desc $ user E.^. UserLastAuthentication
) ]
let users :: Map UserId (Maybe Rational) E.limit $ 2 * btfiCount
users = Map.fromList $ users'' <&> \(E.Value uid', E.Value mDiff) -> (uid', recip . toRational <$> mDiff) return $ user E.^. UserId
chooseUsers :: ConduitT () UserId (Rand Crypto.ChaChaDRG) () genericTake btfiCount <$> shuffleM users
chooseUsers = C.unfoldM chooseUsers' (users, btfiCount)
where chooseUsers' (users', n) = runMaybeT $ do
guard $ n > 0
let getWeighted = MaybeT . Random.weightedMay . mapMaybe (\(uid', mDiff) -> (uid', ) <$> mDiff) $ Map.toList users'
getUnweighted = MaybeT . Random.uniformMay $ Map.keysSet users'
user <- getWeighted <|> getUnweighted
return (user, (Map.delete user users', pred n))
drg <- liftIO Crypto.drgNew
return . flip evalRand drg . runConduit $ chooseUsers .| C.foldMap pure
let let
toTokenFile :: UserId -> DB (Either Void DBFile) toTokenFile :: UserId -> DB (Either Void DBFile)