diff --git a/distributed-fork-aws-lambda.cabal b/distributed-fork-aws-lambda.cabal --- a/distributed-fork-aws-lambda.cabal +++ b/distributed-fork-aws-lambda.cabal @@ -1,11 +1,11 @@ name: distributed-fork-aws-lambda -version: 0.0.1.2 +version: 0.0.2.0 synopsis: AWS Lambda backend for distributed-fork. homepage: https://github.com/utdemir/distributed-fork license: BSD3 license-file: LICENSE author: Utku Demir -maintainer: me@utdemir.com +maintainer: Utku Demir copyright: Utku Demir category: Control build-type: Simple @@ -40,14 +40,15 @@ , base64-bytestring , bytestring , containers + , elf , interpolate , lens , lens-aeson , safe-exceptions + , stm , stratosphere >= 0.15.0 , text , time - , typed-process , unordered-containers , zip-archive diff --git a/src/Control/Concurrent/Throttled.hs b/src/Control/Concurrent/Throttled.hs --- a/src/Control/Concurrent/Throttled.hs +++ b/src/Control/Concurrent/Throttled.hs @@ -1,3 +1,7 @@ +{-# LANGUAGE DeriveGeneric #-} + +-- | +-- A utility module which lets you put a concurrency limit to an IO action. module Control.Concurrent.Throttled ( Throttle , newThrottle @@ -5,26 +9,23 @@ ) where -------------------------------------------------------------------------------- -import Control.Monad -import Control.Exception -import Control.Concurrent +import Control.Monad (when) +import Control.Exception (finally) +import Control.Concurrent.STM -------------------------------------------------------------------------------- newtype Throttle - = Throttle (MVar (MVar ())) + = Throttle (TVar Int) newThrottle :: Int -> IO Throttle -newThrottle maxConcurrent = do - mv <- newEmptyMVar - forM_ [1..maxConcurrent] $ \_ -> forkIO . forever $ do - m <- takeMVar mv - putMVar m () - putMVar m () - return $ Throttle mv +newThrottle lim = + Throttle <$> newTVarIO (max lim 1) throttled :: Throttle -> IO a -> IO a -throttled (Throttle mv) act = do - m <- newEmptyMVar - putMVar mv m - readMVar m - act `finally` takeMVar m +throttled (Throttle tv) act = do + atomically $ do + r <- readTVar tv + when (r == 0) retry + writeTVar tv (r - 1) + finally act $ + atomically $ modifyTVar tv (+ 1) diff --git a/src/Control/Distributed/Fork/Lambda/Internal/Archive.hs b/src/Control/Distributed/Fork/Lambda/Internal/Archive.hs --- a/src/Control/Distributed/Fork/Lambda/Internal/Archive.hs +++ b/src/Control/Distributed/Fork/Lambda/Internal/Archive.hs @@ -1,6 +1,6 @@ -{-# LANGUAGE BinaryLiterals #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE BinaryLiterals #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE ScopedTypeVariables #-} {- This module contains the executables for the Lambda function. @@ -14,18 +14,18 @@ ) where -------------------------------------------------------------------------------- -import Codec.Archive.Zip hiding - (Archive) +import Codec.Archive.Zip hiding + (Archive) import Control.Exception import Control.Monad -import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy as BL +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as BL import Data.Digest.Pure.SHA +import Data.Elf import Data.Function import Data.String.Interpolate -import qualified Data.Text as T -import qualified Data.Text.Encoding as T -import System.Process.Typed +import qualified Data.Text as T +import qualified Data.Text.Encoding as T -------------------------------------------------------------------------------- import Control.Distributed.Fork.Backend import Control.Distributed.Fork.Lambda.Internal.Constants @@ -45,20 +45,28 @@ import boto3 queue_url = os.environ["#{envAnswerQueueUrl}"] -client = boto3.client('sqs') +bucket_url = os.environ["#{envAnswerBucketUrl}"] +sqs = boto3.client('sqs') +s3 = boto3.client('s3') + def handle(event, context): popen = subprocess.Popen( ["./#{hsMainName}", "#{argExecutorMode}"], stdin=subprocess.PIPE, stdout=subprocess.PIPE) (out, _) = popen.communicate(b64decode(event["d"])) - client.send_message( + ret = b64encode(out) + sqs.send_message( QueueUrl=queue_url, - MessageBody=b64encode(out), + MessageBody=ret, MessageAttributes={ "Id": { "DataType": "Number", "StringValue": str(event["i"]) + }, + "AnswerType": { + "DataType": "String", + "StringValue": "inline" } } ) @@ -69,40 +77,52 @@ And we read the current executable. Since it'll run on AWS Lambda, it needs to be a statically linked Linux -executable, so we do a preliminary check here. We're calling the "file" -command here instead of using libmagic, because trying to statically compile -it caused problems on my system. +executable, so we do a preliminary check here. -} mkHsMain :: IO BS.ByteString mkHsMain = do path <- getExecutablePath - checkFile path - BS.readFile path - where - checkFile path = do - ret <- - liftIO $ - T.splitOn ", " . T.decodeUtf8 . BL.toStrict <$> - readProcessStdout_ (proc "file" ["--brief", path]) - let preds :: [(T.Text, T.Text)] - preds = - [ ( "ELF 64-bit LSB executable" - , "file is not an 64 bit ELF executable") - , ( "statically linked" - , "file is not statically linked" - ) - ] - forM_ preds $ \(str, ex) -> - if str `elem` ret - then return () - else throwIO (FileException ex) - return () + contents <- BS.readFile path + assertBinary contents + return contents -newtype FileException - = FileException T.Text - deriving Show +assertBinary :: BS.ByteString -> IO () +assertBinary contents = do + elf <- (return $! parseElf contents) + `catch` (\(_ :: SomeException) -> throwIO FileExceptionNotElf) + unless (elfClass elf == ELFCLASS64) $ + throwIO FileExceptionNot64Bit + when (any (\s -> elfSegmentType s == PT_DYNAMIC) (elfSegments elf)) $ + throwIO FileExceptionNotStatic +data FileException + = FileExceptionNotElf + | FileExceptionNot64Bit + | FileExceptionNotStatic + instance Exception FileException + +instance Show FileException where + show FileExceptionNotElf = [i| + Error: I am not an ELF (Linux) binary. + + The executable will run on AWS environment, because of that + this library currently only supports Linux. + |] + show FileExceptionNot64Bit = [i| + Error: I am not a 64bit executable. + + AWS Lambda currently only runs 64 bit executables. + |] + show FileExceptionNotStatic = [i| + Error: I am not a dynamic executable. + + Since the executable will run on AWS environment, it needs + to be statically linked. + + You can give GHC "-optl-static -optl-pthread -fPIC" flags + to statically compile executables. + |] {- And we're going to put all of them in a zip archive. diff --git a/src/Control/Distributed/Fork/Lambda/Internal/Constants.hs b/src/Control/Distributed/Fork/Lambda/Internal/Constants.hs --- a/src/Control/Distributed/Fork/Lambda/Internal/Constants.hs +++ b/src/Control/Distributed/Fork/Lambda/Internal/Constants.hs @@ -14,3 +14,6 @@ envAnswerQueueUrl :: Text envAnswerQueueUrl = "ANSWER_QUEUE_URL" + +envAnswerBucketUrl :: Text +envAnswerBucketUrl = "ANSWER_BUCKET_URL" diff --git a/src/Control/Distributed/Fork/Lambda/Internal/Invoke.hs b/src/Control/Distributed/Fork/Lambda/Internal/Invoke.hs --- a/src/Control/Distributed/Fork/Lambda/Internal/Invoke.hs +++ b/src/Control/Distributed/Fork/Lambda/Internal/Invoke.hs @@ -182,10 +182,10 @@ withInvoke env stack f = do le <- newLambdaEnv env stack throttle <- newThrottle 128 - let answerT = async $ catchAny (answerThread le) $ \ex -> print ex - deadLetterT = async $ catchAny (deadLetterThread le) $ \ex -> print ex - void . sequence $ replicate 4 answerT ++ replicate 2 deadLetterT - f $ execute le throttle + let answerT = async . forever $ catchAny (answerThread le) $ \ex -> print ex + deadLetterT = async . forever $ catchAny (deadLetterThread le) $ \ex -> print ex + threads <- (++) <$> replicateM 4 answerT <*> replicateM 2 deadLetterT + f (execute le throttle) `finally` mapM_ cancel threads -------------------------------------------------------------------------------- diff --git a/src/Control/Distributed/Fork/Lambda/Internal/Stack.hs b/src/Control/Distributed/Fork/Lambda/Internal/Stack.hs --- a/src/Control/Distributed/Fork/Lambda/Internal/Stack.hs +++ b/src/Control/Distributed/Fork/Lambda/Internal/Stack.hs @@ -1,7 +1,6 @@ -{-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE QuasiQuotes #-} @@ -18,24 +17,24 @@ -------------------------------------------------------------------------------- import Control.Exception.Safe +import Control.Lens import Control.Monad -import Data.Aeson (Value (Object)) -import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy as BL -import qualified Data.HashMap.Strict as HM +import Data.Aeson (Value (Object)) +import Data.Aeson.QQ +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as BL +import qualified Data.HashMap.Strict as HM import Data.List import Data.Maybe import Data.Monoid -import qualified Data.Text as T -import qualified Data.Text.Encoding as T -import Control.Lens -import Network.AWS hiding (environment) +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import Network.AWS hiding (environment) +import Network.AWS.CloudFormation import Network.AWS.Lambda -import qualified Network.AWS.S3 as S3 +import qualified Network.AWS.S3 as S3 import Network.AWS.Waiter -import qualified Stratosphere as S -import Data.Aeson.QQ -import Network.AWS.CloudFormation +import qualified Stratosphere as S -------------------------------------------------------------------------------- import Control.Distributed.Fork.Lambda.Internal.Constants import Control.Distributed.Fork.Lambda.Internal.Types @@ -71,15 +70,15 @@ (deadLetterQueue) for collecting the failures. -} seTemplate :: StackOptions -> S.Template -seTemplate StackOptions{..} = +seTemplate StackOptions{ soLambdaCode = S3Loc (BucketName bucketName) path, soLambdaMemory } = S.template (S.Resources - [ S.resource "func" $ + [ S.resource templateOutputFunc $ S.LambdaFunctionProperties $ S.lambdaFunction (S.lambdaFunctionCode - & S.lfcS3Bucket ?~ S.Ref templateParameterS3Bucket - & S.lfcS3Key ?~ S.Ref templateParameterS3Key) + & S.lfcS3Bucket ?~ S.Literal bucketName + & S.lfcS3Key ?~ S.Literal path) "handler.handle" (S.GetAtt "role" "Arn") (S.Literal S.Python27) @@ -88,20 +87,17 @@ & S.lfDeadLetterConfig ?~ S.LambdaFunctionDeadLetterConfig (Just $ S.GetAtt "deadLetterQueue" "Arn") , S.resource "role" seRole - , S.resource "answerQueue" $ S.SQSQueueProperties S.sqsQueue - , S.resource "deadLetterQueue" $ S.SQSQueueProperties S.sqsQueue + , S.resource templateOutputAnswerQueue $ S.SQSQueueProperties S.sqsQueue + , S.resource templateOutputDeadLetterQueue $ S.SQSQueueProperties S.sqsQueue + , S.resource templateOutputAnswerBucket $ S.S3BucketProperties S.s3Bucket ]) & - S.templateParameters ?~ - S.Parameters - [ S.parameter templateParameterS3Bucket "String" - , S.parameter templateParameterS3Key "String" - ] & S.templateOutputs ?~ - S.Outputs - [ S.output templateOutputFunc (S.Ref "func") - , S.output templateOutputAnswerQueue (S.Ref "answerQueue") - , S.output templateOutputDeadLetterQueue (S.Ref "deadLetterQueue") - ] + S.Outputs + [ S.output templateOutputFunc (S.Ref templateOutputFunc) + , S.output templateOutputAnswerQueue (S.Ref templateOutputAnswerQueue) + , S.output templateOutputDeadLetterQueue (S.Ref templateOutputDeadLetterQueue) + , S.output templateOutputAnswerBucket (S.Ref templateOutputAnswerBucket) + ] seRole :: S.ResourceProperties seRole = @@ -155,12 +151,6 @@ _ -> error "invariant violation" -templateParameterS3Bucket :: T.Text -templateParameterS3Bucket = "s3bucket" - -templateParameterS3Key :: T.Text -templateParameterS3Key = "s3key" - templateOutputFunc :: T.Text templateOutputFunc = "output" @@ -170,6 +160,9 @@ templateOutputDeadLetterQueue :: T.Text templateOutputDeadLetterQueue = "deadLetterQueue" +templateOutputAnswerBucket :: T.Text +templateOutputAnswerBucket = "answerBucket" + -------------------------------------------------------------------------------- {- @@ -214,22 +207,13 @@ } seCreateStack :: StackOptions -> AWS StackInfo -seCreateStack options@StackOptions { soName = StackName stackName - , soLambdaCode = S3Loc (BucketName bucketName) path} = do +seCreateStack options@StackOptions { soName = StackName stackName } = do csrs <- send $ createStack stackName & csTemplateBody ?~ (T.decodeUtf8 . BL.toStrict . S.encodeTemplate $ seTemplate options) & csCapabilities .~ [CapabilityIAM] - & csParameters .~ - [ parameter - & pParameterKey ?~ templateParameterS3Bucket - & pParameterValue ?~ bucketName - , parameter - & pParameterKey ?~ templateParameterS3Key - & pParameterValue ?~ path - ] unless (csrs ^. csrsResponseStatus == 200) $ throwM $ StackException @@ -268,11 +252,17 @@ Nothing -> throwM $ StackException "Could not determine deadLetterQueue URL." Just t -> return t + answerBucket <- case lookupOutput stackRs templateOutputAnswerBucket of + Nothing -> throwM $ StackException "Could not determine answerBucket URL." + Just t -> return t + _ <- send $ updateFunctionConfiguration func & ufcEnvironment ?~ ( environment - & eVariables ?~ HM.fromList [ (envAnswerQueueUrl, answerQueue) ] + & eVariables ?~ HM.fromList [ (envAnswerQueueUrl, answerQueue) + , (envAnswerBucketUrl, answerBucket) + ] ) return $ StackInfo { siId = stackId diff --git a/test/Test.hs b/test/Test.hs --- a/test/Test.hs +++ b/test/Test.hs @@ -10,7 +10,6 @@ import Control.Exception (try) import Data.Monoid ((<>)) import qualified Data.Text as T -import System.Environment import Test.Tasty import Test.Tasty.HUnit -------------------------------------------------------------------------------- @@ -21,12 +20,7 @@ main :: IO () main = do initDistributedFork - tests_ <- lookupEnv "ENABLE_AWS_TESTS" >>= \case - Nothing -> do - putStrLn "ENV[ENABLE_AWS_TESTS] is not set, skipping." - return $ testGroup "No tests" [] - Just _ -> return tests - defaultMain tests_ + defaultMain tests opts :: LambdaBackendOptions opts = lambdaBackendOptions "serverless-batch"