{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module FoundationDB.Layer.Directory.Internal.HCA where

import Control.Concurrent.MVar
import Control.Monad
import Control.Monad.Catch (bracket_)
import Control.Monad.IO.Class (MonadIO (..))
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Serialize.Get (getWord64le, runGet)
import FoundationDB
import FoundationDB.Layer.Directory.Internal.Error
import FoundationDB.Layer.Subspace
import FoundationDB.Layer.Tuple
import FoundationDB.Options.MutationType (add)
import FoundationDB.Options.TransactionOption (nextWriteNoWriteConflictRange)
import System.IO.Unsafe (unsafePerformIO)
import System.Random (getStdRandom, randomR)

oneBytes :: ByteString
oneBytes :: ByteString
oneBytes = [Word8] -> ByteString
BS.pack [Word8
0x01, Word8
0x00, Word8
0x00, Word8
0x00, Word8
0x00, Word8
0x00, Word8
0x00, Word8
0x00]

data HCA = HCA
  { HCA -> Subspace
counters :: Subspace,
    HCA -> Subspace
recent :: Subspace
  }
  deriving (Int -> HCA -> ShowS
[HCA] -> ShowS
HCA -> String
(Int -> HCA -> ShowS)
-> (HCA -> String) -> ([HCA] -> ShowS) -> Show HCA
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [HCA] -> ShowS
$cshowList :: [HCA] -> ShowS
show :: HCA -> String
$cshow :: HCA -> String
showsPrec :: Int -> HCA -> ShowS
$cshowsPrec :: Int -> HCA -> ShowS
Show, HCA -> HCA -> Bool
(HCA -> HCA -> Bool) -> (HCA -> HCA -> Bool) -> Eq HCA
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: HCA -> HCA -> Bool
$c/= :: HCA -> HCA -> Bool
== :: HCA -> HCA -> Bool
$c== :: HCA -> HCA -> Bool
Eq, Eq HCA
Eq HCA
-> (HCA -> HCA -> Ordering)
-> (HCA -> HCA -> Bool)
-> (HCA -> HCA -> Bool)
-> (HCA -> HCA -> Bool)
-> (HCA -> HCA -> Bool)
-> (HCA -> HCA -> HCA)
-> (HCA -> HCA -> HCA)
-> Ord HCA
HCA -> HCA -> Bool
HCA -> HCA -> Ordering
HCA -> HCA -> HCA
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: HCA -> HCA -> HCA
$cmin :: HCA -> HCA -> HCA
max :: HCA -> HCA -> HCA
$cmax :: HCA -> HCA -> HCA
>= :: HCA -> HCA -> Bool
$c>= :: HCA -> HCA -> Bool
> :: HCA -> HCA -> Bool
$c> :: HCA -> HCA -> Bool
<= :: HCA -> HCA -> Bool
$c<= :: HCA -> HCA -> Bool
< :: HCA -> HCA -> Bool
$c< :: HCA -> HCA -> Bool
compare :: HCA -> HCA -> Ordering
$ccompare :: HCA -> HCA -> Ordering
$cp1Ord :: Eq HCA
Ord)

newHCA :: Subspace -> HCA
newHCA :: Subspace -> HCA
newHCA Subspace
s =
  HCA :: Subspace -> Subspace -> HCA
HCA
    { counters :: Subspace
counters = Subspace -> [Elem] -> Subspace
extend Subspace
s [Integer -> Elem
Int Integer
0],
      recent :: Subspace
recent = Subspace -> [Elem] -> Subspace
extend Subspace
s [Integer -> Elem
Int Integer
1]
    }

windowSize :: Int -> Int
windowSize :: Int -> Int
windowSize Int
start
  | Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
255 = Int
64
  | Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
65535 = Int
1024
  | Bool
otherwise = Int
8192

-- | global lock used for allocating. We use this simply because the other
-- clients have it. It appears to exist in order to reduce contention on the HCA
-- counter at the transaction level by pushing some of the contention to the
-- client level instead.
globalAllocateLock :: MVar ()
globalAllocateLock :: MVar ()
globalAllocateLock = IO (MVar ()) -> MVar ()
forall a. IO a -> a
unsafePerformIO (() -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ())
{-# NOINLINE globalAllocateLock #-}

withAllocLock :: Transaction a -> Transaction a
withAllocLock :: Transaction a -> Transaction a
withAllocLock =
  Transaction () -> Transaction () -> Transaction a -> Transaction a
forall (m :: * -> *) a c b. MonadMask m => m a -> m c -> m b -> m b
bracket_
    (IO () -> Transaction ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> Transaction ()) -> IO () -> Transaction ()
forall a b. (a -> b) -> a -> b
$ MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
globalAllocateLock)
    (IO () -> Transaction ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> Transaction ()) -> IO () -> Transaction ()
forall a b. (a -> b) -> a -> b
$ MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
globalAllocateLock ())

