Skip to content

Commit

Permalink
Merge pull request #24 from lafer-fz/main
Browse files Browse the repository at this point in the history
feat: Add trace ID to the response header.
  • Loading branch information
CoderPoet authored Sep 20, 2023
2 parents 6a9c6ca + c2701b3 commit 7f92845
Show file tree
Hide file tree
Showing 5 changed files with 75 additions and 15 deletions.
6 changes: 3 additions & 3 deletions tracing/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,10 @@ require (
github.com/golang/protobuf v1.5.3 // indirect
github.com/henrylee2cn/ameda v1.5.1 // indirect
github.com/klauspost/cpuid/v2 v2.2.3 // indirect
github.com/kr/pretty v0.3.0 // indirect
github.com/kr/pretty v0.1.0 // indirect
github.com/kr/text v0.2.0 // indirect
github.com/nyaruka/phonenumbers v1.1.6 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/rogpeppe/go-internal v1.9.0 // indirect
github.com/tidwall/gjson v1.14.4 // indirect
github.com/tidwall/match v1.1.1 // indirect
github.com/tidwall/pretty v1.2.1 // indirect
Expand All @@ -40,6 +40,6 @@ require (
golang.org/x/sys v0.8.0 // indirect
golang.org/x/text v0.9.0 // indirect
google.golang.org/protobuf v1.30.0 // indirect
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
11 changes: 2 additions & 9 deletions tracing/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,8 @@ github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfV
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
github.com/klauspost/cpuid/v2 v2.2.3 h1:sxCkb+qR91z4vsqw4vGGZlDgPz3G7gjaLyK3V8y70BU=
github.com/klauspost/cpuid/v2 v2.2.3/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
Expand All @@ -62,9 +60,6 @@ github.com/nyaruka/phonenumbers v1.1.6 h1:DcueYq7QrOArAprAYNoQfDgp0KetO4LqtnBtQC
github.com/nyaruka/phonenumbers v1.1.6/go.mod h1:yShPJHDSH3aTKzCbXyVxNpbl2kA+F+Ne5Pun/MvFRos=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d h1:zE9ykElWQ6/NYmHa3jpm/yHnI4xSofP+UP6SpjHcSeM=
github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1:OnSkiWE9lh6wB0YB77sQom3nweQdgAjqCqsofrRNTgc=
github.com/smartystreets/goconvey v1.6.4 h1:fv0U8FUIMPNf1L9lnHLvLhgicrIVChEkdzIKYqbNC9s=
Expand Down Expand Up @@ -132,10 +127,8 @@ google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQ
google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
Expand Down
5 changes: 5 additions & 0 deletions tracing/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -134,5 +134,10 @@ func ServerMiddleware(cfg *Config) app.HandlerFunc {
tc.SetSpan(span)

c.Next(ctx)

if cfg.customResponseHandler != nil {
// execute custom response handler
cfg.customResponseHandler(ctx, c)
}
}
}
49 changes: 49 additions & 0 deletions tracing/middleware_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
// Copyright 2022 CloudWeGo Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tracing

import (
"context"
"net/http"
"testing"
"time"

"github.com/cloudwego/hertz/pkg/common/test/assert"

"github.com/cloudwego/hertz/pkg/app"
"github.com/cloudwego/hertz/pkg/app/server"
"go.opentelemetry.io/otel"
sdktrace "go.opentelemetry.io/otel/sdk/trace"
"go.opentelemetry.io/otel/sdk/trace/tracetest"
oteltrace "go.opentelemetry.io/otel/trace"
)

func TestServerMiddleware(t *testing.T) {
sr := tracetest.NewSpanRecorder()
otel.SetTracerProvider(sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr)))
tracer, cfg := NewServerTracer(WithCustomResponseHandler(func(c context.Context, ctx *app.RequestContext) {
ctx.Header("trace-id", oteltrace.SpanFromContext(c).SpanContext().TraceID().String())
}))
h := server.Default(tracer, server.WithHostPorts("127.0.0.1:6666"))
h.Use(ServerMiddleware(cfg))
h.GET("/ping", func(c context.Context, ctx *app.RequestContext) {
})

go h.Spin()
time.Sleep(100 * time.Millisecond)
resp, err := http.Get("http://127.0.0.1:6666/ping")
assert.Nil(t, err)
assert.True(t, len(resp.Header.Get("trace-id")) != 0)
}
19 changes: 16 additions & 3 deletions tracing/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,9 @@
package tracing

import (
"context"

"github.com/cloudwego/hertz/pkg/app"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/metric"
"go.opentelemetry.io/otel/propagation"
Expand Down Expand Up @@ -45,6 +48,8 @@ type Config struct {
textMapPropagator propagation.TextMapPropagator

recordSourceOperation bool

customResponseHandler app.HandlerFunc
}

func newConfig(opts []Option) *Config {
Expand All @@ -69,9 +74,10 @@ func newConfig(opts []Option) *Config {

func defaultConfig() *Config {
return &Config{
tracerProvider: otel.GetTracerProvider(),
meterProvider: otel.GetMeterProvider(),
textMapPropagator: otel.GetTextMapPropagator(),
tracerProvider: otel.GetTracerProvider(),
meterProvider: otel.GetMeterProvider(),
textMapPropagator: otel.GetTextMapPropagator(),
customResponseHandler: func(c context.Context, ctx *app.RequestContext) {},
}
}

Expand All @@ -88,3 +94,10 @@ func WithTextMapPropagator(p propagation.TextMapPropagator) Option {
cfg.textMapPropagator = p
})
}

// WithCustomResponseHandler configures CustomResponseHandler
func WithCustomResponseHandler(h app.HandlerFunc) Option {
return option(func(cfg *Config) {
cfg.customResponseHandler = h
})
}

0 comments on commit 7f92845

Please sign in to comment.