-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE BangPatterns #-}
-- | Bucketing requests by 'DataSource'.
--
-- When a request is issued by the client via 'dataFetch', it is placed
-- in the 'RequestStore'. When we are ready to fetch the current batch
-- of requests, the 'contents' operation extracts the fetches, bucketed
-- by 'DataSource'.
--
-- This module is provided for access to Haxl internals only; most
-- users should not need to import it.
--
module Haxl.Core.RequestStore
  ( BlockedFetches(..)
  , BlockedFetchInternal(..)
  , RequestStore
  , isEmpty
  , noRequests
  , addRequest
  , contents
  , getSize
  , ReqCountMap(..)
  , emptyReqCounts
  , filterRCMap
  , getMapFromRCMap
  , getSummaryMapFromRCMap
  , addToCountMap
  , subFromCountMap
  ) where

import Haxl.Core.DataSource
import Haxl.Core.Stats
import Data.Map (Map)
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Map.Strict as Map
import Data.Proxy
import Data.Text (Text)
import Data.Kind (Type)
import Data.Typeable
import Unsafe.Coerce

-- | A container for multiple 'BlockedFetch' objects.
newtype RequestStore u = RequestStore (Map TypeRep (BlockedFetches u))
  -- Since we don't know which data sources we will be using, the store
  -- is dynamically-typed.  It maps the TypeRep of the request to the
  -- 'BlockedFetches' for that 'DataSource'.

newtype BlockedFetchInternal = BlockedFetchInternal CallId

-- | A batch of 'BlockedFetch' objects for a single 'DataSource'
data BlockedFetches u =
  forall r. (DataSource u r) =>
        BlockedFetches [BlockedFetch r] [BlockedFetchInternal]

isEmpty :: RequestStore u -> Bool
isEmpty :: RequestStore u -> Bool
isEmpty (RequestStore Map TypeRep (BlockedFetches u)
m) = Map TypeRep (BlockedFetches u) -> Bool
forall k a. Map k a -> Bool
Map.null Map TypeRep (BlockedFetches u)
m

-- | A new empty 'RequestStore'.
noRequests :: RequestStore u
noRequests :: RequestStore u
noRequests = Map TypeRep (BlockedFetches u) -> RequestStore u
forall u. Map TypeRep (BlockedFetches u) -> RequestStore u
RequestStore Map TypeRep (BlockedFetches u)
forall k a. Map k a
Map.empty

-- | Adds a 'BlockedFetch' to a 'RequestStore'.
addRequest
  :: forall u r. (DataSource u r)
  => BlockedFetch r -> BlockedFetchInternal -> RequestStore u -> RequestStore u
addRequest :: BlockedFetch r
-> BlockedFetchInternal -> RequestStore u -> RequestStore u
addRequest BlockedFetch r
bf BlockedFetchInternal
bfi (RequestStore Map TypeRep (BlockedFetches u)
m) =
  Map TypeRep (BlockedFetches u) -> RequestStore u
forall u. Map TypeRep (BlockedFetches u) -> RequestStore u
RequestStore (Map TypeRep (BlockedFetches u) -> RequestStore u)
-> Map TypeRep (BlockedFetches u) -> RequestStore u
forall a b. (a -> b) -> a -> b
$ (BlockedFetches u -> BlockedFetches u -> BlockedFetches u)
-> TypeRep
-> BlockedFetches u
-> Map TypeRep (BlockedFetches u)
-> Map TypeRep (BlockedFetches u)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith BlockedFetches u -> BlockedFetches u -> BlockedFetches u
combine TypeRep
ty ([BlockedFetch r] -> [BlockedFetchInternal] -> BlockedFetches u
forall u (r :: * -> *).
DataSource u r =>
[BlockedFetch r] -> [BlockedFetchInternal] -> BlockedFetches u
BlockedFetches [BlockedFetch r
bf] [BlockedFetchInternal
bfi]) Map TypeRep (BlockedFetches u)
m
 where
  combine :: BlockedFetches u -> BlockedFetches u -> BlockedFetches u
  combine :: BlockedFetches u -> BlockedFetches u -> BlockedFetches u
combine BlockedFetches u
_ (BlockedFetches [BlockedFetch r]
bfs [BlockedFetchInternal]
bfis)
    | r Any -> TypeRep