findStartAndWindow :: HCA -> Bool -> Int -> Transaction (Int, Int)
findStartAndWindow :: HCA -> Bool -> Int -> Transaction (Int, Int)
findStartAndWindow hca :: HCA
hca@HCA {Subspace
recent :: Subspace
counters :: Subspace
recent :: HCA -> Subspace
counters :: HCA -> Subspace
..} Bool
windowAdvanced Int
start = do
  Future (Maybe ByteString)
countFuture <- Transaction (Future (Maybe ByteString))
-> Transaction (Future (Maybe ByteString))
forall a. Transaction a -> Transaction a
withAllocLock Transaction (Future (Maybe ByteString))
getCount
  Maybe ByteString
countStr <- Future (Maybe ByteString) -> Transaction (Maybe ByteString)
forall a. Future a -> Transaction a
await Future (Maybe ByteString)
countFuture
  Int
count <- Maybe ByteString -> Transaction Int
forall a. Num a => Maybe ByteString -> Transaction a
parseCount Maybe ByteString
countStr
  let window :: Int
window = Int -> Int
windowSize Int
start
  if Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
window
    then (Int, Int) -> Transaction (Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
start, Int
window)
    else HCA -> Bool -> Int -> Transaction (Int, Int)
findStartAndWindow HCA
hca Bool
True (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
window)
  where
    getCount :: Transaction (Future (Maybe ByteString))
getCount = do
      let start' :: Integer
start' = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
start
      Bool -> Transaction () -> Transaction ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
windowAdvanced (Transaction () -> Transaction ())
-> Transaction () -> Transaction ()
forall a b. (a -> b) -> a -> b
$ do
        ByteString -> ByteString -> Transaction ()
