GoとRustで学ぶCustom AWS Lambda Runtimes 〜よりLambdaらしいCustom Runtime実装編〜

概要

• GoとRustで学ぶCustom AWS Lambda Runtimes 〜入門編〜の続きです
• いつものようにfunctionには独自型とhandlerだけ書きたい場合は　Custom Runtimeをどう実装すればよいか、Rust版のCustom AWS Lambda Runtimesを見ながら考え実装してみました。
• 不完全な箇所もありつつざっと一通り実装したものなのでGitHubの方のコードは改善していくと思います。

目指すもののおさらい

Custom Runtimeでもこんな風に独自イベントやhandlerを書きたい！

bootstrap.go

package main

import (
	"encoding/json"
	"fmt"
	"log"

	"github.com/toshi0607/go-custom-runtime-sample/runtime"
)

type MyEvent struct {
	Name string `json:"name"`
}

func MyHandler(ctx runtime.Context, event MyEvent) ([]byte, error) {
	return fmt.Sprintf("Hello %s!", event.Name), nil
}

func main() {
	runtime.Start(MyHandler)
}

実装仕方で工夫できるものの、前の記事の実装ではbootstrap.goにイベント取得処理やイベントを待ち受けるループを実装していてこれじゃない感を醸成していました。

runtime-client

前記事でも紹介したRustのCustom Runtimeのおおよその構成を踏襲して実装して見ます。

ただし、ここではすべてruntimeというパッケージの中にbootstrap以外のコードを置きます。もちろん分割するのもありだと思います。

runtime-clientはCustom Runtimeで登場した4つのAPIを愚直にハンドルするためのAPIクライアントです。必要な型も定義しています。

client.go

package runtime

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/pkg/errors"
)

const (
	RUNTIME_API_VERSION = "2018-06-01"
	CONTENT_TYPE        = "application/json"
)

type Client interface {
	NextInvocation() ([]byte, *EventContext, error)
	InvocationResponse(awsRequestId string, content []byte) error
	InvocationError(awsRequestId string, err error) error
	InitializationError(err error) error
}

// https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/runtimes-api.html#runtimes-api-next
type EventContext struct {
	InvokedFunctionArn string
	AwsRequestId       string
	RuntimeTraceId     string
	DeadlineMs         int64
	ClientContext      ClientContext
	CognitoIdentity    CognitoIdentity
}

type CognitoIdentity struct {
	identityId     string
	identityPoolId string
}

type ClientContext struct {
	client      ClientApplication
	custom      map[string]string
	environment map[string]string
}

type ClientApplication struct {
	installationId string
	appTitle       string
	appVersionName string
	appVersionCode string
	appPackageName string
}

type client struct {
	endpoint string
	client   *http.Client
}

func NewClient(endpoint string) Client {
	return &client{
		endpoint: endpoint,
		client:   http.DefaultClient,
	}
}

// http://${AWS_LAMBDA_RUNTIME_API}/2018-06-01/runtime/invocation/next

// https://docs.aws.amazon.com/lambda/latest/dg/runtimes-api.html#runtimes-api-next
func (c *client) NextInvocation() ([]byte, *EventContext, error) {
	url := c.endpoint + "/" + RUNTIME_API_VERSION + "/runtime/invocation/next"
	resp, err := c.client.Get(url)
	if err != nil {
		return nil, nil, errors.New("failed to get a response")
	}
	defer resp.Body.Close()
	ec, err := getEventContext(resp.Header)
	if err != nil {
		return nil, nil, errors.New("failed to get an event context")
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, nil, errors.New("failed to get an event")
	}

	return b, ec, nil
}

// https://docs.aws.amazon.com/lambda/latest/dg/runtimes-api.html#runtimes-api-response
func (c *client) InvocationResponse(awsRequestId string, content []byte) error {
	// /runtime/invocation/AwsRequestId/response
	url := c.endpoint + "/" + RUNTIME_API_VERSION + "/runtime/invocation/" + awsRequestId + "/response"

	resp, err := c.client.Post(url, CONTENT_TYPE, bytes.NewBuffer(content))
	if err != nil {
		return errors.Wrap(err, "failed to post content")
	}
	defer resp.Body.Close()

	return nil
}

// https://docs.aws.amazon.com/lambda/latest/dg/runtimes-api.html#runtimes-api-invokeerror
func (c *client) InvocationError(awsRequestId string, aerr error) error {
	// /runtime/invocation/AwsRequestId/error
	url := c.endpoint + "/" + RUNTIME_API_VERSION + "/runtime/invocation/" + awsRequestId + "/error"
	errs, err := json.Marshal(aerr)
	if err != nil {
		return errors.Wrap(err, "failed to post content")
	}

	resp, err := c.client.Post(url, CONTENT_TYPE, bytes.NewBuffer(errs))
	if err != nil {
		return errors.Wrap(err, "failed to post content")
	}
	defer resp.Body.Close()

	return nil
}

