optimizations: conditionally enable tracing, use worker pool in stream controller

Author: kralicky

clientconn.go

@@ -104,15 +104,17 @@ func (cc *ClientConn) invoke(
 		panic(fmt.Sprintf("[totem] unsupported request type: %T", req))
 	}
 
-	name, attr := spanInfo(method, peerFromCtx(ctx))
-	attr = append(attr, attribute.String("func", "clientConn.Invoke"))
-	ctx, span := cc.tracer.Start(ctx, name,
-		trace.WithSpanKind(trace.SpanKindClient),
-		trace.WithAttributes(attr...),
-	)
-	defer span.End()
-	otel.GetTextMapPropagator().Inject(ctx, &mdSupplier)
-
+	var span trace.Span
+	if TracingEnabled {
+		name, attr := spanInfo(method, peerFromCtx(ctx))
+		attr = append(attr, attribute.String("func", "clientConn.Invoke"))
+		ctx, span = cc.tracer.Start(ctx, name,
+			trace.WithSpanKind(trace.SpanKindClient),
+			trace.WithAttributes(attr...),
+		)
+		defer span.End()
+		otel.GetTextMapPropagator().Inject(ctx, &mdSupplier)
+	}
 	cc.metrics.TrackTxBytes(serviceName, methodName, int64(len(reqMsg)))
 
 	rpc := &RPC{

discovery.go

@@ -15,12 +15,15 @@ func discoverServices(ctx context.Context, ctrl *StreamController, maxHops int32
 		RemainingHops: maxHops,
 	})
 
