Nil checks, simplify GRPC headers

cisco-open · May 10, 2023 · 217d71d · 217d71d
1 parent d5c5461
commit 217d71d
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 29 deletions.
diff --git a/pkg/istio/istio.go b/pkg/istio/istio.go
@@ -215,9 +215,11 @@ func NewIstioIntegrationHandler(config *IstioIntegrationHandlerConfig, logger lo
 
 	s.environment = e
 
-	s.zipkinTracer, err = tracing.SetupZipkinTracing(s.environment)
-	if err != nil {
-		return nil, err
+	if s.environment.ZipkinAddress != "" {
+		s.zipkinTracer, err = tracing.SetupZipkinTracing(s.environment)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	registry := prometheus.NewRegistry()
@@ -307,16 +309,21 @@ func (h *istioIntegrationHandler) GetHTTPTransport(transport http.RoundTripper)
 		return nil, errors.Wrap(err, "could not get stream handler")
 	}
 
-	zipkinWrappedTransport, err := zipkinhttp.NewTransport(h.zipkinTracer, zipkinhttp.RoundTripper(transport))
-	if err != nil {
-		return nil, errors.Wrap(err, "could not get zipkin tracer")
+	if h.zipkinTracer != nil {
+		transport, err = zipkinhttp.NewTransport(h.zipkinTracer, zipkinhttp.RoundTripper(transport))
+		if err != nil {
+			return nil, errors.Wrap(err, "could not get zipkin tracer")
+		}
+
 	}
 
 	logger := h.logger.WithName("http-transport")
-	tp := NewIstioHTTPRequestTransport(zipkinWrappedTransport, h.caClient, h.discoveryClient, logger, h.zipkinTracer)
+	tp := NewIstioHTTPRequestTransport(transport, h.caClient, h.discoveryClient, logger, h.zipkinTracer)
 	httpTransport := pwhttp.NewHTTPTransport(tp, streamHandler, logger)
 
-	httpTransport.AddMiddleware(tracing.NewZipkinHTTPClientTracingMiddleware(h.zipkinTracer))
+	if h.zipkinTracer != nil {
+		httpTransport.AddMiddleware(tracing.NewZipkinHTTPClientTracingMiddleware(h.zipkinTracer))
+	}
 
 	httpTransport.AddMiddleware(middleware.NewEnvoyHTTPHandlerMiddleware())
 	httpTransport.AddMiddleware(NewIstioHTTPHandlerMiddleware())
@@ -332,6 +339,10 @@ func (h *istioIntegrationHandler) GetGRPCDialOptions() ([]grpc.DialOption, error
 
 	grpcDialer := pwgrpc.NewGRPCDialer(h.caClient, streamHandler, h.discoveryClient, h.logger)
 
+	if h.zipkinTracer != nil {
+		grpcDialer.AddMiddleware(tracing.NewZipkinGRPCTracingMiddleware(h.zipkinTracer))
+	}
+
 	grpcDialer.AddMiddleware(middleware.NewEnvoyHTTPHandlerMiddleware())
 	grpcDialer.AddMiddleware(NewIstioHTTPHandlerMiddleware())
 
@@ -355,7 +366,13 @@ func (h *istioIntegrationHandler) ListenAndServe(ctx context.Context, listenAddr
 
 	httpHandler := pwhttp.NewHandler(handler, streamHandler, api.ListenerDirectionInbound)
 
-	httpHandler.AddMiddleware(tracing.NewZipkinHTTPTracingMiddleware(h.zipkinTracer))
+	if h.zipkinTracer != nil {
+		httpHandler.AddMiddleware(tracing.NewZipkinHTTPTracingMiddleware(h.zipkinTracer))
+
+		// TODO(bertab) should it be merged with the http middleware?
+		// if including both will it duplicate spans?
+		// httpHandler.AddMiddleware(tracing.NewZipkinGRPCTracingMiddleware(h.zipkinTracer))
+	}
 
 	httpHandler.AddMiddleware(middleware.NewEnvoyHTTPHandlerMiddleware())
 	httpHandler.AddMiddleware(NewIstioHTTPHandlerMiddleware())

diff --git a/pkg/istio/tracing/zipkin.go b/pkg/istio/tracing/zipkin.go
@@ -156,11 +156,11 @@ func InjectHTTPHeaders(r api.HTTPRequest, opts ...InjectOption) propagation.Inje
 func ExtractGRPCHeaders(req api.HTTPRequest) propagation.Extractor {
 	return func() (*model.SpanContext, error) {
 		var (
-			traceIDHeader      = GetGRPCHeader(req.Header(), b3.TraceID)
-			spanIDHeader       = GetGRPCHeader(req.Header(), b3.SpanID)
-			parentSpanIDHeader = GetGRPCHeader(req.Header(), b3.ParentSpanID)
-			sampledHeader      = GetGRPCHeader(req.Header(), b3.Sampled)
-			flagsHeader        = GetGRPCHeader(req.Header(), b3.Flags)
+			traceIDHeader, _      = req.Header().Get(b3.TraceID)
+			spanIDHeader, _       = req.Header().Get(b3.SpanID)
+			parentSpanIDHeader, _ = req.Header().Get(b3.ParentSpanID)
+			sampledHeader, _      = req.Header().Get(b3.Sampled)
+			flagsHeader, _        = req.Header().Get(b3.Flags)
 		)
 
 		return b3.ParseHeaders(
@@ -177,33 +177,24 @@ func InjectGRPCHeaders(req api.HTTPRequest) propagation.Injector {
 		}
 
 		if sc.Debug {
-			setGRPCHeader(req.Header(), b3.Flags, "1")
+			req.Header().Set(b3.Flags, "1")
 		} else if sc.Sampled != nil {
 			if *sc.Sampled {
-				setGRPCHeader(req.Header(), b3.Sampled, "1")
+				req.Header().Set(b3.Sampled, "1")
 			} else {
-				setGRPCHeader(req.Header(), b3.Sampled, "0")
+				req.Header().Set(b3.Sampled, "0")
 			}
 		}
 
 		if !sc.TraceID.Empty() && sc.ID > 0 {
 			// set identifiers
-			setGRPCHeader(req.Header(), b3.TraceID, sc.TraceID.String())
-			setGRPCHeader(req.Header(), b3.SpanID, sc.ID.String())
+			req.Header().Set(b3.TraceID, sc.TraceID.String())
+			req.Header().Set(b3.SpanID, sc.ID.String())
 			if sc.ParentID != nil {
-				setGRPCHeader(req.Header(), b3.ParentSpanID, sc.ParentID.String())
+				req.Header().Set(b3.ParentSpanID, sc.ParentID.String())
 			}
 		}
 
 		return nil
 	}
 }
-
-func GetGRPCHeader(header api.HeaderMap, key string) string {
-	v, _ := header.Get(key)
-	return v
-}
-
-func setGRPCHeader(header api.HeaderMap, key, value string) {
-	header.Set(key, value)
-}