Improvements for jaegertracing (#62)

* Improvements for jaegertracing

Changes:
- fixed http.status_code value and remove ugly workaround
- improve readability of middleware code
- refactor bodyDumpResponseWriter to private responseDumper with limit for body size
- added tags for client real IP and X-Request-ID header
- replace span.SetTag with span.LogKV for values with big cardinality.

Author: ont

jaegertracing/body_dump_response_writer.go

@@ -22,6 +22,7 @@ package jaegertracing
 
 import (
 	"bytes"
+	"crypto/rand"
 	"errors"
 	"fmt"
 	"io"
@@ -54,6 +55,13 @@ type (
 
 		// add req body & resp body to tracing tags
 		IsBodyDump bool
+
+		// prevent logging long http request bodies
+		LimitHTTPBody bool
+
+		// http body limit size (in bytes)
+		// NOTE: don't specify values larger than 60000 as jaeger can't handle values in span.LogKV larger than 60000 bytes
+		LimitSize int
 	}
 )
 
@@ -63,6 +71,9 @@ var (
 		Skipper:       middleware.DefaultSkipper,
 		ComponentName: defaultComponentName,
 		IsBodyDump:    false,
+
+		LimitHTTPBody: true,
+		LimitSize:     60_000,
 	}
 )
 
@@ -128,82 +139,115 @@ func TraceWithConfig(config TraceConfig) echo.MiddlewareFunc {
 
 			req := c.Request()
 			opname := "HTTP " + req.Method + " URL: " + c.Path()
+			realIP := c.RealIP()
+			requestID := getRequestID(c) // request-id generated by reverse-proxy
+
 			var sp opentracing.Span
-			tr := config.Tracer
-			if ctx, err := tr.Extract(opentracing.HTTPHeaders,
-				opentracing.HTTPHeadersCarrier(req.Header)); err != nil {
-				sp = tr.StartSpan(opname)
+			var err error
+
+			ctx, err := config.Tracer.Extract(
+				opentracing.HTTPHeaders,
+				opentracing.HTTPHeadersCarrier(req.Header),
+			)
+
+			if err != nil {
+				sp = config.Tracer.StartSpan(opname)
 			} else {
-				sp = tr.StartSpan(opname, ext.RPCServerOption(ctx))
+				sp = config.Tracer.StartSpan(opname, ext.RPCServerOption(ctx))
 			}
+			defer sp.Finish()
 
 			ext.HTTPMethod.Set(sp, req.Method)
 			ext.HTTPUrl.Set(sp, req.URL.String())
 			ext.Component.Set(sp, config.ComponentName)
+			sp.SetTag("client_ip", realIP)
+			sp.SetTag("request_id", requestID)
 
 			// Dump request & response body
-			resBody := new(bytes.Buffer)
+			var respDumper *responseDumper
 			if config.IsBodyDump {
 				// request
 				reqBody := []byte{}
-				if c.Request().Body != nil { // Read
+				if c.Request().Body != nil {
 					reqBody, _ = ioutil.ReadAll(c.Request().Body)
-					sp.SetTag("http.req.body", string(reqBody))
+
+					if config.LimitHTTPBody {
+						sp.LogKV("http.req.body", limitString(string(reqBody), config.LimitSize))
+					} else {
+						sp.LogKV("http.req.body", string(reqBody))
+					}
 				}
 
-				req.Body = ioutil.NopCloser(bytes.NewBuffer(reqBody)) // Reset
+				req.Body = ioutil.NopCloser(bytes.NewBuffer(reqBody)) // reset original request body
 
 				// response
-				mw := io.MultiWriter(c.Response().Writer, resBody)
-				writer := &bodyDumpResponseWriter{Writer: mw, ResponseWriter: c.Response().Writer}
-				c.Response().Writer = writer
+				respDumper = newResponseDumper(c.Response())
+				c.Response().Writer = respDumper
 			}
 
+			// setup request context - add opentracing span
 			req = req.WithContext(opentracing.ContextWithSpan(req.Context(), sp))
 			c.SetRequest(req)
 
-			var err error
-			defer func() {
-				committed := c.Response().Committed
-				status := c.Response().Status
-
-				if err != nil {
-					var httpError *echo.HTTPError
-					if errors.As(err, &httpError) {
-						if httpError.Code != 0 {
-							status = httpError.Code
-						}
-						sp.SetTag("error.message", httpError.Message)
-					} else {
-						sp.SetTag("error.message", err.Error())
-					}
-					if status == http.StatusOK {
-						// this is ugly workaround for cases when httpError.code == 0 or error was not httpError and status
-						// in request was 200 (OK). In these cases replace status with something that represents an error
-						// it could be that error handlers or middlewares up in chain will output different status code to
-						// client. but at least we send something better than 200 to jaeger
-						status = http.StatusInternalServerError
-					}
-				}
+			// call next middleware / controller
+			err = next(c)
+			if err != nil {
+				c.Error(err) // call custom registered error handler
+			}
 
-				ext.HTTPStatusCode.Set(sp, uint16(status))
-				if status >= http.StatusInternalServerError || !committed {
-					ext.Error.Set(sp, true)
-				}
+			status := c.Response().Status
+			ext.HTTPStatusCode.Set(sp, uint16(status))
 
-				// Dump response body
-				if config.IsBodyDump {
-					sp.SetTag("http.resp.body", resBody.String())
+			if err != nil {
+				logError(sp, err)
+			}
+
+			// Dump response body
+			if config.IsBodyDump {
+				if config.LimitHTTPBody {
+					sp.LogKV("http.resp.body", limitString(respDumper.GetResponse(), config.LimitSize))
+				} else {
+					sp.LogKV("http.resp.body", respDumper.GetResponse())
 				}
+			}
 
-				sp.Finish()
-			}()
-			err = next(c)
-			return err
+			return nil // error was already processed with ctx.Error(err)
 		}
 	}
 }
 
+func limitString(str string, size int) string {
+	if len(str) > size {
+        return str[:size/2] + "\n---- skipped ----\n" + str[len(str)-size/2:]
+	}
+
+    return str
+}
+
+func logError(span opentracing.Span, err error) {
+	var httpError *echo.HTTPError
+	if errors.As(err, &httpError) {
+		span.LogKV("error.message", httpError.Message)
+	} else {
+		span.LogKV("error.message", err.Error())
+	}
+	span.SetTag("error", true)
+}
+
+func getRequestID(ctx echo.Context) string {
+	requestID := ctx.Request().Header.Get(echo.HeaderXRequestID) // request-id generated by reverse-proxy
+	if requestID == "" {
+		requestID = generateToken() // missed request-id from proxy, we generate it manually
+	}
+	return requestID
+}
+
+func generateToken() string {
+	b := make([]byte, 16)
+	rand.Read(b)
+	return fmt.Sprintf("%x", b)
+}
+
 // TraceFunction wraps funtion with opentracing span adding tags for the function name and caller details
 func TraceFunction(ctx echo.Context, fn interface{}, params ...interface{}) (result []reflect.Value) {
 	// Get function name

jaegertracing/jaegertracing.go

@@ -18,6 +18,7 @@ import (
 type mockSpan struct {
 	tracer   opentracing.Tracer
 	tags     map[string]interface{}
+	logs     map[string]interface{}
 	opName   string
 	finished bool
 }
@@ -26,6 +27,7 @@ func createSpan(tracer opentracing.Tracer) *mockSpan {
 	return &mockSpan{
 		tracer: tracer,
 		tags:   make(map[string]interface{}),
+		logs:   make(map[string]interface{}),
 	}
 }
 
@@ -41,6 +43,10 @@ func (sp *mockSpan) getTag(key string) interface{} {
 	return sp.tags[key]
 }
 
+func (sp *mockSpan) getLog(key string) interface{} {
+	return sp.logs[key]
+}
+
 func (sp *mockSpan) Finish() {
 	sp.finished = true
 }
@@ -60,6 +66,13 @@ func (sp *mockSpan) SetTag(key string, value interface{}) opentracing.Span {
 func (sp *mockSpan) LogFields(fields ...log.Field) {
 }
 func (sp *mockSpan) LogKV(alternatingKeyValues ...interface{}) {
+	for i := 0; i < len(alternatingKeyValues); i += 2 {
+		ikey := alternatingKeyValues[i]
+		value := alternatingKeyValues[i+1]
+		if key, ok := ikey.(string); ok {
+			sp.logs[key] = value
+		}
+	}
 }
 func (sp *mockSpan) SetBaggageItem(restrictedKey, value string) opentracing.Span {
 	return sp
@@ -162,7 +175,7 @@ func TestTraceWithDefaultConfig(t *testing.T) {
 
 		assert.Equal(t, uint16(400), tracer.currentSpan().getTag("http.status_code"))
 		assert.Equal(t, true, tracer.currentSpan().getTag("error"))
-		assert.Equal(t, "baaaad request", tracer.currentSpan().getTag("error.message"))
+		assert.Equal(t, "baaaad request", tracer.currentSpan().getLog("error.message"))
 	})
 
 	t.Run("unknown error", func(t *testing.T) {
@@ -172,7 +185,7 @@ func TestTraceWithDefaultConfig(t *testing.T) {
 
 		assert.Equal(t, uint16(500), tracer.currentSpan().getTag("http.status_code"))
 		assert.Equal(t, true, tracer.currentSpan().getTag("error"))
-		assert.Equal(t, "internal stuff went wrong", tracer.currentSpan().getTag("error.message"))
+		assert.Equal(t, "internal stuff went wrong", tracer.currentSpan().getLog("error.message"))
 	})
 }
 
@@ -215,8 +228,8 @@ func TestTraceWithConfigOfBodyDump(t *testing.T) {
 	assert.Equal(t, true, tracer.currentSpan().isFinished())
 	assert.Equal(t, "EchoTracer", tracer.currentSpan().getTag("component"))
 	assert.Equal(t, "/trace", tracer.currentSpan().getTag("http.url"))
-	assert.Equal(t, `{"name": "Lorem"}`, tracer.currentSpan().getTag("http.req.body"))
-	assert.Equal(t, `Hi`, tracer.currentSpan().getTag("http.resp.body"))
+	assert.Equal(t, `{"name": "Lorem"}`, tracer.currentSpan().getLog("http.req.body"))
+	assert.Equal(t, `Hi`, tracer.currentSpan().getLog("http.resp.body"))
 	assert.Equal(t, uint16(200), tracer.currentSpan().getTag("http.status_code"))
 	assert.Equal(t, nil, tracer.currentSpan().getTag("error"))
 	assert.Equal(t, true, tracer.hasStartSpanWithOption)
@@ -264,3 +277,51 @@ func TestTraceOfNoCurrentSpan(t *testing.T) {
 
 	assert.Equal(t, false, tracer.hasStartSpanWithOption)
 }
+
+func TestTraceWithLimitHTTPBody(t *testing.T) {
+	tracer := createMockTracer()
+
+	e := echo.New()
+	e.Use(TraceWithConfig(TraceConfig{
+		Tracer:        tracer,
+		ComponentName: "EchoTracer",
+		IsBodyDump:    true,
+		LimitHTTPBody: true,
+		LimitSize:     10,
+	}))
+	e.POST("/trace", func(c echo.Context) error {
+		return c.String(200, "Hi 123456789012345678901234567890")
+	})
+
+	req := httptest.NewRequest(http.MethodPost, "/trace", bytes.NewBufferString("123456789012345678901234567890"))
+	rec := httptest.NewRecorder()
+	e.ServeHTTP(rec, req)
+
+	assert.Equal(t, true, tracer.currentSpan().isFinished())
+	assert.Equal(t, "12345\n---- skipped ----\n67890", tracer.currentSpan().getLog("http.req.body"))
+	assert.Equal(t, "Hi 12\n---- skipped ----\n67890", tracer.currentSpan().getLog("http.resp.body"))
+}
+
+func TestTraceWithoutLimitHTTPBody(t *testing.T) {
+	tracer := createMockTracer()
+
+	e := echo.New()
+	e.Use(TraceWithConfig(TraceConfig{
+		Tracer:        tracer,
+		ComponentName: "EchoTracer",
+		IsBodyDump:    true,
+		LimitHTTPBody: false, // disabled
+		LimitSize:     10,
+	}))
+	e.POST("/trace", func(c echo.Context) error {
+		return c.String(200, "Hi 123456789012345678901234567890")
+	})
+
+	req := httptest.NewRequest(http.MethodPost, "/trace", bytes.NewBufferString("123456789012345678901234567890"))
+	rec := httptest.NewRecorder()
+	e.ServeHTTP(rec, req)
+
+	assert.Equal(t, true, tracer.currentSpan().isFinished())
+	assert.Equal(t, "123456789012345678901234567890", tracer.currentSpan().getLog("http.req.body"))
+	assert.Equal(t, "Hi 123456789012345678901234567890", tracer.currentSpan().getLog("http.resp.body"))
+}

jaegertracing/jaegertracing_test.go

@@ -0,0 +1,34 @@
+package jaegertracing
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+
+	"github.com/labstack/echo/v4"
+)
+
+type responseDumper struct {
+	http.ResponseWriter
+
+	mw  io.Writer
+	buf *bytes.Buffer
+}
+
+func newResponseDumper(resp *echo.Response) *responseDumper {
+	buf := new(bytes.Buffer)
+	return &responseDumper{
+		ResponseWriter: resp.Writer,
+
+		mw:  io.MultiWriter(resp.Writer, buf),
+		buf: buf,
+	}
+}
+
+func (d *responseDumper) Write(b []byte) (int, error) {
+	return d.mw.Write(b)
+}
+
+func (d *responseDumper) GetResponse() string {
+	return d.buf.String()
+}