Exemple #1
0
// 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)
	}

}