こんにちは、デジスマチームでエンジニアをやっている堀田です。

先日、社内のTechTalkで「TypeScriptをAWS Lambdaで実行する」というテーマで話をしました。

TSは型で色々できるのだから、Lambdaの実装を型だけでやれたら嬉しいですよね？という内容です。

この記事ではその時の話を紹介します。

0. ゴール設定
1. 普段通りLambdaを使う
2. Lambda関数を型だけで実装する
3. Handler型を動かすために
- Custom Runtime
  - bootstrapの実装
- TypeScript Compiler API
麻雀の点数型をAPIに
まとめ
We are hiring!!

0. ゴール設定

初めにゴール設定をします。

型定義が書かれたTSファイルをアップロードして、
Eventを渡し、
動けばOK

を目標とします。その様子を下図に示します。

1. 普段通りLambdaを使う

比較のために一度、普段通りLambdaを実装してみます。条件は次の通りで、「受け取った2つの数字の足し算結果を返す」Lambdaを考えます。

File: index.ts
Runtime: nodejs20.x
Handler: index.handler

実装は次のようになります。Lambdaにデプロイする際には、TS -> JSにトランスパイルして、JSファイルをアップロードすることになります。

// index.ts

type Event = {
    n1: number;
    n2: number;
}

function sum(n1: number, n2: number): number {
  return n1 + n2;
}

export const handler = async (event: Event) => {
    return sum(event.n1, event.n2);
}

余談ですが、 Node.js 22 からは --experimental-transform-types フラグを付けることでTSファイルのままでも一部実行できるようになります。先日Lambdaでも nodejs22.x Runtime がサポートされましたが、

While you can enable the feature in Lambda, your function entrypoint (i.e. index.mjs or app.cjs) cannot currently be written using TypeScript as the runtime expects a file with a JavaScript extension.

とあるように、現時点ではエントリーポイントにはJSが必須なようです。

2. Lambda関数を型だけで実装する

先ほどの足し算の実装を今度はTSの型だけで行います。

type Num<
  N extends number,
  COUNT extends number[] = []
> = COUNT["length"] extends N ? COUNT : Num<N, [1, ...COUNT]>

type Sum<N1 extends number, N2 extends number> = [
    ...Num<N1>,
    ...Num<N2>
]["length"]

以前、型の足し算についてまとめたことがあるので、詳細はそちらを参照としますが、次のように期待通り動作します。 (この方法では制限がありますが、一旦よしとします。)

let a1: Sum<9, 16> = 25
let a2: Sum<100, 160> = 260
let a3: Sum<999, 999> = 1998

// @ts-expect-error
let b1: Sum<9, 16> = 24
// @ts-expect-error
let b2: Sum<100, 160> = 259
// @ts-expect-error
let b3: Sum<999, 999> = 1997

最終的にLambdaにアップロードする予定なので、Handler型を用意します。

export type Handler<EVENT extends Event> = Sum<EVENT["n1"], EVENT["n2"]>

最初のセクションで作ったhandler関数と同じi/fになるように実装しました。

export const handler = async (event: Event) => {
    return sum(event.n1, event.n2);
}

3. Handler型を動かすために

ここからが本題です。先ほど定義した Handler 型はただの型定義なので、普通にやると動きません。それを動かすために必要なのが、

AWS Lambda Custom Runtime
TypeScript Compiler API

です。

Custom Runtimeは、LambdaのRuntimeを自作するための仕組みで、TypeScript Compiler API はtypescriptの機能が提供されているAPIです。

Custom Runtimeの上でCompiler APIを動かし、アップロードされたHandler型を評価して結果を返すというのが全体の仕組みです。

Custom Runtime

ここからは、Custom Runtimeの詳細を解説します。ポイントは次の通りです。

Lambda側では OS-only Runtime を指定する (今回は provided.al2023)
Lambda実行時のエントリーポイントとなる実行可能ファイル bootstrap を用意する
bootstrap を、関数コードと一緒にアップロードするか、Lambda Layerに組み込んで提供する

今回は、型定義ファイルをアップロードするだけで動く環境を作りたかったので次のようにLambda Layerで提供することにしました。

resource "aws_lambda_layer_version" "tlrt" {
  layer_name = local.runtime_name
  filename   = data.archive_file.lambda.output_path

  source_code_hash = data.archive_file.lambda.output_base64sha256
}

data "archive_file" "lambda" {
  type        = "zip"
  output_path = "${path.module}/${local.runtime_name}.zip"

  source_dir = "${path.module}/runtime"
}

Lambda自体は次のように定義します

