Fix memory leak caused by difference maps.

concordium-consensus/src/Concordium/GlobalState/AccountMap/DifferenceMap.hs

@@ -5,6 +5,8 @@
 --  to disk via 'Concordium.GlobalState.AccountMap.LMDB.insert'.
 module Concordium.GlobalState.AccountMap.DifferenceMap where
 
+import Control.Monad.IO.Class
+import Data.IORef
 import qualified Data.Map.Strict as Map
 import Prelude hiding (lookup)
 
@@ -20,26 +22,26 @@ data DifferenceMap = DifferenceMap
       --  In other words, if the parent block is finalized,
       --  then the parent map is @Notnhing@ as the LMDB account map
       --  should be consulted instead.
-      dmParentMap :: !(Maybe DifferenceMap)
+      dmParentMap :: !(IORef (Maybe DifferenceMap))
     }
-    deriving (Eq, Show)
+    deriving (Eq)
 
 -- | Gather all accounts from the provided 'DifferenceMap' and its parent maps.
 --  Accounts are returned in ascending order of their 'AccountIndex'.
-flatten :: DifferenceMap -> [(AccountAddress, AccountIndex)]
+flatten :: (MonadIO m) => DifferenceMap -> m [(AccountAddress, AccountIndex)]
 flatten dmap = go dmap []
   where
-    go :: DifferenceMap -> [(AccountAddress, AccountIndex)] -> [(AccountAddress, AccountIndex)]
-    go DifferenceMap{dmParentMap = Nothing, ..} !accum =
-        let !listOfAccounts = Map.toList dmAccounts
-        in  listOfAccounts <> accum
-    go DifferenceMap{dmParentMap = Just parentMap, ..} !accum =
-        let !listOfAccounts = Map.toList dmAccounts
-        in  go parentMap $! listOfAccounts <> accum
+    go diffMap !accum = do
+        mParentMap <- liftIO $ readIORef (dmParentMap diffMap)
+        case mParentMap of
+            Nothing -> return collectedAccounts
+            Just parentMap -> go parentMap collectedAccounts
+      where
+        collectedAccounts = Map.toList (dmAccounts diffMap) <> accum
 
--- | Create a new empty 'DifferenceMap' based on the difference map of
+-- | Create a new empty 'DifferenceMap' potentially based on the difference map of
 -- the parent.
-empty :: Maybe DifferenceMap -> DifferenceMap
+empty :: IORef (Maybe DifferenceMap) -> DifferenceMap
 empty mParentDifferenceMap =
     DifferenceMap
         { dmAccounts = Map.empty,
@@ -50,15 +52,19 @@ empty mParentDifferenceMap =
 --  difference maps.
 --  Returns @Just AccountIndex@ if the account is present and
 --  otherwise @Nothing@.
-lookup :: AccountAddress -> DifferenceMap -> Maybe AccountIndex
+lookup :: (MonadIO m) => AccountAddress -> DifferenceMap -> m (Maybe AccountIndex)
 lookup addr = check
   where
-    check DifferenceMap{..} = case Map.lookupGE k dmAccounts of
-        Nothing -> check =<< dmParentMap
+    check diffMap = case Map.lookupGE k (dmAccounts diffMap) of
+        Nothing -> do
+            mParentMap <- liftIO $ readIORef (dmParentMap diffMap)
+            case mParentMap of
+                Nothing -> return Nothing
+                Just parentMap -> check parentMap
         Just (foundAccAddr, accIdx) ->
             if checkEquivalence foundAccAddr
-                then Just accIdx
-                else Nothing
+                then return $ Just accIdx
+                else return Nothing
     k = createAlias addr 0
     checkEquivalence found = accountAddressEmbed k == accountAddressEmbed found
 

concordium-consensus/src/Concordium/GlobalState/Persistent/Accounts.hs

@@ -66,6 +66,7 @@ module Concordium.GlobalState.Persistent.Accounts where
 
 import Control.Monad.Reader
 import Data.Foldable (foldlM)
+import Data.IORef
 import qualified Data.Map.Strict as Map
 import Data.Maybe
 import Data.Serialize
@@ -133,11 +134,11 @@ data Accounts (pv :: ProtocolVersion) = Accounts
       --  Otherwise if the block is not persisted and accounts have been added, then
       --  the 'DiffMap.DifferenceMap' yields the @AccountAddress -> AccountIndex@ mapping for
       --  accounts created in the block.
-      accountDiffMap :: !(Maybe DiffMap.DifferenceMap)
+      accountDiffMap :: !(IORef (Maybe DiffMap.DifferenceMap))
     }
 
 instance (IsProtocolVersion pv) => Show (Accounts pv) where
