Cosmos DB の直接モード利用時に Application Insights で詳細なテレメトリを収集できるようにする

最近は色々なところで Application Insights をインストールして、アプリケーションのテレメトリを継続的にちゃんと取得しましょう的なことを言い続けているのですが、Cosmos DB に関してはデフォルトでは依存関係テレメトリが送信されません。

ドキュメントにも以下のように記載されています。要するに直接モード (TCP / Direct mode) を使っていると、組み込みの依存関係トラッキングの対象外となります。

Azure Cosmos DB
Only tracked automatically if HTTP/HTTPS is used. TCP mode won't be captured by Application Insights.

Dependency tracking in Application Insights - Azure Monitor | Microsoft Learn

Application Insights の依存関係トラッキングが主に HttpClient の DiagnosticSource ベースなので、HTTP 以外は個別対応が必要になるという話です。

単純に Cosmos DB の接続モードを Gateway に変更すれば解決しますが、以下のドキュメントにもあるようにパフォーマンス上のデメリットや、Azure Functions の Consumption Plan では接続数の問題も出てくるので、基本は Direct を使いたいところです。

Cosmos DB SDK v3 ではリクエスト送信前後に独自の処理を追加できる仕組みが用意されているので、それを利用することで Direct でも Application Insights にテレメトリを送信できるようになります。

この辺りは Cosmos DB SDK が公式で Application Insights への統合を用意してくれそうなのですが、時間がかかりそうなのでひとまず独自に実装します。

実装方針含め完全に不明ですが Cosmos DB SDK が DiagnosticSource に対応すると、ユーザーコードを変更することなく対応できるので期待したいところです。

独自のテレメトリを送信する方法はドキュメントに用意されているので、この通り実装すれば簡単です。今回は依存関係を送信したいので DependencyTelemetry を使えばよいです。

以前に Cosmos DB の Request Unit 消費を Application Insights に送信する方法を書きましたが、これは古い v2 SDK を使っていた時のコードなので、今の v3 SDK ではもっと洗練したコードを実現できます。

Cosmos DB SDK v2 は廃止が近いので、まだ使っている場合は早めに移行しておきましょう。

RequestHandler を利用して送信する

Cosmos DB SDK v3 には HttpClient の HttpMessageHandler のような機能として RequestHandler が導入されています。用途はほぼ同じで、メソッドも 1 つしかないので簡単です。

内部で使われるメタデータ系リクエスト以外で呼び出されるようになっているので、これを使って Application Insights にテレメトリを送信する仕組みを用意できます。

実際に作成した RequestHandler の実装は以下の通りになります。これは最低限の情報のみを送信しているので、追加したい情報があれば Properties に入れると良いです。

Application Insights へのテレメトリ送信に必要な TelemetryClient は DI 経由で貰うので、コンストラクタで受け取るようにしておきます。ドキュメントによっては TelemetryConfiguration を使う例があるかも知れませんが、今は TelemetryClient を直接受け取るのが推奨されています。

public class AppInsightsRequestHandler : RequestHandler
{
    public AppInsightsRequestHandler(TelemetryClient telemetryClient)
    {
        _telemetryClient = telemetryClient;
    }

    private readonly TelemetryClient _telemetryClient;

    public override async Task<ResponseMessage> SendAsync(RequestMessage request, CancellationToken cancellationToken)
    {
        // using var で定義しておくと自動で StopOperation が呼び出される
        using var dependency = _telemetryClient.StartOperation<DependencyTelemetry>($"{request.Method} {request.RequestUri.OriginalString}");

        var response = await base.SendAsync(request, cancellationToken);

        var telemetry = dependency.Telemetry;

        // Application Insights 上で Cosmos DB として表示するために必要
        telemetry.Type = "Azure DocumentDB";
        telemetry.Data = request.RequestUri.OriginalString;

        telemetry.ResultCode = ((int)response.StatusCode).ToString();
        telemetry.Success = response.IsSuccessStatusCode;

        // メトリックも同時に送信できるので便利
        telemetry.Metrics["RequestCharge"] = response.Headers.RequestCharge;

        return response;
    }
}

Cosmos DB のプロトコルは Direct でも中身は HTTP ベースなので、特に難しいことはないと思います。Request Charge を送信しておかないと、RU を大量に消費したリクエストを特定出来なくなるので、メトリックとして同時に送信するようにしています。

ペイロードを解析すれば実行されたクエリを同時に送ることも出来ますが、Cosmos DB から Log Analytics に送信する機能を使った方が効率的だと思うので実装していません。

作成した RequestHandler を使うためには、DI を組み合わせて以下のように初期化します。

builder.Services.AddSingleton<AppInsightsRequestHandler>();

builder.Services.AddSingleton(provider =>
{
    var requestHandler = provider.GetRequiredService<AppInsightsRequestHandler>();

    var connectionString = builder.Configuration.GetConnectionString("CosmosConnection");

    return new CosmosClient(connectionString, new CosmosClientOptions
    {
        CustomHandlers = { requestHandler }
    });
});

当然ながら TelemetryClient が DI に登録されている必要があるので、Application Insights SDK のインストールが必要です。Azure Functions の場合は最初から組み込まれているので必要ありません。

