content: module KVStore.Client (runKeyValueClient) where
  
  import Control.Concurrent
  import Control.Exception
  import Control.Monad
  import Data.ByteString (ByteString)
  import Data.ByteString qualified as BS
  import Data.IORef
  import Data.Maybe
  import System.Timeout
  import Text.Printf
  
  import Network.GRPC.Client
  import Network.GRPC.Common
  
  import KVStore.API
  import KVStore.API.JSON qualified as JSON
  import KVStore.API.Protobuf qualified as Protobuf
  import KVStore.Cmdline
  import KVStore.Util.Profiling
  import KVStore.Util.RandomAccessSet (RandomAccessSet)
  import KVStore.Util.RandomAccessSet qualified as RandomAccessSet
  import KVStore.Util.RandomGen qualified as RandomGen
  
  {-------------------------------------------------------------------------------
    Top-level
  -------------------------------------------------------------------------------}
  
  -- | Run the client for the specified time in seconds
  --
  -- Print the number of RPC calls.
  runKeyValueClient :: Cmdline -> IO ()
  runKeyValueClient cmdline@Cmdline{cmdDuration} = do
      statsVar <- newIORef zeroStats
      void $ timeout (cmdDuration * 1_000_000) $ client cmdline statsVar
      putStr . showStats cmdline =<< readIORef statsVar
  
  {-------------------------------------------------------------------------------
    Stats
  -------------------------------------------------------------------------------}
  
  data Stats = Stats {
        statsNumCreate   :: !Int
      , statsNumUpdate   :: !Int
      , statsNumRetrieve :: !Int
      , statsNumDelete   :: !Int
      }
  
  zeroStats :: Stats
  zeroStats = Stats {
        statsNumCreate   = 0
      , statsNumUpdate   = 0
      , statsNumRetrieve = 0
      , statsNumDelete   = 0
      }
  
  statsTotal :: Stats -> Int
  statsTotal stats = sum [
        statsNumCreate   stats
      , statsNumUpdate   stats
      , statsNumRetrieve stats
      , statsNumDelete   stats
      ]
  
  incNumCreate, incNumUpdate, incNumRetrieve, incNumDelete :: Stats -> (Stats, ())
  incNumCreate   stats = (stats{statsNumCreate   = statsNumCreate   stats + 1}, ())
  incNumUpdate   stats = (stats{statsNumUpdate   = statsNumUpdate   stats + 1}, ())
  incNumRetrieve stats = (stats{statsNumRetrieve = statsNumRetrieve stats + 1}, ())
  incNumDelete   stats = (stats{statsNumDelete   = statsNumDelete   stats + 1}, ())
  
  showStats :: Cmdline -> Stats -> String
  showStats Cmdline{cmdDuration} stats = unlines [
        printf "Did %.3f RPCs/s" countPerSec
      , "Totals:"
      , printf "  %d CREATE"   (statsNumCreate   stats)
      , printf "  %d UPDATE"   (statsNumUpdate   stats)
      , printf "  %d RETRIEVE" (statsNumRetrieve stats)
      , printf "  %d DELETE"   (statsNumDelete   stats)
      ]
    where
      countPerSec :: Double
      countPerSec = fromIntegral (statsTotal stats) / fromIntegral cmdDuration
  
  {-------------------------------------------------------------------------------
    Main client
  -------------------------------------------------------------------------------}
  
  -- | Connect to the server and keep calling random RPCs
  --
  -- Caller should create an 'IORef' initialized to 0, then call 'client' in
  -- separate thread, and kill the thread after some amount of time. The number
  -- of RPC calls made can then be read off from the 'IORef'.
  client :: Cmdline -> IORef Stats -> IO ()
  client Cmdline{
             cmdJSON
           , cmdSecure
           , cmdDisableTcpNoDelay
           , cmdPingRateLimit
           , cmdSemaphoreLimit
           } statsVar = do
      knownKeys <- RandomAccessSet.new
      random    <- RandomGen.new
      limiter   <- newQSem (fromMaybe 1 cmdSemaphoreLimit)
  
      withConnection params server $ \conn -> do
        let kvstore :: KVStore
            kvstore
              | cmdJSON   = JSON.client     conn
              | otherwise = Protobuf.client conn
  
        forever $ do
          waitQSem limiter
          forkIO $
            doRandomCRUD random knownKeys kvstore `finally` signalQSem limiter
    where
      params :: ConnParams
      params = def {
            connHTTP2Settings = def {
                http2TcpNoDelay            = not cmdDisableTcpNoDelay
              , http2OverridePingRateLimit = cmdPingRateLimit
              }
          }
  
      server :: Server
      server
        | cmdSecure
        = ServerSecure
            NoServerValidation
            SslKeyLogNone -- Let the server write the log
            Address {
                addressHost      = "127.0.0.1"
              , addressPort      = defaultSecurePort
              , addressAuthority = Nothing
              }
  
        | otherwise
        = ServerInsecure $ Address {
               addressHost      = "127.0.0.1"
             , addressPort      = defaultInsecurePort
             , addressAuthority = Nothing
             }
  
      doRandomCRUD ::
           RandomGen.RandomGen
        -> RandomAccessSet Key
        -> KVStore
        -> IO ()
      doRandomCRUD random knownKeys kvstore = do
          -- Pick a random CRUD action to take
          command <- RandomGen.nextInt random 4
  
          if command == 0 then do
            doCreate kvstore knownKeys
            atomicModifyIORef' statsVar incNumCreate
          else do
            -- If we don't know about any keys, retry with a new random action
            noKnownKeys <- RandomAccessSet.isEmpty knownKeys
            unless noKnownKeys $ do
              case command of
                1 -> do doRetrieve kvstore knownKeys
                        atomicModifyIORef' statsVar incNumRetrieve
                2 -> do doUpdate kvstore knownKeys
                        atomicModifyIORef' statsVar incNumUpdate
                3 -> do doDelete kvstore knownKeys
                        atomicModifyIORef' statsVar incNumDelete
                _ -> error "impossible"
  
  {-------------------------------------------------------------------------------
    Access the various server features
  -------------------------------------------------------------------------------}
  
  -- | Create a random key and value
  doCreate :: KVStore -> RandomAccessSet Key -> IO ()
  doCreate kvstore knownKeys = markRequest "CREATE" $ do
      key   <- createRandomKey knownKeys
      value <- Value <$> randomBytes meanValueSize
  
      let handleGrpcException :: GrpcException -> IO ()
          handleGrpcException e =
              if grpcError e == GrpcAlreadyExists
                then putStrLn "Key already existed"
                else throwIO e
  
      handle handleGrpcException $
        create kvstore (key, value)
  
  -- | Retrieve the value of a random key
  doRetrieve :: KVStore -> RandomAccessSet Key -> IO ()
  doRetrieve kvstore knownKeys = markRequest "RETRIEVE" $ do
      key <- RandomAccessSet.getRandomKey knownKeys
  
      let handleGrpcException :: GrpcException -> IO ()
          handleGrpcException e =
              if grpcError e == GrpcNotFound then do
                RandomAccessSet.remove knownKeys key
                putStrLn "Key not found"
              else
                throwIO e
  
      handle handleGrpcException $ do
        Value res <- retrieve kvstore key
        when (BS.length res < 1) $
          error "Invalid response"
  
  -- | Update a random key with a random value
  doUpdate :: KVStore -> RandomAccessSet Key -> IO ()
  doUpdate kvstore knownKeys = markRequest "UPDATE" $ do
      key   <- RandomAccessSet.getRandomKey knownKeys
      value <- Value <$> randomBytes meanValueSize
  
      let handleGrpcException :: GrpcException -> IO ()
          handleGrpcException e =
              if grpcError e == GrpcNotFound then do
                RandomAccessSet.remove knownKeys key
                putStrLn "Key not found"
              else
                throwIO e
  
      handle handleGrpcException $
        update kvstore (key, value)
  
  -- | Delete the value of a random key
  doDelete :: KVStore -> RandomAccessSet Key -> IO ()
  doDelete kvstore knownKeys = markRequest "DELETE" $ do
      key <- RandomAccessSet.getRandomKey knownKeys
      delete kvstore key
      RandomAccessSet.remove knownKeys key
  
  {-------------------------------------------------------------------------------
    Creating random keys and values
  -------------------------------------------------------------------------------}
  
  meanKeySize, meanValueSize :: Int
  meanKeySize   = 64
  meanValueSize = 65536
  
  -- | Create and add a key to the set of known keys
  createRandomKey :: RandomAccessSet Key -> IO Key
  createRandomKey knownKeys = loop
    where
      loop :: IO Key
      loop = do
        key   <- Key <$> randomBytes meanKeySize
        added <- RandomAccessSet.add knownKeys key
        if added
          then return key
          else loop
  
  -- | Create an exponentially sized byte string with a mean size
  randomBytes :: Int -> IO ByteString
  randomBytes mean = do
      random <- RandomGen.new
      d      <- RandomGen.nextDouble random
      let size :: Int
          size = round (fromIntegral mean * (-1 * log (1 - d)))
      RandomGen.getRandomBytes <$> RandomGen.nextBytes random (1 + size)
  
  {-------------------------------------------------------------------------------
    Debugging
  -------------------------------------------------------------------------------}
  
  markRequest :: String -> IO a -> IO a
  markRequest label =
      bracket_ (markEvent $ "CLIENT start " ++ label)
               (markEvent $ "CLIENT stop "  ++ label)