File Coverage

blib/lib/AWS/Lambda/Bootstrap.pm

Criterion	Covered	Total	%
statement	134	156	85.9
branch	18	32	56.2
condition	13	35	37.1
subroutine	26	29	89.6
pod	0	8	0.0
total	191	260	73.4

package AWS::Lambda::Bootstrap;
use 5.026000;
use utf8;
use strict;
use warnings;
use HTTP::Tiny;
use JSON::XS qw/decode_json encode_json/;
use Try::Tiny;
use AWS::Lambda;
use AWS::Lambda::Context;
use AWS::Lambda::ResponseWriter;
use Scalar::Util qw(blessed);
use Exporter 'import';

our @EXPORT = ('bootstrap');

sub bootstrap {
    my $handler = shift;
    my $bootstrap = AWS::Lambda::Bootstrap->new(
        handler => $handler,
    );
    $bootstrap->handle_events;
}

sub new {
    my $proto = shift;
    my $class = ref $proto || $proto;
    my %args;
    if (@_ == 1 && ref $_[0] eq 'HASH') {
        %args = %{$_[0]};
    } else {
        %args = @_;
    }

    my $api_version = '2018-06-01';
    my $env_handler = $args{handler} // $ENV{'_HANDLER'} // die '$_HANDLER is not found';
    my ($handler, $function) = split(/[.]/, $env_handler, 2);
    my $runtime_api = $args{runtime_api} // $ENV{'AWS_LAMBDA_RUNTIME_API'} // die '$AWS_LAMBDA_RUNTIME_API is not found';
    my $task_root = $args{task_root} // $ENV{'LAMBDA_TASK_ROOT'} // die '$LAMBDA_TASK_ROOT is not found';
    my $max_workers = $args{max_workers} // $ENV{'AWS_LAMBDA_MAX_CONCURRENCY'} // 0;
    die "max_workers must be a non-negative integer, got: $max_workers"
        unless $max_workers =~ /^\d+$/ && $max_workers >= 0;
    my $self = bless +{
        task_root      => $task_root,
        handler        => $handler,
        function_name  => $function,
        runtime_api    => $runtime_api,
        api_version    => $api_version,
        next_event_url => "http://${runtime_api}/${api_version}/runtime/invocation/next",
        max_workers    => $max_workers,
        http           => HTTP::Tiny->new(
            # XXX: I want to disable timeout, but it seems HTTP::Tiny does not support it.
            # So, I set a long timeout.
            timeout => 365*24*60*60, # 365 days
        ),
    }, $class;
    return $self;
}

sub handle_events {
    my $self = shift;
    $self->_init or return;

    if ($self->{max_workers} > 0) {
        require Parallel::Prefork;
        my $pm = Parallel::Prefork->new({
            max_workers => $self->{max_workers},
            trap_signals => {
                TERM => 'TERM',
                HUP  => 'TERM',
            },
        });
        while ($pm->signal_received ne 'TERM') {
            $pm->start and next;
            $self->_handle_events;
            $pm->finish;
        }
        $pm->wait_all_children;
    } else {
        $self->_handle_events
    }
}

sub _handle_events {
    my $self = shift;
    while(1) {
        $self->handle_event;
    }
}

sub _init {
    my $self = shift;
    if (my $func = $self->{function}) {
        return $func;
    }

    my $task_root = $self->{task_root};
    my $handler = $self->{handler};
    my $name = $self->{function_name};
    return try {
        package main;
        require "${task_root}/${handler}.pl";
        my $f = main->can($name) // die "handler $name is not found";
        $self->{function} = $f;
    } catch {
        $self->lambda_init_error($_);
        $self->{function} = sub {};
        undef;
    };
}

