Account for pending XFTP files in storage quota

src/Simplex/FileTransfer/Server.hs

@@ -16,6 +16,7 @@
 module Simplex.FileTransfer.Server
   ( runXFTPServer,
     runXFTPServerBlocking,
+    checkedStorageRelease,
   ) where
 
 import Control.Logger.Simple
@@ -491,7 +492,10 @@ processXFTPRequest HTTP2Body {bodyPart} = \case
         unless (size file `elem` sizes) $ throwE SIZE
         ts <- liftIO getFileTime
         -- TODO validate body empty
-        sId <- ExceptT $ addFileRetry st file 3 ts
+        unlessM (lift $ reserveStorage $ fromIntegral $ size file) $ throwE QUOTA
+        sId <- lift (addFileRetry st file 3 ts) >>= \case
+          Left e -> lift (releaseStorage $ fromIntegral $ size file) >> throwE e
+          Right sId -> pure sId
         rcps <- mapM (ExceptT . addRecipientRetry st 3 sId) rks
         lift $ withFileLog $ \sl -> do
           logAddFile sl sId file ts EntityActive
@@ -535,7 +539,7 @@ processXFTPRequest HTTP2Body {bodyPart} = \case
     receiveServerFile FileRec {senderId, fileInfo = FileInfo {size, digest}, filePath} = case bodyPart of
       Nothing -> pure $ FRErr SIZE
       -- TODO validate body size from request before downloading, once it's populated
-      Just getBody -> skipCommitted $ ifM reserve receive (pure $ FRErr QUOTA)
+      Just getBody -> skipCommitted receive
         where
           -- having a filePath means the file is already uploaded and committed, must not change anything
           skipCommitted = ifM (isJust <$> readTVarIO filePath) (liftIO $ drain $ fromIntegral size)
@@ -548,11 +552,6 @@ processXFTPRequest HTTP2Body {bodyPart} = \case
                   | bs == s -> pure FROk
                   | bs == 0 || bs > s -> pure $ FRErr SIZE
                   | otherwise -> drain (s - bs)
-          reserve = do
-            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)
           receive = do
             path <- asks $ filesPath . config
             let fPath = path </> B.unpack (B64.encode $ unEntityId senderId)
@@ -568,13 +567,9 @@ processXFTPRequest HTTP2Body {bodyPart} = \case
                     liftIO $ atomicModifyIORef'_ (filesSize stats) (+ fromIntegral size)
                     pure FROk
                   Left _e -> do
-                    us <- asks usedStorage
-                    atomically $ modifyTVar' us $ subtract (fromIntegral size)
                     liftIO $ whenM (doesFileExist fPath) (removeFile fPath) `catch` logFileError
                     pure $ FRErr AUTH
               Left e -> do
-                us <- asks usedStorage
-                atomically $ modifyTVar' us $ subtract (fromIntegral size)
                 liftIO $ whenM (doesFileExist fPath) (removeFile fPath) `catch` logFileError
                 pure $ FRErr e
           receiveChunk spec = do
@@ -616,24 +611,22 @@ processXFTPRequest HTTP2Body {bodyPart} = \case
 deleteServerFile_ :: FileStoreClass s => FileRec -> M s (Either XFTPErrorType ())
 deleteServerFile_ fr@FileRec {senderId} = do
   withFileLog (`logDeleteFile` senderId)
-  deleteOrBlockServerFile_ fr filesDeleted (`deleteFile` senderId)
+  deleteOrBlockServerFile_ fr filesDeleted $ \st -> fmap (fmap Just) $ deleteFileSize st senderId
 
 -- this also deletes the file from storage, but doesn't include it in delete statistics
 blockServerFile :: FileStoreClass s => FileRec -> BlockingInfo -> M s (Either XFTPErrorType ())
 blockServerFile fr@FileRec {senderId} info = do
   withFileLog $ \sl -> logBlockFile sl senderId info