clearRange (Subspace -> [Elem] -> ByteString
pack Subspace
counters []) (Subspace -> [Elem] -> ByteString
pack Subspace
counters [Integer -> Elem
Int Integer
start'])
        TransactionOption -> Transaction ()
setOption TransactionOption
nextWriteNoWriteConflictRange
        ByteString -> ByteString -> Transaction ()
clearRange (Subspace -> [Elem] -> ByteString
pack Subspace
recent []) (Subspace -> [Elem] -> ByteString
pack Subspace
recent [Integer -> Elem
Int Integer
start'])
      ByteString -> MutationType -> Transaction ()
atomicOp (Subspace -> [Elem] -> ByteString
pack Subspace
counters [Integer -> Elem
Int Integer
start']) (ByteString -> MutationType
add ByteString
oneBytes)
      Transaction (Future (Maybe ByteString))
-> Transaction (Future (Maybe ByteString))
forall a. Transaction a -> Transaction a
withSnapshot (Transaction (Future (Maybe ByteString))
 -> Transaction (Future (Maybe ByteString)))
-> Transaction (Future (Maybe ByteString))
-> Transaction (Future (Maybe ByteString))
forall a b. (a -> b) -> a -> b
$ ByteString -> Transaction (Future (Maybe ByteString))
get (Subspace -> [Elem] -> ByteString
pack Subspace
counters [Integer -> Elem
Int Integer
start'])

    parseCount :: Maybe ByteString -> Transaction a
parseCount Maybe ByteString
Nothing = a -> Transaction a
forall (m :: * -> *) a. Monad m => a -> m a
return a
0
    parseCount (Just ByteString
bs) =
      case Get Word64 -> ByteString -> Either String Word64
forall a. Get a -> ByteString -> Either String a
runGet Get Word64
getWord64le ByteString
bs of
        Left String
_ -> String -> Transaction a
forall a. String -> Transaction a
throwDirInternalError (String -> Transaction a) -> String -> Transaction a
forall a b. (a -> b) -> a -> b
$ String
"failed to parse count: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ByteString -> String
forall a. Show a => a -> String
show ByteString
bs
        Right Word64
n -> a -> Transaction a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Transaction a) -> a -> Transaction a
forall a b. (a -> b) -> a -> b
$ Word64 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n

findSubspaceLoop ::
  HCA ->
  Subspace ->
  Int ->
  Int ->
  Transaction (Maybe Subspace)
findSubspaceLoop :: HCA -> Subspace -> Int -> Int -> Transaction (Maybe Subspace)
findSubspaceLoop hca :: HCA
hca@HCA {Subspace
recent :: Subspace
counters :: Subspace
recent :: HCA -> Subspace
counters :: HCA -> Subspace
..} Subspace
s Int
start Int
window = do
  Int
candidate <- IO Int -> Transaction Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> Transaction Int) -> IO Int -> Transaction Int
forall a b. (a -> b) -> a -> b
$ (StdGen -> (Int, StdGen)) -> IO Int
forall a. (StdGen -> (a, StdGen)) -> IO a
getStdRandom ((Int, Int) -> StdGen -> (Int, StdGen)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
start, Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
window))
  let key :: ByteString
key = Subspace -> [Elem] -> ByteString
pack Subspace
recent [Integer -> Elem
Int (Integer -> Elem) -> Integer -> Elem
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
candidate]
  (Maybe (ByteString, ByteString)
latestCounter, Future (Maybe ByteString)
candidateValueF) <- Transaction
  (Maybe (ByteString, ByteString), Future (Maybe ByteString))
-> Transaction
     (Maybe (ByteString, ByteString), Future (Maybe ByteString))
forall a. Transaction a -> Transaction a
withAllocLock (Transaction
   (Maybe (ByteString, ByteString), Future (Maybe ByteString))
 -> Transaction
      (Maybe (ByteString, ByteString), Future (Maybe ByteString)))
-> Transaction
     (Maybe (ByteString, ByteString), Future (Maybe ByteString))
-> Transaction
     (Maybe (ByteString, ByteString), Future (Maybe ByteString))
forall a b. (a -> b) -> a -> b
$ do
    Maybe (ByteString, ByteString)
latestCounter <- Transaction (Maybe (ByteString, ByteString))
-> Transaction (Maybe (ByteString, ByteString))
forall a. Transaction a -> Transaction a
withSnapshot (Transaction (Maybe (ByteString, ByteString))
 -> Transaction (Maybe (ByteString, ByteString)))
-> Transaction (Maybe (ByteString, ByteString))
-> Transaction (Maybe (ByteString, ByteString))
forall a b. (a -> b) -> a -> b
$ Subspace -> Transaction (Maybe (ByteString, ByteString))
getLast Subspace
counters
    Future (Maybe ByteString)
candidateValueF <- ByteString -> Transaction (Future (Maybe ByteString))
get ByteString
key
    TransactionOption -> Transaction ()
setOption TransactionOption
nextWriteNoWriteConflictRange
    ByteString -> ByteString -> Transaction ()
set ByteString
key ByteString
""
    (Maybe (ByteString, ByteString), Future (Maybe ByteString))
-> Transaction
     (Maybe (ByteString, ByteString), Future (Maybe ByteString))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (ByteString, ByteString)
latestCounter, Future (Maybe ByteString)
candidateValueF)
  Integer
currentStart <- case Maybe (ByteString, ByteString)
latestCounter of
    Just (ByteString
k, ByteString
_) -> case Subspace -> ByteString -> Either String [Elem]
unpack Subspace
counters ByteString
k of
      Right (Int Integer
x : [Elem]
_) -> Integer -> Transaction Integer
forall (m :: * -> *) a. Monad m => a -> m a
return Integer
x
      Either String [Elem]
_ -> String -> Transaction Integer
forall a. String -> Transaction a
throwDirInternalError (String -> Transaction Integer) -> String -> Transaction Integer
forall a b. (a -> b) -> a -> b
$ String
"bad counter format: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ByteString -> String
forall a. Show a => a -> String
show ByteString
k
    Maybe (ByteString, ByteString)
_ -> String -> Transaction Integer
forall a. String -> Transaction a
throwDirInternalError String
"failed to find latestCounter"
  if Integer
currentStart Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
start
    then Maybe Subspace -> Transaction (Maybe Subspace)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Subspace
forall a. Maybe a
Nothing
    else
      Future (Maybe ByteString) -> Transaction (Maybe ByteString)
forall a. Future a -> Transaction a
await Future (Maybe ByteString)
candidateValueF Transaction (Maybe ByteString)
-> (Maybe ByteString -> Transaction (Maybe Subspace))
-> Transaction (Maybe Subspace)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just ByteString
_ -> HCA -> Subspace -> Int -> Int -> Transaction (Maybe Subspace)
findSubspaceLoop HCA
hca Subspace
s Int
start Int
window
        Maybe ByteString
Nothing -> do
          ByteString -> ByteString -> FDBConflictRangeType -> Transaction ()
addConflictRange ByteString
key (ByteString
key ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"0x00") FDBConflictRangeType
ConflictRangeTypeWrite
          Maybe Subspace -> Transaction (Maybe Subspace)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Subspace -> Transaction (Maybe Subspace))
-> Maybe Subspace -> Transaction (Maybe Subspace)
forall a b. (a -> b) -> a -> b
$ Subspace -> Maybe Subspace
forall a. a -> Maybe a
Just (Subspace -> Maybe Subspace) -> Subspace -> Maybe Subspace
forall a b. (a -> b) -> a -> b
$ Subspace -> [Elem] -> Subspace
extend Subspace
s [Integer -> Elem
Int (Integer -> Elem) -> Integer -> Elem
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
candidate]

initStart :: HCA -> Transaction Int
initStart :: HCA -> Transaction Int
initStart HCA {Subspace
recent :: Subspace
counters :: Subspace
recent :: HCA -> Subspace
counters :: HCA -> Subspace
..} = do
  Maybe (ByteString, ByteString)
mkv <- Transaction (Maybe (ByteString, ByteString))
-> Transaction (Maybe (ByteString, ByteString))
forall a. Transaction a -> Transaction a
withSnapshot (Transaction (Maybe (ByteString, ByteString))
 -> Transaction (Maybe (ByteString, ByteString)))
-> Transaction (Maybe (ByteString, ByteString))
-> Transaction (Maybe (ByteString, ByteString))
forall a b. (a -> b) -> a -> b
$ Subspace -> Transaction (Maybe (ByteString, ByteString))
getLast Subspace
counters
  case Maybe (ByteString, ByteString)
mkv of
    Just (ByteString
k, ByteString
_) -> case Subspace -> ByteString -> Either String [Elem]
unpack Subspace
counters ByteString
k of
      Right (Int Integer
start : [Elem]
_) -> Int -> Transaction Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Transaction Int) -> Int -> Transaction Int
forall a b. (a -> b) -> a -> b
$ Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
start
      Either String [Elem]
_ -> Int -> Transaction Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
    Maybe (ByteString, ByteString)
Nothing -> Int -> Transaction Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0

allocate :: HCA -> Subspace -> Transaction Subspace
allocate :: HCA -> Subspace -> Transaction Subspace
allocate HCA
hca Subspace
s = do
  Int
start <- HCA -> Transaction Int
initStart HCA
hca
  (Int
start', Int
window) <- HCA -> Bool -> Int -> Transaction (Int, Int)
findStartAndWindow HCA
hca Bool
False Int
start
  Maybe Subspace
msub <- HCA -> Subspace -> Int -> Int -> Transaction (Maybe Subspace)
findSubspaceLoop HCA
hca Subspace
s Int
start' Int
window
  case Maybe Subspace
msub of
    Just Subspace
sub -> Subspace -> Transaction Subspace
forall (m :: * -> *) a. Monad m => a -> m a
return Subspace
sub
    Maybe Subspace
Nothing -> HCA -> Subspace -> Transaction Subspace
allocate HCA
hca Subspace
s