// https://docs.aws.amazon.com/lambda/latest/dg/runtimes-api.html#runtimes-api-initerror
func (c *client) InitializationError(aerr error) error {
	// /runtime/init/error
	url := c.endpoint + "/" + RUNTIME_API_VERSION + "/runtime/init/error"
	errs, err := json.Marshal(aerr)
	if err != nil {
		return errors.Wrap(err, "failed to post content")
	}

	resp, err := c.client.Post(url, CONTENT_TYPE, bytes.NewBuffer(errs))
	if err != nil {
		return errors.Wrap(err, "failed to post content")
	}
	defer resp.Body.Close()

	return nil
}

func getEventContext(header http.Header) (*EventContext, error) {
	invokedFunctionArn := header.Get("Lambda-Runtime-Invoked-Function-Arn")
	awsRequestId := header.Get("Lambda-Runtime-Aws-Request-Id")
	runtimeTraceId := header.Get("Lambda-Runtime-Trace-Id")

	runtimeDeadlineMs, err := strconv.ParseInt(header.Get("Lambda-Runtime-Deadline-Ms"), 10, 64)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse Lambda-Runtime-Deadline-Ms")
	}

	ev := &EventContext{
		InvokedFunctionArn: invokedFunctionArn,
		AwsRequestId:       awsRequestId,
		RuntimeTraceId:     runtimeTraceId,
		DeadlineMs:         runtimeDeadlineMs,
	}

	runtimeClientContext := header.Get("Lambda-Runtime-Client-Context")
	if runtimeClientContext != "" {
		var clientContext ClientContext
		if err := json.Unmarshal([]byte(runtimeClientContext), &clientContext); err != nil {
			return nil, errors.Wrap(err, "failed to unmarshal ClientContext")
		}

		ev.ClientContext = clientContext
	}

	runtimeCognitoIdentity := header.Get("Lambda-Runtime-Cognito-Identity") // 型
	if runtimeCognitoIdentity != "" {
		var cognitoIdentity CognitoIdentity
		if err := json.Unmarshal([]byte(runtimeClientContext), &cognitoIdentity); err != nil {
			return nil, errors.Wrap(err, "failed to unmarshal ClientContext")
		}

		ev.CognitoIdentity = cognitoIdentity
	}

	return ev, nil
}

runtime

runtime-clientを利用してイベントループを回すruntimeそのものです。

env.goは実行時に環境変数からハンドラーに渡すコンテキスト情報（APIで取得するもの以外）を集めるためのインターフェースとしてConfigProviderを定義しています。

Rust実装では同ファイルにテストで差し替える用のエンドポイントや仮の環境変数がセットできるようになっています。

env.go

package runtime

import (
	"os"
	"strconv"

	"github.com/pkg/errors"
)

type ConfigProvider interface {
	GetFunctionSettings() (*FunctionSettings, error)
	GetRuntimeApiEndpoint() string
}

type FunctionSettings struct {
	functionName string
	memorySize   int32
	version      string
	logStream    string
	logGroup     string
}

type configProvider struct{}

func NewConfigProvider() ConfigProvider {
	return &configProvider{}
}

func (c configProvider) GetFunctionSettings() (*FunctionSettings, error) {
	functionName := os.Getenv("AWS_LAMBDA_FUNCTION_NAME")
	version := os.Getenv("AWS_LAMBDA_FUNCTION_VERSION")
	logStream := os.Getenv("AWS_LAMBDA_LOG_STREAM_NAME")
	logGroup := os.Getenv("AWS_LAMBDA_LOG_GROUP_NAME")
	memoryStr := os.Getenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
	memorySize, err := strconv.Atoi(memoryStr)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to parse memoryStr: %s to int", memoryStr)
	}

	return &FunctionSettings{
		functionName: functionName,
		memorySize:   int32(memorySize),
		version:      version,
		logStream:    logStream,
		logGroup:     logGroup,
	}, nil
}

func (c configProvider) GetRuntimeApiEndpoint() string {
	return "http://" + os.Getenv("AWS_LAMBDA_RUNTIME_API")
}

runtime.goにメインの処理がまとまっています。ハンドラーを定義するファイルからruntime.start(MyHandler)のように呼び出すstart関数もここで定義しています。

runtime.go

package runtime

import (
	"encoding/json"
	"log"
)

