content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  {-# LANGUAGE CPP #-}
  module Glean.Server.Sharding (
    shardManagerConfig,
    withShardsUpdater,
    waitForTerminateSignalsAndGracefulShutdown,
    -- for testing
    dbUpdateNotifierThread) where
  
  import Control.Concurrent (
    modifyMVar_,
    newMVar,
    MVar,
    newEmptyMVar,
    tryPutMVar,
    takeMVar)
  import Control.Concurrent.Async ( withAsync )
  import Control.Exception
  import Control.Monad ( when, void, unless )
  import Data.HashSet (HashSet, toList)
  import qualified Data.HashSet as HashSet
  import Data.List (sort)
  #if GLEAN_FACEBOOK
  import Data.Maybe
  #endif
  import qualified Data.Set as Set
  import GHC.Conc (unsafeIOToSTM)
  import Glean.Backend.Types (dbShard)
  import qualified Glean.Database.Catalog as Catalog
  import Glean.Database.Catalog.Filter as Catalog
  import Glean.Database.Types ( Env(..) )
  #if GLEAN_FACEBOOK
  import Glean.Impl.ShardManager
  #endif
  import Glean.Server.Config (Config, cfgPort, cfgPublishShards)
  import Glean.Server.PublishShards ( ShardKey, getShardKey, updateShards )
  import qualified Glean.ServerConfig.Types as ServerConfig
  import Glean.Types ( Repo )
  import Glean.Util.Observed (Observed)
  import qualified Glean.Util.Observed as Observed
  import Glean.Util.Periodic ( doPeriodically )
  import Glean.Util.ShardManager
  import Glean.Util.ThriftService (DbShard)
  import Util.Time ( seconds )
  import System.Exit (die)
  import System.Posix.Signals (installHandler, sigTERM, Handler (Catch), sigINT)
  import System.Time.Extra ( showDuration, sleep, Seconds )
  import Util.Control.Exception ( swallow )
  import Util.EventBase ( EventBaseDataplane )
  import Util.Log ( vlog, logInfo )
  import Util.STM
    (retry, atomically, writeTVar, STM, TVar, readTVar, registerDelay)
  
  
  type PortNumber = Int
  
  shardManagerConfig ::
    -- | Application port
    Maybe PortNumber ->
    Catalog.Catalog ->
    Observed ServerConfig.Config ->
    (SomeShardManager -> IO b) ->
    IO b
  shardManagerConfig _mbPort _catalog smCfgServerConfig callback = do
    config <- Observed.get smCfgServerConfig
    case ServerConfig.config_sharding config of
      ServerConfig.ShardingPolicy_no_shards {} ->
        callback $ SomeShardManager noSharding
      ServerConfig.ShardingPolicy_static_assignment {} ->
        callback $
          SomeShardManager $
            shardByRepo $ do
              config <- Observed.get smCfgServerConfig
              case ServerConfig.config_sharding config of
                ServerConfig.ShardingPolicy_static_assignment assignment ->
                  return $
                    Just $
                      Set.toList $
                        ServerConfig.staticShardsPolicy_shards assignment
                _ ->
                  return Nothing
  #if GLEAN_FACEBOOK
      ServerConfig.ShardingPolicy_shard_manager policy -> do
        let smCliArgs = ShardManagerClientArgs
              { serviceName = ServerConfig.shardManagerPolicy_service_name policy
              , applicationPortNumber = fromMaybe 0 _mbPort
              , numberOfShards = fromIntegral $
                  ServerConfig.shardManagerPolicy_nshards policy
              , defaultDomainName =
                ServerConfig.shardManagerPolicy_default_domain_id policy
              }
        withShardManager smCliArgs $ \sm -> callback $ SomeShardManager sm
      ServerConfig.ShardingPolicy_shard_manager_most_recent policy -> do
        let ServerConfig.ShardManagerMostRecentPolicy{
              ServerConfig.shardManagerMostRecentPolicy_shard_manager_policy =
                ServerConfig.ShardManagerPolicy{..},
              .. } = policy
            smCliArgs = ShardManagerClientArgs
              { serviceName = shardManagerPolicy_service_name
              , applicationPortNumber = fromMaybe 0 _mbPort
              , numberOfShards = fromIntegral shardManagerPolicy_nshards
              , defaultDomainName = shardManagerPolicy_default_domain_id
              }
            getMostRecent =
              Set.fromList . map itemRepo <$> Catalog.listMostRecent _catalog
  
        withShardManagerForMostRecent
          smCliArgs
          shardManagerMostRecentPolicy_most_recent_domain_id
          getMostRecent $ \sm -> callback $ SomeShardManager sm
  #endif
      other ->
        error $ "Unsupported sharding policy: " <> show other
  
  databasesUpdatedCallback
    :: EventBaseDataplane
    -> ShardKey
    -> MVar (Maybe [DbShard])
    -> HashSet Repo
    -> IO ()
  databasesUpdatedCallback evb shardKey currentShards dbs = swallow $ do
    modifyMVar_ currentShards $ \prevShards -> do
      let newShards = sort $ map dbShard $ HashSet.toList dbs
  
      if Just newShards == prevShards then do
        logInfo $ "no change in shards: " <> show newShards
        return prevShards
      else do
        updateShards evb shardKey newShards
        return (Just newShards)
  
  -- The 'dbUpdateNotifierThread' will sit in a loop waiting for changes
  -- to the local databases (using STM retry to detect changes).  When
  -- changes are detected, it waits 1s so that multiple changes are
  -- processed in a single batch, and then invokes the callback.
  --
  -- The doPeriodically on the outside is just a fallback in case something
  -- goes wrong; it ensures that the exception is caught and logged, and we
  -- don't immediately retry in a loop.
  --
  dbUpdateNotifierThread
    :: Env
    -> Seconds
    -> STM Bool  -- ^ is the server gracefullly shutting down?
    -> (HashSet Repo -> IO ())
    -> IO ()
  dbUpdateNotifierThread env delay terminating callback =
    doPeriodically (seconds 30) $ do
      initial <- updated
      go initial
      where
      go prev = do
        atomically $ do
          termSignal <- terminating
          dbs <- fetchCurrentShardList env termSignal
          when (dbs == prev) retry
  
        vlog 1 $ "DB update detected, waiting " <> showDuration delay
        sleep delay -- wait 1s so we can batch updates
        current <- updated
        go current
  
      updated = do
        termSignal <- atomically terminating
        dbs <- atomically $ fetchCurrentShardList env termSignal
        vlog 1 "DB update notification"
        callback dbs
        return dbs
  
  withShardsUpdater
    :: EventBaseDataplane
    -> Config
    -> Env
    -> Seconds
    -> STM Bool  -- ^ is the server gracefullly shutting down?
    -> IO a
    -> IO a
  withShardsUpdater evb cfg env delay terminating action
    | cfgPublishShards cfg = do
      port <- case cfgPort cfg of
        Just port -> return port
        Nothing -> die "--publish-shards requires --port"
      currentShards <- newMVar Nothing
      key <- getShardKey evb port
      withAsync
        (dbUpdateNotifierThread env delay terminating
          $ databasesUpdatedCallback evb key currentShards)
        $ const action
    | otherwise = action
  
  
  waitForTerminateSignalsAndGracefulShutdown
    :: Env
    -> TVar Bool -- ^ broadcast channel for initiating the timeout
    -> Seconds -- ^ amount of time to wait before forcing a shutdown
    -> IO ()
  waitForTerminateSignalsAndGracefulShutdown env terminating timeout = do
    -- To wait in Haskell-land while the server is taking requests,
    -- use an mvar that gets filled when the right signals are read
    mvar <- newEmptyMVar
    let sigHandler = void $ tryPutMVar mvar ()
    withSignalHandler sigTERM sigHandler $ \_ ->
      withSignalHandler sigINT sigHandler $ \_ -> do
        -- Haskell will wait here until being instructed to stop
        takeMVar mvar
        logInfo "SIGTERM/SIGINT received"
  
        -- stop publishing complete shards
        atomically $ writeTVar terminating True
  
        -- start the timeout (if any)
        timeoutElapsedSTM <- if timeout > 0
          then readTVar <$> registerDelay (floor $ timeout * 1000000)
          else return $ pure True
  
        -- block until we do not advertise any shards anymore or run out of time
        atomically $ do
          dbs <- list
          timeoutElapsed <- timeoutElapsedSTM
  
          -- terminate when either the list is empty or we exceed the timeout
          unless (null dbs || timeoutElapsed) $ do
            unsafeIOToSTM $ logInfo $
              "Waiting for incomplete dbs: " <> show (toList dbs)
            retry
  
        logInfo "Shutting down"
  
    where
      withSignalHandler sig h = bracket
        (installHandler sig (Catch h) Nothing)
        (\old -> installHandler sig old Nothing)
  
      list = fetchCurrentShardList env True
  
  
  fetchCurrentShardList :: Env -> Bool -> STM (HashSet Repo)
  fetchCurrentShardList Env{..} terminating =
      let filter = if terminating
                    then incompleteQueryableF
                    else queryableF
      in
      HashSet.fromList . map Catalog.itemRepo
        <$> Catalog.list envCatalog [Catalog.Local] filter
  