mismi-s3: Smarter uploadRecursive

Uploads of multiple small files is now concurrent.

Author: erikd-ambiata

mismi-s3/mismi-s3.cabal

@@ -35,6 +35,7 @@ library
                      , filepath                        >= 1.3        && < 1.5
                      , http-client                     >= 0.4.18     && < 0.5
                      , http-types                      == 0.8.*
+                     , lifted-async                    == 0.9.*
                      , mtl                             >= 2.1        && < 2.3
                      , process                         >= 1.2        && < 1.5
                      , resourcet                       == 1.1.*

mismi-s3/src/Mismi/S3/Commands.hs

@@ -70,7 +70,7 @@ module Mismi.S3.Commands (
   ) where
 
 import           Control.Arrow ((***))
-
+import           Control.Concurrent.Async.Lifted (mapConcurrently_)
 import           Control.Exception (ioError)
 import qualified Control.Exception as CE
 import           Control.Lens ((.~), (^.), to, view)
@@ -124,7 +124,6 @@ import           System.Directory (createDirectoryIfMissing, doesFileExist, getD
 import           System.FilePath (FilePath, (</>), takeDirectory)
 import           System.Posix.IO (OpenMode(..), openFd, closeFd, fdSeek, defaultFileFlags)
 import           System.Posix.Files (fileSize, getFileStatus, isDirectory, isRegularFile)
-import           System.Posix.Types (FileOffset)
 import qualified "unix-bytestring" System.Posix.IO.ByteString as UBS
 
 import           System.Timeout.Lifted (timeout)
@@ -372,13 +371,13 @@ upload :: FilePath -> Address -> EitherT UploadError AWS ()
 upload =
   uploadWithMode Fail
 
-uploadRecursive :: FilePath -> Address -> EitherT UploadError AWS ()
+uploadRecursive :: FilePath -> Address -> Int -> EitherT UploadError AWS ()
 uploadRecursive =
   uploadRecursiveWithMode Fail
 
-uploadRecursiveOrFail :: FilePath -> Address -> AWS ()
-uploadRecursiveOrFail f a =
-  eitherT hoistUploadError pure $ uploadRecursive f a
+uploadRecursiveOrFail :: FilePath -> Address -> Int -> AWS ()
+uploadRecursiveOrFail f a i =
+  eitherT hoistUploadError pure $ uploadRecursive f a i
 
 uploadOrFail :: FilePath -> Address -> AWS ()
 uploadOrFail f a =
@@ -388,9 +387,9 @@ uploadWithModeOrFail :: WriteMode -> FilePath -> Address -> AWS ()
 uploadWithModeOrFail w f a =
   eitherT hoistUploadError pure $ uploadWithMode w f a
 
-uploadRecursiveWithModeOrFail :: WriteMode -> FilePath -> Address -> AWS ()
-uploadRecursiveWithModeOrFail w f a =
-  eitherT hoistUploadError pure $ uploadRecursiveWithMode w f a
+uploadRecursiveWithModeOrFail :: WriteMode -> FilePath -> Address -> Int -> AWS ()
+uploadRecursiveWithModeOrFail w f a i =
+  eitherT hoistUploadError pure $ uploadRecursiveWithMode w f a i
 
 hoistUploadError :: UploadError -> AWS ()
 hoistUploadError e =
@@ -439,11 +438,11 @@ uploadSingle file a = do
   void . send $ f' A.putObject a rq & A.poServerSideEncryption .~ pure sse
 
 multipartUpload :: FilePath -> Address -> Integer -> Integer -> Int -> EitherT UploadError AWS ()
-multipartUpload file a fileSize chunk fork = do
+multipartUpload file a fSize chunk fork = do
   e <- ask
   mpu <- lift $ createMultipartUpload a
 
-  let chunks = calculateChunksCapped (fromInteger fileSize) (fromInteger chunk) 4096 -- max 4096 prts returned
+  let chunks = calculateChunksCapped (fromInteger fSize) (fromInteger chunk) 4096 -- max 4096 prts returned
 
   r <- liftIO $
     consume (forM_ chunks . writeQueue) fork $ multipartUploadWorker e mpu file a
@@ -481,40 +480,46 @@ multipartUploadWorker e mpu file a (o, c, i) =
         pure $! Right $! PartResponse i m
 
 
-uploadRecursiveWithMode :: WriteMode -> FilePath -> Address -> EitherT UploadError AWS ()
-uploadRecursiveWithMode m src (Address buck ky) = do
+uploadRecursiveWithMode :: WriteMode -> FilePath -> Address -> Int -> EitherT UploadError AWS ()
+uploadRecursiveWithMode mode src (Address buck ky) fork = do
   es <- tryIO $ getFileStatus src
   case es of
     Left _ -> left $ UploadSourceMissing src
     Right st -> unless (isDirectory st) . left $ UploadSourceNotDirectory src
-  files <- fmap fst <$> liftIO (listRecursivelyLocal src)
-  let outputAddrs = fmap (\fp -> Address buck (ky // Key (T.pack $ L.drop prefixLen fp))) files
-  mapM_ (uncurry (uploadWithMode m)) $ L.zip files outputAddrs
+  files <- liftIO (listRecursivelyLocal src)
+  mapM_ uploadFiles $ chunkFilesBySize fork (fromIntegral bigChunkSize) files
   where
+    uploadFiles :: [(FilePath, Int64)] -> EitherT UploadError AWS ()
+    uploadFiles [] = pure ()
+    uploadFiles [(f,s)]
+      | fromIntegral s < bigChunkSize = lift . uploadSingle f $ uploadAddress f
+      | otherwise = uploadWithMode mode f $ uploadAddress f
+    uploadFiles xs =
+      mapConcurrently_ (\ (f, _) -> lift . uploadSingle f $ uploadAddress f) xs
+
+
     prefixLen = L.length (src </> "a") - 1
 
-    -- uploadAddress :: FilePath -> Address
-    -- uploadAddress fp = Address buck (ky // Key (T.pack $ L.drop prefixLen fp))
+    uploadAddress :: FilePath -> Address
+    uploadAddress fp = Address buck (ky // Key (T.pack $ L.drop prefixLen fp))
 
--- Take a list of files and their sizes, and return a list of list of files
--- where the total size of of the files in the sub list is less than `bigChunkSize`
+-- Take a list of files and their sizes, and convert it to a list of tests
+-- where the total size of of the files in the sub list is less than `maxSize`
 -- and the length of the sub lists is <= `maxCount`.
-chunkFilesBySize :: Int -> Int64 -> [(FilePath, Int64)] -> [[FilePath]]
+chunkFilesBySize :: Int -> Int64 -> [(FilePath, Int64)] -> [[(FilePath, Int64)]]
 chunkFilesBySize maxCount maxSize =
   takeFiles 0 [] . L.sortOn snd
   where
-    takeFiles :: Int64 -> [FilePath] -> [(FilePath, Int64)] -> [[FilePath]]
+    takeFiles :: Int64 -> [(FilePath, Int64)] -> [(FilePath, Int64)] -> [[(FilePath, Int64)]]
     takeFiles _ acc [] = [acc]
     takeFiles current acc ((x, s):xs) =
       if current + s < maxSize && L.length acc < maxCount
-        then takeFiles (current + s) (x:acc) xs
-        else (x:acc) : takeFiles 0 [] xs
-
--- bigChunkSize
+        then takeFiles (current + s) ((x, s):acc) xs
+        else ((x, s):acc) : takeFiles 0 [] xs
 
 -- | Like `listRecursively` but for the local filesystem.
 -- Also returns
-listRecursivelyLocal :: MonadIO m => FilePath -> m [(FilePath, FileOffset)]
+listRecursivelyLocal :: MonadIO m => FilePath -> m [(FilePath, Int64)]
 listRecursivelyLocal topdir = do
   entries <- liftIO $ listDirectory topdir
   (dirs, files) <- liftIO . partitionDirsFilesWithSizes $ fmap (topdir </>) entries
@@ -530,16 +535,17 @@ listDirectory path =
     f filename =
       filename /= "." && filename /= ".."
 
-partitionDirsFilesWithSizes :: MonadIO m => [FilePath] -> m ([FilePath], [(FilePath, FileOffset)])
+partitionDirsFilesWithSizes :: MonadIO m => [FilePath] -> m ([FilePath], [(FilePath, Int64)])
 partitionDirsFilesWithSizes =
   pworker ([], [])
   where
     pworker (dirs, files) [] = pure (dirs, files)
     pworker (dirs, files) (x:xs) = do
       xstat <- liftIO $ getFileStatus x
-      pworker
-        (if isDirectory xstat then x : dirs else dirs, if isRegularFile xstat then (x, fileSize xstat) : files else files)
-        xs
+      let xsize = fromIntegral $ fileSize xstat
+          newDirs = if isDirectory xstat then x : dirs else dirs
+          newFiles = if isRegularFile xstat then (x, xsize) : files else files
+      pworker (newDirs, newFiles) xs
 
 write :: Address -> Text -> AWS WriteResult
 write =

mismi-s3/test/Test/IO/Mismi/S3/Commands.hs

@@ -532,7 +532,7 @@ prop_upload_recursive = once . testAWS $ do
 
   addr <- withKey (// Key "top") <$> newAddress
 
-  eitherT (fail . show) pure $ uploadRecursive tmpdir addr
+  eitherT (fail . show) pure $ uploadRecursive tmpdir addr 2
 
   a <- read (withKey (// Key "a") addr)
   c <- read (withKey (// Key "b/c") addr)

mismi-s3/test/Test/Mismi/S3/Commands.hs

@@ -47,7 +47,7 @@ prop_chunk_files_by_size =
   QC.forAll (QC.choose (1000, 10000)) $ \ maxChunkSize ->
   QC.forAll (fileNameSizePairs fileCount) $ \ pairs ->
     let chunks = chunkFilesBySize maxFilesPerChunk maxChunkSize pairs
-        chunkSizes = DL.map (multiChunkSum (DM.fromList pairs)) chunks
+        chunkSizes = DL.map (multiChunkSum (DM.fromList pairs) . DL.map fst) chunks
     in
       DL.filter (> maxChunkSize) chunkSizes === []