resource "aws_lambda_function" "sample" {
  function_name = "sample"
  role          = aws_iam_role.sample.arn

  runtime  = "provided.al2023"
  handler  = "index.Handler"
  filename         = data.archive_file.sample.output_path
  source_code_hash = data.archive_file.sample.output_base64sha256

  layers = [
    aws_lambda_layer_version.tlrt.arn,
  ]

  timeout = 600
}

export const handler が export type Handler に変わったので、 handler変数には "index.Handler" を設定するようにしています。

この辺り、成果物とソースコードはGitHubに上げています。

bootstrapの実装

Custom Runtimeを動かすための最低限の実装を紹介します。そのために必要なのは、

実行時に渡されたEventを取得する
結果を計算し、それを返却する

です。

Lambdaでは、Custom Runtimeを実装するための Runtime API が提供されていて、 bootstrapからこのAPIを呼べば、両方クリアできます。それ以外にもエラーを報告するAPIなどもありますが、ここでは省略します。

APIのendpointも、 AWS_LAMBDA_RUNTIME_API の環境変数でLambda側から提供されます。bootstrapは実行可能ファイルであればなんでも良いので、今回はgoで実装しました。

package cmd

import (
    "io"
    "net/http"
    "os"
)

type LambdaRuntimeApi struct {
}

// AWS_LAMBDA_RUNTIME_API is an environment variable that is set by the Lambda service.
var baseUrl = "http://" + os.Getenv("AWS_LAMBDA_RUNTIME_API") + "/2018-06-01"

// 1. 実行時に渡されたEventを取得する
func (api *LambdaRuntimeApi) GetInvocation() (*http.Response, error) {
    path := "/runtime/invocation/next"
    return http.Get(baseUrl + path)
}

// 2. 結果を計算し、それを返却する
func (api *LambdaRuntimeApi) PostResponse(requestId string, body io.Reader) (*http.Response, error) {
    path := "/runtime/invocation/" + requestId + "/response"
    return http.Post(baseUrl+path, "application/json", body)
}

これで、入力と出力が可能になりました。次に必要なのは、型を評価する処理です。これにはTypeScript Compiler APIで実装したTSコード (をbundle & minifyしたJSコード)を使います。

bootstrapからは次のように呼び出しています。

nodeRes, err := exec.Command("/opt/node", "/opt/compute-type-value.js", handler, event, lambdaTaskRoot).Output()

ここで登場した、 lambdaTaskRoot は Lambdaにアップロードしたコードへのパスが指定されている変数です。Lambdaから LAMBDA_TASK_ROOT 環境変数として提供されているのでそれを使っています。

Custom Runtimeにnodejsを乗せているのは、Compiler APIを使うためにどうしても必要だったからです。*1

TypeScript Compiler API

型を評価する実装を書いていきます。Compiler APIを使いますが、私が把握している限りAPIの仕様をまとめた公式ドキュメントはありません。ただし、公式wiki があり、そこには豊富な例と使い方のヒントとなるような文章が書かれています。

例えば、このAPIを使う時、ProgramやHost、 Sourceなどの用語が出てきますが、これらは

- A Program which is the TypeScript terminology for your whole application
- A CompilerHost which represents the users' system, with an API for reading files, checking directories and case sensitivity etc.
- Many SourceFiles which represent each source file in the application, hosting both the text and TypeScript AST

として紹介されています。個人的には、「ファイル書き込みなどシステム関連の操作をしたい時はHost、AST取得や型チェックなどTypeScript関連の処理をやりたい場合はProgram、tsファイルを読み込んだ後はSourceとして扱うんだな」というかなり大雑把な理解で扱っています。

今回、このAPIを使ってやりたいことは、アップロードされたHandler型に実行時に渡されたEventデータをマッピングして評価することです。例えば、下のようなEventが渡されるとしましょう。

const event = {
  n1: 9,
  n2: 16
}

この場合、次のようにHandler型に {n1: 9, n2: 16} をマッピングした、新しいMain型を実行時に作り、それを評価するという流れになります。

import type { Handler } from "/path/to/uploaded_handler"

type Main = Handler<{n1: 9, n2: 16}>

Eventをマッピングして新しくMain型を作る処理を実装します。全てを載せるとかなり長くなるので、一部省いていますが、次のようにCompiler APIの factory を使うことで、TSコードを組み立てられます。

import {...} from "typescript" 