sub handle_event {
    my $self = shift;
    $self->_init or return;
    my ($payload, $context) = $self->lambda_next;
    my $response = try {
        local $AWS::Lambda::context = $context;
        local $ENV{_X_AMZN_TRACE_ID} = $context->{trace_id};
        $self->{function}->($payload, $context);
    } catch {
        my $err = $_;
        print STDERR "$err";
        $self->lambda_error($err, $context);
        bless {}, 'AWS::Lambda::ErrorSentinel';
    };
    my $ref = ref($response);
    if ($ref eq 'AWS::Lambda::ErrorSentinel') {
        return;
    }
    if ($ref eq 'CODE') {
        $self->lambda_response_streaming($response, $context);
    } else {
        $self->lambda_response($response, $context);
    }
    return 1;
}

sub lambda_next {
    my $self = shift;
    my $resp = $self->{http}->get($self->{next_event_url});
    if (!$resp->{success}) {
        die "failed to retrieve the next event: $resp->{status} $resp->{reason}";
    }
    my $h = $resp->{headers};
    my $payload = decode_json($resp->{content});
    return $payload, AWS::Lambda::Context->new(
        deadline_ms          => $h->{'lambda-runtime-deadline-ms'},
        aws_request_id       => $h->{'lambda-runtime-aws-request-id'},
        invoked_function_arn => $h->{'lambda-runtime-invoked-function-arn'},
        trace_id             => $h->{'lambda-runtime-trace-id'},
        tenant_id            => $h->{'lambda-runtime-aws-tenant-id'},
    );
}

sub lambda_response {
    my $self = shift;
    my ($response, $context) = @_;
    my $runtime_api = $self->{runtime_api};
    my $api_version = $self->{api_version};
    my $request_id = $context->aws_request_id;
    my $url = "http://${runtime_api}/${api_version}/runtime/invocation/${request_id}/response";
    my $resp = $self->{http}->post($url, {
        content => encode_json($response),
    });
    if (!$resp->{success}) {
        die "failed to response of execution: $resp->{status} $resp->{reason}";
    }
}

sub lambda_response_streaming {
    my $self = shift;
    my ($response, $context) = @_;
    my $runtime_api = $self->{runtime_api};
    my $api_version = $self->{api_version};
    my $request_id = $context->aws_request_id;
    my $url = "http://${runtime_api}/${api_version}/runtime/invocation/${request_id}/response";
    my $writer = undef;
    try {
        $response->(sub {
            my $content_type = shift;
            $writer = AWS::Lambda::ResponseWriter->new(
                response_url => $url,
                http         => $self->{http},
            );
            $writer->_request($content_type);
            return $writer;
        });
    } catch {
        my $err = $_;
        print STDERR "$err";
        if ($writer) {
            $writer->_close_with_error($err);
        } else {
            $self->lambda_error($err, $context);
        }
    };
    if ($writer) {
        my $response = $writer->_handle_response;
        if (!$response->{success}) {
            die "failed to response of execution: $response->{status} $response->{reason}";
        }
    }
}

sub lambda_error {
    my $self = shift;
    my ($error, $context) = @_;
    my $runtime_api = $self->{runtime_api};
    my $api_version = $self->{api_version};
    my $request_id = $context->aws_request_id;
    my $url = "http://${runtime_api}/${api_version}/runtime/invocation/${request_id}/error";
    my $type = blessed($error) // "Error";
    my $resp = $self->{http}->post($url, {
        content => encode_json({
            errorMessage => "$error",
            errorType => "$type",
        }),
    });
    if (!$resp->{success}) {
        die "failed to send error of execution: $resp->{status} $resp->{reason}";
    }
}

sub lambda_init_error {
    my $self = shift;
    my $error = shift;
    my $runtime_api = $self->{runtime_api};
    my $api_version = $self->{api_version};
    my $url = "http://${runtime_api}/${api_version}/runtime/init/error";
    my $type = blessed($error) // "Error";
    my $resp = $self->{http}->post($url, {
        content => encode_json({
            errorMessage => "$error",
            errorType => "$type",
        }),
    });
    if (!$resp->{success}) {
        die "failed to send error of execution: $resp->{status} $resp->{reason}";
    }
}

1;
__END__