-  deleteOrBlockServerFile_ fr filesBlocked $ \st -> blockFile st senderId info True
+  deleteOrBlockServerFile_ fr filesBlocked $ \st -> fmap (fmap $ const Nothing) $ blockFile st senderId info True
 
-deleteOrBlockServerFile_ :: FileStoreClass s => FileRec -> (FileServerStats -> IORef Int) -> (s -> IO (Either XFTPErrorType ())) -> M s (Either XFTPErrorType ())
+deleteOrBlockServerFile_ :: FileRec -> (FileServerStats -> IORef Int) -> (s -> IO (Either XFTPErrorType (Maybe Word32))) -> M s (Either XFTPErrorType ())
 deleteOrBlockServerFile_ FileRec {filePath, fileInfo} stat storeAction = runExceptT $ do
   path <- readTVarIO filePath
   stats <- asks serverStats
   ExceptT $ first (\(_ :: SomeException) -> FILE_IO) <$> try (forM_ path $ \p -> whenM (doesFileExist p) (removeFile p >> deletedStats stats))
   st <- asks fileStore
-  ExceptT $ liftIO $ storeAction st
-  forM_ path $ \_ -> do
-    us <- asks usedStorage
-    atomically $ modifyTVar' us $ subtract (fromIntegral $ size fileInfo)
+  released <- ExceptT $ liftIO $ storeAction st
+  forM_ released $ lift . releaseStorage . fromIntegral
   lift $ incFileStat stat
   where
     deletedStats stats = do
@@ -651,26 +644,42 @@ expireServerFiles itemDelay expCfg = do
   old <- liftIO $ expireBeforeEpoch expCfg
   filesCount <- liftIO $ getFileCount st
   logNote $ "Expiration check: " <> tshow filesCount <> " files"
-  expireLoop st us old
+  expireLoop st old
   usedEnd <- readTVarIO us
   logNote $ "Used " <> mbs usedStart <> " -> " <> mbs usedEnd <> ", " <> mbs (usedStart - usedEnd) <> " reclaimed."
   where
     mbs bs = tshow (bs `div` 1048576) <> "mb"
-    expireLoop st us old = do
+    expireLoop st old = do
       expired <- liftIO $ expiredFiles st old 10000
-      forM_ expired $ \(sId, filePath_, fileSize) -> do
+      forM_ expired $ \(_sId, filePath_, _fileSize) -> do
         mapM_ threadDelay itemDelay
         forM_ filePath_ $ \fp ->
           whenM (doesFileExist fp) $
             removeFile fp `catch` \(e :: SomeException) -> logError $ "failed to remove expired file " <> tshow fp <> ": " <> tshow e
-        forM_ filePath_ $ \_ ->
-          atomically $ modifyTVar' us $ subtract (fromIntegral fileSize)
         incFileStat filesExpired
       let sIds = map (\(sId, _, _) -> sId) expired
       unless (null sIds) $ do
         withFileLog $ \sl -> mapM_ (logDeleteFile sl) sIds
-        liftIO $ deleteFiles st sIds
-        expireLoop st us old
+        deletedSizes <- liftIO $ deleteFilesSizes st sIds
+        forM_ deletedSizes $ releaseStorage . fromIntegral
+        expireLoop st old
+
+reserveStorage :: Int64 -> M s Bool
+reserveStorage size = do
+  us <- asks usedStorage
+  quota <- asks $ fromMaybe maxBound . fileSizeQuota . config
+  atomically . stateTVar us $ \used ->
+    if used <= quota && size <= quota - used then (True, used + size) else (False, used)
+
+releaseStorage :: Int64 -> M s ()
+releaseStorage size = do
+  us <- asks usedStorage
+  released <- atomically . stateTVar us $ \used ->
+    maybe (False, used) (True,) $ checkedStorageRelease used size
+  unless released $ logError "File storage reservation underflow"
+
+checkedStorageRelease :: Int64 -> Int64 -> Maybe Int64
+checkedStorageRelease used size = (used - size) <$ guard (0 <= size && size <= used)
 
 randomId :: Int -> M s ByteString
 randomId n = atomically . C.randomBytes n =<< asks random

