エムスリーテックブログ

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

Server Functions っぽい仕組みを自作して Lambda 関数呼び出しに適用してみた

デジスマチームの池奥です。 新卒2年目のソフトウェアエンジニアです!

先週、人生で初めて外部モニターを買いました。快適な作業環境を手に入れたと思ったのですが、モニターのある環境に慣れておらずすっかり持て余しています。今のところ YouTube ビュワーとして活躍しています。

このブログは、デジスマチームブログリレーの 1 日目の記事です 🎉

イントロ

普段、私はNode.jsでバックエンドを実装することが多いのですが、一部の処理を AWS Lambda のような別の実行環境に切り出したい、というシーンに時々遭遇します。このような場合、Lambda の実装は別のパッケージに置き、元のサーバーからはSDK経由で呼び出す形にすることが多いと思います。

しかし、皆さんは「もし、Lambda関数を普通の関数のように呼び出せたら」と思ったことはありませんか?

別のパッケージに実装を置いてしまうと、呼び出し元のアプリケーションと Lambda 関数でビジネスロジックが分断され、コードの見通しが悪くなりがちです。さらに、データの受け渡しにはシリアライズ・デシリアライズの処理が必要となり、TypeScript の強力な型システムの恩恵を受けることが難しくなるという課題があります。

そこで今回、発想を転換してみます!ロジックを単一コードベース上に書き、ビルド時にランタイムごとに処理を分離するというアプローチです。あたかもローカル関数を呼び出すように Lambda 関数を使えるようになれば、開発体験も向上するはずです。

このアプローチのヒントになったのが、React Server Functionsです。React Server Functions は、コード上では単なる関数呼び出しに見える部分が、実行時には HTTP 通信に置き換わる仕組みです。呼び出し元はフロントエンド、呼び出された側はサーバーサイドで動作しますが、これをシームレスに記述できます。

まさに今回の要件にぴったりだと感じ、今回はこの仕組みをNode.jsバックエンドとLambda関数に応用するプロトタイプを作成しました!なお、あくまで趣味の実験として作ったものですので、一部の実装が雑な作りになっていますがご容赦ください。また、今回は AWS Lambda にデプロイすることを想定しています。

作ったもの

早速ですが、作ったツールの動作を説明します。

次のサンプルコードはAPIサーバー上に/addエンドポイントが定義されていて、受け取ったパラメータを calcAddition 関数で処理して、結果を返すというシンプルな実装です。今回、calcAdditionを AWS Lambda 上で実行したいとしましょう。

// src/action.ts (AWS Lambda で実行される)
"use lambda";
export const calcAddition = (a: number, b: number): number => {
  return a + b;
};

// src/route.ts (サーバーで実行される)
import { calcAddition } from "./action";

app.post("/add", (req, res) => {
  const { a, b } = req.body;
  const result = calcAddition(a, b);
  res.json({ result });
});

まずcalcAddition 関数には "use lambda" というディレクティブが付されています。これは今回新しく追加したもので、バンドラに「AWS Lambda で実行して欲しい」というメッセージを伝える特別な符号のようなものです。

バンドラはこの符号を見つけると、その関数を AWS Lambda にデプロイできるようなコードにビルドします。一方、route.ts のコードは単に calcAddition()として関数を呼び出していますが、こちらも実際にビルドを行うと、AWS Lambda を Invoke するコードに変換されます。

実際に動かしたい方は https://github.com/yuta-ike/use-lambda-prototype をご参照ください。なお、あくまでパパッと作ったプロトタイプなのでプラグインは非常に雑な実装です。ご容赦ください。

Rolldown プラグインの実装

React Server Functions は、ビルド時のメタプログラミングによって実現されています。今回はRolldownというバンドラのプラグインを作成し、その中でコードを動的に変換していきたいと思います。

Rolldown について

Rolldown は、Rust 製の JavaScript/TypeScript バンドラーです。元々RollupというJavaScript製のバンドラがあり、これと互換性を持つインタフェースを持っていることが特徴です。なお、プラグインはJavaScriptで実装できます。 Vite でもバージョン 7 から採用されています。

rolldown.rs