function createMain(
  entryFile: string,
  entryType: string,
  events: Record<string, string | number>,
): string {
  const source = createSourceFile("", "", ScriptTarget.ES2020);
  const printer = createPrinter({ newLine: NewLineKind.LineFeed });

  return printer.printNode(
    EmitHint.Unspecified,
    factory.createSourceFile(
      [
        factory.createImportDeclaration(
          undefined,
          factory.createImportClause(
            true,
            undefined,
            factory.createNamedImports([
              factory.createImportSpecifier(
                false,
                ...

「空のファイルの先頭にimport文を書いて、そのimport文はimport typeで、アップロードされたファイルからHandler型をimportして、type Mainを宣言して、Handlerの型パラメータにEventの値を渡す」というコードを愚直に書いていってます。

次にMain型を読み込んで、型を評価する実装をします。ここではHostとProgramを使用します。まず、 createCompilerHost を使ってHostを作成し、上で作ったMain型のファイルを出力(host.writeFile)します。

その次に、 createProgram を使ってProgramを作成しますが、この時にLambdaにアップロードされたファイルへのパス *2 (entryFile)と出力したMain型のファイルへのパス(mainFileName)を渡します。

export function evaluate(
  entryFile: string,
  entryType: string,
  events: Record<string, string | number>,
): LiteralType["value"] {
  const options = {
    module: ModuleKind.ES2020,
    target: ScriptTarget.ES2020,
    strict: true,
  };

  const mainTs = createMain(entryFile, entryType, events);
  const mainFileName = `/tmp/main.ts`;
  const host = createCompilerHost(options);
  host.writeFile(mainFileName, mainTs, false);

  const program = createProgram([mainFileName, entryFile], options, host);
}

最後にProgramから型チェッカー (checker) を取得し、評価対象のMain型を探し出し、その結果を返します。

  const checker = program.getTypeChecker();

  const source = program.getSourceFile(mainFileName);

  let evaluatedType: LiteralType["value"] | undefined;
  source?.forEachChild((node) => {
    if (isTypeAliasDeclaration(node) && node.name.escapedText === "Main") {
      const type = checker.getTypeAtLocation(node);

      if (type.isLiteral()) {
        evaluatedType = type.value;
        return;
      }
    }
  });

  if (evaluatedType === undefined) {
    throw new Error("RuntimeError: No Main type found in /tmp/main.ts");
  }

  return evaluatedType;

host.writeFileで書き込む必要があるのか？や他にスマートな方法無いのか？などは考えましたが、とりあえずこれでやりたいことを実現できたのでOKとします。

麻雀の点数型をAPIに

以上の仕組みによって、型定義ファイルをアップロードするだけでLambdaを実行できます。冒頭で紹介した足し算Lambdaも動きます。

実は私は2年前に、簡単な麻雀の点数型を作ったことがあります。これを使って、麻雀点数APIを作ってみました。

https://github.com/horita-yuya/mahjong-kun

import type { MahjongKun, Tile } from "mahjong-kun";

export type Handler<
  EVENTS extends {
    t1: Tile;
    t2: Tile;
    t3: Tile;
    t4: Tile;
    t5: Tile;
    t6: Tile;
    t7: Tile;
    t8: Tile;
    t9: Tile;
    t10: Tile;
    t11: Tile;
    t12: Tile;
    t13: Tile;
    t14: Tile;
    winning: Tile;
    style: "ron" | "tumo";
    position: "parent" | "child";
  },
> = MahjongKun<
  [
    EVENTS["t1"],
    EVENTS["t2"],
    EVENTS["t3"],
    EVENTS["t4"],
    EVENTS["t5"],
    EVENTS["t6"],
    EVENTS["t7"],
    EVENTS["t8"],
    EVENTS["t9"],
    EVENTS["t10"],
    EVENTS["t11"],
    EVENTS["t12"],
    EVENTS["t13"],
    EVENTS["t14"],
  ],
  EVENTS["winning"],
  [],
  EVENTS["style"],
  EVENTS["position"]
>;

麻雀の点数型を読み込んで、Handler型を作っただけです。アップロードする際は、bundleされたd.tsファイルが出力されるように、ビルドスクリプトを次のようにしました。

tsup index.ts --dts-resolve

自分が子で、ロン上がり、タンヤオのみを想定してAPIを呼ぶと

1300点が返ってきました。

まとめ

TypeScriptは型だけで色々表現できます。それらを実世界で活用するために、Type Level Runtimeを作った話をしました。 2年前に作った麻雀の点数型、型の世界で閉じていましたがこのRuntimeのおかげでAPIとして提供できるようになりました。

We are hiring!!

エムスリーでは絶賛エンジニアを募集中です！今回紹介した技術スタック・アーキテクチャ以外にも様々な構成のプロダクトがありますので、ご興味ある方は是非こちらからお願いします！

jobs.m3.com

*1:Rustで書かれたstcを使ってみる選択もありましたが、そこはtscにこだわりました。

*2:_HANDLER環境変数で取得可能

エムスリーテックブログ

エムスリー(m3)のエンジニア・開発メンバーによる技術ブログです

Type Level Runtime for AWS Lambda