src/Simplex/FileTransfer/Server/Store.hs

@@ -26,7 +26,8 @@ import Control.Monad
 import qualified Data.Attoparsec.ByteString.Char8 as A
 import Data.Int (Int64)
 import qualified Data.Map.Strict as M
-import Data.Maybe (catMaybes, isJust)
+import Data.Either (rights)
+import Data.Maybe (catMaybes)
 import Data.Set (Set)
 import qualified Data.Set as S
 import Data.Word (Word32)
@@ -77,9 +78,13 @@ class FileStoreClass s where
   addFile :: s -> SenderId -> FileInfo -> RoundedFileTime -> ServerEntityStatus -> IO (Either XFTPErrorType ())
   setFilePath :: s -> SenderId -> FilePath -> IO (Either XFTPErrorType ())
   addRecipient :: s -> SenderId -> FileRecipient -> IO (Either XFTPErrorType ())
+  deleteFileSize :: s -> SenderId -> IO (Either XFTPErrorType Word32)
   deleteFile :: s -> SenderId -> IO (Either XFTPErrorType ())
+  deleteFile s = fmap (fmap $ const ()) . deleteFileSize s
+  deleteFilesSizes :: s -> [SenderId] -> IO [Word32]
+  deleteFilesSizes s = fmap rights . mapM (deleteFileSize s)
   deleteFiles :: s -> [SenderId] -> IO ()
-  deleteFiles s = mapM_ (void . deleteFile s)
+  deleteFiles s = void . deleteFilesSizes s
   blockFile :: s -> SenderId -> BlockingInfo -> Bool -> IO (Either XFTPErrorType ())
   deleteRecipient :: s -> RecipientId -> FileRec -> IO ()
   getFile :: s -> SFileParty p -> XFTPFileId -> IO (Either XFTPErrorType (FileRec, C.APublicAuthKey))
@@ -135,11 +140,11 @@ instance FileStoreClass STMFileStore where
           TM.insert rId (senderId, rKey) recipients
           pure $ Right ()
 
-  deleteFile STMFileStore {files, recipients} senderId = atomically $ do
+  deleteFileSize STMFileStore {files, recipients} senderId = atomically $ do
     TM.lookupDelete senderId files >>= \case
-      Just FileRec {recipientIds} -> do
+      Just FileRec {fileInfo = FileInfo {size}, recipientIds} -> do
         readTVar recipientIds >>= mapM_ (`TM.delete` recipients)
-        pure $ Right ()
+        pure $ Right size
       _ -> pure $ Left AUTH
 
   blockFile st senderId info _deleted = atomically $
@@ -177,8 +182,7 @@ instance FileStoreClass STMFileStore where
 
   getUsedStorage STMFileStore {files} = foldM addSize 0 =<< readTVarIO files
     where
-      addSize acc FileRec {fileInfo = FileInfo {size}, filePath} =
-        ifM (isJust <$> readTVarIO filePath) (pure $! acc + fromIntegral size) (pure acc)
+      addSize acc FileRec {fileInfo = FileInfo {size}} = pure $! acc + fromIntegral size
 
   getFileCount STMFileStore {files} = M.size <$> readTVarIO files
 

src/Simplex/FileTransfer/Server/Store/Postgres.hs

@@ -110,15 +110,20 @@ instance FileStoreClass PostgresFileStore where
         >>= either handleDuplicate (pure . Right)
     withLog "addRecipient" st $ \s -> logAddRecipients s senderId (pure $ FileRecipient rId rKey)
 
-  deleteFile st sId = E.uninterruptibleMask_ $ runExceptT $ do
-    assertUpdated $ withDB' "deleteFile" st $ \db ->
-      DB.execute db "DELETE FROM files WHERE sender_id = ?" (Only sId)
+  deleteFileSize st sId = E.uninterruptibleMask_ $ runExceptT $ do
+    sizes <- withDB' "deleteFile" st $ \db ->
+      DB.query db "DELETE FROM files WHERE sender_id = ? RETURNING file_size" (Only sId)
+    size <- case sizes of
+      [Only (n :: Int32)] -> pure $ fromIntegral n
+      _ -> throwE AUTH
     withLog "deleteFile" st $ \s -> logDeleteFile s sId
