{-# OPTIONS_GHC -Wno-orphans #-}

-- | Handling of client or server disconnections occurring with ongoing RPCs on
-- a shared connection.
--
-- When a server disconnects, we expect:
--
-- 1. All current calls fail with 'Client.ServerDisconnected'
-- 2. Future calls (after reconnection) succeed
--
-- When a client disconnects, we expect:
--
-- 1. The handlers dealing with that client (i.e. on that connection) should
--    fail with 'Server.ClientDisconnected'
-- 2. Future calls (after reconnection) succeed
module Test.Sanity.Disconnect (tests) where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Data.ByteString.Lazy qualified as Lazy (ByteString)
import Data.IORef
import Data.Word
import Foreign.C.Types (CInt(..))
import Network.Socket
import System.Posix
import Test.Tasty
import Test.Tasty.HUnit
import Text.Read hiding (step)

import Network.GRPC.Client qualified as Client
import Network.GRPC.Client.Binary qualified as Binary
import Network.GRPC.Common
import Network.GRPC.Server qualified as Server
import Network.GRPC.Server.Binary qualified as Binary
import Network.GRPC.Server.Run

import Proto.API.Trivial

import Test.Util

{-------------------------------------------------------------------------------
  Top-level
-------------------------------------------------------------------------------}

tests :: TestTree
tests = testGroup "Test.Sanity.Disconnect" [
      testCase "client" test_clientDisconnect
    , testCase "server" test_serverDisconnect
    ]

{-------------------------------------------------------------------------------
  Disconnecting clients
-------------------------------------------------------------------------------}

-- | Two separate clients make many concurrent calls, one of them disconnects.
test_clientDisconnect :: Assertion
test_clientDisconnect = do
    -- Create the server
    disconnectCounter1 <- newIORef 0
    disconnectCounter2 <- newIORef 0
    server <-
      Server.mkGrpcServer def [
          Server.someRpcHandler $
              Server.mkRpcHandler @RPC1 $ echoHandler (Just disconnectCounter1)
        , Server.someRpcHandler $
              Server.mkRpcHandler @RPC2 $ echoHandler (Just disconnectCounter2)
        ]

    -- Start server
    let serverConfig = ServerConfig {
            serverInsecure = Just $ InsecureConfig {
                insecureHost = Just "127.0.0.1"
              , insecurePort = 0
              }
          , serverSecure = Nothing
          }
    portSignal <- newEmptyMVar
    void $ forkIO $ forkServer def serverConfig server $ \runningServer -> do
      putMVar portSignal =<< getServerPort runningServer
      waitServer runningServer

    -- Wait for the server to signal its port
    serverPort <- readMVar portSignal
    let serverAddress =
          Client.ServerInsecure Client.Address {
              addressHost      = "127.0.0.1"
            , addressPort      = serverPort
            , addressAuthority = Nothing
            }


    -- Start a client in a separate process
    let numCalls = 10
    dyingChild <- forkProcess $
      Client.withConnection def serverAddress $ \conn -> do
        inLockstep conn (Proxy @RPC1) numCalls NeverTerminate $ \results _getFinal -> do
          -- Wait until we are sure that all clients have started their RPC,
          -- then kill the process. This avoids race conditions and guarantees
          -- that the server will see @numCalls@ clients disconnecting.
          _ <- waitForHistoryOfMinLen results 1
          c_exit 1

    -- Start two more clients; these will not disconnect
    let numSteps = 5
    (result1, result2) <- concurrently
      ( Client.withConnection def serverAddress $ \conn -> do
          inLockstep conn (Proxy @RPC1) numCalls (TerminateAfter numSteps) $ \_results getFinal ->
            getFinal
      )
      ( Client.withConnection def serverAddress $ \conn -> do
          inLockstep conn (Proxy @RPC2) numCalls (TerminateAfter numSteps) $ \_results getFinal ->
            getFinal
      )

    -- Wait for the forked process to terminate
    _status <- getProcessStatus True False dyingChild

    -- All calls by clients in /this/ process (not the ones we killed) should
    -- have finished normally
    let expectedResult = [
            replicate numCalls (StepOk i)
          | i <- reverse [1 .. numSteps]
          ]
    assertEqual "" expectedResult result1
    assertEqual "" expectedResult result2

    -- We should also see only @numCalls@ client disconnects for the first
    -- handler and none for the second
    clientDisconnects1 <- readIORef disconnectCounter1
    clientDisconnects2 <- readIORef disconnectCounter2
    assertEqual "" numCalls clientDisconnects1
    assertEqual "" 0        clientDisconnects2

-- We need to use this to properly simulate the execution environment crashing
-- in an unrecoverable way. In particular, we don't want to give the program a
-- chance to do any of its normal exception handling/cleanup behavior.
foreign import ccall unsafe "exit" c_exit :: CInt -> IO ()

{-------------------------------------------------------------------------------
  Disconnecting servers
-------------------------------------------------------------------------------}

-- | Client makes many concurrent calls, server disconnects
test_serverDisconnect :: Assertion
test_serverDisconnect = withTemporaryFile $ \ipcFile -> do
    -- We use a temporary file as a very rudimentary means of inter-process
    -- communication so the server (which runs in a separate process) can make
    -- the client aware of the port it is assigned by the OS.
    let ipcWrite :: PortNumber -> IO ()
        ipcWrite port = do
          writeFile ipcFile (show port)

        ipcRead :: IO PortNumber
        ipcRead = do
          fmap (readMaybe @PortNumber) (readFile ipcFile) >>= \case
            Nothing -> do
              ipcRead
            Just p -> do
              writeFile ipcFile ""
              return p

    -- Create the server
    server <-
      Server.mkGrpcServer def [
          Server.someRpcHandler $
              Server.mkRpcHandler @RPC1 $ echoHandler Nothing
        ]

    let serverConfig = ServerConfig {
            serverInsecure = Just $ InsecureConfig {
                insecureHost = Just "127.0.0.1"
              , insecurePort = 0
              }
          , serverSecure = Nothing
          }

        -- Starts the server in a new process. Gives back an action that kills
        -- the created server process.
        startServer :: IO (IO ())
        startServer = do
          serverPid <-
            forkProcess $
              forkServer def serverConfig server $ \runningServer -> do
                ipcWrite =<< getServerPort runningServer
                waitServer runningServer
          return $ signalProcess sigKILL serverPid

    -- Start server, get the initial port
    killServer    <- startServer
    port1         <- ipcRead
    signalRestart <- newEmptyMVar
    let serverAddress port =
          Client.ServerInsecure Client.Address {
              addressHost      = "127.0.0.1"
            , addressPort      = port
            , addressAuthority = Nothing
            }

        reconnectPolicy :: Client.ReconnectPolicy
        reconnectPolicy = go 0
          where
            go :: Int -> Client.ReconnectPolicy
            go n
              | n == 5
              = Client.ReconnectPolicy $ do
                  killRestarted <- startServer
                  port2 <- ipcRead
                  putMVar signalRestart killRestarted
                  return $ Client.DoReconnect Client.Reconnect {
                      Client.nextPolicy =
                        Client.ReconnectPolicy $
                          pure $ Client.DoReconnect Client.Reconnect {
                                Client.reconnectTo =
                                  Client.ReconnectToNew $ serverAddress port2
                              , Client.nextPolicy =
                                  Client.ReconnectPolicy $ pure Client.DontReconnect
                              , Client.onReconnect = Nothing
                              }
                      , Client.reconnectTo = Client.ReconnectToOriginal
                      , Client.onReconnect = def
                    }
              | otherwise
              = Client.ReconnectPolicy $ do
                  threadDelay 10000
                  return $
                    Client.DoReconnect Client.Reconnect {
                        reconnectTo = Client.ReconnectToOriginal
                      , onReconnect = def
                      , nextPolicy  = go (n + 1)
                      }

        connParams :: Client.ConnParams
        connParams = def { Client.connReconnectPolicy = reconnectPolicy }

    Client.withConnection connParams (serverAddress port1) $ \conn -> do
      let numCalls = 10
      results <-
        inLockstep conn (Proxy @RPC1) numCalls NeverTerminate $ \results getFinal -> do
          -- Once all clients have started their RPC, kill the server
          _ <- waitForHistoryOfMinLen results 1
          killServer
          getFinal

      -- All calls should have failed (but we don't know in which step)
      assertEqual "" numCalls $ length $ filter stepFailed (concat results)

      -- New calls should succeed (after reconnection)
      killRestarted <- takeMVar signalRestart
      result <-
        inLockstep conn (Proxy @RPC1) numCalls (TerminateAfter 1) $ \_results getFinal ->
          getFinal

      let expectedResult = [replicate numCalls $ StepOk 1]
      assertEqual "" expectedResult result

      -- Do not leave the server process hanging around
      killRestarted

{-------------------------------------------------------------------------------
  Auxiliary: echo handler
-------------------------------------------------------------------------------}

-- | Echos any input
echoHandler ::
     TrivialRpc rpc
  => Maybe (IORef Int)
  -> Server.Call rpc -> IO ()
echoHandler disconnectCounter call =
    trackDisconnects disconnectCounter $ loop
  where
    loop :: IO ()
    loop = do
        inp <- Binary.recvInput @Word64 call
        case inp of
          StreamElem n   -> Binary.sendNextOutput @Word64 call n >> loop
          FinalElem  n _ -> Binary.sendFinalOutput @Word64 call (n, NoMetadata)
          NoMoreElems  _ -> Server.sendTrailers call NoMetadata

    trackDisconnects :: Maybe (IORef Int) -> IO () -> IO ()
    trackDisconnects Nothing        = id
    trackDisconnects (Just counter) =
        handle $ \(_e :: Server.ClientDisconnected) ->
          atomicModifyIORef' counter $ \n -> (n + 1, ())

{-------------------------------------------------------------------------------
  Bunch of clients all executing in lockstep
-------------------------------------------------------------------------------}

data NumSteps = TerminateAfter Int | NeverTerminate

data Results = Results {
      -- | Results for the current step
      resultsCurr :: TVar [StepResult]

      -- | Number of the current step
    , resultsStep :: Int

      -- | Previous results (in reverse order)
    , resultsHist :: [[StepResult]]
    }

data StepResult = StepOk Int | StepFailed SomeException
  deriving stock (Show)

stepFailed :: StepResult -> Bool
stepFailed StepOk{}     = False
stepFailed StepFailed{} = True

instance Eq StepResult where
  StepOk     i == StepOk     i' = i == i'
  StepFailed _ == StepFailed _  = True -- the exception is merely for debugging
  StepOk     _ == StepFailed _  = False
  StepFailed _ == StepOk     _  = False

initResults :: IO (TVar Results)
initResults = do
    resultsCurr <- newTVarIO []
    newTVarIO Results{
        resultsCurr
      , resultsStep = 1
      , resultsHist = []
      }

-- | Keep collecting results (never terminates)
collectResults :: Int -> TVar Results -> IO a
collectResults numClients results =
    forever $
      atomically $ do
        Results{resultsCurr, resultsStep, resultsHist} <- readTVar results
        current <- readTVar resultsCurr
        if length current < numClients
          then retry
          else do
            current' <- newTVar []
            writeTVar results Results{
                resultsCurr = current'
              , resultsStep = succ resultsStep
              , resultsHist = current : resultsHist
              }

-- | Get the 'TVar' for the specified step, blocking until that step is reached
--
-- This is executed by each client on each step. As a result, we can assume that
-- the required step can never be /before/ the current step (because all clients
-- must deliver their result for the current step before the step advances).
waitForStep :: TVar Results -> Int -> IO (TVar [StepResult])
waitForStep results step = atomically $ do
    Results{resultsCurr, resultsStep} <- readTVar results
    if resultsStep < step
      then retry
      else return resultsCurr

-- | Wait until a history of at least the specified length is ready
waitForHistoryOfMinLen :: TVar Results -> Int -> IO [[StepResult]]
waitForHistoryOfMinLen results numSteps = atomically $ do
    Results{resultsHist} <- readTVar results
    if length resultsHist < numSteps
      then retry
      else return resultsHist

inLockstep :: forall rpc a.
     TrivialRpc rpc
  => Client.Connection -- ^ Server to connect to
  -> Proxy rpc         -- ^ Method to call
  -> Int               -- ^ Number of clients
  -> NumSteps          -- ^ How many steps each client should take
  -> (TVar Results -> IO [[StepResult]] -> IO a)
     -- ^ Monitor the results
     --
     -- This is also passed a function to get the /final/ results, after all
     -- clients have terminated. If some clients never terminate, this function
     -- will block indefinitely.
  -> IO a
inLockstep conn rpc numClients numSteps monitor = do
    results <- initResults
    withAsync (collectResults numClients results) $ \_ ->
      withAsync (runClients results) $ \clients ->
        monitor results (wait clients)
  where
    runClients :: TVar Results -> IO [[StepResult]]
    runClients results = do
        replicateConcurrently_ numClients $
          Client.withRPC conn def rpc (client results)
        resultsHist <$> readTVarIO results

    client :: TVar Results -> Client.Call rpc -> IO ()
    client results call = loop 1
      where
        loop :: Int -> IO ()
        loop n = do
            current <- waitForStep results n
            handle (recordException current) $
              case numSteps of
                TerminateAfter n' | n == n' -> do
                  Binary.sendFinalInput call n
                  (resp, NoMetadata) <- Binary.recvFinalOutput call
                  atomically $ modifyTVar current (StepOk resp:)
                _otherwise -> do
                  Binary.sendNextInput call n
                  resp <- Binary.recvNextOutput call
                  atomically $ modifyTVar current (StepOk resp:)
                  loop (succ n)

        recordException :: TVar [StepResult] -> SomeException -> IO ()
        recordException current e =
            atomically $ modifyTVar current (StepFailed e:)

{-------------------------------------------------------------------------------
  Auxiliary: trivial RPCs

  We want two distinct handler so we have two trivial RPCs.
-------------------------------------------------------------------------------}

type TrivialRpc rpc = (
    SupportsClientRpc        rpc
  , Input                    rpc ~ Lazy.ByteString
  , Output                   rpc ~ Lazy.ByteString
  , RequestMetadata          rpc ~ NoMetadata
  , ResponseInitialMetadata  rpc ~ NoMetadata
  , ResponseTrailingMetadata rpc ~ NoMetadata
  )

type RPC1 = Trivial' "rpc1"
type RPC2 = Trivial' "rpc2"