Rolldown によるコード変換はいくつかのフェーズに分かれていますが、フェーズごとに hook という形で任意の処理を挟むことができる仕組みが用意されています。よく利用されるものとして、ファイル単位でコードの書き換えを行える transform や、モジュールの解決先を操作するための resolveId などの hook があります。Rolldown プラグインはこれらの hook を組み合わせて実現されています。

なお、Rolldownのプラグイン関連のドキュメントはまだWIP段階のものが多く、同じインタフェースを持つRollupのドキュメントを参照すると良いです。

プラグインの概要

まずは、変換後のコードの全体像をお見せします。 次のサンプルコードは、メインとなる Node.js のサーバーのコードと、サーバーから呼び出される Lambda 関数のコードの 2 種類にビルドされます。これは単に Rolldown の設定を 2 通り用意して、2 回ビルドを実施しています。

// src/action.ts (Lambda関数で実行される)
"use lambda";
export const calcAddition = (a: number, b: number): number => {
  return a + b;
};

// src/route.ts (サーバーで実行される)
import { calcAddition } from "./action";

app.post("/add", (req, res) => {
  const { a, b } = req.body;
  const result = calcAddition(a, b);
  res.json({ result });
});

便宜上、以降は Node.js のサーバーのコードをメインアプリケーションと呼びます。メインアプリケーションのコードは、dist/server ディレクトリに、Lambda 関数のコードは dist/lambda ディレクトリに出力されます。

メインアプリケーションの出力は次のようになります。actions.js のコードがまるっと消えており、Lambda 関数を呼び出すコードに差し替えられています。

// dist/server/actions.js
export const calcAddition = (...input) => {
  // 中身が差し代わる
  const { Payload: payload } = await lambdaClient.send(
    new InvokeCommand({
      FunctionName: "action-calcAddition",
      InvocationType: "RequestResponse",
      Payload: JSON.stringify(input),
    })
  );
};

// dist/server/route.js
import { calcAddition } from "./action";

app.post("/add", (req, res) => {
  const { a, b } = req.body;
  const result = calcAddition(a, b);
  res.json({ result });
});

Lambda 関数側は次のようになります。actions.js で定義されている calcAddition 関数がそのままバンドルされていますが、Lambda 関数のエントリポイントとなる handler 関数が追加されています。

// dist/lambda/action-calcAddition.js
// 新しく追加された handler 関数
export const handler = async (event) => {
  const result = await calcAddition(...event);
  return JSON.stringify(result);
};

const calcAddition = (a, b) => {
  return a + b;
};

このように 2 種類のビルドを行うことで、1つのコードベースをメインアプリケーション向けと Lambda 関数向けの2つのランタイム向けに分離します。

メインアプリケーション側のプラグイン実装

メインアプリケーション用のプラグインの実装を紹介します。Rolldown の transform hook を利用します。

主な処理は次のようになります。

  1. コードを AST の形式にパースし、 "use lambda" というディレクティブが存在するかを確認する
  2. 存在する場合は、AST を走査し、ExportNamedDeclaration(名前付き Export)を抜き出す
  3. Export する関数を差し替えたコードを生成する
transform(code, id) {
  // 1-1. コードをパース -----------------------------------------------------
  const parsed = this.parse(code, {
    lang: "ts",
    astType: "ts",
  });

  const firstStatement = parsed.body.shift();
  const isUseLambdaDirective =
    firstStatement != null &&
    firstStatement.type === "ExpressionStatement" &&
    firstStatement.expression.type === "Literal" &&
    firstStatement.expression.value === "use lambda";

  // 1-2. "use lambda" ディレクティブが存在しない場合は何もしない -----------------
  if (!isUseLambdaDirective) {
    return code;
  }

  // 2. ExportNamedDeclaration を抽出
  const exports = parsed.body.filter(
    (statement) => statement.type === "ExportNamedDeclaration"
  );

  const identifiers = exports.flatMap((node) =>
    node.declaration?.type === "VariableDeclaration"
      ? node.declaration.declarations.flatMap((d) =>
          d.id.type === "Identifier" ? d.id.name : []
        )
      : node.declaration?.type === "FunctionDeclaration"
      ? node.declaration.id != null
        ? [node.declaration.id.name]
        : []
      : []
  );
  lambdaInfo.set(relative(`${import.meta.dirname}`, id), identifiers);

  // 3-1. 差し替えたコードを生成-----------------------------------------------
  const newCode = identifiers
    .map(
      (identifier) => `export const ${identifier} = async (...input) => {
          try {
            const command = new InvokeCommand({
              FunctionName: "action-${identifier}",
              InvocationType: "RequestResponse",
              Payload: JSON.stringify(input),
            });
            const { Payload: payload } = await lambdaClient.send(command);
            const result = JSON.parse(Buffer.from(payload).toString("utf8"));
            console.log(result)
            const parsed = JSON.parse(result);
            return parsed;
          } catch (error) {
            console.error("Error invoking Lambda function:", error);
            throw error;
          }
        };`
    )
    .join("\n\n");

  // 3-2. import 文などを追加 -----------------------------------------------
  const header = [
    `import { LambdaClient, InvokeCommand } from "@aws-sdk/client-lambda";`,
    `const lambdaClient = new LambdaClient({ region: "ap-northeast-1" });`,
  ].join("\n");

  return {
    code: `${header}\n\n${newCode}`,
  };
}

