Merge pull request #224 from lolepezy/fix/erase-all-maps

Erase every map from LMDB when cleaning cache

Author: lolepezy

src/RPKI/Store/Base/LMDB.hs

@@ -9,7 +9,7 @@
 
 module RPKI.Store.Base.LMDB where
 
-import Control.Monad (forM, forever)
+import Control.Monad
 import Control.Concurrent.STM
 
 import qualified Data.ByteString as BS
@@ -35,7 +35,6 @@ import qualified Lmdb.Map as LMap
 import qualified Lmdb.Multimap as LMMap
 import qualified Lmdb.Types as Lmdb
 
-
 import Pipes
 
 type Env = Lmdb.Environment 'Lmdb.ReadWrite
@@ -54,7 +53,7 @@ data LmdbStore (name :: Symbol) = LmdbStore {
 }
 
 data LmdbMultiStore (name :: Symbol) = LmdbMultiStore { 
-    db :: Lmdb.MultiDatabase BS.ByteString BS.ByteString,
+    db  :: Lmdb.MultiDatabase BS.ByteString BS.ByteString,
     env :: LmdbEnv
 }
 
@@ -314,5 +313,24 @@ getMapNames tx db =
         void $ runEffect $ LMap.firstForward c >-> do
             forever $ do
                 Lmdb.KeyValue name _ <- await
-                lift $ modifyIORef' maps ([name] <>)
-        readIORef maps     
+                lift $ modifyIORef' maps (name :)
+        readIORef maps     
+
+
+eraseEnv :: Env -> Tx LmdbStorage 'RW -> IO [BS.ByteString]
+eraseEnv env (LmdbTx tx) = do     
+    db <- openDatabase tx Nothing defaultDbSettings    
+    mapNames <- getMapNames tx db        
+    forM_ mapNames $ \mapName -> do 
+        -- first open it as is
+        m <- openDatabase tx (Just $ convert mapName) defaultDbSettings
+        isMulti <- isMultiDatabase tx m                            
+        if isMulti 
+            then do 
+                -- close and reopen as multi map
+                closeDatabase env m
+                m' <- openMultiDatabase tx (Just $ convert mapName) defaultMultiDbSettngs                        
+                LMMap.clear tx m'
+            else 
+                LMap.clear tx m
+    pure mapNames

src/RPKI/Store/Base/Map.hs

@@ -18,9 +18,6 @@ data SMap (name :: Symbol) s k v where
 instance Storage s => WithStorage s (SMap name s k v) where
     storage (SMap s _) = s
 
-instance WithTx s => CanErase s (SMap name s k v) where
-    erase tx (SMap _ s) = S.clear tx s
-
 put :: (AsStorable k, AsStorable v) =>
         Tx s 'RW -> SMap name s k v -> k -> v -> IO ()
 put tx (SMap _ s) k v = S.put tx s (storableKey k) (storableValue v)    

src/RPKI/Store/Base/MultiMap.hs

@@ -14,9 +14,6 @@ data SMultiMap (name :: Symbol) s k v where
 instance Storage s => WithStorage s (SMultiMap name s k v) where
     storage (SMultiMap s _) = s
 
-instance WithTx s => CanErase s (SMultiMap name s k v) where
-    erase tx (SMultiMap _ s) = S.clearMu tx s
-
 put :: (AsStorable k, AsStorable v) =>
         Tx s 'RW -> SMultiMap name s k v -> k -> v -> IO ()
 put tx (SMultiMap _ s) k v = S.putMu tx s (storableKey k) (storableValue v)    

src/RPKI/Store/Base/Storage.hs

@@ -63,7 +63,3 @@ rwTxT :: (MonadIO m, WithStorage s ws)
 rwTxT tdb f = liftIO $ do 
     db <- readTVarIO tdb
     rwTx db $ \tx -> f tx db
-
-
-class WithTx s => CanErase s a where
-    erase :: Tx s 'RW -> a -> IO ()

src/RPKI/Store/Database.hs

@@ -69,17 +69,14 @@ import           RPKI.Time
 -- It is brittle and inconvenient, but so far seems to be 
 -- the only realistic option.
 currentDatabaseVersion :: Integer
-currentDatabaseVersion = 35
+currentDatabaseVersion = 36
 
 -- Some constant keys
 databaseVersionKey, forAsyncFetchKey, validatedByVersionKey :: Text
 databaseVersionKey = "database-version"
 forAsyncFetchKey  = "for-async-fetch"
 validatedByVersionKey  = "validated-by-version-map"
 
