Go 言語で AWS S3 のダウンロードに時間がかかったとき中断する

こんにちは、 rightgo09 です。

以下は、Go 言語で AWS の S3 からファイルをダウンロードするコードです。

package main

import (
	"io"
	"os"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

func main() {
	sess := session.Must(session.NewSession())

	s3svc := s3.New(sess)

	srcObject, err := s3svc.GetObject(&s3.GetObjectInput{
		Bucket: aws.String("mybucket"),
		Key:    aws.String("my/picture/thumbnail-938423.jpg"),
	})
	if err != nil {
		panic(err)
	}
	defer srcObject.Body.Close()

	f, err := os.Create("thumbnail-938423.jpg")
	if err != nil {
		panic(err)
	}
	defer f.Close()
	if _, err := io.Copy(f, srcObject.Body); err != nil {
		panic(err)
	}
}

context.Context で、ダウンロードに 100 ミリ秒かかったらエラーとなるように、タイムアウト処理を追加してみます。

package main

import (
	"context"
	"io"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

func main() {
	sess := session.Must(session.NewSession())
	s3svc := s3.New(sess)

	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	srcObject, err := s3svc.GetObjectWithContext(ctx, &s3.GetObjectInput{
		Bucket: aws.String("mybucket"),
		Key:    aws.String("my/picture/thumbnail-938423.jpg"),
	})
	if err != nil {
		panic(err)
	}
	defer srcObject.Body.Close()

	f, err := os.Create("thumbnail-938423.jpg")
	if err != nil {
		panic(err)
	}
	defer f.Close()
	if _, err := io.Copy(f, srcObject.Body); err != nil {
		panic(err)
	}
}

もし初期化やダウンロードに時間がかかると、 s3.GetObjectWithContext や io.Copy の err で以下のエラーが格納されます。

panic: context deadline exceeded

AWS Lambda の場合

Lambda では実行時に context.Context を受け取ることができます。

これを元に context.Context を作って、同じく S3 からのダウンロードに 100 ミリ秒以上かかったら失敗としてみます。

package main

import (
	"context"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

func main() {
	lambda.Start(run)
}

func run(ctx context.Context) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
	defer cancel()

	errCh := make(chan error)
	go func() {
		errCh <- download(ctx)
	}()

	select {
	case err := <-errCh:
		if err != nil {
			return "", err
		}
	}

	return "ok", nil
}

func download(ctx context.Context) error {
	sess := session.Must(session.NewSession())
	s3svc := s3.New(sess)

	errCh := make(chan error)

	go func() {
		srcObject, err := s3svc.GetObjectWithContext(ctx, &s3.GetObjectInput{
			Bucket: aws.String("mybucket"),
			Key:    aws.String("my/picture/thumbnail-938423.jpg"),
		})
		if err != nil {
			errCh <- err
			return
		}
		defer srcObject.Body.Close()

		f, err := os.Create("/tmp/thumbnail-938423.jpg")
		if err != nil {
			errCh <- err
			return
		}
		defer f.Close()

		if _, err := io.Copy(f, srcObject.Body); err != nil {
			errCh <- err
			return
		}

		errCh <- nil
	}()

	select {
	case err := <-errCh:
		if err != nil {
			return err
		}
	case <-ctx.Done():
		return ctx.Err()
	}

	return nil
}

成功した場合のレスポンス

"ok"

ダウンロードに時間がかかり、失敗した場合のレスポンス

{
  "errorMessage": "context deadline exceeded",
  "errorType": "deadlineExceededError"
}

ちなみにこの context.Context は「Lambda が起動した時間＋Lambdaのタイムアウト秒」を DeadLine として、既に設定されています。

Lambda 設定 タイムアウト： 15秒
起動時間： 2019-07-25 02:05:43

func run(ctx context.Context) (string, error) {
	fmt.Println("now: ", time.Now())
	deadline, _ := ctx.Deadline()
	fmt.Println("deadline: ", deadline)

	return "ok", nil
}

now:  2019-07-25 02:05:43.876590122 +0000 UTC m=+0.068351685
deadline:  2019-07-25 02:05:58.875287691 +0000 UTC

これは、「以降の処理があと５秒はかかるのがわかっているのに、 Deadline があと 2 秒後に迫っている」というような状況のとき、その時点で諦める、という使い方ができます。これはリソースの節約につながります。

deadline, _ := ctx.Deadline()
if deadline.Sub(time.Now().Add(5 * time.Second)) < 0 {
	return context.DeadlineExceeded
}

以上です。