-    show accts = "Accounts: " <> show (accountTable accts) <> "DiffMap: " <> show (accountDiffMap accts)
+    show accts = "Accounts: " <> show (accountTable accts)
 
 -- | A constraint that ensures a monad @m@ supports the persistent account operations.
 --  This essentially requires that the monad support 'MonadBlobStore', and 'MonadCache' for
@@ -162,16 +163,20 @@ instance (SupportsPersistentAccount pv m) => BlobStorable m (Accounts pv) where
     storeUpdate Accounts{..} = do
         (pTable, accountTable') <- storeUpdate accountTable
         (pRegIdHistory, regIdHistory') <- storeUpdate accountRegIdHistory
-        case accountDiffMap of
+        mAccDiffMap <- liftIO $ readIORef accountDiffMap
+        case mAccDiffMap of
             Nothing -> return ()
-            Just accDiffMap -> LMDBAccountMap.insert $! DiffMap.flatten accDiffMap
+            Just accDiffMap -> do
+                flattenedAccounts <- liftIO $ DiffMap.flatten accDiffMap
+                LMDBAccountMap.insert flattenedAccounts
+                -- The difference map is set to @Nothing@ as any potential new accounts have been written to the
+                -- lmdb account map.
+                liftIO $ modifyIORef' accountDiffMap (const Nothing)
         let newAccounts =
                 Accounts
                     { accountTable = accountTable',
                       accountRegIdHistory = regIdHistory',
-                      -- The difference map is set to @Nothing@ as any potential new accounts have been written to the
-                      -- lmdb account map.
-                      accountDiffMap = Nothing
+                      ..
                     }
 
         -- put an empty 'OldMap.PersistentAccountMap'.
@@ -191,7 +196,7 @@ instance (SupportsPersistentAccount pv m) => BlobStorable m (Accounts pv) where
         return $ do
             accountTable <- maccountTable
             accountRegIdHistory <- mrRIH
-            let accountDiffMap = Nothing
+            accountDiffMap <- liftIO $ newIORef Nothing
             return $ Accounts{..}
 
 instance (SupportsPersistentAccount pv m, av ~ AccountVersionFor pv) => Cacheable1 m (Accounts pv) (PersistentAccount av) where
@@ -200,9 +205,10 @@ instance (SupportsPersistentAccount pv m, av ~ AccountVersionFor pv) => Cacheabl
         return
             accts{accountTable = acctTable}
 
-emptyAccounts :: Maybe DiffMap.DifferenceMap -> Accounts pv
-emptyAccounts Nothing = Accounts L.empty Trie.empty Nothing
-emptyAccounts successorDiffMap = Accounts L.empty Trie.empty $ Just $ DiffMap.empty successorDiffMap
+emptyAccounts :: (MonadIO m) => m (Accounts pv)
+emptyAccounts = do
+    accountDiffMap <- liftIO $ newIORef Nothing
+    return $ Accounts L.empty Trie.empty accountDiffMap
 
 -- | Add a new account. Returns @Just idx@ if the new account is fresh, i.e., the address does not exist,
 --  or @Nothing@ in case the account already exists. In the latter case there is no change to the accounts structure.
@@ -213,14 +219,21 @@ putNewAccount !acct a0@Accounts{..} = do
         True -> return (Nothing, a0)
         False -> do
             (accIdx, newAccountTable) <- L.append acct accountTable
-            let newDiffMap = case accountDiffMap of
-                    Nothing -> DiffMap.insert addr accIdx $ DiffMap.empty Nothing
-                    Just accDiffMap -> DiffMap.insert addr accIdx accDiffMap
-            return (Just accIdx, a0{accountTable = newAccountTable, accountDiffMap = Just newDiffMap})
+            mAccountDiffMap <- liftIO $ readIORef accountDiffMap
+            newDiffMap <- case mAccountDiffMap of
+                Nothing -> do
+                    freshDifferenceMap <- liftIO $ newIORef (Nothing :: Maybe DiffMap.DifferenceMap)
+                    return $ DiffMap.insert addr accIdx $ DiffMap.empty freshDifferenceMap
+                Just accDiffMap -> do
+                    return $ DiffMap.insert addr accIdx accDiffMap
+            liftIO $ modifyIORef' accountDiffMap (const $ Just newDiffMap)
+            return (Just accIdx, a0{accountTable = newAccountTable})
 
 -- | Construct an 'Accounts' from a list of accounts. Inserted in the order of the list.
 fromList :: (SupportsPersistentAccount pv m) => [PersistentAccount (AccountVersionFor pv)] -> m (Accounts pv)
-fromList = foldlM insert $ emptyAccounts Nothing
+fromList accs = do
+    accum <- emptyAccounts
+    foldlM insert accum accs
   where
     insert accounts account = snd <$> putNewAccount account accounts
 
@@ -238,12 +251,14 @@ getAccountByCredId cid accs@Accounts{..} =
 
 -- | Get the account at a given index (if any).
 getAccountIndex :: (SupportsPersistentAccount pv m) => AccountAddress -> Accounts pv -> m (Maybe AccountIndex)
-getAccountIndex addr Accounts{..} =
-    case accountDiffMap of
+getAccountIndex addr Accounts{..} = do
+    mAccountDiffMap <- liftIO $ readIORef accountDiffMap
+    case mAccountDiffMap of
         Nothing -> lookupDisk
-        Just accDiffMap -> case DiffMap.lookup addr accDiffMap of
-            Just accIdx -> return $ Just accIdx
-            Nothing -> lookupDisk
+        Just accDiffMap ->
+            DiffMap.lookup addr accDiffMap >>= \case
+                Just accIdx -> return $ Just accIdx
+                Nothing -> lookupDisk
   where
     -- Lookup the 'AccountIndex' in the lmdb backed account map.
     lookupDisk =
@@ -343,9 +358,12 @@ updateAccountsAtIndex' fupd ai = fmap snd . updateAccountsAtIndex fupd' ai
 allAccounts :: (SupportsPersistentAccount pv m) => Accounts pv -> m [(AccountAddress, AccountIndex)]
 allAccounts accounts = do
     persistedAccs <- LMDBAccountMap.all
-    case accountDiffMap accounts of
+    mDiffMap <- liftIO $ readIORef (accountDiffMap accounts)
+    case mDiffMap of
         Nothing -> return persistedAccs
-        Just accDiffMap -> return $! persistedAccs <> DiffMap.flatten accDiffMap
+        Just accDiffMap -> do
+            flattenedDiffMapAccounts <- DiffMap.flatten accDiffMap
+            return $! persistedAccs <> flattenedDiffMapAccounts
 
 -- | Get a list of all account addresses.
 accountAddresses :: (SupportsPersistentAccount pv m) => Accounts pv -> m [AccountAddress]
@@ -405,9 +423,10 @@ migrateAccounts migration Accounts{..} = do
     -- The account registration IDs are not cached. There is a separate cache
     -- that is purely in-memory and just copied over.
     newAccountRegIds <- Trie.migrateUnbufferedTrieN return accountRegIdHistory
+    emptyAccountDiffMap <- liftIO $ newIORef (Nothing :: Maybe DiffMap.DifferenceMap)
     return $!
         Accounts
             { accountTable = newAccountTable,
               accountRegIdHistory = newAccountRegIds,
-              accountDiffMap = Nothing
+              accountDiffMap = emptyAccountDiffMap
             }

concordium-consensus/src/Concordium/GlobalState/Persistent/BlockState.hs

@@ -933,10 +933,11 @@ emptyBlockState bspBirkParameters cryptParams keysCollection chainParams = do
     bspUpdates <- refMake =<< initialUpdates keysCollection chainParams
     bspReleaseSchedule <- emptyReleaseSchedule
     bspRewardDetails <- emptyBlockRewardDetails
+    bspAccounts <- Accounts.emptyAccounts
     bsp <-
         makeBufferedRef $
             BlockStatePointers
-                { bspAccounts = Accounts.emptyAccounts Nothing,
+                { bspAccounts = bspAccounts,
                   bspInstances = Instances.emptyInstances,
                   bspModules = modules,
                   bspBank = makeHashed Rewards.emptyBankStatus,
@@ -3710,9 +3711,14 @@ doThawBlockState ::
 doThawBlockState HashedPersistentBlockState{..} = do
     -- This load is cheap as the underlying block state is retained in memory as we're building from it, so it must be the "best" block.
     bsp@BlockStatePointers{bspAccounts = a0@Accounts.Accounts{..}} <- loadPBS hpbsPointers
-    let newDiffMap = case accountDiffMap of
-            Nothing -> Nothing
-            Just diffMap -> Just $ DiffMap.empty (Just diffMap)
+    mDiffMap <- liftIO $ readIORef accountDiffMap
+    newDiffMap <- case mDiffMap of
+        -- reuse the reference pointing to @Nothing@.
+        Nothing -> return accountDiffMap
+        Just _ -> do
+            -- create a new reference pointing to
+            -- a new difference map which inherits the parent difference map.
+            liftIO $ newIORef $ Just (DiffMap.empty accountDiffMap)
     let bsp' = bsp{bspAccounts = a0{Accounts.accountDiffMap = newDiffMap}}
     liftIO $ newIORef =<< makeBufferedRef bsp'
 

concordium-consensus/src/Concordium/GlobalState/Persistent/Genesis.hs

@@ -110,18 +110,20 @@ data AccumGenesisState pv = AccumGenesisState
 --------- Helper functions ----------
 
 -- | The initial value for accumulating data from genesis data accounts.
-initialAccumGenesisState :: AccumGenesisState pv
-initialAccumGenesisState =
-    AccumGenesisState
-        { agsAllAccounts = Accounts.emptyAccounts Nothing,
-          agsBakerIds = Trie.empty,
-          agsBakerKeys = Trie.empty,
-          agsTotal = 0,
-          agsStakedTotal = 0,
-          agsBakerInfoRefs = Vec.empty,
-          agsBakerStakes = Vec.empty,
-          agsBakerCapitals = Vec.empty
-        }
+initialAccumGenesisState :: (MTL.MonadIO m) => m (AccumGenesisState pv)
+initialAccumGenesisState = do
+    emptyAccs <- Accounts.emptyAccounts
+    return $
+        AccumGenesisState
+            { agsAllAccounts = emptyAccs,
+              agsBakerIds = Trie.empty,
+              agsBakerKeys = Trie.empty,
+              agsTotal = 0,
+              agsStakedTotal = 0,
+              agsBakerInfoRefs = Vec.empty,
+              agsBakerStakes = Vec.empty,
+              agsBakerCapitals = Vec.empty
+            }
 
 -- | Construct a hashed persistent block state from the data in genesis.
 -- The result is immediately flushed to disc and cached.
@@ -132,8 +134,9 @@ buildGenesisBlockState ::
     GenesisData.GenesisState pv ->
     MTL.ExceptT String m (BS.HashedPersistentBlockState pv, TransactionTable.TransactionTable)
 buildGenesisBlockState vcgp GenesisData.GenesisState{..} = do
+    initState <- initialAccumGenesisState
     -- Iterate the accounts in genesis once and accumulate all relevant information.
-    AccumGenesisState{..} <- Vec.ifoldM' accumStateFromGenesisAccounts initialAccumGenesisState genesisAccounts
+    AccumGenesisState{..} <- Vec.ifoldM' accumStateFromGenesisAccounts initState genesisAccounts
 
     -- Birk parameters
     persistentBirkParameters :: BS.PersistentBirkParameters pv <- do