// The tracer middleware emits TraceEntry objects to the supplied Collector whenever the
// server processes an incoming request.
//
// Two Trace entries will be emitted for each request, one for the incoming request
// and one for the outgoing response.
//
// The middleware will inject a traceId for each incoming request into the context
// that gets passed to the request handler. To ensure that any further RPC requests
// that occur inside the wrapped handler are associated with the current request, the
// handler should pass its context to any performed RPC client requests.
//
// This function is designed to emit events in non-blocking mode. If the Collector does
// not have enough capacity to store a generated TraceEntry then it will be silently dropped.
func Tracer(collector *tracePkg.Collector) usrv.EndpointOption {
return func(ep *usrv.Endpoint) error {
hostname, err := os.Hostname()
if err != nil {
return err
}
// Wrap original method
originalHandler := ep.Handler
ep.Handler = usrv.HandlerFunc(func(ctx context.Context, responseWriter usrv.ResponseWriter, request *usrv.Message) {
var traceId string
// Check if the request contains a trace id. If no trace is
// available allocate a new traceId and inject it in the
// request context that gets passed to the handler
trace := request.Headers.Get(CtxTraceId)
if trace == nil {
traceId = uuid.New()
ctx = context.WithValue(ctx, CtxTraceId, traceId)
} else {
traceId = trace.(string)
}
// Inject trace into outgoing message
responseWriter.Header().Set(CtxTraceId, traceId)
// Trace incoming request. This call is non-blocking
collector.Add(&tracePkg.Record{
Timestamp: time.Now(),
TraceId: traceId,
CorrelationId: request.CorrelationId,
Type: tracePkg.Request,
From: request.From,
To: request.To,
Host: hostname,
})
// Trace response when the handler returns
defer func(start time.Time) {
var errMsg string
errVal := responseWriter.Header().Get("error")
if errVal != nil {
errMsg = errVal.(string)
}
// Trace response. This call is non-blocking
collector.Add(&tracePkg.Record{
Timestamp: time.Now(),
TraceId: traceId,
CorrelationId: request.CorrelationId,
Type: tracePkg.Response,
From: request.To, // when responding we switch From/To
To: request.From,
Host: hostname,
Duration: time.Since(start).Nanoseconds(),
Error: errMsg,
})
}(time.Now())
// Invoke the original handler
originalHandler.Serve(ctx, responseWriter, request)
})
return nil
}
}
func TestTracerWithError(t *testing.T) {
var err error
processedChan := make(chan struct{})
storage := storage.Memory
defer storage.Close()
storage.AfterStore(func() {
processedChan <- struct{}{}
})
collector, err := tracePkg.NewCollector(storage, 1000, time.Hour)
if err != nil {
t.Fatalf("Error creating collector: %v", err)
}
ep := usrv.Endpoint{
Name: "traceTest",
Handler: usrv.HandlerFunc(func(ctx context.Context, rw usrv.ResponseWriter, req *usrv.Message) {
rw.WriteError(errors.New("I cannot allow you to do that Dave"))
}),
}
err = Tracer(collector)(&ep)
if err != nil {
t.Fatalf("Error applying Tracer() to endpoint: %v", err)
}
msg := &usrv.Message{
From: "sender",
To: "recipient",
CorrelationId: "123",
}
// Send request
w := usrvtest.NewRecorder()
ep.Handler.Serve(context.Background(), w, msg)
traceId := w.Header().Get(CtxTraceId)
if traceId == nil {
t.Fatalf("Expected middleware to set response writer header %s", CtxTraceId)
}
// Block till both entries are processed
<-processedChan
<-processedChan
// Fetch trace
traceLog, err := storage.GetTrace(traceId.(string))
if err != nil {
t.Fatalf("Error retrieving trace with id %s: %v", traceId, err)
}
if len(traceLog) != 2 {
t.Fatalf("Expected trace len to be 2; got %d", len(traceLog))
}
traceEntryOut := traceLog[1]
// Validate RES trace
if traceEntryOut.Type != tracePkg.Response {
t.Fatalf("Expected trace to be of type %v; got %v", tracePkg.Response, traceEntryOut.Type)
}
if traceEntryOut.CorrelationId != msg.CorrelationId {
t.Fatalf("Expected trace CorrelationId to be %s; got %s", msg.CorrelationId, traceEntryOut.CorrelationId)
}
if traceEntryOut.TraceId != traceId {
t.Fatalf("Expected trace TraceId to be %s; got %s", traceId, traceEntryOut.TraceId)
}
if traceEntryOut.Error != "I cannot allow you to do that Dave" {
t.Fatalf("Expected trace Error to be 'I cannot allow you to do that Dave'; got %v", traceEntryOut.Error)
}
// Out trace should reverse From and To
if traceEntryOut.From != msg.To {
t.Fatalf("Expected trace From to be %s; got %s", msg.To, traceEntryOut.From)
}
if traceEntryOut.To != msg.From {
t.Fatalf("Expected trace To to be %s; got %s", msg.From, traceEntryOut.To)
}
}