forall (t :: * -> *) a. Typeable t => t a -> TypeRep
typeOf1 ([BlockedFetch r] -> r Any
forall (r1 :: * -> *) a. [BlockedFetch r1] -> r1 a
getR [BlockedFetch r]
bfs) TypeRep -> TypeRep -> Bool
forall a. Eq a => a -> a -> Bool
== TypeRep
ty = [BlockedFetch r] -> [BlockedFetchInternal] -> BlockedFetches u
forall u (r :: * -> *).
DataSource u r =>
[BlockedFetch r] -> [BlockedFetchInternal] -> BlockedFetches u
BlockedFetches (BlockedFetch r -> BlockedFetch r
forall a b. a -> b
unsafeCoerce BlockedFetch r
bfBlockedFetch r -> [BlockedFetch r] -> [BlockedFetch r]
forall a. a -> [a] -> [a]
:[BlockedFetch r]
bfs) (BlockedFetchInternal
bfiBlockedFetchInternal
-> [BlockedFetchInternal] -> [BlockedFetchInternal]
forall a. a -> [a] -> [a]
:[BlockedFetchInternal]
bfis)
    | Bool
otherwise                = [Char] -> BlockedFetches u
forall a. HasCallStack => [Char] -> a
error [Char]
"RequestStore.insert"
         -- the dynamic type check here should be unnecessary, but if
         -- there are bugs in `Typeable` or `Map` then we'll get an
         -- error instead of a crash.  The overhead is negligible.

  -- a type conversion only, so we can get the type of the reqeusts from
  -- the list of BlockedFetch.
  getR :: [BlockedFetch r1] -> r1 a
  getR :: [BlockedFetch r1] -> r1 a
getR [BlockedFetch r1]
_ = r1 a
forall a. HasCallStack => a
undefined

  -- The TypeRep of requests for this data source
  ty :: TypeRep
  !ty :: TypeRep
ty = r Any -> TypeRep
forall (t :: * -> *) a. Typeable t => t a -> TypeRep
typeOf1 (forall a. r a
forall a. HasCallStack => a
undefined :: r a)

-- | Retrieves the whole contents of the 'RequestStore'.
contents :: RequestStore u -> [BlockedFetches u]
contents :: RequestStore u -> [BlockedFetches u]
contents (RequestStore Map TypeRep (BlockedFetches u)
m) = Map TypeRep (BlockedFetches u) -> [BlockedFetches u]
forall k a. Map k a -> [a]
Map.elems Map TypeRep (BlockedFetches u)
m

getSize :: RequestStore u -> Int
getSize :: RequestStore u -> Int
getSize (RequestStore Map TypeRep (BlockedFetches u)
m) = Map TypeRep (BlockedFetches u) -> Int
forall k a. Map k a -> Int
Map.size Map TypeRep (BlockedFetches u)
m

-- A counter to keep track of outgone requests. Entries are added to this
-- map as we send requests to datasources, and removed as these fetches
-- are completed.
-- This is a 2 level map: the 1st level stores requests for a particular
-- datasource, the 2nd level stores count of requests per type.
newtype ReqCountMap = ReqCountMap (Map Text (Map TypeRep Int))
  deriving (Int -> ReqCountMap -> ShowS
[ReqCountMap] -> ShowS
ReqCountMap -> [Char]
(Int -> ReqCountMap -> ShowS)
-> (ReqCountMap -> [Char])
-> ([ReqCountMap] -> ShowS)
-> Show ReqCountMap
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [ReqCountMap] -> ShowS
$cshowList :: [ReqCountMap] -> ShowS
show :: ReqCountMap -> [Char]
$cshow :: ReqCountMap -> [Char]
showsPrec :: Int -> ReqCountMap -> ShowS
$cshowsPrec :: Int -> ReqCountMap -> ShowS
Show)

emptyReqCounts :: ReqCountMap
emptyReqCounts :: ReqCountMap
emptyReqCounts = Map Text (Map TypeRep Int) -> ReqCountMap
ReqCountMap Map Text (Map TypeRep Int)
forall k a. Map k a
Map.empty

addToCountMap
  :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
  => Proxy r
  -> Int -- type and number of requests
  -> ReqCountMap
  -> ReqCountMap
addToCountMap :: Proxy r -> Int -> ReqCountMap -> ReqCountMap
addToCountMap = (Int -> Int -> Int) -> Proxy r -> Int -> ReqCountMap -> ReqCountMap
forall (r :: * -> *).
(DataSourceName r, Typeable r) =>
(Int -> Int -> Int) -> Proxy r -> Int -> ReqCountMap -> ReqCountMap
updateCountMap Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)

subFromCountMap
  :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
  => Proxy r
  -> Int -- type and number of requests
  -> ReqCountMap
  -> ReqCountMap
subFromCountMap :: Proxy r -> Int -> ReqCountMap -> ReqCountMap
subFromCountMap = (Int -> Int -> Int) -> Proxy r -> Int -> ReqCountMap -> ReqCountMap
forall (r :: * -> *).
(DataSourceName r, Typeable r) =>
(Int -> Int -> Int) -> Proxy r -> Int -> ReqCountMap -> ReqCountMap
updateCountMap (-)

