{-
  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.
-}

module Glean.Database.Close (
  closeDatabase,
  closeDatabases,
  closeIdleDatabase,
  closeIdleDatabases,
  closeOpenDB,
) where

import Control.Exception hiding(handle)
import Control.Monad.Extra
import qualified Data.HashMap.Strict as HashMap
import Data.Maybe
import qualified Data.Text.Encoding as Text

import ServiceData.GlobalStats as Stats
import Util.STM

import Glean.Database.Open
import qualified Glean.Database.Storage as Storage
import Glean.Database.Types
import Glean.Database.Writes
import qualified Glean.RTS.Foreign.LookupCache as LookupCache
import Glean.Types hiding (Database)
import qualified Glean.Types as Thrift
import Glean.Util.Mutex
import Util.Time


closeDatabases :: Env -> IO ()
closeDatabases env@Env{..} = do
  dbs <- readTVarIO envActive
  forM_ (HashMap.keys dbs) $ closeDatabase env

isIdle :: (TimePoint -> Bool) -> DB s -> OpenDB s -> STM Bool
isIdle long_enough db odb = and <$> sequence
  [ (== 1) <$> readTVar (dbUsers db)  -- we are the only user
  , long_enough <$> readTVar (odbIdleSince odb)
  , case odbWriting odb of
      Just Writing{..} -> do
        writeQueueSize <- readTVar (writeQueueSize wrQueue)
        commit <- readTVar wrCommit
        return $ writeQueueSize == 0 && isNothing commit
      Nothing -> return True
  ]

closeIf
  :: (forall s . DB s -> DBState s -> STM (Maybe (OpenDB s)))
  -> Env
  -> Repo
  -> IO ()
closeIf should_close env repo = usingActiveDatabase env repo $ \r ->
  forM_ r $ \db -> mask_ $ do
    r <- atomically $ do
      state <- readTVar $ dbState db
      o <- should_close db state
      case o of
        Just odb -> do
          idle <- isIdle (const True) db odb
          when (not idle) retry
          writeTVar (dbState db) Closing
          return $ Just (db,odb)
        Nothing -> return Nothing

    forM_ r $ \(db, odb) ->
      -- the actual closing of the DB must be uninterruptible
      uninterruptibleMask_ (closeOpenDB env odb)
        `finally` atomically (writeTVar (dbState db) Closed)

closeDatabase :: Env -> Repo -> IO ()
closeDatabase env = closeIf
  (\_ state -> case state of
    Opening -> retry
    Open odb -> do
      deleteWriteQueues env odb
      case odbWriting odb of
        Just Writing{..} -> do
          r <- readTVar wrCommit
          when (isJust r) retry
        Nothing -> return ()
      return $ Just odb
    Closing -> retry
    Closed -> return Nothing)
  env

-- | Synchronously close a database if it has been idle for more than
-- `duration`.
closeIdleDatabase :: Env -> Repo -> DiffTimePoints -> IO ()
closeIdleDatabase env repo duration = do
  now <- getTimePoint
  closeIf
    (\db state -> case state of
        Opening -> return Nothing
        Open odb -> do
          idle <- isIdle
            (\last_use -> diffTimePoints last_use now >= duration)
            db
            odb
          return $ if idle then Just odb else Nothing
        Closing -> return Nothing
        Closed -> return Nothing)
    env
    repo

closeIdleDatabases :: Env -> DiffTimePoints -> [Repo] -> IO ()
closeIdleDatabases env@Env{..} duration blocklist = do
  dbs <- readTVarIO envActive
  let notBlocklisted = filter (not . (`elem` blocklist)) (HashMap.keys dbs)
  forM_ notBlocklisted $ \repo -> closeIdleDatabase env repo duration
  exportOpenDBStats env

-- | set a counter glean.db.<repo>.open to the number of currently
-- open DBs for that particular repo name.
exportOpenDBStats :: Env -> IO ()
exportOpenDBStats Env{..} = do
  opens <- atomically $ do
    dbs <- readTVar envActive
    forM (HashMap.toList dbs) $ \(repo, db) -> do
      state <- readTVar (dbState db)
      case state of
        Open{} -> return [(Thrift.repo_name repo, 1)]
        Opening{} -> return [(Thrift.repo_name repo, 1)]
        _ -> return []
  let repoOpenCounts = HashMap.fromListWith (+) (concat opens)
  forM_ (HashMap.toList repoOpenCounts) $ \(repoNm,count) -> do
    setCounter ("glean.db." <> Text.encodeUtf8 repoNm <> ".open") count

closeOpenDB :: Storage.Storage s => Env -> OpenDB s -> IO ()
closeOpenDB env OpenDB{..} = do
  case odbWriting of
    Just Writing{..} -> do
      -- free memory and update counters
      withMutex wrLock $ const $ LookupCache.clear wrLookupCache
      updateLookupCacheStats env
    Nothing -> return ()
  Storage.close odbHandle