diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2016 Alex Mason + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/Main.hs b/Main.hs new file mode 100644 --- /dev/null +++ b/Main.hs @@ -0,0 +1,56 @@ + +{-# LANGUAGE OverloadedStrings #-} +module Main where + + +import Data.Conduit (($$)) +import Data.Conduit.Binary (sourceHandle) +import Data.Text (pack) +import Network.AWS +import Network.AWS.Data.Text (fromText) +import Network.AWS.S3.CreateMultipartUpload +import Network.AWS.S3.StreamingUpload +import System.Environment +import System.IO (BufferMode (BlockBuffering), + hSetBuffering, stdin) + +main :: IO () +main = do + args <- getArgs + case args of + (region:profile:credfile:bucket:key:file:_) -> + case (,,,) <$> (FromFile <$> fromText (pack profile) <*> pure credfile) + <*> fromText (pack region) + <*> fromText (pack bucket) + <*> fromText (pack key) + of + Right (env',reg,buck,ky) -> do + env <- newEnv reg env' + hSetBuffering stdin (BlockBuffering Nothing) + res <- runResourceT . runAWS env $ case file of + -- Stream data from stdin + "-" -> sourceHandle stdin $$ streamUpload (createMultipartUpload buck ky) + -- Read from a file + _ -> concurrentUpload (FP file) $ createMultipartUpload buck ky + print res + Left err -> print err >> usage + ("abort":region:profile:credfile:bucket:_) -> + case (,,) <$> (FromFile <$> fromText (pack profile) <*> pure credfile) + <*> fromText (pack region) + <*> fromText (pack bucket) + of + Right (env',reg,buck) -> do + env <- newEnv reg env' + res <- runResourceT . runAWS env . abortAllUploads $ buck + print res + Left err -> print err >> usage + + _ -> usage + +usage :: IO () +usage = putStrLn "\nUsage: \n\n\ + \ Upload file:\n\ + \ s3upload \n\ + \ Abort all unfinished uploads for bucket:\n\ + \ s3upload abort \n\n\ + \all arguments must be supplied" diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,2 @@ +# amazonka-s3-streaming +Provides a conduit based interface to uploading data to S3 using the Multipart API diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/amazonka-s3-streaming.cabal b/amazonka-s3-streaming.cabal new file mode 100644 --- /dev/null +++ b/amazonka-s3-streaming.cabal @@ -0,0 +1,62 @@ +name: amazonka-s3-streaming +version: 0.1.0.0 +cabal-version: >=1.10 +build-type: Simple +license: BSD3 +license-file: LICENSE +copyright: Copyright (c) 2016 Commonwealth Scientific and Industrial Research Organisation (CSIRO) +maintainer: Alex.Mason@data61.csiro.au +homepage: https://github.com/Axman6/amazonka-s3-streaming#readme +synopsis: Provides conduits to upload data to S3 using the Multipart API +description: + Please see README.md +category: Network, AWS, Cloud, Distributed Computing +author: Alex Mason +extra-source-files: + README.md + +source-repository head + type: git + location: https://github.com/Axman6/amazonka-s3-streaming + +library + + if impl(ghc <7.11) + build-depends: + semigroups >=0.12 && <0.19, + transformers >=0.4 && <0.6 + exposed-modules: + Network.AWS.S3.StreamingUpload + build-depends: + base >=4.7 && <5, + amazonka >=1.3.7 && <1.5, + amazonka-core >=1.3.7 && <1.5, + amazonka-s3 >=1.3.7 && <1.5, + resourcet >=1.1.7.4 && <1.2, + conduit >=1.2.6.6 && <1.3, + bytestring >=0.10.6.0 && <0.11, + mmorph >=1.0.6 && <1.1, + lens >=4.13 && <5.0, + mtl >=2.2.1 && <2.3, + exceptions >=0.8.2.1 && <0.9, + dlist >=0.8.0.2 && <0.9, + lifted-async >=0.9.0 && <0.10, + mmap >=0.5.9 && <0.6 + default-language: Haskell2010 + hs-source-dirs: src + +executable s3upload + main-is: Main.hs + build-depends: + base >=4.7 && <5, + amazonka >=1.4.3 && <1.5, + amazonka-core >=1.4.3 && <1.5, + amazonka-s3 >=1.4.3 && <1.5, + amazonka-s3-streaming >=0.1.0.0 && <0.2, + conduit-extra >=1.1.15 && <1.2, + conduit >=1.2.8 && <1.3, + bytestring >=0.10.8.1 && <0.11, + text >=1.2.2.1 && <1.3 + default-language: Haskell2010 + ghc-options: -threaded -rtsopts -with-rtsopts=-N + diff --git a/src/Network/AWS/S3/StreamingUpload.hs b/src/Network/AWS/S3/StreamingUpload.hs new file mode 100644 --- /dev/null +++ b/src/Network/AWS/S3/StreamingUpload.hs @@ -0,0 +1,229 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ParallelListComp #-} +{-# LANGUAGE PatternGuards #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Network.AWS.S3.StreamingUpload +( + streamUpload + , UploadLocation(..) + , concurrentUpload + , abortAllUploads + , module Network.AWS.S3.CreateMultipartUpload + , module Network.AWS.S3.CompleteMultipartUpload + , chunkSize +) where + +import Network.AWS (Error, HasEnv (..), + LogLevel (..), + MonadAWS, getFileSize, + hashedBody, send, + toBody) + +import Control.Monad.Trans.AWS (AWSConstraint) +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.IO.Class (MonadIO, liftIO) +import Control.Monad.Morph (lift) +import Control.Monad.Trans.Resource (MonadBaseControl, + MonadResource, throwM) + +import Data.Conduit (Sink, await) +import Data.Conduit.List (sourceList) + +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import Data.ByteString.Builder (stringUtf8) +import System.IO.MMap (mmapFileByteString) + +import qualified Data.DList as D +import Data.List (unfoldr) +import Data.List.NonEmpty (nonEmpty) + +import Control.Exception.Lens (catching, handling) +import Control.Lens + +import Text.Printf (printf) + +import Control.Concurrent.Async.Lifted (forConcurrently) + +-- | Minimum size of data which will be sent in a single part, currently 6MB +chunkSize :: Int +chunkSize = 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 'chunkSize' and return +the 'CompleteMultipartUploadResponse'. + +If uploading of any parts fails, an attempt is made to abort the Multipart +upload, but it is likely that if an upload fails this abort may also fail. +'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 'Error' +-} +streamUpload :: (MonadResource m, MonadAWS m, AWSConstraint r m) + => CreateMultipartUpload + -> Sink ByteString m CompleteMultipartUploadResponse +streamUpload cmu = do + logger <- lift $ view envLogger + let logStr :: MonadIO m => String -> m () + logStr = liftIO . logger Info . stringUtf8 + + cmur <- lift (send cmu) + when (cmur ^. cmursResponseStatus /= 200) $ + fail "Failed to create upload" + + logStr "\n**** Created upload\n" + + let Just upId = cmur ^. cmursUploadId + bucket = cmu ^. cmuBucket + key = cmu ^. cmuKey + -- go :: Text -> Builder -> Int -> Int -> Sink ByteString m () + go !bss !bufsize !ctx !partnum !completed = Data.Conduit.await >>= \mbs -> case mbs of + 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 $ partUploader 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) + go empty 0 hashInit (partnum+1) $ D.snoc completed part + + Nothing -> lift $ do + rs <- partUploader 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) + prts = nonEmpty =<< sequence allParts + + send $ completeMultipartUpload bucket key upId + & cMultipartUpload ?~ set cmuParts prts completedMultipartUpload + + + partUploader :: MonadAWS m => Int -> Int -> Digest SHA256 -> D.DList ByteString -> m UploadPartResponse + partUploader pnum size digest = + D.toList + >>> sourceList + >>> hashedBody digest (fromIntegral size) + >>> toBody + >>> uploadPart bucket key pnum upId + >>> send + >=> checkUpload + + checkUpload :: (Monad m) => UploadPartResponse -> m UploadPartResponse + checkUpload upr = do + when (upr ^. uprsResponseStatus /= 200) $ fail "Failed to upload piece" + return upr + + catching id (go D.empty 0 hashInit 1 D.empty) $ \e -> + lift (send (abortMultipartUpload bucket key upId)) >> throwM e + -- Whatever happens, we abort the upload and rethrow + + +-- | Specifies whether to upload a file or 'ByteString +data UploadLocation + = FP FilePath -- ^ A file to be uploaded + | BS ByteString -- ^ A strict 'ByteString' + + -- IO (Int -> IO (Maybe ByteString)) (Either (IO ()) (IO ())) + -- part number as input, may be called many times until Nothing is returned, + -- and a function to close either this part or all parts + +{-| +Allows a file or 'ByteString' to be uploaded concurrently, using the +async library. 'ByteString's are split into 'chunkSize' chunks +and uploaded directly. + +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 `Error`, or `IOError`. +-} +concurrentUpload :: (MonadAWS m, MonadBaseControl IO m) + => UploadLocation -> CreateMultipartUpload -> m CompleteMultipartUploadResponse +concurrentUpload ud cmu = do + cmur <- send cmu + when (cmur ^. cmursResponseStatus /= 200) $ + fail "Failed to create upload" + let Just upId = cmur ^. cmursUploadId + bucket = cmu ^. cmuBucket + key = cmu ^. cmuKey + -- hndlr :: SomeException -> m CompleteMultipartUploadResponse + hndlr e = send (abortMultipartUpload bucket key upId) >> throwM e + + handling id hndlr $ do + umrs <- case ud of + BS bs -> forConcurrently (zip [1..] $ chunksOf chunkSize bs) $ \(partnum, b) -> do + umr <- send . uploadPart bucket key partnum upId . toBody $ b + pure $ completedPart partnum <$> (umr ^. uprsETag) + + FP fp -> do + fsize <- liftIO $ getFileSize fp + let (count,lst) = divMod (fromIntegral fsize) chunkSize + params = [(partnum, chunkSize*offset, size) + | partnum <- [1..] + | offset <- [0..count] + | size <- (chunkSize <$ [0..count-1]) ++ [lst] + ] + + forConcurrently params $ \(partnum,off,size) -> do + b <- liftIO $ mmapFileByteString fp (Just (fromIntegral off,size)) + umr <- send . uploadPart bucket key partnum upId . toBody $ b + pure $ completedPart partnum <$> (umr ^. uprsETag) + + let prts = nonEmpty =<< sequence umrs + send $ completeMultipartUpload bucket key upId + & cMultipartUpload ?~ set cmuParts prts 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 + case mki of + Nothing -> pure () + Just (key,uid) -> send (abortMultipartUpload bucket key uid) >> pure () + + + +-- 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))