🍃
Aurora DSQL のために、マルチリージョン対応の DataSource を実装してみる

2025/01/06に公開
AWS
 可用性 99.999% の javax.sql.DataSourceマルチリージョン構成の Amazon Aurora DSQL の可用性は99.999%です。アクティブ-アクティブ構成で常にデータ同期されているため、アプリケーションはどちらのリージョンに接続しても問題なく処理を実行することが出来ます。もちろん同じリージョンのエンドポイントに接続したほうが早く応答が返ってくるのですが、遠いほうのリージョンに接続したとしても100-200ミリ秒程度の追加時間で済むのであれば許容範囲であるケースもあるかと思います。
というわけで、以下のようなアーキテクチャを実現する独自実装の DataSource を作ってみました、という話です。


 独自実装 DataSource の特徴試しに作った DataSource の実装は以下のような特徴を持っています。
org.springframework.jdbc.datasource.AbstractDataSource を継承しています。
接続プールの機能はありません。接続プールの機能は HikariCP でラップすることにより実現しています。
Aurora DSQL のエンドポイントからラウンドロビンで次の接続の取得先を決定します。
取得した接続に対して、 isValid() メソッドを呼び出して正常性を確認します。
接続や正常性検査に失敗した場合は SQLException をスローして例外通知します。
SQLException をスローした場合、メソッド全体を再試行します。
再試行時は例外をスローしたエンドポイントとは別のエンドポイントを利用します。（ラウンドロビンなので）
実装が必要になる java.sql.DataSource#getConnection() メソッドはこんな感じのコードになります。リトライロジックは@Retryアノテーションの付与のみで実現しています。
    /**
     * {@link javax.sql.DataSource#getConnection(String, String)} の実装
     */
    @Retryable(retryFor = SQLException.class, maxAttempts = 4, backoff = @Backoff(delay = 0))
    @Override
    public Connection getConnection(final String username, final String password) throws SQLException {

        // DSQLのエンドポイントからラウンドロビンで次の接続の取得先を決定
        final int index = this.counter.getAndUpdate(v -> v == Integer.MAX_VALUE ? 0 : v + 1);
        final String region = this.regions[index % this.regions.length];
        final String endpoint = regionalEndpoints.get(region);

        try {
            // 接続に利用するトークンをキャッシュから取得。有効期限切れの場合は新たに生成
            final String token = this.tokenCache.compute(endpoint,
                    (k, v) -> v == null || v.isExpired(System.currentTimeMillis())
                            ? new CacheEntry(
                                    this.generateAuthToken(endpoint, region, this.dsqlTokenExpirationSeconds),
                                    this.dsqlTokenExpirationSeconds * 900 + System.currentTimeMillis())
                            : v).getToken();

            // 選択されたエンドポイントとトークンを使って新しい接続を取得
            final Connection connection = regionalDataSources.get(region).getConnection(username, token);
            if (!connection.isValid(this.connectionTimeout)) {
                // 取得した接続が正常ではない場合はSQLExceptionをスローして例外通知
                // 例外発生時はこのメソッド全体がリトライされ、異なるリージョンからの接続の取得が試行される
                throw new SQLException("The connection obtained from the DataSource is invalid.");
            }

            // 取得した接続が正常な場合は返却
            this.logger.debug("get a new connection from: " + region + ", connection=" + connection.toString());
            return connection;
        } catch (SQLException e) {
            this.logger.debug("a sql exception occurred at: " + region + ", message=" + e.getMessage());
            throw e;
        }
    }

 接続用トークンの取得とキャッシュ現状、Aurora DSQL では固定されたパスワードではなく、APIを使って有効期限のあるトークンを取得して接続します。トークン取得用のAPIはリージョンごとに個別に呼び出す必要があるため、DataSourceとして単一のパスワードを保持するのではなく、接続を取得するたびにトークンを取得する実装としてみました。
ただし、本当に接続するたびに新しいトークンを発行するのは負荷が高いので、トークンの有効期限内はメモリ上にキャッシュしたトークンを再利用する仕組みも取り入れています。

 独自 DataSource をラップする HikariCP の設定接続プールの機能は独自 DataSource には持たせずに、springframework の既定の接続プール機能である HikariCP に移譲しています。
    @Bean
    @Primary
    public DataSource dataSource(@Qualifier("rawDataSource") final DataSource rawDataSource) {

        final HikariDataSource ds = new HikariDataSource();
        ds.setDataSource(rawDataSource); // 独自実装のDataSourceをラップ
        ds.setExceptionOverrideClassName(DsqlExceptionOverride.class.getName()); // Aurora DSQL用に例外ハンドリングを最適化

        return ds;
    }
プロパティで設定できないので見落とされがちかもしれませんが、実はHikariCPは接続を取得する DataSource を外部から渡してやることが可能です。新しい接続の取得は外部の DataSource に任せて、HikariCP自身は取得した接続をプールするロジックに集中する、といった役割分担が出来るようになっています。
また、SQLExceptionの派生クラスがスローされた際に、一部のケースでは接続は閉じずにリトライを行いたいので HikariCP の例外ハンドリングを上書きし、Aurora DSQL 向けに最適化しています。詳しい実装は以下を参照してください。

https://docs.aws.amazon.com/ja_jp/aurora-dsql/latest/userguide/SECTION_program-with-java-apps.html#SECTION_program-with-java-apps-exceptions

 リトライロジック接続取得を再試行するためのロジックは一切書いていません。唯一、@Retryableというアノテーションを付与したのみです。このアノテーションは Spring Retry というプロジェクトから提供されています。このライブラリは Spring Batch や Spring Integration などの Springframework の他のライブラリの中で利用されている実績があるものなので、一定の安心感とともに採用してもよいかと思います。

https://github.com/spring-projects/spring-retry
Spring Retry と Aurora DSQL の組み合わせをつかったお試し実装はこちらでも行っています。

https://zenn.dev/akring/articles/1fc39ead4f6604

 まとめという感じで、非常にお手軽にシームレスに複数のリージョンにまたがってエンドポイントを選択して利用可能な DataSource の実装を作ることが出来ました。接続先の選択ロジックを、自身が稼働しているリージョンのエンドポイントを優先するとか、一定期間エラーが継続するエンドポイントは選択候補から外すとか、もっと賢くする余地は多分にありますが、とりあえずの実装としては十分そうな気がします。
こんなに簡単に実装できてしまって大丈夫なのか不安な点はありますが、従来の、書き込み処理の有無で選択可能なエンドポイントが変わる、データ同期の遅延を意識しないとならない、などの問題から解放されたおかげで実装が非常に簡易になったのでは？と思っています。なにしろ、接続プールの中にあるどの接続を選択しても全く同じようにSQLを実行してデータの読み書きが可能なので、ラウンドロビンのような雑な選択方法でも動くものが作れてしまったので。 Aurora DSQL なかなかに強力で面白いですね。
可用性 99.999% の javax.sql.DataSource

独自実装 DataSource の特徴

接続用トークンの取得とキャッシュ

独自 DataSource をラップする HikariCP の設定

リトライロジック

まとめ

Discussion
