109 lines
3.1 KiB
C#
109 lines
3.1 KiB
C#
using Microsoft.AspNetCore.Http;
|
||
using Microsoft.Extensions.Logging;
|
||
using System.Diagnostics;
|
||
using TakeoutSaaS.Shared.Abstractions.Diagnostics;
|
||
using TakeoutSaaS.Shared.Abstractions.Ids;
|
||
|
||
namespace TakeoutSaaS.Shared.Web.Middleware;
|
||
|
||
/// <summary>
|
||
/// 统一 TraceId/CorrelationId，贯穿日志与响应。
|
||
/// </summary>
|
||
public sealed class CorrelationIdMiddleware(RequestDelegate next, ILogger<CorrelationIdMiddleware> logger, IIdGenerator idGenerator)
|
||
{
|
||
private const string TraceHeader = "X-Trace-Id";
|
||
private const string SpanHeader = "X-Span-Id";
|
||
private const string RequestHeader = "X-Request-Id";
|
||
|
||
/// <summary>
|
||
/// 管道入口，确保 TraceId/SpanId 贯穿请求。
|
||
/// </summary>
|
||
/// <param name="context">HTTP 上下文。</param>
|
||
public async Task InvokeAsync(HttpContext context)
|
||
{
|
||
// 1. 确保活动存在并启动
|
||
var ownsActivity = Activity.Current is null;
|
||
var activity = Activity.Current ?? new Activity("TakeoutSaaS.Request");
|
||
|
||
if (activity.Id is null)
|
||
{
|
||
activity.SetIdFormat(ActivityIdFormat.W3C);
|
||
activity.Start();
|
||
}
|
||
|
||
// 2. 生成/解析 TraceId、SpanId
|
||
var traceId = activity.TraceId.ToString();
|
||
var spanId = activity.SpanId.ToString();
|
||
|
||
if (string.IsNullOrWhiteSpace(traceId))
|
||
{
|
||
traceId = ResolveTraceId(context);
|
||
}
|
||
|
||
// 3. 写入上下文与响应头
|
||
context.TraceIdentifier = traceId;
|
||
TraceContext.TraceId = traceId;
|
||
TraceContext.SpanId = spanId;
|
||
|
||
context.Response.OnStarting(() =>
|
||
{
|
||
context.Response.Headers[TraceHeader] = traceId;
|
||
context.Response.Headers[SpanHeader] = spanId;
|
||
return Task.CompletedTask;
|
||
});
|
||
|
||
// 4. 带 Scope 调用后续中间件
|
||
using (logger.BeginScope(new Dictionary<string, object>
|
||
{
|
||
["TraceId"] = traceId,
|
||
["SpanId"] = spanId
|
||
}))
|
||
{
|
||
try
|
||
{
|
||
await next(context);
|
||
}
|
||
finally
|
||
{
|
||
// 5. 清理上下文与活动
|
||
TraceContext.Clear();
|
||
if (ownsActivity)
|
||
{
|
||
activity.Stop();
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
private string ResolveTraceId(HttpContext context)
|
||
{
|
||
if (TryGetHeader(context, TraceHeader, out var traceId))
|
||
{
|
||
return traceId;
|
||
}
|
||
|
||
if (TryGetHeader(context, RequestHeader, out var requestId))
|
||
{
|
||
return requestId;
|
||
}
|
||
|
||
return idGenerator.NextId().ToString();
|
||
}
|
||
|
||
private static bool TryGetHeader(HttpContext context, string headerName, out string value)
|
||
{
|
||
if (context.Request.Headers.TryGetValue(headerName, out var values))
|
||
{
|
||
var headerValue = values.ToString();
|
||
if (!string.IsNullOrWhiteSpace(headerValue))
|
||
{
|
||
value = headerValue;
|
||
return true;
|
||
}
|
||
}
|
||
|
||
value = string.Empty;
|
||
return false;
|
||
}
|
||
}
|