type Handler interface {
	// ここでイベントの型に[]byteを指定してしまうと各ハンドラーでイベントをデシリアライズしないといけないので
	// interface{}型などにして型チェックを行いながら共通のデシリアライズ処理を実装すると
	// より柔軟になりそうです。
	Run(ctx Context, event []byte) ([]byte, error)
}

type handler func(ctx Context, event []byte) ([]byte, error)

// func型に名前を付けてRunメソッドを定義することでHandler interfaceを実装し、
// 自分で定義したハンドラーを渡せるようにしているところがポイントです。
// 本家aws-lambda-goでもそのように実装してあり参考にしました。
func (h handler) Run(ctx Context, event []byte) ([]byte, error) {
	response, err := h(ctx, event)
	if err != nil {
		return nil, err
	}

	responseBytes, err := json.Marshal(response)
	if err != nil {
		return nil, err
	}

	return responseBytes, nil
}

func Start(h handler) {
	startWithConfig(h, NewConfigProvider())
}

func startWithConfig(h handler, config ConfigProvider) {
	endpoint := config.GetRuntimeApiEndpoint()
	settings, err := config.GetFunctionSettings()
	if endpoint == "" || err != nil {
		log.Fatal("failed to init lambda")
	}

	startWithRuntimeClient(h, *settings, NewClient(endpoint))

}

func startWithRuntimeClient(h handler, s FunctionSettings, c Client) {
	NewRuntime(c, h, s).Start()
}

type Runtime struct {
	client   Client
	handler  Handler
	settings FunctionSettings
}

type Context struct {
	MemorySize         int32
	FunctionName       string
	FunctionVersion    string
	InvokedFunctionArn string
	AwsRequestId       string
	XrayTraceId        string
	LogStreamName      string
	LogGroupName       string
	ClientContext      ClientContext
	Identity           CognitoIdentity
	Deadline           int64
}

func NewRuntime(client Client, handler Handler, settings FunctionSettings) *Runtime {
	return &Runtime{
		client:   client,
		handler:  handler,
		settings: settings,
	}
}

func (r *Runtime) Start() {
	for {
		ctx, ev := r.getNextEvent()
		requestId := ctx.AwsRequestId
		result, err := r.handler.Run(ctx, ev)
		if err != nil {
			err := r.client.InvocationError(requestId, err)
			if err != nil {
				log.Fatal("failed to invoke an error")
			}
		}
		err = r.client.InvocationResponse(requestId, result)
		if err != nil {
			err := r.client.InvocationError(requestId, err)
			if err != nil {
				log.Fatal("failed to invoke an error")
			}
		}
	}
}

func (r *Runtime) getNextEvent() (Context, []byte) {
	ev, ctx, err := r.client.NextInvocation()
	if err != nil {
		r.client.InitializationError(err)
	}
	return Context{
		MemorySize:         r.settings.memorySize,
		FunctionName:       r.settings.functionName,
		FunctionVersion:    r.settings.version,
		LogStreamName:      r.settings.logStream,
		LogGroupName:       r.settings.logGroup,
		InvokedFunctionArn: ctx.InvokedFunctionArn,
		AwsRequestId:       ctx.AwsRequestId,
		XrayTraceId:        ctx.RuntimeTraceId,
		ClientContext:      ctx.ClientContext,
		Identity:           ctx.CognitoIdentity,
		Deadline:           ctx.DeadlineMs,
	}, ev
}

bootstrap

以上のようにruntime（とruntime-client）を実装すると、bootstrapを次のように実装できるようになります。

bootstrap.go

package main

import (
	"encoding/json"
	"fmt"
	"log"

	"github.com/toshi0607/go-custom-runtime-sample/runtime"
)

type MyEvent struct {
	Name string `json:"name"`
}

func MyHandler(ctx runtime.Context, event []byte) ([]byte, error) {
	// runtime.Startの引数を[]byteとして実装しているので`event MyEvent`とは書けません…
	log.Printf("request id: %s\n", ctx.AwsRequestId)

	var me MyEvent
	json.Unmarshal(event, &me)

	str := fmt.Sprintf("Hello %s!", me.Name)
	log.Println(str)
	return []byte(str), nil
}

func main() {
	runtime.Start(MyHandler)
}

これでいつものLambda（に前の記事のよりは近いが不完全なもの）を実装することができました。

Rust版ではリトライ処理が入っていたり、aws-lambda-go（Custom Runtimeでなく公式サポートされているGo向けのライブラリ）ではハンドラーの型チェックが入っていたり、実装を色々比較しながら自分で実装してみると気づきがあって面白かったです。

ただ冒頭でも書いた通りいろいろと省略してしまっているので改善したい気持ちが強いです。