fix: fix trace ctx should not be overwrite (#648)

Author: fatelei

internal/core/local_runtime/environment_python.go

@@ -10,7 +10,7 @@ import (
 
 func (p *LocalPluginRuntime) InitPythonEnvironment() error {
 	// root span for python env init
-_, span := p.startSpan("python.init_env", attribute.String("plugin.identity", p.Config.Identity()))
+	_, span := p.startSpan("python.init_env", attribute.String("plugin.identity", p.Config.Identity()))
 	defer span.End()
 
 	// prepare uv environment
@@ -25,7 +25,10 @@ _, span := p.startSpan("python.init_env", attribute.String("plugin.identity", p.
 	case ErrVirtualEnvironmentInvalid:
 		// remove the venv and rebuild it
 		log.Warn("virtual environment for %s is invalid; deleting and recreating", p.Config.Identity())
-		p.deleteVirtualEnvironment()
+		err = p.deleteVirtualEnvironment()
+		if err != nil {
+			return err
+		}
 
 		// create virtual environment
 		venv, err = p.createVirtualEnvironment(uvPath)

internal/core/local_runtime/otel_test.go

@@ -12,14 +12,12 @@ import (
 )
 
 func TestPythonInitSpansReuseUpstreamTrace(t *testing.T) {
-	// setup tracer provider with span recorder
 	sr := tracetest.NewSpanRecorder()
 	tp := sdktrace.NewTracerProvider()
 	tp.RegisterSpanProcessor(sr)
 	gootel.SetTracerProvider(tp)
 	gootel.SetTextMapPropagator(propagation.TraceContext{})
 
-	// create parent span with known trace id
 	ctx := context.Background()
 	tr := gootel.Tracer("test")
 	ctx, parent := tr.Start(ctx, "parent")
@@ -34,11 +32,44 @@ func TestPythonInitSpansReuseUpstreamTrace(t *testing.T) {
 	spans := sr.Ended()
 	require.Len(t, spans, 2)
 
-	// ensure child trace id matches parent trace id
 	parentSpan := spans[0]
 	childSpan := spans[1]
 	if parentSpan.Name() == "python.init_env" {
 		parentSpan, childSpan = childSpan, parentSpan
 	}
 	require.Equal(t, parentSpan.SpanContext().TraceID(), childSpan.SpanContext().TraceID())
 }
+
+func TestStartSpanDoesNotAffectSubsequentSpans(t *testing.T) {
+	sr := tracetest.NewSpanRecorder()
+	tp := sdktrace.NewTracerProvider()
+	tp.RegisterSpanProcessor(sr)
+	gootel.SetTracerProvider(tp)
+	gootel.SetTextMapPropagator(propagation.TraceContext{})
+
+	ctx := context.Background()
+	tr := gootel.Tracer("test")
+	ctx, parent := tr.Start(ctx, "parent")
+
+	runtime := &LocalPluginRuntime{}
+	runtime.SetTraceContext(ctx)
+
+	ctx1, span1 := runtime.startSpan("span1")
+	span1.End()
+
+	ctx2, span2 := runtime.startSpan("span2")
+
+	require.False(t, ctx1.Err() != nil, "ctx1 should not be canceled yet")
+	require.False(t, ctx2.Err() != nil, "ctx2 should not be canceled")
+
+	span2.End()
+	parent.End()
+
+	spans := sr.Ended()
+	require.Len(t, spans, 3)
+
+	traceID := spans[0].SpanContext().TraceID()
+	for _, span := range spans {
+		require.Equal(t, traceID, span.SpanContext().TraceID(), "all spans should share the same trace ID")
+	}
+}

internal/core/local_runtime/setup_python_environment.go

@@ -44,16 +44,14 @@ func (p *LocalPluginRuntime) ensureTraceCtx() context.Context {
 }
 
 func (p *LocalPluginRuntime) startSpan(name string, attrs ...attribute.KeyValue) (context.Context, trace.Span) {
-	ctx := p.ensureTraceCtx()
-	ctx, sp := p.otelTracer().Start(ctx, name)
+	rootCtx := p.ensureTraceCtx()
+	ctx, sp := p.otelTracer().Start(rootCtx, name)
 	if id, ok := log.IdentityFromContext(ctx); ok && id.TenantID != "" {
 		sp.SetAttributes(attribute.String("tenant_id", id.TenantID))
 	}
 	if len(attrs) > 0 {
 		sp.SetAttributes(attrs...)
 	}
-	// keep last context for potential child spans
-	p.traceCtx = ctx
 	return ctx, sp
 }
 
@@ -205,11 +203,15 @@ func (p *LocalPluginRuntime) installDependencies(
 
 	defer func() {
 		if cmd.Process != nil {
-			cmd.Process.Kill()
+			err := cmd.Process.Kill()
+			if err != nil {
+				log.Warn("failed to kill python process", "error", err)
+			}
 		}
 	}()
 
 	var errMsg strings.Builder
+	var errMsgMu sync.Mutex
 	var wg sync.WaitGroup
 	wg.Add(2)
 
@@ -220,14 +222,12 @@ func (p *LocalPluginRuntime) installDependencies(
 		routinepkg.RoutineLabelKeyMethod: "InitPythonEnvironment",
 	}, func() {
 		defer wg.Done()
-		// read stdout
 		buf := make([]byte, 1024)
 		for {
 			n, err := stdout.Read(buf)
 			if err != nil {
 				break
 			}
-			// FIXME: move the log to separated layer
 			log.Info("installing plugin", "plugin", p.Config.Identity(), "output", string(buf[:n]))
 			lastActiveAt = time.Now()
 		}
@@ -238,20 +238,23 @@ func (p *LocalPluginRuntime) installDependencies(
 		routinepkg.RoutineLabelKeyMethod: "InitPythonEnvironment",
 	}, func() {
 		defer wg.Done()
-		// read stderr
 		buf := make([]byte, 1024)
 		for {
 			n, err := stderr.Read(buf)
-			if err != nil && err != os.ErrClosed {
+			if err != nil && !errors.Is(err, os.ErrClosed) {
 				lastActiveAt = time.Now()
+				errMsgMu.Lock()
 				errMsg.WriteString(string(buf[:n]))
+				errMsgMu.Unlock()
 				break
-			} else if err == os.ErrClosed {
+			} else if errors.Is(err, os.ErrClosed) {
 				break
 			}
 
 			if n > 0 {
+				errMsgMu.Lock()
 				errMsg.WriteString(string(buf[:n]))
+				errMsgMu.Unlock()
 				lastActiveAt = time.Now()
 			}
 		}
@@ -271,11 +274,16 @@ func (p *LocalPluginRuntime) installDependencies(
 			if time.Since(lastActiveAt) > time.Duration(
 				p.appConfig.PythonEnvInitTimeout,
 			)*time.Second {
-				cmd.Process.Kill()
+				err := cmd.Process.Kill()
+				errMsgMu.Lock()
+				if err != nil {
+					errMsg.WriteString(fmt.Sprintf("failed to kill python process: %s", err))
+				}
 				errMsg.WriteString(fmt.Sprintf(
 					"init process exited due to no activity for %d seconds",
 					p.appConfig.PythonEnvInitTimeout,
 				))
+				errMsgMu.Unlock()
 				break
 			}
 		}