From d703cfae8724b6d41de07b69ec8c8e07205936d4 Mon Sep 17 00:00:00 2001 From: shum Date: Wed, 1 Apr 2026 12:59:48 +0000 Subject: [PATCH] refactor: move usedStorage from FileStore to XFTPEnv --- src/Simplex/FileTransfer/Server.hs | 16 ++++++++++------ src/Simplex/FileTransfer/Server/Env.hs | 5 +++-- src/Simplex/FileTransfer/Server/Store.hs | 20 +++++++------------- 3 files changed, 20 insertions(+), 21 deletions(-) diff --git a/src/Simplex/FileTransfer/Server.hs b/src/Simplex/FileTransfer/Server.hs index 6e0a9735a..75d16e310 100644 --- a/src/Simplex/FileTransfer/Server.hs +++ b/src/Simplex/FileTransfer/Server.hs @@ -549,7 +549,7 @@ processXFTPRequest HTTP2Body {bodyPart} = \case | bs == 0 || bs > s -> pure $ FRErr SIZE | otherwise -> drain (s - bs) reserve = do - us <- asks $ usedStorage . store + us <- asks usedStorage quota <- asks $ fromMaybe maxBound . fileSizeQuota . config atomically . stateTVar us $ \used -> let used' = used + fromIntegral size in if used' <= quota then (True, used') else (False, used) @@ -566,7 +566,7 @@ processXFTPRequest HTTP2Body {bodyPart} = \case liftIO $ atomicModifyIORef'_ (filesSize stats) (+ fromIntegral size) pure FROk Left e -> do - us <- asks $ usedStorage . store + us <- asks usedStorage atomically $ modifyTVar' us $ subtract (fromIntegral size) liftIO $ whenM (doesFileExist fPath) (removeFile fPath) `catch` logFileError pure $ FRErr e @@ -624,6 +624,9 @@ deleteOrBlockServerFile_ FileRec {filePath, fileInfo} stat storeAction = runExce ExceptT $ first (\(_ :: SomeException) -> FILE_IO) <$> try (forM_ path $ \p -> whenM (doesFileExist p) (removeFile p >> deletedStats stats)) st <- asks store void $ atomically $ storeAction st + forM_ path $ \_ -> do + us <- asks usedStorage + atomically $ modifyTVar' us $ subtract (fromIntegral $ size fileInfo) lift $ incFileStat stat where deletedStats stats = do @@ -636,7 +639,8 @@ getFileTime = getRoundedSystemTime expireServerFiles :: Maybe Int -> ExpirationConfig -> M () expireServerFiles itemDelay expCfg = do st <- asks store - usedStart <- readTVarIO $ usedStorage st + us <- asks usedStorage + usedStart <- readTVarIO us old <- liftIO $ expireBeforeEpoch expCfg files' <- readTVarIO (files st) logNote $ "Expiration check: " <> tshow (M.size files') <> " files" @@ -644,7 +648,7 @@ expireServerFiles itemDelay expCfg = do mapM_ threadDelay itemDelay atomically (expiredFilePath st sId old) >>= mapM_ (maybeRemove $ delete st sId) - usedEnd <- readTVarIO $ usedStorage st + usedEnd <- readTVarIO us logNote $ "Used " <> mbs usedStart <> " -> " <> mbs usedEnd <> ", " <> mbs (usedStart - usedEnd) <> " reclaimed." where mbs bs = tshow (bs `div` 1048576) <> "mb" @@ -691,9 +695,9 @@ restoreServerStats = asks (serverStatsBackupFile . config) >>= mapM_ restoreStat liftIO (strDecode <$> B.readFile f) >>= \case Right d@FileServerStatsData {_filesCount = statsFilesCount, _filesSize = statsFilesSize} -> do s <- asks serverStats - FileStore {files, usedStorage} <- asks store + FileStore {files} <- asks store _filesCount <- M.size <$> readTVarIO files - _filesSize <- readTVarIO usedStorage + _filesSize <- readTVarIO =<< asks usedStorage liftIO $ setFileServerStats s d {_filesCount, _filesSize} renameFile f $ f <> ".bak" logNote "server stats restored" diff --git a/src/Simplex/FileTransfer/Server/Env.hs b/src/Simplex/FileTransfer/Server/Env.hs index d4c58df66..f38cc5e9d 100644 --- a/src/Simplex/FileTransfer/Server/Env.hs +++ b/src/Simplex/FileTransfer/Server/Env.hs @@ -91,6 +91,7 @@ defaultInactiveClientExpiration = data XFTPEnv = XFTPEnv { config :: XFTPServerConfig, store :: FileStore, + usedStorage :: TVar Int64, storeLog :: Maybe (StoreLog 'WriteMode), random :: TVar ChaChaDRG, serverIdentity :: C.KeyHash, @@ -115,7 +116,7 @@ newXFTPServerEnv config@XFTPServerConfig {storeLogFile, fileSizeQuota, xftpCrede store <- newFileStore storeLog <- mapM (`readWriteFileStore` store) storeLogFile used <- countUsedStorage <$> readTVarIO (files store) - atomically $ writeTVar (usedStorage store) used + usedStorage <- newTVarIO used forM_ fileSizeQuota $ \quota -> do logNote $ "Total / available storage: " <> tshow quota <> " / " <> tshow (quota - used) when (quota < used) $ logWarn "WARNING: storage quota is less than used storage, no files can be uploaded!" @@ -123,7 +124,7 @@ newXFTPServerEnv config@XFTPServerConfig {storeLogFile, fileSizeQuota, xftpCrede httpServerCreds <- mapM loadServerCredential httpCredentials Fingerprint fp <- loadFingerprint xftpCredentials serverStats <- newFileServerStats =<< getCurrentTime - pure XFTPEnv {config, store, storeLog, random, tlsServerCreds, httpServerCreds, serverIdentity = C.KeyHash fp, serverStats} + pure XFTPEnv {config, store, usedStorage, storeLog, random, tlsServerCreds, httpServerCreds, serverIdentity = C.KeyHash fp, serverStats} countUsedStorage :: M.Map k FileRec -> Int64 countUsedStorage = M.foldl' (\acc FileRec {fileInfo = FileInfo {size}} -> acc + fromIntegral size) 0 diff --git a/src/Simplex/FileTransfer/Server/Store.hs b/src/Simplex/FileTransfer/Server/Store.hs index eec481a21..e3860eae6 100644 --- a/src/Simplex/FileTransfer/Server/Store.hs +++ b/src/Simplex/FileTransfer/Server/Store.hs @@ -25,7 +25,6 @@ module Simplex.FileTransfer.Server.Store where import Control.Concurrent.STM -import Control.Monad import qualified Data.Attoparsec.ByteString.Char8 as A import Data.Int (Int64) import Data.Set (Set) @@ -43,8 +42,7 @@ import Simplex.Messaging.Util (ifM, ($>>=)) data FileStore = FileStore { files :: TMap SenderId FileRec, - recipients :: TMap RecipientId (SenderId, RcvPublicAuthKey), - usedStorage :: TVar Int64 + recipients :: TMap RecipientId (SenderId, RcvPublicAuthKey) } data FileRec = FileRec @@ -72,8 +70,7 @@ newFileStore :: IO FileStore newFileStore = do files <- TM.emptyIO recipients <- TM.emptyIO - usedStorage <- newTVarIO 0 - pure FileStore {files, recipients, usedStorage} + pure FileStore {files, recipients} addFile :: FileStore -> SenderId -> FileInfo -> RoundedFileTime -> ServerEntityStatus -> STM (Either XFTPErrorType ()) addFile FileStore {files} sId fileInfo createdAt status = @@ -91,9 +88,8 @@ newFileRec senderId fileInfo createdAt status = do setFilePath :: FileStore -> SenderId -> FilePath -> STM (Either XFTPErrorType ()) setFilePath st sId fPath = - withFile st sId $ \FileRec {fileInfo, filePath} -> do + withFile st sId $ \FileRec {filePath} -> do writeTVar filePath (Just fPath) - modifyTVar' (usedStorage st) (+ fromIntegral (size fileInfo)) pure $ Right () addRecipient :: FileStore -> SenderId -> FileRecipient -> STM (Either XFTPErrorType ()) @@ -110,19 +106,17 @@ addRecipient st@FileStore {recipients} senderId (FileRecipient rId rKey) = -- this function must be called after the file is deleted from the file system deleteFile :: FileStore -> SenderId -> STM (Either XFTPErrorType ()) -deleteFile FileStore {files, recipients, usedStorage} senderId = do +deleteFile FileStore {files, recipients} senderId = do TM.lookupDelete senderId files >>= \case - Just FileRec {fileInfo, recipientIds} -> do + Just FileRec {recipientIds} -> do readTVar recipientIds >>= mapM_ (`TM.delete` recipients) - modifyTVar' usedStorage $ subtract (fromIntegral $ size fileInfo) pure $ Right () _ -> pure $ Left AUTH -- this function must be called after the file is deleted from the file system blockFile :: FileStore -> SenderId -> BlockingInfo -> Bool -> STM (Either XFTPErrorType ()) -blockFile st@FileStore {usedStorage} senderId info deleted = - withFile st senderId $ \FileRec {fileInfo, fileStatus} -> do - when deleted $ modifyTVar' usedStorage $ subtract (fromIntegral $ size fileInfo) +blockFile st senderId info _deleted = + withFile st senderId $ \FileRec {fileStatus} -> do writeTVar fileStatus $! EntityBlocked info pure $ Right ()