-	ctx, span := Tracer().Start(ctx, "totem.discoverServices")
-	defer span.End()
-
-	lg := ctrl.logger.With(
-		zap.String("traceID", span.SpanContext().TraceID().String()),
-	)
+	lg := ctrl.logger
+	var span trace.Span
+	if TracingEnabled {
+		ctx, span = Tracer().Start(ctx, "totem.discoverServices")
+		defer span.End()
+		lg = lg.With(
+			zap.String("traceID", span.SpanContext().TraceID().String()),
+		)
+	}
 	lg.Debug("starting service discovery")
 
 	respC := ctrl.Request(ctx, &RPC{
@@ -47,9 +50,11 @@ func discoverServices(ctx context.Context, ctrl *StreamController, maxHops int32
 		return nil, err
 	}
 
-	span.AddEvent("Results", trace.WithAttributes(
-		attribute.StringSlice("methods", infoMsg.MethodNames()),
-	))
+	if TracingEnabled {
+		span.AddEvent("Results", trace.WithAttributes(
+			attribute.StringSlice("methods", infoMsg.MethodNames()),
+		))
+	}
 
 	lg.With(
 		zap.Any("methods", infoMsg.MethodNames()),

examples/simple/main.go

@@ -68,10 +68,6 @@ type helloServer struct {
 // There is nothing special about the Hello implementation - it is just a
 // regular unary RPC.
 func (h *helloServer) Hello(ctx context.Context, req *HelloRequest) (*HelloResponse, error) {
-	// We can check if we're inside a Totem stream using the context:
-	if totem.CheckContext(ctx) {
-		// inside a Totem stream (not a real grpc server)
-	}
 	return &HelloResponse{
 		Message: "Hello, " + req.GetName(),
 	}, nil

go.mod

@@ -3,6 +3,7 @@ module github.com/kralicky/totem
 go 1.21
 
 require (
+	github.com/alitto/pond v1.8.3
 	github.com/charmbracelet/lipgloss v0.7.1
 	github.com/google/uuid v1.3.0
 	github.com/jhump/protoreflect v1.15.1
@@ -22,7 +23,6 @@ require (
 	go.uber.org/atomic v1.11.0
 	go.uber.org/zap v1.25.0
 	golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63
-	golang.org/x/sync v0.3.0
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20230803162519-f966b187b2e5
 	google.golang.org/grpc v1.57.0
 	google.golang.org/protobuf v1.31.0
@@ -56,6 +56,7 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/mod v0.12.0 // indirect
 	golang.org/x/net v0.14.0 // indirect
+	golang.org/x/sync v0.3.0 // indirect
 	golang.org/x/sys v0.11.0 // indirect
 	golang.org/x/text v0.12.0 // indirect
 	golang.org/x/tools v0.12.1-0.20230815132531-74c255bcf846 // indirect

go.sum

@@ -2,6 +2,8 @@ cloud.google.com/go v0.110.7 h1:rJyC7nWRg2jWGZ4wSJ5nY65GTdYJkg0cd/uXb+ACI6o=
 cloud.google.com/go/compute v1.23.0 h1:tP41Zoavr8ptEqaW6j+LQOnyBBhO7OkOMAGrgLopTwY=
 cloud.google.com/go/compute v1.23.0/go.mod h1:4tCnrn48xsqlwSAiLf1HXMQk8CONslYbdiEZc9FEIbM=
 cloud.google.com/go/compute/metadata v0.2.3 h1:mg4jlk7mCAj6xXp9UJ4fjI9VUI5rubuGBW5aJ7UnBMY=
+github.com/alitto/pond v1.8.3 h1:ydIqygCLVPqIX/USe5EaV/aSRXTRXDEI9JwuDdu+/xs=
+github.com/alitto/pond v1.8.3/go.mod h1:CmvIIGd5jKLasGI3D87qDkQxjzChdKMmnXMg3fG6M6Q=
 cloud.google.com/go/compute/metadata v0.2.3/go.mod h1:VAV5nSsACxMJvgaAuX6Pk2AawlZn8kiOGuCv6gTkwuA=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=

invokers.go

@@ -61,18 +61,21 @@ func (l *localServiceInvoker) Invoke(ctx context.Context, req *RPC) ([]byte, err
 		zap.Strings("md", req.GetMetadata().Keys()),
 	).Debug("invoking method using local service")
 
-	attrs := []attribute.KeyValue{
-		attribute.String("func", "localServiceInvoker.Invoke"),
-		attribute.String("name", l.service.ServiceName),
-	}
+	var span trace.Span
+	if TracingEnabled {
+		attrs := []attribute.KeyValue{
+			attribute.String("func", "localServiceInvoker.Invoke"),
+			attribute.String("name", l.service.ServiceName),
+		}
 
-	ctx, span := Tracer().Start(ctx, "Invoke/Local: "+req.QualifiedMethodName(),
-		trace.WithAttributes(attrs...))
-	defer span.End()
+		ctx, span = Tracer().Start(ctx, "Invoke/Local: "+req.QualifiedMethodName(),
+			trace.WithAttributes(attrs...))
+		defer span.End()
+	}
 
 	if m, ok := l.methods[req.MethodName]; ok {
 		startTime := time.Now()
-		resp, err := m.Handler(l.serviceImpl, addTotemToContext(ctx), func(v any) error {
+		resp, err := m.Handler(l.serviceImpl, ctx, func(v any) error {
 			reqBytes := req.GetRequest()
 			l.metrics.TrackRxBytes(serviceName, methodName, int64(len(reqBytes)))
 			return proto.Unmarshal(reqBytes, protoimpl.X.ProtoMessageV2Of(v))
@@ -129,14 +132,17 @@ func (r *streamControllerInvoker) Invoke(ctx context.Context, req *RPC) ([]byte,
 		ctx = metadata.NewOutgoingContext(ctx, md)
 	}
 
-	attrs := []attribute.KeyValue{
-		attribute.String("func", "streamControllerInvoker.Invoke"),
-		attribute.String("name", r.controller.name),
-	}
+	var span trace.Span
+	if TracingEnabled {
+		attrs := []attribute.KeyValue{
+			attribute.String("func", "streamControllerInvoker.Invoke"),
+			attribute.String("name", r.controller.name),
+		}
 
-	ctx, span := Tracer().Start(ctx, "Invoke/Stream: "+req.QualifiedMethodName(),
-		trace.WithAttributes(attrs...))
-	defer span.End()
+		ctx, span = Tracer().Start(ctx, "Invoke/Stream: "+req.QualifiedMethodName(),
+			trace.WithAttributes(attrs...))
+		defer span.End()
+	}
 
 	r.controller.metrics.TrackTxBytes(serviceName, methodName, int64(len(req.GetRequest())))
 	rc := r.controller.Request(ctx, req)

server.go

@@ -165,19 +165,25 @@ func (r *Server) Splice(stream Stream, opts ...StreamControllerOption) error {
 		}
 	}()
 
-	ctx, span := Tracer().Start(stream.Context(), "Server.Splice/Discovery",
-		trace.WithAttributes(
-			attribute.String("name", r.name),
-		),
-	)
+	ctx := stream.Context()
+	var span trace.Span
+	if TracingEnabled {
+		ctx, span = Tracer().Start(ctx, "Server.Splice/Discovery",
+			trace.WithAttributes(
+				attribute.String("name", r.name),
+			),
+		)
+	}
 	info, err := discoverServices(ctx, ctrl, r.discoveryHopLimit)
 	if err != nil {
 		err := fmt.Errorf("service discovery failed: %w", err)
 		span.RecordError(err)
 		span.SetStatus(otelcodes.Error, err.Error())
 		return err
 	}
-	span.End()
+	if TracingEnabled {
+		span.End()
+	}
 	r.logger.With(
 		zap.Any("methods", info.MethodNames()),
 	).Debug("splicing stream")
@@ -253,13 +259,19 @@ func (r *Server) Serve() (grpc.ClientConnInterface, <-chan error) {
 		ch <- runErr
 	}()
 
-	ctx, span := Tracer().Start(r.Context(), "Server.Serve/Discovery",
-		trace.WithAttributes(
-			attribute.String("name", r.name),
-		),
-	)
+	ctx := context.Background()
+	var span trace.Span
+	if TracingEnabled {
+		ctx, span = Tracer().Start(r.Context(), "Server.Serve/Discovery",
+			trace.WithAttributes(
+				attribute.String("name", r.name),
+			),
+		)
+	}
 	info, err := discoverServices(ctx, r.controller, r.discoveryHopLimit)
-	span.End()
+	if TracingEnabled {
+		span.End()
+	}
 
 	if err != nil {
 		r.controller.Kick(fmt.Errorf("service discovery failed: %w", err))