Lambda 関数側のプラグイン実装

次に、Lambda 関数側の実装を見ていきます。まずは、"use lambda"ディレクティブが付されたファイルで export されている関数を収集します。そして 1 つの関数につき、1 つの Lambda 関数を生成します。

まず、具体例を見てみます。このような関数が定義されているとします。

// calc.ts
"use lambda";
export const calcAddition = (a: number, b: number): number => {
  return a + b;
};

この関数に対して、新しいエントリポイントファイル(action-calc.js)を作成し、calc.ts の calcAddition 関数を import して呼び出します。action-calc.jsは実際にファイルを作成するのではなく、Rolldown の load hook を利用して、仮想的なモジュールとして扱います。

// action-calc.js (仮想的なモジュール)
import { calcAddition } from "./calc.js";
export const handler = async (event) => {
  const result = await calcAddition(...event);
  return JSON.stringify(result);
};

// calc.js
export const calcAddition = (a: number, b: number): number => {
  return a + b;
};

プラグインの実装は次のようになります。まず、エントリポイントとして、本来のファイル名ではなく、action-calc.js?lambda-entry=trueというクエリパラメータを付けた架空ファイル名を利用します。

このままでは実際のこのファイルを読み込もうとしたときにエラーになってしまいます。そこで、load hook を利用します。load hook は、モジュールの読み込み時に呼び出されるhookです。与えられた ID(ファイルパスのようなもの) に対してどのようなコードを返すかを定義できます。ここで仮想的な ID に対して適当なコードを返すことで、仮想的なモジュールに任意の処理を注入できます。

注入するコードは次のようなものです。originalEntryPathは実際にユーザーが定義している関数のファイルパスが入っており、この関数を呼び出して実行する内容となっています。

const code = `
  import { ${functionName} } from "./${originalEntryPath}.js";
  export const handler = async (event) => {
    const result = await ${functionName}(...event)
    return JSON.stringify(result)
  }
`;

プラグイン全体はこのようになります。

{
  input: `${entry}?lambda-entry=true`,
  plugins: [
    {
      resolveId(source) {
        if (source.endsWith("?lambda-entry=true")) {
          return source;
        }
      },
      load(id) {
        if (id.endsWith("?lambda-entry=true")) {
          const functionName = basename(
            id.replace("?lambda-entry=true", ""),
            ".ts"
          ).replace(/action-/, "");
          const originalEntryPath = basename(
            functionNameToDefinedPath[functionName],
            ".ts"
          );
          const code = `
            import { ${functionName} } from "./${originalEntryPath}.js";
            export const handler = async (event) => {
              const result = await ${functionName}(...event)
              return JSON.stringify(result)
            }
          `;
          return code;
        }
      },
    },
  ],
}

以上の実装により、Lambda 関数に準拠したエントリポイントを持ちつつ、実態としてユーザーが定義した関数を呼び出すコードを生成できました。

Lambda 関数の実装箇所の特定

先ほどの説明では省きましたが、Lambda 関数側には固定のエントリポイントがありません。"use lambda"ディレクティブが付されたファイルを見つけて、そこから export されている関数がエントリポイントとなります。

