diff --git a/Changelog.md b/Changelog.md --- a/Changelog.md +++ b/Changelog.md @@ -1,5 +1,8 @@ # Changelog - amazonka-s3-streaming +## 2.0.0.0 + - Update all the things for Amazonka-2.0 - Thanks to @endgame, @domenkozar, @koterpillar, @mpickering, @jhrcek, @roberth for your patience and input to the process (over the years...) + ## 1.1.0.0 - Adds MonadFail constraints, thanks @utdemir diff --git a/Main.hs b/Main.hs --- a/Main.hs +++ b/Main.hs @@ -1,67 +1,60 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE OverloadedStrings #-} -module Main where +module Main where import Data.Conduit ( runConduit, (.|) ) import Data.Conduit.Binary ( sourceHandle ) +import Data.Functor ( (<&>) ) import Data.Text ( pack ) -import Network.AWS -import Network.AWS.Data.Text ( fromText ) -import Network.AWS.S3.CreateMultipartUpload -import Network.AWS.S3.StreamingUpload +import Amazonka.S3.StreamingUpload + ( UploadLocation(FP), abortAllUploads, concurrentUpload, streamUpload ) -import Control.Monad.IO.Class ( liftIO ) -import System.Environment -import System.IO ( BufferMode(BlockBuffering), hSetBuffering, stdin ) +import Amazonka.Auth ( discover ) +import Amazonka.Env ( newEnv ) +import Amazonka.S3.CreateMultipartUpload ( newCreateMultipartUpload ) +import Amazonka.S3.Types ( BucketName(..), ObjectKey(..) ) -#if !MIN_VERSION_base(4,8,0) -import Control.Applicative ( pure, (<$>), (<*>) ) -#endif +import Control.Monad.IO.Class ( liftIO ) +import Control.Monad.Trans.Resource ( runResourceT ) +import System.Environment ( getArgs ) +import System.IO ( BufferMode(BlockBuffering), hSetBuffering, stdin ) main :: IO () main = do args <- getArgs + + env <- newEnv discover + case args of - (region:profile:credfile:bucket:key:file:_) -> - case (,,,) <$> (FromFile <$> fromText (pack profile) <*> pure credfile) - <*> (fromText (pack region) :: Either String Region) - <*> fromText (pack bucket) - <*> fromText (pack key) - of - Right (creds,_reg,buck,ky) -> do - env <- newEnv creds - hSetBuffering stdin (BlockBuffering Nothing) - res <- runResourceT . runAWS env $ case file of - "-" -> runConduit (sourceHandle stdin .| streamUpload Nothing (createMultipartUpload buck ky)) - >>= liftIO . either print print - _ -> concurrentUpload Nothing Nothing (FP file) (createMultipartUpload buck ky) - >>= liftIO . print + ("upload":bucket:key:file:_) -> do + let buck = BucketName $ pack bucket + ky = ObjectKey $ pack key + hSetBuffering stdin (BlockBuffering Nothing) + res <- runResourceT $ case file of + "-" -> runConduit (sourceHandle stdin .| streamUpload env Nothing (newCreateMultipartUpload buck ky)) + >>= liftIO . either print print + _ -> concurrentUpload env Nothing Nothing (FP file) (newCreateMultipartUpload buck ky) + >>= liftIO . print - print res - Left err -> print err >> usage - ("abort":region:profile:credfile:bucket:_) -> - case (,,) <$> (FromFile <$> fromText (pack profile) <*> pure credfile) - <*> (fromText (pack region) :: Either String Region) - <*> fromText (pack bucket) - of - Right (creds,_reg,buck) -> do - env <- newEnv creds - res <- runResourceT . runAWS env . abortAllUploads $ buck - print res - Left err -> print err >> usage + print res + ("abort":bucket:_) -> do + res <- runResourceT $ abortAllUploads env (BucketName $ pack bucket) + print res _ -> usage usage :: IO () usage = putStrLn . unlines $ - [ "Usage: \n" + [ "Usage:" + , "" , " Upload file:" - , " s3upload " + , " s3upload upload OR \"-\" for stdin" + , "" , " Abort all unfinished uploads for bucket:" - , " s3upload abort \n" - , "all arguments must be supplied - the region will be obtained from the AWS_REGION env var" - , "if compiled with amazonka > 1.4.4, but must still be supplied (making an option parsing library" - , "a dependency of this package seemed overkill)" + , " s3upload abort " + , "" + , "Uses `newEnv discover` to make the Amazonka environment, so it wil look at" + , "appropriate env vars, or ~/.aws/credentials, etc." ] diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# amazonka-s3-streaming [![Build Status](https://travis-ci.org/axman6/amazonka-s3-streaming.svg?branch=master)](https://travis-ci.org/axman6/amazonka-s3-streaming) +# amazonka-s3-streaming [![Haskell-CI](https://github.com/axman6/amazonka-s3-streaming/actions/workflows/haskell-ci.yml/badge.svg)](https://github.com/axman6/amazonka-s3-streaming/actions/workflows/haskell-ci.yml) Provides a conduit based streaming interface and a concurrent interface to uploading data to S3 using the Multipart API. Also provides method to upload files or bytestrings of known size in parallel. diff --git a/amazonka-s3-streaming.cabal b/amazonka-s3-streaming.cabal --- a/amazonka-s3-streaming.cabal +++ b/amazonka-s3-streaming.cabal @@ -1,5 +1,5 @@ name: amazonka-s3-streaming -version: 1.1.0.0 +version: 2.0.0.0 synopsis: Provides conduits to upload data to S3 using the Multipart API description: Provides a conduit based streaming interface and a concurrent interface to uploading data to S3 using the Multipart API. Also provides method to upload @@ -8,31 +8,41 @@ license: BSD3 license-file: LICENSE author: Alex Mason -maintainer: axman6+hackage@gmail.com -copyright: Alex Mason, Copyright (c) 2016 Commonwealth Scientific and Industrial Research Organisation (CSIRO) +maintainer: amazonka-s3-streaming@me.axman6.com +copyright: Copyright © 2023 Alex Mason, Copyright © 2016 Commonwealth Scientific and Industrial Research Organisation (CSIRO) category: Network, AWS, Cloud, Distributed Computing build-type: Simple extra-source-files: README.md, Changelog.md cabal-version: >=1.10 -tested-with: GHC == 8.0.* || == 8.2.2 || == 8.4.* || == 8.6.* || == 8.8.* +tested-with: GHC ==9.6.2 || ==9.4.5 || ==9.2.8 || ==9.0.2 || ==8.10.7 + library hs-source-dirs: src - exposed-modules: Network.AWS.S3.StreamingUpload + exposed-modules: Amazonka.S3.StreamingUpload default-language: Haskell2010 - build-depends: base >= 4.9 && < 5 - , amazonka >= 1.6 && < 1.7 - , amazonka-core >= 1.6 && < 1.7 - , amazonka-s3 >= 1.6 && < 1.7 + ghc-options: -Wall + -Wincomplete-uni-patterns + -Wincomplete-record-updates + -Wcompat + -Widentities + -Wredundant-constraints + -Wpartial-fields + -Wno-unrecognised-pragmas + -Wunused-packages + build-depends: base >= 4.14 && < 5 + , amazonka >= 2.0 && < 2.1 + , amazonka-s3 >= 2.0 && < 2.1 , conduit >= 1.3 && < 1.4 - , bytestring >= 0.10.8.0 && < 0.11 - , mmorph >= 1.0.6 && < 1.2 - , lens >= 4.13 && < 5.0 - , mtl >= 2.2.1 && < 2.3 - , exceptions >= 0.8.2.1 && < 0.11 - , dlist >= 0.8 && < 0.9 , async >= 2 && < 2.3 - , http-client >= 0.4 && < 0.7 + , bytestring >= 0.10.8.0 && < 0.13 + , deepseq >= 1.4.4 && < 1.6 + , exceptions >= 0.8.2.1 && < 0.11 + , http-client >= 0.4 && < 0.8 + , http-client-tls >= 0.3 && < 0.4 + , transformers >= 0.5 && < 0.7 + , text >= 1.2.4 && < 1.3 || >= 2.0 && < 2.2 + , resourcet >= 1.2.0 && < 1.4 flag s3upload-exe Description: Whether to build the s3upload executable for uploading files using this library. @@ -45,18 +55,17 @@ executable s3upload main-is: Main.hs - ghc-options: -threaded -rtsopts -with-rtsopts=-N + ghc-options: -threaded -rtsopts "-with-rtsopts=-N -qg" default-language: Haskell2010 if flag(s3upload-exe) - buildable: True + buildable: True else - buildable: False + buildable: False build-depends: base , amazonka - , amazonka-core , amazonka-s3 , amazonka-s3-streaming , conduit-extra , conduit , text - + , resourcet diff --git a/src/Amazonka/S3/StreamingUpload.hs b/src/Amazonka/S3/StreamingUpload.hs new file mode 100644 --- /dev/null +++ b/src/Amazonka/S3/StreamingUpload.hs @@ -0,0 +1,332 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE DuplicateRecordFields #-} +{-# LANGUAGE ImportQualifiedPost #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ParallelListComp #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Amazonka.S3.StreamingUpload + ( streamUpload + , ChunkSize + , minimumChunkSize + , NumThreads + , concurrentUpload + , UploadLocation(..) + , abortAllUploads + , module Amazonka.S3.CreateMultipartUpload + , module Amazonka.S3.CompleteMultipartUpload + ) where + +import Amazonka ( HashedBody(..), LogLevel(..), getFileSize, hashedFileRange, send, toBody ) +import Amazonka.Crypto ( hash ) +import Amazonka.Env ( Env, logger, manager ) + +import Amazonka.S3.AbortMultipartUpload ( AbortMultipartUploadResponse, newAbortMultipartUpload ) +import Amazonka.S3.CompleteMultipartUpload + ( CompleteMultipartUpload(..), CompleteMultipartUploadResponse, newCompleteMultipartUpload ) +import Amazonka.S3.CreateMultipartUpload + ( CreateMultipartUpload(..), CreateMultipartUploadResponse(..) ) +import Amazonka.S3.ListMultipartUploads + ( ListMultipartUploadsResponse(..), newListMultipartUploads, uploads ) +import Amazonka.S3.Types + ( BucketName, CompletedMultipartUpload(..), CompletedPart, MultipartUpload(..), + newCompletedMultipartUpload, newCompletedPart ) +import Amazonka.S3.UploadPart ( UploadPartResponse(..), newUploadPart ) + +import Network.HTTP.Client ( managerConnCount, newManager ) +import Network.HTTP.Client.TLS ( tlsManagerSettings ) + +import Control.Monad.Catch ( Exception, MonadCatch, onException ) +import Control.Monad.IO.Class ( MonadIO, liftIO ) +import Control.Monad.Trans.Class ( lift ) +import Control.Monad.Trans.Resource ( MonadResource, runResourceT ) + +import Conduit ( MonadUnliftIO(..) ) +import Data.Conduit ( ConduitT, Void, await, handleC, yield, (.|) ) +import Data.Conduit.Combinators ( sinkList ) +import Data.Conduit.Combinators qualified as CC + +import Data.ByteString qualified as BS +import Data.ByteString.Builder ( stringUtf8 ) +import Data.ByteString.Builder.Extra ( byteStringCopy, runBuilder ) +import Data.ByteString.Internal ( ByteString(PS) ) + +import Data.List ( unfoldr ) +import Data.List.NonEmpty ( fromList, nonEmpty ) +import Data.Text ( Text ) + +import Control.Concurrent ( newQSem, signalQSem, waitQSem ) +import Control.Concurrent.Async ( forConcurrently ) +import Control.Exception.Base ( SomeException, bracket_ ) + +import Foreign.ForeignPtr ( ForeignPtr, mallocForeignPtrBytes, plusForeignPtr ) +import Foreign.ForeignPtr.Unsafe ( unsafeForeignPtrToPtr ) + +import Control.DeepSeq ( rwhnf ) +import Data.Foldable ( for_, traverse_ ) +import Data.Typeable ( Typeable ) +import Data.Word ( Word8 ) +import Control.Monad ((>=>)) + + +type ChunkSize = Int +type NumThreads = Int + +-- | Minimum size of data which will be sent in a single part, currently 6MB +minimumChunkSize :: ChunkSize +minimumChunkSize = 6*1024*1024 -- Making this 5MB+1 seemed to cause AWS to complain + + +data StreamingError + = UnableToCreateMultipartUpload CreateMultipartUploadResponse + | FailedToUploadPiece UploadPartResponse + | Other String + deriving stock (Show, Eq, Typeable) + +instance Exception StreamingError + + + +{- | +Given a 'CreateMultipartUpload', creates a 'Sink' which will sequentially +upload the data streamed in in chunks of at least 'minimumChunkSize' and return either +the 'CompleteMultipartUploadResponse', or if an exception is thrown, +`AbortMultipartUploadResponse` and the exception as `SomeException`. If aborting +the upload also fails then the exception caused by the call to abort will be thrown. + +'Amazonka.S3.ListMultipartUploads' can be used to list any pending +uploads - it is important to abort multipart uploads because you will +be charged for storage of the parts until it is completed or aborted. +See the AWS documentation for more details. + +Internally, a single @chunkSize@d buffer will be allocated and reused between +requests to avoid holding onto incoming @ByteString@s. + +May throw 'Amazonka.Error' +-} +streamUpload :: forall m. (MonadUnliftIO m, MonadResource m) + => Env + -> Maybe ChunkSize -- ^ Optional chunk size + -> CreateMultipartUpload -- ^ Upload location + -> ConduitT ByteString Void m (Either (AbortMultipartUploadResponse, SomeException) CompleteMultipartUploadResponse) +streamUpload env mChunkSize multiPartUploadDesc@CreateMultipartUpload'{bucket = buck, key = k} = do + buffer <- liftIO $ allocBuffer chunkSize + unsafeWriteChunksToBuffer buffer + .| enumerateConduit + .| startUpload buffer + where + chunkSize :: ChunkSize + chunkSize = maybe minimumChunkSize (max minimumChunkSize) mChunkSize + + logStr :: String -> m () + logStr msg = do + liftIO $ logger env Debug $ stringUtf8 msg + + startUpload :: Buffer + -> ConduitT (Int, BufferResult) Void m + (Either (AbortMultipartUploadResponse, SomeException) + CompleteMultipartUploadResponse) + startUpload buffer = do + CreateMultipartUploadResponse'{uploadId = upId} <- lift $ send env multiPartUploadDesc + lift $ logStr "\n**** Created upload\n" + + handleC (cancelMultiUploadConduit upId) $ + CC.mapM (multiUpload buffer upId) + .| finishMultiUploadConduit upId + + multiUpload :: Buffer -> Text -> (Int, BufferResult) -> m (Maybe CompletedPart) + multiUpload buffer upId (partnum, result) = do + let !bs = bufferToByteString buffer result + !bsHash = hash bs + UploadPartResponse'{eTag} <- send env $! newUploadPart buck k partnum upId $! toBody $! HashedBytes bsHash bs + let !_ = rwhnf eTag + logStr $ "\n**** Uploaded part " <> show partnum + return $! newCompletedPart partnum <$> eTag + + -- collect all the parts + finishMultiUploadConduit :: Text + -> ConduitT (Maybe CompletedPart) Void m + (Either (AbortMultipartUploadResponse, SomeException) CompleteMultipartUploadResponse) + finishMultiUploadConduit upId = do + parts <- sinkList + res <- lift $ send env $ (newCompleteMultipartUpload buck k upId) + { multipartUpload = + Just $ newCompletedMultipartUpload {parts = sequenceA $ fromList parts} + } + + return $ Right res + + -- in case of an exception, return Left + cancelMultiUploadConduit :: Text -> SomeException + -> ConduitT i Void m + (Either (AbortMultipartUploadResponse, SomeException) CompleteMultipartUploadResponse) + cancelMultiUploadConduit upId exc = do + res <- lift $ send env $ newAbortMultipartUpload buck k upId + return $ Left (res, exc) + + -- count from 1 + enumerateConduit :: ConduitT a (Int, a) m () + enumerateConduit = loop 1 + where + loop i = await >>= maybe (return ()) (go i) + go i x = do + yield (i, x) + loop (i + 1) + {-# INLINE enumerateConduit #-} + +-- The number of bytes remaining in a buffer, and the pointer that backs it. +data Buffer = Buffer {remaining :: !Int, _fptr :: !(ForeignPtr Word8)} + +data PutResult + = Ok Buffer -- Didn't fill the buffer, updated buffer. + | Full ByteString -- Buffer is full, the unwritten remaining string. + +data BufferResult = FullBuffer | Incomplete Int + +-- Accepts @ByteString@s and writes them into @Buffer@. When the buffer is full, +-- @FullBuffer@ is emitted. If there is no more input, @Incomplete@ is emitted with +-- the number of bytes remaining in the buffer. +unsafeWriteChunksToBuffer :: MonadIO m => Buffer -> ConduitT ByteString BufferResult m () +unsafeWriteChunksToBuffer buffer0 = awaitLoop buffer0 where + awaitLoop buf = + await >>= maybe (yield $ Incomplete $ remaining buf) + (liftIO . putBuffer buf >=> \case + Full next -> yield FullBuffer *> chunkLoop buffer0 next + Ok buf' -> awaitLoop buf' + ) + -- Handle inputs which are larger than the chunkSize + chunkLoop buf = liftIO . putBuffer buf >=> \case + Full next -> yield FullBuffer *> chunkLoop buffer0 next + Ok buf' -> awaitLoop buf' + +bufferToByteString :: Buffer -> BufferResult -> ByteString +bufferToByteString (Buffer bufSize fptr) FullBuffer = PS fptr 0 bufSize +bufferToByteString (Buffer bufSize fptr) (Incomplete remaining) = PS fptr 0 (bufSize - remaining) + +allocBuffer :: Int -> IO Buffer +allocBuffer chunkSize = Buffer chunkSize <$> mallocForeignPtrBytes chunkSize + +putBuffer :: Buffer -> ByteString -> IO PutResult +putBuffer buffer bs + | BS.length bs <= remaining buffer = + Ok <$> unsafeWriteBuffer buffer bs + | otherwise = do + let (remainder,rest) = BS.splitAt (remaining buffer) bs + Full rest <$ unsafeWriteBuffer buffer remainder + +-- The length of the bytestring must be less than or equal to the number +-- of bytes remaining. +unsafeWriteBuffer :: Buffer -> ByteString -> IO Buffer +unsafeWriteBuffer (Buffer remaining fptr) bs = do + let ptr = unsafeForeignPtrToPtr fptr + len = BS.length bs + _ <- runBuilder (byteStringCopy bs) ptr remaining + pure $ Buffer (remaining - len) (plusForeignPtr fptr len) + + +-- | Specifies whether to upload a file or 'ByteString'. +data UploadLocation + = FP FilePath -- ^ A file to be uploaded + | BS ByteString -- ^ A strict 'ByteString' + +{-| +Allows a file or 'ByteString' to be uploaded concurrently, using the +async library. The chunk size may optionally be specified, but will be at least +`minimumChunkSize`, and may be made larger than if the `ByteString` or file +is larger enough to cause more than 10,000 chunks. + +Files are mmapped into 'chunkSize' chunks and each chunk is uploaded in parallel. +This considerably reduces the memory necessary compared to reading the contents +into memory as a strict 'ByteString'. The usual caveats about mmaped files apply: +if the file is modified during this operation, the data may become corrupt. + +May throw `Amazonka.Error`, or `IOError`; an attempt is made to cancel the +multipart upload on any error, but this may also fail if, for example, the network +connection has been broken. See `abortAllUploads` for a crude cleanup method. +-} +concurrentUpload :: forall m. + (MonadResource m, MonadCatch m) + => Env + -> Maybe ChunkSize -- ^ Optional chunk size + -> Maybe NumThreads -- ^ Optional number of threads to upload with + -> UploadLocation -- ^ Whether to upload a file on disk or a `ByteString` that's already in memory. + -> CreateMultipartUpload -- ^ Description of where to upload. + -> m CompleteMultipartUploadResponse +concurrentUpload env' mChunkSize mNumThreads uploadLoc + multiPartUploadDesc@CreateMultipartUpload'{bucket = buck, key = k} + = do + CreateMultipartUploadResponse'{uploadId = upId} <- send env' multiPartUploadDesc + + let logStr :: MonadIO n => String -> n () + logStr = liftIO . logger env' Info . stringUtf8 + + calculateChunkSize :: Int -> Int + calculateChunkSize len = + let chunkSize' = maybe minimumChunkSize (max minimumChunkSize) mChunkSize + in if len `div` chunkSize' >= 10000 then len `div` 9999 else chunkSize' + + mConnCount = managerConnCount tlsManagerSettings + nThreads = maybe mConnCount (max 1) mNumThreads + + env <- if maybe False (> mConnCount) mNumThreads + then do + mgr' <- liftIO $ newManager tlsManagerSettings{managerConnCount = nThreads} + pure env'{manager = mgr'} + else pure env' + flip onException (send env (newAbortMultipartUpload buck k upId)) $ do + sem <- liftIO $ newQSem nThreads + uploadResponses <- case uploadLoc of + BS bytes -> + let chunkSize = calculateChunkSize $ BS.length bytes + in liftIO $ forConcurrently (zip [1..] $ chunksOf chunkSize bytes) $ \(partnum, chunk) -> + bracket_ (waitQSem sem) (signalQSem sem) $ do + logStr $ "Starting part: " ++ show partnum + UploadPartResponse'{eTag} <- runResourceT $ send env . newUploadPart buck k partnum upId . toBody $ chunk + logStr $ "Finished part: " ++ show partnum + pure $ newCompletedPart partnum <$> eTag + + FP filePath -> do + fsize <- liftIO $ getFileSize filePath + let chunkSize = calculateChunkSize $ fromIntegral fsize + (count,lst) = fromIntegral fsize `divMod` chunkSize + params = [(partnum, chunkSize*offset, size) + | partnum <- [1..] + | offset <- [0..count] + | size <- (chunkSize <$ [0..count-1]) ++ [lst] + ] + + liftIO $ forConcurrently params $ \(partnum,off,size) -> + bracket_ (waitQSem sem) (signalQSem sem) $ do + logStr $ "Starting file part: " ++ show partnum + chunkStream <- hashedFileRange filePath (fromIntegral off) (fromIntegral size) + UploadPartResponse'{eTag} <- runResourceT $ + send env . newUploadPart buck k partnum upId . toBody $ chunkStream + logStr $ "Finished file part: " ++ show partnum + pure $ newCompletedPart partnum <$> eTag + + let parts = nonEmpty =<< sequence uploadResponses + send env $ (newCompleteMultipartUpload buck k upId) + { multipartUpload = Just $ newCompletedMultipartUpload { parts } } + +-- | Aborts all uploads in a given bucket - useful for cleaning up. +abortAllUploads :: MonadResource m => Env -> BucketName -> m () +abortAllUploads env buck = do + ListMultipartUploadsResponse' {uploads} <- send env $ newListMultipartUploads buck + flip (traverse_ . traverse_) uploads $ \MultipartUpload'{key, uploadId} -> do + let mki = (,) <$> key <*> uploadId + for_ mki $ \(key',uid) -> send env (newAbortMultipartUpload buck key' uid) + + + +-- http://stackoverflow.com/questions/32826539/chunksof-analog-for-bytestring +justWhen :: (a -> Bool) -> (a -> b) -> a -> Maybe b +justWhen f g a = if f a then Just (g a) else Nothing + +nothingWhen :: (a -> Bool) -> (a -> b) -> a -> Maybe b +nothingWhen f = justWhen (not . f) + +chunksOf :: Int -> BS.ByteString -> [BS.ByteString] +chunksOf x = unfoldr (nothingWhen BS.null (BS.splitAt x)) diff --git a/src/Network/AWS/S3/StreamingUpload.hs b/src/Network/AWS/S3/StreamingUpload.hs deleted file mode 100644 --- a/src/Network/AWS/S3/StreamingUpload.hs +++ /dev/null @@ -1,262 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE ParallelListComp #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TupleSections #-} - -module Network.AWS.S3.StreamingUpload - ( streamUpload - , ChunkSize - , minimumChunkSize - , NumThreads - , concurrentUpload - , UploadLocation(..) - , abortAllUploads - , module Network.AWS.S3.CreateMultipartUpload - , module Network.AWS.S3.CompleteMultipartUpload - ) where - -import Network.AWS - ( AWS, HasEnv(..), LogLevel(..), MonadAWS, getFileSize, hashedBody, hashedFileRange, - liftAWS, runAWS, runResourceT, send, toBody ) - -import Network.AWS.Data.Crypto ( Digest, SHA256, hashFinalize, hashInit, hashUpdate ) - -import Network.AWS.S3.AbortMultipartUpload -import Network.AWS.S3.CompleteMultipartUpload -import Network.AWS.S3.CreateMultipartUpload -import Network.AWS.S3.ListMultipartUploads -import Network.AWS.S3.Types - ( BucketName, cmuParts, completedMultipartUpload, completedPart, muKey, muUploadId ) -import Network.AWS.S3.UploadPart - -import Control.Applicative -import Control.Category ( (>>>) ) -import Control.Monad ( forM_, when, (>=>) ) -import Control.Monad.Fail ( MonadFail ) -import Control.Monad.IO.Class ( MonadIO, liftIO ) -import Control.Monad.Morph ( lift ) -import Control.Monad.Reader.Class ( local ) - -import Conduit ( MonadUnliftIO(..) ) -import Data.Conduit ( ConduitT, Void, await, catchC ) -import Data.Conduit.List ( sourceList ) - -import Data.ByteString ( ByteString ) -import qualified Data.ByteString as BS -import Data.ByteString.Builder ( stringUtf8 ) - -import qualified Data.DList as D -import Data.List ( unfoldr ) -import Data.List.NonEmpty ( nonEmpty ) - -import Control.Lens ( set, view ) -import Control.Lens.Operators - -import Text.Printf ( printf ) - -import Control.Concurrent ( newQSem, signalQSem, waitQSem ) -import Control.Concurrent.Async ( forConcurrently ) -import Control.Exception.Base ( SomeException, bracket_ ) -import Control.Monad.Catch ( onException ) -import System.Mem ( performGC ) - -import Network.HTTP.Client ( defaultManagerSettings, managerConnCount, newManager ) - -type ChunkSize = Int -type NumThreads = Int - --- | Minimum size of data which will be sent in a single part, currently 6MB -minimumChunkSize :: ChunkSize -minimumChunkSize = 6*1024*1024 -- Making this 5MB+1 seemed to cause AWS to complain - - -{- | -Given a 'CreateMultipartUpload', creates a 'Sink' which will sequentially -upload the data streamed in in chunks of at least 'minimumChunkSize' and return either -the 'CompleteMultipartUploadResponse', or if an exception is thrown, -`AbortMultipartUploadResponse` and the exception as `SomeException`. If aborting -the upload also fails then the exception caused by the call to abort will be thrown. - -'Network.AWS.S3.ListMultipartUploads' can be used to list any pending -uploads - it is important to abort multipart uploads because you will -be charged for storage of the parts until it is completed or aborted. -See the AWS documentation for more details. - -May throw 'Network.AWS.Error' --} -streamUpload :: (MonadUnliftIO m, MonadAWS m, MonadFail m) - => Maybe ChunkSize -- ^ Optional chunk size - -> CreateMultipartUpload -- ^ Upload location - -> ConduitT ByteString Void m (Either (AbortMultipartUploadResponse, SomeException) CompleteMultipartUploadResponse) -streamUpload mChunkSize multiPartUploadDesc = do - logger <- lift $ liftAWS $ view envLogger - let logStr :: MonadIO m => String -> m () - logStr = liftIO . logger Debug . stringUtf8 - chunkSize = maybe minimumChunkSize (max minimumChunkSize) mChunkSize - - multiPartUpload <- lift $ send multiPartUploadDesc - when (multiPartUpload ^. cmursResponseStatus /= 200) $ - fail "Failed to create upload" - - logStr "\n**** Created upload\n" - - let Just upId = multiPartUpload ^. cmursUploadId - bucket = multiPartUploadDesc ^. cmuBucket - key = multiPartUploadDesc ^. cmuKey - -- go :: DList ByteString -> Int -> Context SHA256 -> Int -> DList (Maybe CompletedPart) -> Sink ByteString m () - go !bss !bufsize !ctx !partnum !completed = await >>= \case - Just bs - | l <- BS.length bs, bufsize + l <= chunkSize - -> go (D.snoc bss bs) (bufsize + l) (hashUpdate ctx bs) partnum completed - | otherwise -> do - rs <- lift $ performUpload partnum (bufsize + BS.length bs) - (hashFinalize $ hashUpdate ctx bs) - (D.snoc bss bs) - - logStr $ printf "\n**** Uploaded part %d size %d\n" partnum bufsize - let !part = completedPart partnum <$> (rs ^. uprsETag) - liftIO performGC - go empty 0 hashInit (partnum+1) . D.snoc completed $! part - - Nothing -> lift $ do - prts <- if bufsize > 0 - then do - rs <- performUpload partnum bufsize (hashFinalize ctx) bss - logStr $ printf "\n**** Uploaded (final) part %d size %d\n" partnum bufsize - let allParts = D.toList $ D.snoc completed $ completedPart partnum <$> (rs ^. uprsETag) - pure $ nonEmpty =<< sequence allParts - else do - logStr $ printf "\n**** No final data to upload\n" - pure $ nonEmpty =<< sequence (D.toList completed) - - send $ completeMultipartUpload bucket key upId - & cMultipartUpload ?~ set cmuParts prts completedMultipartUpload - - - performUpload :: (MonadAWS m, MonadFail m) => Int -> Int -> Digest SHA256 -> D.DList ByteString -> m UploadPartResponse - performUpload pnum size digest = - D.toList - >>> sourceList - >>> hashedBody digest (fromIntegral size) - >>> toBody - >>> uploadPart bucket key pnum upId - >>> send - >=> checkUpload - - checkUpload :: (Monad m, MonadFail m) => UploadPartResponse -> m UploadPartResponse - checkUpload upr = do - when (upr ^. uprsResponseStatus /= 200) $ fail "Failed to upload piece" - return upr - - (Right <$> go D.empty 0 hashInit 1 D.empty) `catchC` \(except :: SomeException) -> - Left . (,except) <$> lift (send (abortMultipartUpload bucket key upId)) - -- Whatever happens, we abort the upload and return the exception - - --- | Specifies whether to upload a file or 'ByteString -data UploadLocation - = FP FilePath -- ^ A file to be uploaded - | BS ByteString -- ^ A strict 'ByteString' - -{-| -Allows a file or 'ByteString' to be uploaded concurrently, using the -async library. The chunk size may optionally be specified, but will be at least -`minimumChunkSize`, and may be made larger than if the `ByteString` or file -is larger enough to cause more than 10,000 chunks. - -Files are mmapped into 'chunkSize' chunks and each chunk is uploaded in parallel. -This considerably reduces the memory necessary compared to reading the contents -into memory as a strict 'ByteString'. The usual caveats about mmaped files apply: -if the file is modified during this operation, the data may become corrupt. - -May throw `Network.AWS.Error`, or `IOError`; an attempt is made to cancel the -multipart upload on any error, but this may also fail if, for example, the network -connection has been broken. See `abortAllUploads` for a crude cleanup method. --} -concurrentUpload :: (MonadAWS m, MonadFail m) - => Maybe ChunkSize -- ^ Optional chunk size - -> Maybe NumThreads -- ^ Optional number of threads to upload with - -> UploadLocation -- ^ Whether to upload a file on disk or a `ByteString` that's already in memory. - -> CreateMultipartUpload -- ^ Description of where to upload. - -> m CompleteMultipartUploadResponse -concurrentUpload mChunkSize mNumThreads uploadLoc multiPartUploadDesc = do - env <- liftAWS $ view environment - cmur <- send multiPartUploadDesc - when (cmur ^. cmursResponseStatus /= 200) $ - fail "Failed to create upload" - - let logStr :: MonadIO m => String -> m () - logStr = liftIO . (env ^. envLogger) Info . stringUtf8 - bucket = multiPartUploadDesc ^. cmuBucket - key = multiPartUploadDesc ^. cmuKey - Just upId = cmur ^. cmursUploadId - - calculateChunkSize :: Int -> Int - calculateChunkSize len = - let chunkSize' = maybe minimumChunkSize (max minimumChunkSize) mChunkSize - in if len `div` chunkSize' >= 10000 then len `div` 9999 else chunkSize' - - mConnCount = managerConnCount defaultManagerSettings - nThreads = maybe mConnCount (max 1) mNumThreads - - exec :: MonadAWS m => AWS a -> m a - exec act = if maybe False (> mConnCount) mNumThreads - then do - mgr' <- liftIO $ newManager defaultManagerSettings{managerConnCount = nThreads} - liftAWS $ local (envManager .~ mgr') act - else liftAWS act - exec $ flip onException (send (abortMultipartUpload bucket key upId)) $ do - sem <- liftIO $ newQSem nThreads - uploadResponses <- case uploadLoc of - BS bytes -> - let chunkSize = calculateChunkSize $ BS.length bytes - in liftIO $ forConcurrently (zip [1..] $ chunksOf chunkSize bytes) $ \(partnum, chunk) -> - bracket_ (waitQSem sem) (signalQSem sem) $ do - logStr $ "Starting part: " ++ show partnum - umr <- runResourceT $ runAWS env $ send . uploadPart bucket key partnum upId . toBody $ chunk - logStr $ "Finished part: " ++ show partnum - pure $ completedPart partnum <$> (umr ^. uprsETag) - - FP filePath -> do - fsize <- liftIO $ getFileSize filePath - let chunkSize = calculateChunkSize $ fromIntegral fsize - (count,lst) = fromIntegral fsize `divMod` chunkSize - params = [(partnum, chunkSize*offset, size) - | partnum <- [1..] - | offset <- [0..count] - | size <- (chunkSize <$ [0..count-1]) ++ [lst] - ] - - liftIO $ forConcurrently params $ \(partnum,off,size) -> - bracket_ (waitQSem sem) (signalQSem sem) $ do - logStr $ "Starting file part: " ++ show partnum - chunkStream <- hashedFileRange filePath (fromIntegral off) (fromIntegral size) - uploadResp <- runResourceT $ runAWS env $ - send . uploadPart bucket key partnum upId . toBody $ chunkStream - logStr $ "Finished file part: " ++ show partnum - pure $ completedPart partnum <$> (uploadResp ^. uprsETag) - - let parts = nonEmpty =<< sequence uploadResponses - send $ completeMultipartUpload bucket key upId - & cMultipartUpload ?~ set cmuParts parts completedMultipartUpload - --- | Aborts all uploads in a given bucket - useful for cleaning up. -abortAllUploads :: (MonadAWS m) => BucketName -> m () -abortAllUploads bucket = do - rs <- send (listMultipartUploads bucket) - forM_ (rs ^. lmursUploads) $ \mu -> do - let mki = (,) <$> mu ^. muKey <*> mu ^. muUploadId - forM_ mki $ \(key,uid) -> send (abortMultipartUpload bucket key uid) - - - --- http://stackoverflow.com/questions/32826539/chunksof-analog-for-bytestring -justWhen :: (a -> Bool) -> (a -> b) -> a -> Maybe b -justWhen f g a = if f a then Just (g a) else Nothing - -nothingWhen :: (a -> Bool) -> (a -> b) -> a -> Maybe b -nothingWhen f = justWhen (not . f) - -chunksOf :: Int -> BS.ByteString -> [BS.ByteString] -chunksOf x = unfoldr (nothingWhen BS.null (BS.splitAt x))