updateCountMap
  :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
  => (Int -> Int -> Int)
  -> Proxy r
  -> Int -- type and number of requests
  -> ReqCountMap
  -> ReqCountMap
updateCountMap :: (Int -> Int -> Int) -> Proxy r -> Int -> ReqCountMap -> ReqCountMap
updateCountMap Int -> Int -> Int
op Proxy r
p Int
n (ReqCountMap Map Text (Map TypeRep Int)
m) = Map Text (Map TypeRep Int) -> ReqCountMap
ReqCountMap (Map Text (Map TypeRep Int) -> ReqCountMap)
-> Map Text (Map TypeRep Int) -> ReqCountMap
forall a b. (a -> b) -> a -> b
$ (Map TypeRep Int -> Map TypeRep Int -> Map TypeRep Int)
-> Text
-> Map TypeRep Int
-> Map Text (Map TypeRep Int)
-> Map Text (Map TypeRep Int)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith
  ((Map TypeRep Int -> Map TypeRep Int -> Map TypeRep Int)
-> Map TypeRep Int -> Map TypeRep Int -> Map TypeRep Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> Int -> Int)
-> Map TypeRep Int -> Map TypeRep Int -> Map TypeRep Int
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Int -> Int -> Int
op)) -- flip is important as "op" is not commutative
  (Proxy r -> Text
forall (req :: * -> *). DataSourceName req => Proxy req -> Text
dataSourceName Proxy r
p) (TypeRep -> Int -> Map TypeRep Int
forall k a. k -> a -> Map k a
Map.singleton TypeRep
ty Int
n)
  Map Text (Map TypeRep Int)
m
  where
    -- The TypeRep of requests for this data source
    -- The way this is implemented, all elements in the 2nd level map will be
    -- mapped to the same key, as all requests to a datasource have the same
    -- "type". It will be more beneficial to be able to instead map requests
    -- to their names (ie, data constructor) - but there's no cheap way of doing
    -- that.
    ty :: TypeRep
    !ty :: TypeRep
ty = r Any -> TypeRep
forall (t :: * -> *) a. Typeable t => t a -> TypeRep
typeOf1 (forall a. r a
forall a. HasCallStack => a
undefined :: r a)

-- Filter all keys with 0 fetches. Since ReqCountMap is a 2-level map, we need
-- nested filter operations.
filterRCMap :: ReqCountMap -> ReqCountMap
filterRCMap :: ReqCountMap -> ReqCountMap
filterRCMap (ReqCountMap Map Text (Map TypeRep Int)
m) = Map Text (Map TypeRep Int) -> ReqCountMap
ReqCountMap (Map Text (Map TypeRep Int) -> ReqCountMap)
-> Map Text (Map TypeRep Int) -> ReqCountMap
forall a b. (a -> b) -> a -> b
$
  (Map TypeRep Int -> Bool)
-> Map Text (Map TypeRep Int) -> Map Text (Map TypeRep Int)
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Int -> Bool)
-> (Map TypeRep Int -> Int) -> Map TypeRep Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map TypeRep Int -> Int
forall k a. Map k a -> Int
Map.size) ((Int -> Bool) -> Map TypeRep Int -> Map TypeRep Int
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Map TypeRep Int -> Map TypeRep Int)
-> Map Text (Map TypeRep Int) -> Map Text (Map TypeRep Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Text (Map TypeRep Int)
m)

-- Filters the ReqCountMap by default
getMapFromRCMap :: ReqCountMap -> Map Text (Map TypeRep Int)
getMapFromRCMap :: ReqCountMap -> Map Text (Map TypeRep Int)
getMapFromRCMap ReqCountMap
r
  | ReqCountMap Map Text (Map TypeRep Int)
m <- ReqCountMap -> ReqCountMap
filterRCMap ReqCountMap
r = Map Text (Map TypeRep Int)
m

getSummaryMapFromRCMap :: ReqCountMap -> HashMap.HashMap Text Int
getSummaryMapFromRCMap :: ReqCountMap -> HashMap Text Int
getSummaryMapFromRCMap (ReqCountMap Map Text (Map TypeRep Int)
m) = [(Text, Int)] -> HashMap Text Int
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList
  [ (Text
k, Int
s)
  | (Text
k, Map TypeRep Int
v) <- Map Text (Map TypeRep Int) -> [(Text, Map TypeRep Int)]
forall k a. Map k a -> [(k, a)]
Map.toList Map Text (Map TypeRep Int)
m
  , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Map TypeRep Int -> Bool
forall k a. Map k a -> Bool
Map.null Map TypeRep Int
v
  , let s :: Int
s = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Map TypeRep Int -> [Int]
forall k a. Map k a -> [a]
Map.elems Map TypeRep Int
v
  , Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
  ]