content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PatternSynonyms #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE ViewPatterns #-}
  
  -- | This module defines data structures and operations pertaining to registered
  -- servers using registered calls; for unregistered support, see
  -- `Network.GRPC.LowLevel.Server.Unregistered`.
  module Network.GRPC.LowLevel.Server where
  
  import Control.Concurrent (
    ThreadId,
    forkFinally,
    killThread,
    myThreadId,
   )
  import Control.Concurrent.STM (
    atomically,
    check,
   )
  import Control.Concurrent.STM.TVar (
    TVar,
    modifyTVar',
    newTVarIO,
    readTVar,
    readTVarIO,
    writeTVar,
   )
  import Control.Exception (bracket)
  import Control.Monad
  import Control.Monad.IO.Class
  import Control.Monad.Trans.Except
  import Data.ByteString (ByteString)
  import qualified Data.ByteString as B
  import qualified Data.Set as S
  import Network.GRPC.LowLevel.Call
  import Network.GRPC.LowLevel.CompletionQueue (
    CompletionQueue,
    createCompletionQueue,
    pluck,
    serverRegisterCompletionQueue,
    serverRequestCall,
    serverShutdownAndNotify,
    shutdownCompletionQueue,
   )
  import Network.GRPC.LowLevel.GRPC
  import Network.GRPC.LowLevel.Op
  import qualified Network.GRPC.Unsafe as C
  import qualified Network.GRPC.Unsafe.ChannelArgs as C
  import qualified Network.GRPC.Unsafe.Op as C
  import qualified Network.GRPC.Unsafe.Security as C
  
  -- | Wraps various gRPC state needed to run a server.
  data Server = Server
    { serverGRPC :: GRPC
    , unsafeServer :: C.Server
    , listeningPort :: Port
    , serverCQ :: CompletionQueue
    -- ^ CQ used for receiving new calls.
    , serverCallCQ :: CompletionQueue
    -- ^ CQ for running ops on calls. Not used to receive new calls.
    , normalMethods :: [RegisteredMethod 'Normal]
    , sstreamingMethods :: [RegisteredMethod 'ServerStreaming]
    , cstreamingMethods :: [RegisteredMethod 'ClientStreaming]
    , bidiStreamingMethods :: [RegisteredMethod 'BiDiStreaming]
    , serverConfig :: ServerConfig
    , outstandingForks :: TVar (S.Set ThreadId)
    , serverShuttingDown :: TVar Bool
    }
  
  -- TODO: should we make a forkGRPC function instead? I am not sure if it would
  -- be safe to let the call handlers threads keep running after the server stops,
  -- so I'm taking the more conservative route of ensuring the server will
  -- stay alive. Experiment more when time permits.
  
  -- | Fork a thread from the server (presumably for a handler) with the guarantee
  -- that the server won't shut down while the thread is alive.
  -- Returns true if the fork happens successfully, and false if the server is
  -- already shutting down (in which case the function to fork is never executed).
  -- If the thread stays alive too long while the server is trying to shut down,
  -- the thread will be killed with 'killThread'.
  -- The purpose of this is to prevent memory access
  -- errors at the C level of the library, not to ensure that application layer
  -- operations in user code complete successfully.
  forkServer :: Server -> IO () -> IO Bool
  forkServer Server{..} f = do
    shutdown <- readTVarIO serverShuttingDown
    case shutdown of
      True -> return False
      False -> do
        -- NB: The spawned thread waits on 'ready' before running 'f' to ensure
        -- that its ThreadId is inserted into outstandingForks before the cleanup
        -- function deletes it from there. Not doing this can lead to stale thread
        -- ids in the set when handlers cleanup ahead of the insertion, and a
        -- subsequent deadlock in stopServer. We can use a dead-list instead if we
        -- need something more performant.
        ready <- newTVarIO False
        tid <- forkFinally (atomically (check =<< readTVar ready) >> f) cleanup
  
        atomically $ do
          modifyTVar' outstandingForks (S.insert tid)
          modifyTVar' ready (const True)
  
        debug
  
        return True
    where
      cleanup _ = do
        tid <- myThreadId
        atomically $ modifyTVar' outstandingForks (S.delete tid)
  
  #ifdef DEBUG
      -- This is intentionally moved outside of the 'do' block so that this file
      -- can be auto-formatted by fourmolu. Fourmolu's support for formatting CPP
      -- fails inside of 'do' blocks.
      debug = do
        tids <- readTVarIO outstandingForks
        grpcDebug ("after fork and bookkeeping: outstandingForks=" ++ show tids)
  # else
      debug = pure ()
  #endif
  
  -- | Configuration for SSL.
  data ServerSSLConfig = ServerSSLConfig
    { clientRootCert :: Maybe FilePath
    , serverPrivateKey :: FilePath
    , serverCert :: FilePath
    , clientCertRequest :: C.SslClientCertificateRequestType
    -- ^ Whether to request a certificate from the client, and what to do with it
    -- if received.
    , customMetadataProcessor :: Maybe C.ProcessMeta
    }
  
  -- | Configuration needed to start a server.
  data ServerConfig = ServerConfig
    { host :: Host
    -- ^ Name of the host the server is running on. Not sure how this is
    -- used. Setting to "localhost" works fine in tests.
    , port :: Port
    -- ^ Port on which to listen for requests.
    , methodsToRegisterNormal :: [MethodName]
    -- ^ List of normal (non-streaming) methods to register.
    , methodsToRegisterClientStreaming :: [MethodName]
    , methodsToRegisterServerStreaming :: [MethodName]
    , methodsToRegisterBiDiStreaming :: [MethodName]
    , serverArgs :: [C.Arg]
    -- ^ Optional arguments for setting up the channel on the server. Supplying an
    -- empty list will cause the channel to use gRPC's default options.
    , sslConfig :: Maybe ServerSSLConfig
    -- ^ Server-side SSL configuration. If 'Nothing', the server will use an
    -- insecure connection.
    }
  
  serverEndpoint :: ServerConfig -> Endpoint
  serverEndpoint ServerConfig{..} = endpoint host port
  
  addPort :: C.Server -> ServerConfig -> IO Int
  addPort server conf@ServerConfig{..} =
    case sslConfig of
      Nothing -> C.withInsecureServerCredentials $ C.grpcServerAddHttp2Port server e
      Just ServerSSLConfig{..} ->
        do
          crc <- mapM B.readFile clientRootCert
          spk <- B.readFile serverPrivateKey
          sc <- B.readFile serverCert
          C.withServerCredentials crc spk sc clientCertRequest $ \creds -> do
            case customMetadataProcessor of
              Just p -> C.setMetadataProcessor creds p
              Nothing -> return ()
            C.grpcServerAddHttp2Port server e creds
    where
      e = unEndpoint $ serverEndpoint conf
  
  startServer :: GRPC -> ServerConfig -> IO Server
  startServer grpc conf@ServerConfig{..} =
    C.withChannelArgs serverArgs $ \args -> do
      let e = serverEndpoint conf
      server <- C.grpcServerCreate args C.reserved
      actualPort <- addPort server conf
      when (unPort port > 0 && actualPort /= unPort port) $
        error $
          "Unable to bind port: " ++ show port
      cq <- createCompletionQueue grpc
      grpcDebug $ "startServer: server CQ: " ++ show cq
      serverRegisterCompletionQueue server cq
  
      -- Register methods according to their GRPCMethodType kind. It's a bit ugly
      -- to partition them this way, but we get very convenient phantom typing
      -- elsewhere by doing so.
      -- TODO: change order of args so we can eta reduce.
      ns <-
        mapM
          (\nm -> serverRegisterMethodNormal server nm e)
          methodsToRegisterNormal
      ss <-
        mapM
          (\nm -> serverRegisterMethodServerStreaming server nm e)
          methodsToRegisterServerStreaming
      cs <-
        mapM
          (\nm -> serverRegisterMethodClientStreaming server nm e)
          methodsToRegisterClientStreaming
      bs <-
        mapM
          (\nm -> serverRegisterMethodBiDiStreaming server nm e)
          methodsToRegisterBiDiStreaming
      C.grpcServerStart server
      forks <- newTVarIO S.empty
      shutdown <- newTVarIO False
      ccq <- createCompletionQueue grpc
      return $
        Server
          grpc
          server
          (Port actualPort)
          cq
          ccq
          ns
          ss
          cs
          bs
          conf
          forks
          shutdown
  
  stopServer :: Server -> IO ()
  -- TODO: Do method handles need to be freed?
  stopServer Server{unsafeServer = s, ..} = do
    grpcDebug "stopServer: calling shutdownNotify."
    shutdownNotify serverCQ
    grpcDebug "stopServer: cancelling all calls."
    C.grpcServerCancelAllCalls s
    cleanupForks
    grpcDebug "stopServer: call grpc_server_destroy."
    C.grpcServerDestroy s
    grpcDebug "stopServer: shutting down CQ."
    shutdownCQ serverCQ
    shutdownCQ serverCallCQ
    where
      shutdownCQ scq = do
        shutdownResult <- shutdownCompletionQueue scq
        case shutdownResult of
          Left _ -> do
            putStrLn "Warning: completion queue didn't shut down."
            putStrLn "Trying to stop server anyway."
          Right _ -> return ()
      shutdownNotify scq = do
        let shutdownTag = C.tag 0
        serverShutdownAndNotify s scq shutdownTag
        grpcDebug "called serverShutdownAndNotify; plucking."
        shutdownEvent <- pluck scq shutdownTag (Just 30)
        grpcDebug $ "shutdownNotify: got shutdown event" ++ show shutdownEvent
        case shutdownEvent of
          -- This case occurs when we pluck but the queue is already in the
          -- 'shuttingDown' state, implying we already tried to shut down.
          Left GRPCIOShutdown -> error "Called stopServer twice!"
          Left _ -> error "Failed to stop server."
          Right _ -> return ()
      cleanupForks = do
        atomically $ writeTVar serverShuttingDown True
        liveForks <- readTVarIO outstandingForks
        grpcDebug $ "Server shutdown: killing threads: " ++ show liveForks
        mapM_ killThread liveForks
        -- wait for threads to shut down
        grpcDebug "Server shutdown: waiting until all threads are dead."
        atomically $ check . (== 0) . S.size =<< readTVar outstandingForks
        grpcDebug "Server shutdown: All forks cleaned up."
  
  -- Uses 'bracket' to safely start and stop a server, even if exceptions occur.
  withServer :: GRPC -> ServerConfig -> (Server -> IO a) -> IO a
  withServer grpc cfg = bracket (startServer grpc cfg) stopServer
  
  -- | Less precisely-typed registration function used in
  -- 'serverRegisterMethodNormal', 'serverRegisterMethodServerStreaming',
  -- 'serverRegisterMethodClientStreaming', and
  -- 'serverRegisterMethodBiDiStreaming'.
  serverRegisterMethod ::
    C.Server ->
    MethodName ->
    Endpoint ->
    GRPCMethodType ->
    IO C.CallHandle
  serverRegisterMethod s nm e mty =
    C.grpcServerRegisterMethod
      s
      (unMethodName nm)
      (unEndpoint e)
      (payloadHandling mty)
  
  {-
  TODO: Consolidate the register functions below.
  
  It seems like we'd need true dependent types to use only one
    registration function. Ideally we'd want a type like
    serverRegisterMethod :: C.Server
                            -> MethodName
                            -> Endpoint
                            -> (t :: GRPCMethodType)
                            -> IO (RegisteredMethod (Lifted t))
  
  where `Lifted t` is the type in the t data kind that corresponds to the data
  constructor t the function was given.
  
  -}
  
  -- | Register a method on a server. The 'RegisteredMethod' type can then be used
  -- to wait for a request to arrive. Note: gRPC claims this must be called before
  -- the server is started, so we do it during startup according to the
  -- 'ServerConfig'.
  serverRegisterMethodNormal ::
    C.Server ->
    -- | method name, e.g. "/foo"
    MethodName ->
    -- | Endpoint name name, e.g. "localhost:9999". I have no
    -- idea why this is needed since we have to provide these
    -- parameters to start a server in the first place. It
    -- doesn't seem to have any effect, even if it's filled
    -- with nonsense.
    Endpoint ->
    IO (RegisteredMethod 'Normal)
  serverRegisterMethodNormal internalServer meth e = do
    h <- serverRegisterMethod internalServer meth e Normal
    return $ RegisteredMethodNormal meth e h
  
  serverRegisterMethodClientStreaming ::
    C.Server ->
    -- | method name, e.g. "/foo"
    MethodName ->
    -- | Endpoint name name, e.g. "localhost:9999". I have no
    -- idea why this is needed since we have to provide these
    -- parameters to start a server in the first place. It
    -- doesn't seem to have any effect, even if it's filled
    -- with nonsense.
    Endpoint ->
    IO (RegisteredMethod 'ClientStreaming)
  serverRegisterMethodClientStreaming internalServer meth e = do
    h <- serverRegisterMethod internalServer meth e ClientStreaming
    return $ RegisteredMethodClientStreaming meth e h
  
  serverRegisterMethodServerStreaming ::
    C.Server ->
    -- | method name, e.g. "/foo"
    MethodName ->
    -- | Endpoint name name, e.g. "localhost:9999". I have no
    -- idea why this is needed since we have to provide these
    -- parameters to start a server in the first place. It
    -- doesn't seem to have any effect, even if it's filled
    -- with nonsense.
    Endpoint ->
    IO (RegisteredMethod 'ServerStreaming)
  serverRegisterMethodServerStreaming internalServer meth e = do
    h <- serverRegisterMethod internalServer meth e ServerStreaming
    return $ RegisteredMethodServerStreaming meth e h
  
  serverRegisterMethodBiDiStreaming ::
    C.Server ->
    -- | method name, e.g. "/foo"
    MethodName ->
    -- | Endpoint name name, e.g. "localhost:9999". I have no
    -- idea why this is needed since we have to provide these
    -- parameters to start a server in the first place. It
    -- doesn't seem to have any effect, even if it's filled
    -- with nonsense.
    Endpoint ->
    IO (RegisteredMethod 'BiDiStreaming)
  serverRegisterMethodBiDiStreaming internalServer meth e = do
    h <- serverRegisterMethod internalServer meth e BiDiStreaming
    return $ RegisteredMethodBiDiStreaming meth e h
  
  -- | Create a 'Call' with which to wait for the invocation of a registered
  -- method.
  serverCreateCall ::
    Server ->
    RegisteredMethod mt ->
    IO (Either GRPCIOError (ServerCall (MethodPayload mt)))
  serverCreateCall Server{..} rm =
    serverRequestCall rm unsafeServer serverCQ serverCallCQ
  
  withServerCall ::
    Server ->
    RegisteredMethod mt ->
    (ServerCall (MethodPayload mt) -> IO (Either GRPCIOError a)) ->
    IO (Either GRPCIOError a)
  withServerCall s rm f =
    bracket (serverCreateCall s rm) cleanup $ \case
      Left e -> return (Left e)
      Right c -> do
        debugServerCall c
        f c
    where
      cleanup (Left _) = pure ()
      cleanup (Right c) = do
        grpcDebug "withServerCall(R): destroying."
        destroyServerCall c
  
  --------------------------------------------------------------------------------
  -- serverReader (server side of client streaming mode)
  
  type ServerReaderHandlerLL =
    ServerCall (MethodPayload 'ClientStreaming) ->
    StreamRecv ByteString ->
    IO (Maybe ByteString, MetadataMap, C.StatusCode, StatusDetails)
  
  serverReader ::
    Server ->
    RegisteredMethod 'ClientStreaming ->
    -- | Initial server metadata
    MetadataMap ->
    ServerReaderHandlerLL ->
    IO (Either GRPCIOError ())
  serverReader s rm initMeta f =
    withServerCall s rm (\sc -> serverReader' s sc initMeta f)
  
  serverReader' ::
    Server ->
    ServerCall (MethodPayload 'ClientStreaming) ->
    -- | Initial server metadata
    MetadataMap ->
    ServerReaderHandlerLL ->
    IO (Either GRPCIOError ())
  serverReader' _ sc@ServerCall{unsafeSC = c, callCQ = ccq} initMeta f =
    runExceptT $ do
      (mmsg, trailMeta, st, ds) <- liftIO $ f sc (streamRecvPrim c ccq)
      void $
        runOps'
          c
          ccq
          ( OpSendInitialMetadata initMeta
              : OpSendStatusFromServer trailMeta st ds
              : maybe [] ((: []) . OpSendMessage) mmsg
          )
  
  --------------------------------------------------------------------------------
  -- serverWriter (server side of server streaming mode)
  
  type ServerWriterHandlerLL =
    ServerCall (MethodPayload 'ServerStreaming) ->
    StreamSend ByteString ->
    IO (MetadataMap, C.StatusCode, StatusDetails)
  
  -- | Wait for and then handle a registered, server-streaming call.
  serverWriter ::
    Server ->
    RegisteredMethod 'ServerStreaming ->
    -- | Initial server metadata
    MetadataMap ->
    ServerWriterHandlerLL ->
    IO (Either GRPCIOError ())
  serverWriter s rm initMeta f =
    withServerCall s rm (\sc -> serverWriter' s sc initMeta f)
  
  serverWriter' ::
    Server ->
    ServerCall (MethodPayload 'ServerStreaming) ->
    MetadataMap ->
    ServerWriterHandlerLL ->
    IO (Either GRPCIOError ())
  serverWriter' _ sc@ServerCall{unsafeSC = c, callCQ = ccq} initMeta f =
    runExceptT $ do
      sendInitialMetadata c ccq initMeta
      st <- liftIO $ f sc (streamSendPrim c ccq)
      sendStatusFromServer c ccq st
  
  --------------------------------------------------------------------------------
  -- serverRW (bidirectional streaming mode)
  
  type ServerRWHandlerLL =
    ServerCall (MethodPayload 'BiDiStreaming) ->
    StreamRecv ByteString ->
    StreamSend ByteString ->
    IO (MetadataMap, C.StatusCode, StatusDetails)
  
  serverRW ::
    Server ->
    RegisteredMethod 'BiDiStreaming ->
    -- | initial server metadata
    MetadataMap ->
    ServerRWHandlerLL ->
    IO (Either GRPCIOError ())
  serverRW s rm initMeta f =
    withServerCall s rm (\sc -> serverRW' s sc initMeta f)
  
  serverRW' ::
    Server ->
    ServerCall (MethodPayload 'BiDiStreaming) ->
    MetadataMap ->
    ServerRWHandlerLL ->
    IO (Either GRPCIOError ())
  serverRW' _ sc@ServerCall{unsafeSC = c, callCQ = ccq} initMeta f =
    runExceptT $ do
      sendInitialMetadata c ccq initMeta
      st <- liftIO $ f sc (streamRecvPrim c ccq) (streamSendPrim c ccq)
      sendStatusFromServer c ccq st
  
  --------------------------------------------------------------------------------
  -- serverHandleNormalCall (server side of normal request/response)
  
  -- | A handler for a registered server call; bytestring parameter is request
  -- body, with the bytestring response body in the result tuple. The first
  -- metadata parameter refers to the request metadata, with the two metadata
  -- values in the result tuple being the initial and trailing metadata
  -- respectively. We pass in the 'ServerCall' so that the server can call
  -- 'serverCallCancel' on it if needed.
  type ServerHandlerLL =
    ServerCall (MethodPayload 'Normal) ->
    IO (ByteString, MetadataMap, C.StatusCode, StatusDetails)
  
  -- | Wait for and then handle a normal (non-streaming) call.
  serverHandleNormalCall ::
    Server ->
    RegisteredMethod 'Normal ->
    -- | Initial server metadata
    MetadataMap ->
    ServerHandlerLL ->
    IO (Either GRPCIOError ())
  serverHandleNormalCall s rm initMeta f =
    withServerCall s rm go
    where
      go sc@ServerCall{unsafeSC = c, callCQ = ccq} = do
        (rsp, trailMeta, st, ds) <- f sc
        void
          <$> runOps
            c
            ccq
            [ OpSendInitialMetadata initMeta
            , OpRecvCloseOnServer
            , OpSendMessage rsp
            , OpSendStatusFromServer trailMeta st ds
            ]
  