この関数は全ファイルを走査して見つけても良いのですが、1 つ前に実装したメインアプリケーション側の処理を少し工夫します。メインアプリケーション側の変換を行なっている際に、実際に呼び出す lambda 関数のリストを得ることができます。これを manifest.jsonという JSON ファイルに出力しておきます。Lambda 関数側のビルド時にこの JSON ファイルを読み込むことで、簡単に Lambda 関数の実装箇所を特定できます。

const lambdaInfo = new Map<string, string[]>();
const useLambdaPlugin = () => {
  const lambdaInfo = new Map<string, string[]>();

  return {
    transform(code, id) {
      // ...
      const exports = parsed.body.filter(
        (statement) => statement.type === "ExportNamedDeclaration"
      );

      const identifiers = exports.flatMap((node) =>
        node.declaration?.type === "VariableDeclaration"
          ? node.declaration.declarations.flatMap((d) =>
              d.id.type === "Identifier" ? d.id.name : []
            )
          : node.declaration?.type === "FunctionDeclaration"
          ? node.declaration.id != null
            ? [node.declaration.id.name]
            : []
          : []
      );

      // 1. 収集した関数名とファイルパスをマップに保存 ----------------------------
      lambdaInfo.set(relative(`${import.meta.dirname}`, id), identifiers);
      // ...
    },
    generateBundle() {
      const outputData = Object.fromEntries(lambdaInfo);
      // 2. manifest.json として出力 ----------------------------------------
      this.emitFile({
        type: "asset",
        fileName: "manifest.json",
        source: JSON.stringify(outputData, null, 2),
      });
    },
  };
};

manifest.jsonは次のようなシンプルな JSON ファイルになります。lambda 関数のエントリポイントとなるファイル名と、そこから export される関数名のリストが記述されています。

{
  "src/action.ts": ["lambdaAction"],
  "src/routes/action.ts": ["calcAddition", "calcSubtraction"]
}

課題

今回は趣味の実験として作っていましたが、この仕組みを実戦投入することを考えてみると、たくさん課題があります。

例えば今回の Lambda 関数は同期的に呼び出すことを想定しているため、メインアプリケーション側では Lambda 関数からのレスポンスを受け取ることができます。しかし非同期的な呼び出しを行う場合はレスポンスを受け取ることができません。つまり、非同期呼び出しを行う場合は、"use lambda" が付された関数はレスポンスを返せないことになります。これは一般的な関数呼び出しとは異なるトリッキーな挙動になってしまいます。

また、シリアライズできない値(コールバック関数など)を渡せないことや、巨大なデータを引数に与えてしまうとLambda 関数の制限に引っかかる可能性もあります。

余談になりますが、これらの課題のいくつかは React Server Functions にも同様に存在しています。React Server Functions も今回試作した仕組みも、決して HTTP 通信や Lambda 関数の呼び出しを気にしなくて良くなるわけではない点に注意が必要です。表面上は単なる関数呼び出しでも、実際にはランタイムを跨いだ通信が発生していることを、実装者が常に意識する必要があります。

まとめ

今回は、Rolldown プラグインを利用して、Lambda 関数の呼び出しをシームレスに行う仕組みを実装しました。この仕組みが実運用に耐えうるのかはさておき、メタプログラミングができて楽しかったです。

今回の実装は、React Server Functions に対応しているフレームワークの実装を参考にしました。実際のフレームワークでも、地道なコード変換によって Server Functions が実現されています。今回は諦めてしまいましたが、React Server Functions にはインラインでの Server Functions というものがあります。これは各フレームワークの実装でも一番複雑になっていると感じる部分ですが、AST ベースでパターンマッチをしていく面白い実装になっています。興味のある方はぜひ調べてみてください。

export const ProfileForm = () => {
  const handleSubmit = async (user: User) => {
    "use server";
    // インラインで定義された Server Functions
    await sql`INSERT INTO users (name) VALUES (${user.name})`;
  };

  return <form onSubmit={handleSubmit}>{/* フォームの内容 */}</form>;
};

We are Hiring!

エムスリーでは、フロントエンド・バックエンドに関わらず、新しい技術に興味のあるエンジニアを募集しています。新卒もお待ちしております!

エンジニア採用ページはこちら

jobs.m3.com

エンジニア新卒採用サイト! !

fresh.m3recruit.com

カジュアル面談! !

jobs.m3.com