組み込んだアプリケーションを実行して、Cosmos DB へリクエストを投げる処理を実行すると、Application Insights の Dependencies にテレメトリが表示されるようになるはずです。

テレメトリを選択すると詳細が表示されるので、同時に送信した Request Charge が確認できます。応答時間や成功したかどうかも E2E タイムライン上で確認できるので、追跡が格段に行いやすくなります。

同時に送信した Request Charge は Application Insights の Log-based Metrics から確認できるので、グラフ表示も自由に行えます。もちろんアラート作成やダッシュボードに表示することも出来ます。

これで直接モードを使っていた場合にブラックボックスになってしまっていた Cosmos DB のリクエストを、Application Insights 上で他のテレメトリと同様に可視化出来るようになりました。

送信するテレメトリをさらに詳細化する

Cosmos DB の接続モードを Gateway で利用した場合に送信されるテレメトリは、行われた処理名が具体的に表示されるようになっていてわかりやすいですが、今回作成した送信処理では HTTP ベースのメソッドと URL が表示されるだけでした。

特にクエリ系は POST として表示されず分かりにくいので、Application Insights で使われている処理を利用してテレメトリの詳細化を行います。実装は以下に存在しています。

やっていることはメソッドと URL から具体的な処理を特定して、その名前をテレメトリとして送信しているだけですが、URL の正規化やマッピングが必要になるので少し手間です。

Application Insights SDK は MIT ライセンスで公開されているので、既存の実装をそのまま利用してしまうのが簡単です。以下のコードでは変更した部分のみ載せています。

public class AppInsightsRequestHandler : RequestHandler
{
    public AppInsightsRequestHandler(TelemetryClient telemetryClient)
    {
        _telemetryClient = telemetryClient;
    }

    private readonly TelemetryClient _telemetryClient;

    public override async Task<ResponseMessage> SendAsync(RequestMessage request, CancellationToken cancellationToken)
    {
        using var dependency = _telemetryClient.StartOperation<DependencyTelemetry>("Cosmos");

        var response = await base.SendAsync(request, cancellationToken);

        var telemetry = dependency.Telemetry;

        var resourcePath = HttpParsingHelper.ParseResourcePath(request.RequestUri.OriginalString);

        // populate properties
        foreach (var (key, value) in resourcePath)
        {
            if (value is not null)
            {
                var propertyName = GetPropertyNameForResource(key);
                if (propertyName is not null)
                {
                    telemetry.Properties[propertyName] = value;
                }
            }
        }

        var operation = HttpParsingHelper.BuildOperationMoniker(request.Method.ToString(), resourcePath);
        var operationName = GetOperationName(operation);

        telemetry.Type = "Azure DocumentDB";
        telemetry.Name = operationName;
        telemetry.Data = request.RequestUri.OriginalString;

        telemetry.ResultCode = ((int)response.StatusCode).ToString();
        telemetry.Success = response.IsSuccessStatusCode;

        telemetry.Metrics["RequestCharge"] = response.Headers.RequestCharge;

        return response;
    }

    private static readonly Dictionary<string, string> OperationNames = new()
    {
        // Database operations
        ["POST /dbs"] = "Create database",
        ["GET /dbs"] = "List databases",
        ["GET /dbs/*"] = "Get database",
        ["DELETE /dbs/*"] = "Delete database",

        // Collection operations
        ["POST /dbs/*/colls"] = "Create collection",
        ["GET /dbs/*/colls"] = "List collections",
        ["POST /dbs/*/colls/*"] = "Query documents",
        ["GET /dbs/*/colls/*"] = "Get collection",
        ["DELETE /dbs/*/colls/*"] = "Delete collection",
        ["PUT /dbs/*/colls/*"] = "Replace collection",

        // Document operations
        ["POST /dbs/*/colls/*/docs"] = "Create document",
        ["GET /dbs/*/colls/*/docs"] = "List documents",
        ["GET /dbs/*/colls/*/docs/*"] = "Get document",
        ["PUT /dbs/*/colls/*/docs/*"] = "Replace document",
        ["DELETE /dbs/*/colls/*/docs/*"] = "Delete document"
    };

    private static string? GetPropertyNameForResource(string resourceType)
    {
        // ignore high cardinality resources (documents, attachments, etc.)
        return resourceType switch
        {
            "dbs" => "Database",
            "colls" => "Collection",
            _ => null
        };
    }

    private static string GetOperationName(string operation)
    {
        return OperationNames.TryGetValue(operation, out var operationName) ? operationName : operation;
    }
}

長いコードですが、変わっていることは operationName の組み立てだけで、それ以外は前のコードと同じです。クエリに関しては Gateway と Direct で API が変わった感があるので、マッピングは修正しています。

この実装を使って Application Insights にテレメトリを送信すると、Application Insights 組み込みの実装と同じように実際の処理名が表示されるようになります。

同時に Properties に Database 名と Container 名を含めるようになっているので、カスタムプロパティを使ってグルーピングやフィルタリングが簡単に行えるようになっています。

やはり公式で Direct 対応の実装が入って、何もしなくても有効になる世界になって欲しいと思いますが、現実問題として Cosmos DB の監視は非常に重要になるので独自の実装でも入れておくのが良いです。

同時に Cosmos DB から Log Analytics への転送も行っておくと、パフォーマンスの問題に対応しやすいです。