APM: wait to shutdown trace-agent components until all payloads are processed to avoid lost data and panics

Author: ajgajg1134

pkg/trace/agent/agent.go

@@ -128,6 +128,8 @@ type Agent struct {
 	ctx context.Context
 
 	firstSpanMap sync.Map
+
+	processWg *sync.WaitGroup
 }
 
 // SpanModifier is an interface that allows to modify spans while they are
@@ -168,6 +170,7 @@ func NewAgent(ctx context.Context, conf *config.AgentConfig, telemetryCollector
 		DebugServer:           api.NewDebugServer(conf),
 		Statsd:                statsd,
 		Timing:                timing,
+		processWg:             &sync.WaitGroup{},
 	}
 	agnt.SamplerMetrics.Add(agnt.PrioritySampler, agnt.ErrorsSampler, agnt.NoPrioritySampler, agnt.RareSampler)
 	agnt.Receiver = api.NewHTTPReceiver(conf, dynConf, in, agnt, telemetryCollector, statsd, timing)
@@ -203,6 +206,7 @@ func (a *Agent) Run() {
 	workers := max(runtime.GOMAXPROCS(0), 1)
 
 	log.Infof("Processing Pipeline configured with %d workers", workers)
+	a.processWg.Add(workers)
 	for i := 0; i < workers; i++ {
 		go a.work()
 	}
@@ -241,24 +245,29 @@ func (a *Agent) UpdateAPIKey(oldKey, newKey string) {
 }
 
 func (a *Agent) work() {
+	defer a.processWg.Done()
 	for {
 		p, ok := <-a.In
 		if !ok {
 			return
 		}
 		a.Process(p)
 	}
-
 }
 
 func (a *Agent) loop() {
 	<-a.ctx.Done()
 	log.Info("Exiting...")
 
-	a.OTLPReceiver.Stop() // Stop OTLPReceiver before Receiver to avoid sending to closed channel
+	// Stop the receiver first so we do not receive more payloads and start to shut down the workers
 	if err := a.Receiver.Stop(); err != nil {
 		log.Error(err)
 	}
+
+	//Wait to process any leftover payloads in flight before closing components that might be needed
+	a.processWg.Wait()
+
+	a.OTLPReceiver.Stop() // Stop OTLPReceiver before Receiver to avoid sending to closed channel
 	for _, stopper := range []interface{ Stop() }{
 		a.Concentrator,
 		a.ClientStatsAggregator,
@@ -383,6 +392,7 @@ func (a *Agent) Process(p *api.Payload) {
 			if a.SpanModifier != nil {
 				a.SpanModifier.ModifySpan(chunk, span)
 			}
+			time.Sleep(100 * time.Millisecond)
 			a.obfuscateSpan(span)
 			a.Truncate(span)
 			if p.ClientComputedTopLevel {

pkg/trace/agent/agent_test.go

@@ -129,6 +129,46 @@ func TestFormatTrace(t *testing.T) {
 	assert.Contains(result.Meta["sql.query"], "SELECT name FROM people WHERE age = ?")
 }
 
+func TestStopWaits(t *testing.T) {
+	cfg := config.New()
+	cfg.Endpoints[0].APIKey = "test"
+	cfg.Obfuscation.Cache.Enabled = true
+	cfg.Obfuscation.Cache.MaxSize = 1_000
+	ctx, cancel := context.WithCancel(context.Background())
+	agnt := NewTestAgent(ctx, cfg, telemetry.NewNoopCollector())
+
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		agnt.Run()
+	}()
+
+	now := time.Now()
+	span := &pb.Span{
+		TraceID:  1,
+		SpanID:   1,
+		Resource: "SELECT name FROM people WHERE age = 42 AND extra = 55",
+		Type:     "sql",
+		Start:    now.Add(-time.Second).UnixNano(),
+		Duration: (500 * time.Millisecond).Nanoseconds(),
+	}
+
+	agnt.In <- &api.Payload{
+		TracerPayload: testutil.TracerPayloadWithChunk(testutil.TraceChunkWithSpan(span)),
+		Source:        info.NewReceiverStats().GetTagStats(info.Tags{}),
+	}
+
+	cancel()
+	wg.Wait() // Wait for agent to completely exit
+
+	mtw := agnt.TraceWriter.(*mockTraceWriter)
+
+	assert := assert.New(t)
+	assert.Len(mtw.payloads, 1)
+	assert.Equal("SELECT name FROM people WHERE age = ? AND extra = ?", mtw.payloads[0].TracerPayload.Chunks[0].Spans[0].Meta["sql.query"])
+}
+
 func TestProcess(t *testing.T) {
 	t.Run("Replacer", func(t *testing.T) {
 		// Ensures that for "sql" type spans:

releasenotes/notes/APMWaitShutdown-d35a8b0cedec0fe8.yaml

@@ -0,0 +1,11 @@
+# Each section from every release note are combined when the
+# CHANGELOG.rst is rendered. So the text needs to be worded so that
+# it does not depend on any information only available in another
+# section. This may mean repeating some details, but each section
+# must be readable independently of the other.
+#
+# Each section note must be formatted as reStructuredText.
+---
+fixes:
+  - |
+    APM: Fix an issue where the trace-agent could panic during shutdown trying to obfuscate a sql payload.