-data EraseWrapper s where 
-    EraseWrapper :: forall t s . (Storage s, CanErase s t) => t -> EraseWrapper s
-
 -- All of the stores of the application in one place
 data DB s = DB {
     taStore          :: TAStore s, 
@@ -96,8 +93,7 @@ data DB s = DB {
     slurmStore       :: SlurmStore s,
     jobStore         :: JobStore s,
     sequences        :: SequenceMap s,
-    metadataStore    :: MetadataStore s,
-    erasables        :: [EraseWrapper s]
+    metadataStore    :: MetadataStore s
 } deriving stock (Generic)
 
 instance Storage s => WithStorage s (DB s) where
@@ -964,10 +960,6 @@ getRtrPayloads tx db worldVersion =
 -- Get all SStats and `<>` them
 totalStats :: StorageStats -> SStats
 totalStats (StorageStats s) = mconcat $ Map.elems s
-
-emptyDBMaps :: (MonadIO m, Storage s) => Tx s 'RW -> DB s -> m ()
-emptyDBMaps tx DB {..} = liftIO $ 
-    forM_ erasables $ \(EraseWrapper t) -> erase tx t
 
 
 -- Utilities to have storage transaction in ValidatorT monad.

src/RPKI/Store/MakeLmdb.hs

@@ -8,7 +8,6 @@ module RPKI.Store.MakeLmdb where
 import Control.Lens
 import Control.Concurrent.STM
 
-import           Data.IORef
 import           Data.String.Interpolate.IsString
 
 import           GHC.TypeLits
@@ -35,9 +34,8 @@ data DbCheckResult = WasIncompatible | WasCompatible | DidntHaveVersion
 
 createDatabase :: LmdbEnv -> AppLogger -> IncompatibleDbCheck -> IO (DB LmdbStorage, DbCheckResult)
 createDatabase env logger checkAction = do 
-
-    erasables <- newIORef []
-    db <- doCreateDb erasables
+    
+    db <- doCreateDb
 
     case checkAction of     
         CheckVersion -> 
@@ -51,24 +49,31 @@ createDatabase env logger checkAction = do
             dbVersion <- getDatabaseVersion tx db
             case dbVersion of 
                 Nothing -> do
-                    logInfo logger [i|Cache version is not set, setting it to #{currentDatabaseVersion}, cleaning up the cache.|]
-                    (_, ms) <- timedMS $ emptyDBMaps tx db
-                    logDebug logger  [i|Erasing cache took #{ms}ms.|]
+                    logInfo logger [i|Cache version is not set, will set the version to #{currentDatabaseVersion} and clean up the cache.|]                    
+                    ms <- eraseCache tx
+                    logDebug logger [i|Erased cache in #{ms}ms.|]
                     saveCurrentDatabaseVersion tx db
                     pure DidntHaveVersion
-                Just version -> 
-                    if version /= currentDatabaseVersion then do
+                Just version
+                    | version == currentDatabaseVersion -> 
+                        pure WasCompatible
+                    | otherwise -> do
                         -- We are seeing incompatible storage. The only option 
                         -- now is to erase all the maps and start from scratch.
-                        logInfo logger [i|Persisted cache version is #{version} and expected version is #{currentDatabaseVersion}, dropping the cache.|]    
-                        (_, ms) <- timedMS $ emptyDBMaps tx db
-                        logDebug logger [i|Erasing cache took #{ms}ms.|]                        
+                        logInfo logger [i|Persisted cache version is #{version} and expected version is #{currentDatabaseVersion}, will drop the cache.|]    
+                        -- NOTE: We erase every map in the cache, including metadata, but that's not an problem, 
+                        -- since we'll set new DB version here
+                        ms <- eraseCache tx
+                        logDebug logger [i|Erased cache in #{ms}ms.|]                  
                         saveCurrentDatabaseVersion tx db
-                        pure WasIncompatible
-                    else
-                        pure WasCompatible
+                        pure WasIncompatible                        
 
-    doCreateDb erasablesRef = do 
+    eraseCache tx = do 
+        nativeEnv <- atomically $ getNativeEnv env
+        (_, ms) <- timedMS $ eraseEnv nativeEnv tx        
+        pure ms
+
+    doCreateDb = do 
         sequences        <- createMap
         taStore          <- TAStore <$> createMap        
         validationsStore <- ValidationsStore <$> createMap
@@ -83,10 +88,7 @@ createDatabase env logger checkAction = do
         jobStore         <- JobStore <$> createMap        
         metadataStore    <- MetadataStore <$> createMap          
         repositoryStore  <- createRepositoryStore
-        objectStore      <- createObjectStore sequences            
-
-        erasables <- readIORef erasablesRef
-
+        objectStore      <- createObjectStore sequences
         pure DB {..}
       where
 
@@ -112,16 +114,10 @@ createDatabase env logger checkAction = do
         lmdb = LmdbStorage env
 
         createMap :: forall k v name . (KnownSymbol name) => IO (SMap name LmdbStorage k v)
-        createMap = do 
-            sm <- SMap lmdb <$> createLmdbStore env
-            modifyIORef' erasablesRef (EraseWrapper sm :)
-            pure sm
+        createMap = SMap lmdb <$> createLmdbStore env            
 
         createMultiMap :: forall k v name . (KnownSymbol name) => IO (SMultiMap name LmdbStorage k v)
-        createMultiMap = do 
-            sm <- SMultiMap lmdb <$> createLmdbMultiStore env
-            modifyIORef' erasablesRef (EraseWrapper sm :)
-            pure sm
+        createMultiMap = SMultiMap lmdb <$> createLmdbMultiStore env            
 
 
 mkLmdb :: FilePath -> Config -> IO LmdbEnv