Skip to content

Commit

Permalink
Merge pull request #11 from bstasyszyn/correlation-id
Browse files Browse the repository at this point in the history
feat: Correlation ID OpenTelemetry middleware
  • Loading branch information
bstasyszyn authored Nov 1, 2024
2 parents 76049ff + fc17b0c commit 4f33656
Show file tree
Hide file tree
Showing 7 changed files with 272 additions and 11 deletions.
13 changes: 11 additions & 2 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@ module github.com/trustbloc/logutil-go
go 1.19

require (
github.com/stretchr/testify v1.8.1
github.com/labstack/echo/v4 v4.12.0
github.com/stretchr/testify v1.8.4
go.opentelemetry.io/otel v1.12.0
go.opentelemetry.io/otel/sdk v1.12.0
go.opentelemetry.io/otel/trace v1.12.0
Expand All @@ -18,9 +19,17 @@ require (
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/go-logr/logr v1.2.3 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/labstack/gommon v0.4.2 // indirect
github.com/mattn/go-colorable v0.1.13 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/valyala/bytebufferpool v1.0.0 // indirect
github.com/valyala/fasttemplate v1.2.2 // indirect
go.uber.org/atomic v1.7.0 // indirect
go.uber.org/multierr v1.6.0 // indirect
golang.org/x/sys v0.0.0-20220919091848-fb04ddd9f9c8 // indirect
golang.org/x/crypto v0.22.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
34 changes: 25 additions & 9 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -8,17 +8,26 @@ github.com/go-logr/logr v1.2.3/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbV
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
github.com/labstack/echo/v4 v4.12.0 h1:IKpw49IMryVB2p1a4dzwlhP1O2Tf2E0Ir/450lH+kI0=
github.com/labstack/echo/v4 v4.12.0/go.mod h1:UP9Cr2DJXbOK3Kr9ONYzNowSh7HP0aG0ShAyycHSJvM=
github.com/labstack/gommon v0.4.2 h1:F8qTUNXgG1+6WQmqoUWnz8WiEU60mXVVw0P4ht1WRA0=
github.com/labstack/gommon v0.4.2/go.mod h1:QlUFxVM+SNXhDL/Z7YhocGIBYOiwB0mXm1+1bAPHPyU=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
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/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
github.com/valyala/fasttemplate v1.2.2 h1:lxLXG0uE3Qnshl9QyaK6XJxMXlQZELvChBOCmQD0Loo=
github.com/valyala/fasttemplate v1.2.2/go.mod h1:KHLXt3tVN2HBp8eijSv/kGJopbvo7S+qRAEEKiv+SiQ=
go.opentelemetry.io/otel v1.12.0 h1:IgfC7kqQrRccIKuB7Cl+SRUmsKbEwSGPr0Eu+/ht1SQ=
go.opentelemetry.io/otel v1.12.0/go.mod h1:geaoz0L0r1BEOR81k7/n9W4TCXYCJ7bPO7K374jQHG0=
go.opentelemetry.io/otel/sdk v1.12.0 h1:8npliVYV7qc0t1FKdpU08eMnOjgPFMnriPhn0HH4q3o=
Expand All @@ -32,10 +41,17 @@ go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
go.uber.org/zap v1.23.0 h1:OjGQ5KQDEUawVHxNwQgPpiypGHOxo2mNZsOqTak4fFY=
go.uber.org/zap v1.23.0/go.mod h1:D+nX8jyLsMHMYrln8A0rJjFt/T/9/bGgIhAqxv5URuY=
golang.org/x/sys v0.0.0-20220919091848-fb04ddd9f9c8 h1:h+EGohizhe9XlX18rfpa8k8RAc5XyaeamM+0VHRd4lc=
golang.org/x/sys v0.0.0-20220919091848-fb04ddd9f9c8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/crypto v0.22.0 h1:g1v0xeRhjcugydODzvb3mEM9SQ0HGp9s/nh3COQ/C30=
golang.org/x/crypto v0.22.0/go.mod h1:vr6Su+7cTlO45qkww3VDJlzDn0ctJvRgYbC2NvXHt+M=
golang.org/x/net v0.24.0 h1:1PcaxkF854Fu3+lvBIx5SYn9wRlBzzcnHZSiaFFAb0w=
golang.org/x/net v0.24.0/go.mod h1:2Q7sJY5mzlzWjKtYUEXSlBWCdyaioyXzRB2RtU8KVE8=
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.19.0 h1:q5f1RH2jigJ1MoAWp2KTp3gm5zAGFUTarQZ5U386+4o=
golang.org/x/sys v0.19.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
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=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
15 changes: 15 additions & 0 deletions pkg/otel/api/api.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
/*
Copyright Gen Digital Inc. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package api

const (
// CorrelationIDHeader is the HTTP header key for the correlation ID.
CorrelationIDHeader = "X-Correlation-ID"

// CorrelationIDAttribute is the Open Telemetry span attribute key for the correlation ID.
CorrelationIDAttribute = "dts.correlation_id"
)
30 changes: 30 additions & 0 deletions pkg/otel/correlationidecho/correlationecho.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
/*
Copyright Gen Digital Inc. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package correlationidecho

import (
"github.com/labstack/echo/v4"
"github.com/trustbloc/logutil-go/pkg/otel/api"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
)

// Middleware reads the X-Correlation-Id header and, if found, sets the
// dts.correlation_id attribute on the current span.
func Middleware() echo.MiddlewareFunc {
return func(next echo.HandlerFunc) echo.HandlerFunc {
return func(c echo.Context) error {
correlationID := c.Request().Header.Get(api.CorrelationIDHeader)
if correlationID != "" {
span := trace.SpanFromContext(c.Request().Context())
span.SetAttributes(attribute.String(api.CorrelationIDAttribute, correlationID))
}

return next(c)
}
}
}
38 changes: 38 additions & 0 deletions pkg/otel/correlationidecho/correlationecho_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
/*
Copyright Gen Digital Inc. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package correlationidecho

import (
"net/http"
"net/http/httptest"
"testing"

"github.com/labstack/echo/v4"
"github.com/stretchr/testify/require"
)

func TestMiddleware(t *testing.T) {
const correlationID1 = "correlationID1"

m := Middleware()

handler := m(func(c echo.Context) error {
return nil
})
require.NotNil(t, handler)

e := echo.New()
req := httptest.NewRequest(http.MethodGet, "/", nil)
req.Header.Set("X-Correlation-Id", correlationID1)

rec := httptest.NewRecorder()

ctx := e.NewContext(req, rec)

err := handler(ctx)
require.NoError(t, err)
}
92 changes: 92 additions & 0 deletions pkg/otel/correlationidtransport/correlationidtransport.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
/*
Copyright Gen Digital Inc. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package correlationidtransport

import (
"crypto/rand"
"crypto/sha256"
"encoding/hex"
"fmt"
"net/http"
"strings"

"go.opentelemetry.io/otel/trace"

"github.com/trustbloc/logutil-go/pkg/otel/api"
)

const (
nilTraceID = "00000000000000000000000000000000"
defaultCorrelationIDLength = 8
)

// Transport is an http.RoundTripper that adds a correlation ID to the request.
type Transport struct {
defaultTransport http.RoundTripper
correlationIDLength int
}

type Opt func(*Transport)

// WithCorrelationIDLength sets the length of the correlation ID.
func WithCorrelationIDLength(length int) Opt {
return func(t *Transport) {
t.correlationIDLength = length
}
}

// New creates a new Transport.
func New(defaultTransport http.RoundTripper, opts ...Opt) *Transport {
t := &Transport{
defaultTransport: defaultTransport,
correlationIDLength: defaultCorrelationIDLength,
}

for _, opt := range opts {
opt(t)
}

return t
}

// RoundTrip executes a single HTTP transaction.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
var correlationID string

span := trace.SpanFromContext(req.Context())

traceID := span.SpanContext().TraceID().String()
if traceID == "" || traceID == nilTraceID {
var err error
correlationID, err = t.generateID()
if err != nil {
return nil, fmt.Errorf("generate correlation ID: %w", err)
}
} else {
correlationID = t.shortenID(traceID)
}

clonedReq := req.Clone(req.Context())
clonedReq.Header.Add(api.CorrelationIDHeader, correlationID)

return t.defaultTransport.RoundTrip(clonedReq)
}

func (t *Transport) generateID() (string, error) {
bytes := make([]byte, t.correlationIDLength/2) //nolint:gomnd

if _, err := rand.Read(bytes); err != nil {
return "", err
}

return strings.ToUpper(hex.EncodeToString(bytes)), nil
}

func (t *Transport) shortenID(id string) string {
hash := sha256.Sum256([]byte(id))
return strings.ToUpper(hex.EncodeToString(hash[:t.correlationIDLength/2])) //nolint:gomnd
}
61 changes: 61 additions & 0 deletions pkg/otel/correlationidtransport/correlationidtransport_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
/*
Copyright Gen Digital Inc. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package correlationidtransport

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

"github.com/stretchr/testify/require"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/sdk/trace"

"github.com/trustbloc/logutil-go/pkg/otel/api"
)

func TestTransport_RoundTrip(t *testing.T) {
var rt mockRoundTripperFunc = func(req *http.Request) (*http.Response, error) {
correlationID := req.Header.Get(api.CorrelationIDHeader)

require.Len(t, correlationID, 8)
return &http.Response{}, nil
}

transport := New(rt, WithCorrelationIDLength(8))

t.Run("No span", func(t *testing.T) {
req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://example.com", nil)
require.NoError(t, err)

resp, err := transport.RoundTrip(req)
require.NoError(t, err)
require.NotNil(t, resp)
})

t.Run("With span", func(t *testing.T) {
tp := trace.NewTracerProvider()

otel.SetTracerProvider(tp)

ctx, span := tp.Tracer("test").Start(context.Background(), "test")
require.NotNil(t, span)

req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.com", nil)
require.NoError(t, err)

resp, err := transport.RoundTrip(req)
require.NoError(t, err)
require.NotNil(t, resp)
})
}

type mockRoundTripperFunc func(*http.Request) (*http.Response, error)

func (fn mockRoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return fn(req)
}

0 comments on commit 4f33656

Please sign in to comment.