-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #11 from bstasyszyn/correlation-id
feat: Correlation ID OpenTelemetry middleware
- Loading branch information
Showing
7 changed files
with
272 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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" | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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
61
pkg/otel/correlationidtransport/correlationidtransport_test.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | ||
} |