+    pure size
 
-  deleteFiles st sIds = E.uninterruptibleMask_ $ do
-    withTransaction (dbStore st) $ \db ->
-      DB.execute db "DELETE FROM files WHERE sender_id IN ?" (Only (In sIds))
+  deleteFilesSizes st sIds = E.uninterruptibleMask_ $ do
+    sizes <- withTransaction (dbStore st) $ \db ->
+      DB.query db "DELETE FROM files WHERE sender_id IN ? RETURNING file_size" (Only (In sIds)) :: IO [Only Int32]
     withLog "deleteFiles" st $ \s -> mapM_ (logDeleteFile s) sIds
+    pure $ map (fromIntegral . fromOnly) sizes
 
   blockFile st sId info _deleted = E.uninterruptibleMask_ $ runExceptT $ do
     assertUpdated $ withDB' "blockFile" st $ \db ->
@@ -164,7 +169,7 @@ instance FileStoreClass PostgresFileStore where
 
   getUsedStorage st =
     withTransaction (dbStore st) $ \db -> do
-      [Only total] <- DB.query_ db "SELECT COALESCE(SUM(file_size::BIGINT), 0)::BIGINT FROM files WHERE file_path IS NOT NULL"
+      [Only total] <- DB.query_ db "SELECT COALESCE(SUM(file_size::BIGINT), 0)::BIGINT FROM files"
       pure total
 
   getFileCount st =

tests/CoreTests/XFTPStoreTests.hs

@@ -18,14 +18,17 @@ import Simplex.Messaging.Server.StoreLog (openWriteStoreLog)
 import Simplex.Messaging.SystemTime (RoundedSystemTime (..))
 import System.Directory (doesFileExist, removeFile)
 import Test.Hspec hiding (fit, it)
+import UnliftIO.Async (concurrently)
 import UnliftIO.STM
 import Util
 import XFTPClient (testXFTPPostgresCfg)
 
 xftpStoreTests :: Spec
 xftpStoreTests = do
-  describe "STMFileStore operations" $
-    it "should compute committed used storage and file count" testSTMStorageAndCount
+  describe "STMFileStore operations" $ do
+    it "should compute reserved storage and file count" testSTMStorageAndCount
+    it "should return one reservation owner across racing deletes" testSTMDeletionOwnership
+    it "should retain blocked reservations until deletion" testSTMBlockedReservation
   describe "PostgresFileStore operations" $ do
     it "should add and get file by sender" testAddGetFileSender
     it "should add and get file by recipient" testAddGetFileRecipient
@@ -37,7 +40,9 @@ xftpStoreTests = do
     it "should block file and update status" testBlockFile
     it "should ack file reception" testAckFile
     it "should return expired files with limit" testExpiredFiles
-    it "should compute committed used storage and file count" testStorageAndCount
+    it "should compute reserved storage and file count" testStorageAndCount
+    it "should return one reservation owner across racing deletes" testDeletionOwnership
+    it "should retain blocked reservations until deletion" testBlockedReservation
 
 xftpMigrationTests :: Spec
 xftpMigrationTests = describe "XFTP migration round-trip" $ do
@@ -212,6 +217,18 @@ testSTMStorageAndCount = do
   testStorageAndCountForStore st
   closeFileStore st
 
+testSTMDeletionOwnership :: Expectation
+testSTMDeletionOwnership = do
+  st <- newFileStore () :: IO STMFileStore
+  testDeletionOwnershipForStore st
+  closeFileStore st
+
+testSTMBlockedReservation :: Expectation
+testSTMBlockedReservation = do
+  st <- newFileStore () :: IO STMFileStore
+  testBlockedReservationForStore st
+  closeFileStore st
+
 testStorageAndCountForStore :: FileStoreClass s => s -> Expectation
 testStorageAndCountForStore st = do
   g <- C.newRandom
