47 lines
1.4 KiB
Go
47 lines
1.4 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// rockTraceKey 用于把采购追踪器挂到 context 上。
|
||
// 采用 context 传递而不是客户端全局字段，避免并发采购时互相串数据。
|
||
type rockTraceKey struct{}
|
||
|
||
// RockTrace 收集一次采购流程中所有上游 HTTP 调用的原始请求与响应。
|
||
// OnStage 在客户端每次请求结束后同步调用，由采购日志记录器落库。
|
||
type RockTrace struct {
|
||
mu sync.Mutex
|
||
OnStage func(stage, method, url string, reqBody []byte, status int, respBody []byte, dur time.Duration, err error)
|
||
}
|
||
|
||
// WithRockTrace 返回带追踪器的 context；trace 为 nil 时原样返回。
|
||
func WithRockTrace(ctx context.Context, trace *RockTrace) context.Context {
|
||
if trace == nil {
|
||
return ctx
|
||
}
|
||
return context.WithValue(ctx, rockTraceKey{}, trace)
|
||
}
|
||
|
||
func traceFromContext(ctx context.Context) *RockTrace {
|
||
if ctx == nil {
|
||
return nil
|
||
}
|
||
trace, _ := ctx.Value(rockTraceKey{}).(*RockTrace)
|
||
return trace
|
||
}
|
||
|
||
// traceRockRequest 记录一次上游调用；没有追踪器时是空操作。
|
||
func traceRockRequest(ctx context.Context, stage, method, url string, reqBody []byte, status int, respBody []byte, dur time.Duration, err error) {
|
||
trace := traceFromContext(ctx)
|
||
if trace == nil || trace.OnStage == nil {
|
||
return
|
||
}
|
||
trace.mu.Lock()
|
||
handler := trace.OnStage
|
||
trace.mu.Unlock()
|
||
handler(stage, method, url, reqBody, status, respBody, dur, err)
|
||
}
|