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 TypeApplications #-}
  module LifecycleTest (main) where
  
  import Control.Concurrent.Async
  import Control.Concurrent.MVar
  import Util.STM
  import Control.Exception hiding (assert)
  import Control.Monad
  import Data.Default
  import qualified Data.HashMap.Strict as HashMap
  import Data.List
  import Data.Proxy
  import GHC.Stack (HasCallStack)
  import Test.HUnit
  import Test.QuickCheck.IO ()
  
  import TestRunner
  
  import Glean (fillDatabase, enqueueJsonBatch)
  import qualified Glean.Database.Catalog as Catalog
  import Glean.Database.Catalog.Filter
  import Glean.Database.Catalog.Test
  import Glean.Database.Config
  import Glean.Database.Close
  import Glean.Database.Create
  import Glean.Database.Open
  import Glean.Database.Delete
  import Glean.Database.Storage.Memory (newStorage)
  import Glean.Database.Test
  import Glean.Database.Types
  import Glean.Init
  import Glean.Internal.Types
  import Glean.Test.HUnit
  import Glean.Test.Mock
  import Glean.Types hiding (Exception)
  
  data TEnv = TEnv
    { tEnv :: Env
    , tStore :: MockStore
    }
  
  withTEnv :: (TEnv -> IO a) -> IO a
  withTEnv f = do
    catalog <- memStore
    storage <- newStorage
    withTestEnv
      [\cfg -> cfg
          { cfgDataStore = DataStore
              { withDataStore = \_ f -> f catalog storage
              , dataStoreTag = "test"
              }
          }]
      $ \env -> do
        x <- f $ TEnv env catalog
        checkActive env False
        return x
  
  data TestError = TestError
    deriving(Eq,Show)
  
  instance Exception TestError
  
  data InconsistentDBError = InconsistentDBError Repo String
    deriving(Show)
  
  instance Exception InconsistentDBError
  
  mkDB :: HasCallStack => Env -> Repo -> IO ()
  mkDB env repo =
    fillDatabase env repo Nothing
      (fail "database already exists") (return ())
  
  checkActive :: HasCallStack => Env -> Bool -> IO ()
  checkActive Env{..} allow_deleting = do
    checkConsistency envCatalog
    atomically $ do
      active <- readTVar envActive
      forM_ (HashMap.toList active) $ \(repo, DB{..}) -> do
        ex <- Catalog.exists envCatalog [Local] repo
        when (not ex) $ throwSTM $ InconsistentDBError repo
          "in envActive but not in Catalog"
        users <- readTVar dbUsers
        when (users == 0) $ do
          state <- readTVar dbState
          case state of
            Open{} -> return ()
            Closed -> do
              meta <- Catalog.readMeta envCatalog repo
              case metaCompleteness meta of
                Incomplete{} -> return ()
                _ -> throwSTM $ InconsistentDBError repo
                  "in envActive but not active"
            Opening -> throwSTM $ InconsistentDBError repo "no users but Opening"
            Closing -> throwSTM $ InconsistentDBError repo "no users but Closing"
      deleting <- readTVar envDeleting
      if allow_deleting
        then case HashMap.keys $ HashMap.intersection active deleting of
          [] -> return ()
          repo : _ -> throwSTM $ InconsistentDBError repo
            "in envActive and envDeleting"
        else case HashMap.keys deleting of
          [] -> return ()
          repo : _ -> throwSTM $ InconsistentDBError repo
            "in envDeleting"
  
  kickOff :: HasCallStack => Env -> Repo -> IO Bool
  kickOff env repo = do
    KickOffResponse ex <- kickOffDatabase env def
      { kickOff_repo = repo
      }
    return ex
  
  repo1 :: Repo
  repo1 = Repo "foo" "1"
  
  repo2 :: Repo
  repo2 = Repo "foo" "2"
  
  hasActiveDatabase :: Env -> Repo -> IO Bool
  hasActiveDatabase Env{..} repo = HashMap.member repo <$> readTVarIO envActive
  
  kickOffTwice :: Test
  kickOffTwice = TestCase $ withTEnv $ \TEnv{..} -> do
    assert $ not <$> kickOff tEnv repo1
    assert $ hasActiveDatabase tEnv repo1
    assert $ kickOff tEnv repo1
    assert $ hasActiveDatabase tEnv repo1
  
  kickOffFail :: Test
  kickOffFail = TestCase $ withTEnv $ \TEnv{..} -> do
    prepare (storeCreate tStore) [ doBefore $ throwIO TestError ]
    assertThrows "" TestError $ kickOff tEnv repo1
    assert $ fmap not $ hasActiveDatabase tEnv repo1
  
  kickOffRace :: Test
  kickOffRace = TestCase $ withTEnv $ \TEnv{..} -> do
    v1 <- newEmptyMVar
    v2 <- newEmptyMVar
    prepare (storeCreate tStore) [ doBefore $ do putMVar v1 () ; takeMVar v2 ]
    withAsync (kickOff tEnv repo1) $ \k -> do
      takeMVar v1
      assert $ kickOff tEnv repo1
      assert $ fmap not $ hasActiveDatabase tEnv repo1
      putMVar v2 ()
      assert $ not <$> wait k
    assert $ hasActiveDatabase tEnv repo1
  
  kickOffConcurrently :: Test
  kickOffConcurrently = TestCase $ withTEnv $ \TEnv{..} -> do
    v1 <- newEmptyMVar
    v2 <- newEmptyMVar
    prepare (storeCreate tStore) [ doBefore $ do putMVar v1 () ; takeMVar v2 ]
    withAsync (kickOff tEnv repo1) $ \k -> do
      takeMVar v1
      assert $ not <$> kickOff tEnv repo2
      assert $ fmap not $ hasActiveDatabase tEnv repo1
      assert $ hasActiveDatabase tEnv repo2
      putMVar v2 ()
      assert $ not <$> wait k
    assert $ hasActiveDatabase tEnv repo1
  
  closeIdle :: Test
  closeIdle = TestCase $ withTEnv $ \TEnv{..} -> do
    mkDB tEnv repo1
    withOpenDatabase tEnv repo1 $ \_ -> return ()
    assert $ not <$> isDatabaseClosed tEnv repo1
    -- make sure we release the use count correctly
    assertThrows "" TestError $
      withOpenDatabase tEnv repo1 (\_ -> throwIO TestError)
  
    closeIdleDatabase tEnv repo1 0
    assert $ isDatabaseClosed tEnv repo1
  
  closeUsed :: Test
  closeUsed = TestCase $ withTEnv $ \TEnv{..} -> do
    mkDB tEnv repo1
    withOpenDatabase tEnv repo1 $ \_ -> return ()
    assert $ not <$> isDatabaseClosed tEnv repo1
  
    withOpenDatabase tEnv repo1 $ \_ -> closeIdleDatabase tEnv repo1 0
    assert $ not <$> isDatabaseClosed tEnv repo1
  
  deleteOpen :: Test
  deleteOpen = TestCase $ withTEnv $ \TEnv{..} -> do
    mkDB tEnv repo1
    withOpenDatabase tEnv repo1 $ \_ -> return ()
    assert $ not <$> isDatabaseClosed tEnv repo1
  
    deleteDatabase tEnv repo1
    atomically $ do
      ex <- Catalog.exists (envCatalog tEnv) [Local] repo1
      when ex retry
    assertThrowsType "" (Proxy @UnknownDatabase) $
      withOpenDatabase tEnv repo1 $ \_ -> return ()
  
  deleteWhileUsing :: Test
  deleteWhileUsing = TestCase $ withTEnv $ \TEnv{..} -> do
    mkDB tEnv repo1
    del <- withOpenDatabase tEnv repo1 $ \_ -> do
      del <- asyncDeleteDatabase tEnv repo1
      atomically $ do
        deleting <- readTVar (envDeleting tEnv)
        when (not $ repo1 `HashMap.member` deleting) retry
      checkActive tEnv True
      return del
    wait del
  
  kickOffWhileDeleting :: Test
  kickOffWhileDeleting = TestCase $ withTEnv $ \TEnv{..} -> do
    mkDB tEnv repo1
    withOpenDatabase tEnv repo1 $ \_ -> return ()
    assert $ not <$> isDatabaseClosed tEnv repo1
  
    v1 <- newEmptyMVar
    v2 <- newEmptyMVar
  
    prepare (storeDelete tStore) [ doBefore $ do putMVar v1 () ; takeMVar v2 ]
    concurrently_ (deleteDatabase tEnv repo1) $ do
      takeMVar v1
      True <- kickOff tEnv repo1
      putMVar v2 ()
  
  useWhileDeleting :: Test
  useWhileDeleting = TestCase $ withTEnv $ \TEnv{..} -> do
    mkDB tEnv repo1
    withOpenDatabase tEnv repo1 $ \_ -> return ()
    assert $ not <$> isDatabaseClosed tEnv repo1
  
    v1 <- newEmptyMVar
    v2 <- newEmptyMVar
  
    prepare (storeDelete tStore) [ doBefore $ do putMVar v1 () ; takeMVar v2 ]
    concurrently_ (deleteDatabase tEnv repo1) $ do
      takeMVar v1
      assertThrowsType "" (Proxy @UnknownDatabase) $
        withOpenDatabase tEnv repo1 $ \_ -> return ()
      putMVar v2 ()
  
  writeAfterFinished :: Test
  writeAfterFinished = TestCase $ withTEnv $ \TEnv{..} -> do
    mkDB tEnv repo1
    r <- try $ enqueueJsonBatch tEnv repo1 (def { sendJsonBatch_batches = [] })
    assertBool "writeAfterFinished" $ case r of
      Left err -> "read-only" `isInfixOf` show (err :: SomeException)
      _ -> False
  
  stacked :: Repo -> Dependencies
  stacked (Repo name hash) = Dependencies_stacked $
    Stacked name hash Nothing
  
  kickOffStacked :: Test
  kickOffStacked = TestCase $ withTEnv $ \TEnv{..} -> do
    mkDB tEnv repo1
    waitUntilComplete tEnv repo1
    KickOffResponse ex <- kickOffDatabase tEnv def
      { kickOff_repo = repo2
      , kickOff_dependencies = Just $ stacked repo1
      }
    assert $ not ex
    withOpenDatabase tEnv repo2 $ \_ -> return ()
  
  kickOffStackedIncomplete :: Test
  kickOffStackedIncomplete = TestCase $ withTEnv $ \TEnv{..} -> do
    assert $ not <$> kickOff tEnv repo1
    assertThrowsType "" (Proxy :: Proxy InvalidDependency)
      $ kickOffDatabase tEnv def
        { kickOff_repo = repo2
        , kickOff_dependencies = Just $ stacked repo1
        }
  
  main :: IO ()
  main = withUnitTest $ testRunner $ TestList
    [ TestLabel "kickOffTwice" kickOffTwice
    , TestLabel "kickOffFail" kickOffFail
    , TestLabel "kickOffRace" kickOffRace
    , TestLabel "kickOffConcurrently" kickOffConcurrently
    , TestLabel "closeIdle" closeIdle
    , TestLabel "closeUsed" closeUsed
    , TestLabel "deleteOpen" deleteOpen
    , TestLabel "deleteWhileUsing" deleteWhileUsing
    , TestLabel "kickOffWhileDeleting" kickOffWhileDeleting
    , TestLabel "useWhileDeleting" useWhileDeleting
    , TestLabel "writeAfterFinished" writeAfterFinished
    , TestLabel "kickOffStacked" kickOffStacked
    , TestLabel "kickOffStackedIncomplete" kickOffStackedIncomplete
    ]
  