@@ -225,12 +242,55 @@ testStorageAndCountForStore st = do
   addFile st fileA fileInfoA testCreatedAt EntityActive `shouldReturn` Right ()
   addFile st fileB fileInfoB testCreatedAt EntityActive `shouldReturn` Right ()
   getFileCount st `shouldReturn` 2
-  getUsedStorage st `shouldReturn` 0
+  getUsedStorage st `shouldReturn` 192000
   setFilePath st fileA "/tmp/file_a" `shouldReturn` Right ()
-  getUsedStorage st `shouldReturn` 128000
+  getUsedStorage st `shouldReturn` 192000
   setFilePath st fileB "/tmp/file_b" `shouldReturn` Right ()
   getUsedStorage st `shouldReturn` 192000
 
+testDeletionOwnership :: Expectation
+testDeletionOwnership = withPgStore testDeletionOwnershipForStore
+
+testDeletionOwnershipForStore :: FileStoreClass s => s -> Expectation
+testDeletionOwnershipForStore st = do
+  g <- C.newRandom
+  (sndKey, _) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
+  let fileInfo = testFileInfo sndKey
+      bulkFirst = EntityId "bulk_first______"
+      singleFirst = EntityId "single_first____"
+      add sId = addFile st sId fileInfo (RoundedSystemTime 100000) EntityActive `shouldReturn` Right ()
+  add bulkFirst
+  bulkSnapshot <- map (\(sId, _, _) -> sId) <$> expiredFiles st 500000 100
+  deleteFilesSizes st bulkSnapshot `shouldReturn` [128000]
+  deleteFileSize st bulkFirst `shouldReturn` Left AUTH
+  add singleFirst
+  singleSnapshot <- map (\(sId, _, _) -> sId) <$> expiredFiles st 500000 100
+  deleteFileSize st singleFirst `shouldReturn` Right 128000
+  deleteFilesSizes st singleSnapshot `shouldReturn` []
+  add testSenderId
+  (single, bulk) <- concurrently (deleteFileSize st testSenderId) (deleteFilesSizes st [testSenderId])
+  let released = either (const []) pure single <> bulk
+  released `shouldBe` [128000]
+  deleteFileSize st testSenderId `shouldReturn` Left AUTH
+  deleteFilesSizes st [testSenderId] `shouldReturn` []
+
+testBlockedReservation :: Expectation
+testBlockedReservation = withPgStore testBlockedReservationForStore
+
+testBlockedReservationForStore :: FileStoreClass s => s -> Expectation
+testBlockedReservationForStore st = do
+  g <- C.newRandom
+  (sndKey, _) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
+  addFile st testSenderId (testFileInfo sndKey) testCreatedAt EntityActive `shouldReturn` Right ()
+  let info = BlockingInfo {reason = BRContent, notice = Nothing}
+  blockFile st testSenderId info True `shouldReturn` Right ()
+  blockFile st testSenderId info True `shouldReturn` Right ()
+  getUsedStorage st `shouldReturn` 128000
+  snapshot <- map (\(sId, _, _) -> sId) <$> expiredFiles st 500000 100
+  (single, bulk) <- concurrently (deleteFileSize st testSenderId) (deleteFilesSizes st snapshot)
+  (either (const []) pure single <> bulk) `shouldBe` [128000]
+  getUsedStorage st `shouldReturn` 0
+
 -- Migration round-trip test
 
 testMigrationRoundTrip :: Expectation
@@ -268,6 +328,7 @@ testMigrationRoundTrip = do
   stmStore2 <- newFileStore () :: IO STMFileStore
   sl2 <- readWriteFileStore storeLogPath2 stmStore2
   closeStoreLog sl2
+  getUsedStorage stmStore2 `shouldReturn` 192000
   -- Verify file 1
   result1 <- getFile stmStore2 SFSender sId1
   case result1 of