Honeycomb Go GRPC Middleware

grpc-middleware.go

// HoneycombUnary returns an interceptor that adds Honeycomb metrics to any incoming unary requests.
func HoneycombUnary() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, r interface{}, i *grpc.UnaryServerInfo, h grpc.UnaryHandler) (interface{}, error) {
		ctx, t, s := setupTrace(ctx, r, i.FullMethod, "unary")
		defer t.Send()

		o, err := h(ctx, r)

		result := "success"
		if err != nil {
			result = "failure"
			s.AddField("grpc.error", err.Error())
		}
		s.AddField("grpc.result", result)
		wrapRuntimeStats(s)

		return o, err
	}
}

// HoneycombStream returns an interceptor that adds Honeycomb metrics to any incoming stream requests.
func HoneycombStream() grpc.StreamServerInterceptor {
	return func(srv interface{}, ss grpc.ServerStream, i *grpc.StreamServerInfo, h grpc.StreamHandler) error {
		ctx, t, s := setupTrace(ss.Context(), nil, i.FullMethod, "stream")
		defer t.Send()

		w := ssWrap{ss, ctx}

		err := h(srv, w)
		result := "success"
		if err != nil {
			result = "failure"
			s.AddField("grpc.error", err.Error())
		}
		s.AddField("grpc.result", result)
		wrapRuntimeStats(s)

		return err
	}
}

func wrapRuntimeStats(span *trace.Span) {
	timer := timer.Start()
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	span.AddField("runtime.read_mem_stats_ms", timer.Finish())

	span.AddField("runtime.alloc_mb", bToMb(m.Alloc))
	span.AddField("runtime.sys_mb", bToMb(m.Sys))
	span.AddField("runtime.live_objs", m.Mallocs-m.Frees)
	span.AddField("runtime.num_gc", m.NumGC)
	span.AddField("runtime.last_gc", m.LastGC)
}

type ssWrap struct {
	grpc.ServerStream
	ctx context.Context
}

// Context ...
func (sw ssWrap) Context() context.Context {
	return sw.ctx
}

func setupTrace(ctx context.Context, r interface{}, method, grpcType string) (context.Context, *trace.Trace, *trace.Span) {
	t := trace.GetTraceFromContext(ctx)
	if t == nil {
		h := getBeelineHeader(ctx)
		ctx, t = trace.NewTrace(ctx, h)
	}
	t.AddField("grpc.method", method)

	span := trace.GetSpanFromContext(ctx)
	span.AddField("name", method)
	span.AddField("meta.type", "grpc")
	span.AddField("meta.subtype", grpcType)
	span.AddField("grpc.input", r)

	return ctx, t, span
}

func bToMb(b uint64) uint64 {
	return b / 1024 / 1024
}

func GetBeelineHeader(ctx context.Context) string {
	md, ok := metadata.FromIncomingContext(ctx)
	if ok {
		if tmp, ok := md["trace"]; ok && len(tmp) >= 1 {
			return tmp[0]
		}
	}
	return ""
}