Merge pull request #1472 from trheyi/main
Refactor Trace module to mitigate goroutine leaks and reduce memory g…
This commit is contained in:
commit
eab87539e5
6 changed files with 531 additions and 93 deletions
121
trace/BUGFIX.md
Normal file
121
trace/BUGFIX.md
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
# Trace Module Bug Analysis & Fix
|
||||
|
||||
## 1. Problem Summary
|
||||
|
||||
`yao start` crashes with `SIGSEGV` when Agent `All()` runs concurrent operations.
|
||||
Root cause: three inter-related race conditions in the trace module's channel
|
||||
lifecycle management, triggered when `agent/context.Release()` calls
|
||||
`trace.MarkCancelled/MarkComplete` followed by `trace.Release`.
|
||||
|
||||
## 2. Bug Chain
|
||||
|
||||
```
|
||||
context.Release()
|
||||
-> MarkCancelled/MarkComplete (many safeSend calls to stateCmdChan)
|
||||
-> stateMarkCompleted triggers state.completed = true
|
||||
-> state worker starts 100ms drain then exits (BUG 1)
|
||||
-> trace.Release(traceID)
|
||||
-> close(mgr.stateCmdChan) (BUG 3: races with concurrent safeSend)
|
||||
-> if stateExecuteSpaceOp also writing (BUG 2: bare channel send)
|
||||
-> panic on closed channel
|
||||
-> in CGO/V8 callback stack, recover() may fail -> SIGSEGV
|
||||
```
|
||||
|
||||
### BUG 1: State worker premature exit (state.go)
|
||||
|
||||
`startStateWorker` exits after `state.completed = true` + 100ms drain, but the
|
||||
channel remains open. Subsequent safeSend calls write to an unconsumed channel.
|
||||
If the buffer (100) fills up, safeSend blocks forever. Commands with response
|
||||
channels (e.g. stateMarkCompleted) deadlock permanently.
|
||||
|
||||
### BUG 2: stateExecuteSpaceOp bare channel write (state.go)
|
||||
|
||||
```go
|
||||
m.stateCmdChan <- &cmdSpaceKVOp{...} // no safeSend, panics on closed channel
|
||||
```
|
||||
|
||||
### BUG 3: safeSend vs close race (state.go + trace.go)
|
||||
|
||||
Even with `defer/recover`, the window between `select` choosing the send case
|
||||
and the actual send allows a concurrent `close()` to trigger a panic that
|
||||
cannot be recovered in CGO callback stacks.
|
||||
|
||||
## 3. Triggering Scenario (All() + context.Release)
|
||||
|
||||
```
|
||||
Parent ctx (llm/process.go or caller/process.go)
|
||||
├── Fork child ctx1 -> goroutine 1 (Orchestrator.All)
|
||||
├── Fork child ctx2 -> goroutine 2
|
||||
└── All returns, defer ctx.Release()
|
||||
|
||||
Key: Fork sets trace=nil, but ForkParent.TraceID points to same traceID.
|
||||
All goroutines share one trace manager, writing to one stateCmdChan.
|
||||
|
||||
ctx.Release():
|
||||
1. MarkCancelled/MarkComplete -> many safeSend calls
|
||||
2. trace.Release -> close(stateCmdChan) + cancel()
|
||||
|
||||
If child goroutines still have residual operations:
|
||||
-> safeSend to closed channel -> panic -> SIGSEGV
|
||||
```
|
||||
|
||||
## 4. Fix Applied
|
||||
|
||||
### FIX 1: State worker uses for-range (state.go)
|
||||
|
||||
Removed the premature exit after `state.completed`. Worker now uses idiomatic
|
||||
`for cmd := range m.stateCmdChan` which only exits when the channel is closed
|
||||
by `Release()`. Go guarantees that `for range` drains all buffered commands
|
||||
before returning.
|
||||
|
||||
### FIX 2: stateExecuteSpaceOp uses safeSend (state.go)
|
||||
|
||||
Replaced bare `m.stateCmdChan <- cmd` with `m.safeSend(cmd)`. Returns a
|
||||
descriptive error when the state worker has stopped.
|
||||
|
||||
### FIX 3: Three-step safe shutdown (manager.go + state.go + trace.go)
|
||||
|
||||
Added `closed int32` atomic flag to `manager` struct. Release now follows a
|
||||
strict three-step sequence:
|
||||
|
||||
1. `atomic.StoreInt32(&mgr.closed, 1)` — blocks new `safeSend` calls
|
||||
2. `mgr.cancel()` — unblocks any `safeSend` stuck in `select` via `ctx.Done`
|
||||
3. `close(mgr.stateCmdChan)` — terminates state worker (drains buffer first)
|
||||
|
||||
`safeSend` checks the atomic flag before touching the channel, providing a
|
||||
fast-path rejection that is safe even in CGO callback stacks where `recover()`
|
||||
may not work.
|
||||
|
||||
### FIX 4: context.Release timing (no code change needed)
|
||||
|
||||
`MarkCancelled/MarkComplete` calls are synchronous (wait for resp channel).
|
||||
The only fire-and-forget call is `stateAddUpdate` inside `addUpdateAndBroadcast`.
|
||||
Under FIX 3 protection, this call returns `false` instead of panicking if
|
||||
Release has already started. Losing one update during shutdown is acceptable.
|
||||
|
||||
## 5. Files Modified
|
||||
|
||||
| File | Change |
|
||||
|------|--------|
|
||||
| trace/state.go | FIX 1 (for range) + FIX 2 (safeSend for SpaceOp) + FIX 3 (atomic check in safeSend) |
|
||||
| trace/manager.go | FIX 3 (closed int32 field) |
|
||||
| trace/trace.go | FIX 3 (three-step Release shutdown) |
|
||||
| trace/trace_lifecycle_test.go | NEW: 7 boundary condition tests |
|
||||
|
||||
## 6. Test Coverage
|
||||
|
||||
New tests in `trace_lifecycle_test.go`:
|
||||
|
||||
- **TestReleaseWhileWriting** — 20 writers + concurrent Release
|
||||
- **TestReleaseDuringSpaceOp** — space KV ops + concurrent Release
|
||||
- **TestReleaseAfterMarkComplete** — MarkComplete -> Release -> late operations
|
||||
- **TestConcurrentReleaseAndMarkCancelled** — MarkCancelled and Release race
|
||||
- **TestSafeSendAfterClosed** — operations after closed flag is set
|
||||
- **TestRapidCreateReleaseLoop** — 100x create/release stress test
|
||||
- **TestConcurrentAllPattern** — simulates real All() fork + parent Release
|
||||
|
||||
All existing tests pass unchanged (interfaces not modified).
|
||||
|
||||
---
|
||||
|
||||
_Last updated: February 2026_
|
||||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
This document tracks known issues in the Trace module that are scheduled for refactoring.
|
||||
|
||||
## Goroutine Leak
|
||||
## Goroutine Leak (Mitigated)
|
||||
|
||||
### Symptom
|
||||
|
||||
|
|
@ -11,28 +11,22 @@ Each trace creation starts 2 goroutines that accumulate during rapid iterations:
|
|||
1. `trace/pubsub.(*PubSub).forward()` - PubSub event forwarding
|
||||
2. `trace.(*manager).startStateWorker()` - State machine worker
|
||||
|
||||
### Evidence
|
||||
### Current Status
|
||||
|
||||
```
|
||||
Goroutine growth by function:
|
||||
Function Initial Final Growth
|
||||
--------------------------------------------------------------------------
|
||||
github.com/yaoapp/yao/trace. 0 10 10
|
||||
github.com/yaoapp/yao/trace/pubsub. 0 10 10
|
||||
```
|
||||
**Fixed in Feb 2026**: State worker now uses `for range` on the command channel,
|
||||
which exits cleanly when `Release()` closes the channel. The three-step shutdown
|
||||
sequence (atomic flag -> cancel -> close) ensures:
|
||||
|
||||
### Root Cause
|
||||
- No new commands are accepted after the flag is set
|
||||
- In-flight `safeSend` calls unblock via `ctx.Done`
|
||||
- Worker drains remaining buffered commands before exiting
|
||||
|
||||
- Goroutines exit when `Release()` closes their channels
|
||||
- Exit is **asynchronous** (goroutine needs to reach select statement)
|
||||
- Go runtime needs time to schedule and cleanup
|
||||
- In rapid iterations, new goroutines are created before old ones fully exit
|
||||
### Residual Behavior
|
||||
|
||||
### Current Behavior
|
||||
|
||||
- **NOT a true leak**: Goroutines eventually exit (channels are closed)
|
||||
- **No unbounded growth**: They will be GC'd eventually
|
||||
- **Typical pattern**: Async cleanup in Go
|
||||
- PubSub `forward()` goroutine still exits asynchronously after `Stop()` closes its channel
|
||||
- In rapid create/release loops, there may be a brief overlap where old goroutines
|
||||
haven't exited before new ones start. This is normal Go async cleanup behavior.
|
||||
- **NOT a true leak**: goroutines eventually exit (channels are closed)
|
||||
|
||||
### Impact on Tests
|
||||
|
||||
|
|
@ -40,45 +34,44 @@ Memory leak tests use a 20KB/iteration threshold to accommodate this overhead:
|
|||
|
||||
| Test | Actual Growth | Threshold |
|
||||
| ----------------- | -------------- | --------- |
|
||||
| StandardMode | ~11-15 KB/iter | 20 KB |
|
||||
| BusinessScenarios | ~13-16 KB/iter | 20 KB |
|
||||
| StandardMode | ~15 bytes/iter | 20 KB |
|
||||
| BusinessScenarios | ~80 bytes/iter | 20 KB |
|
||||
| NestedCalls | ~13 KB/iter | 20 KB |
|
||||
|
||||
## Memory Growth
|
||||
## Memory Growth (Reduced)
|
||||
|
||||
### Symptom
|
||||
|
||||
Linear memory growth during trace operations:
|
||||
Linear memory growth during trace operations.
|
||||
|
||||
```
|
||||
Batch | Iterations | HeapAlloc (MB) | Growth/iter (bytes)
|
||||
------|------------|----------------|--------------------
|
||||
1 | 1000 | 23.28 | 12014.42
|
||||
2 | 2000 | 37.06 | 13229.02
|
||||
3 | 3000 | 50.63 | 13562.39
|
||||
4 | 4000 | 64.54 | 13819.49
|
||||
5 | 5000 | 78.15 | 13910.01
|
||||
```
|
||||
### Current Status
|
||||
|
||||
### Root Cause
|
||||
**Improved in Feb 2026**: The state worker lifecycle fix eliminates the scenario
|
||||
where the worker exits prematurely while the channel remains open, which could
|
||||
cause command objects to accumulate in the buffer without being consumed.
|
||||
|
||||
Trace-related objects are not fully released during `ctx.Release()`:
|
||||
The three-step shutdown ensures all buffered commands are processed before the
|
||||
worker exits, reducing memory retention from unconsumed channel entries.
|
||||
|
||||
- State machine data
|
||||
- PubSub subscriptions
|
||||
- Trace node references
|
||||
### Residual Growth Sources
|
||||
|
||||
- PubSub subscription objects (cleaned up on Stop)
|
||||
- Driver I/O buffers (transient, GC-eligible)
|
||||
- Trace node references in memory state (released on worker exit)
|
||||
|
||||
### Workaround
|
||||
|
||||
The 20KB threshold in memory leak tests accommodates this known overhead while still detecting severe leaks (50KB+ growth would indicate a real problem).
|
||||
The 20KB threshold in memory leak tests accommodates known overhead while still
|
||||
detecting severe leaks (50KB+ growth would indicate a real problem).
|
||||
|
||||
## Planned Refactoring
|
||||
|
||||
The Trace module is scheduled for refactoring to address:
|
||||
The Trace module is scheduled for further refactoring:
|
||||
|
||||
1. **Synchronous cleanup**: Ensure goroutines exit before `Release()` returns
|
||||
2. **Memory management**: Properly release all trace-related objects
|
||||
3. **Resource pooling**: Consider reusing trace resources to reduce allocation overhead
|
||||
1. **Global Event Service**: Decouple event broadcasting from trace manager into
|
||||
a process-level daemon (separate plan)
|
||||
2. **Resource pooling**: Consider reusing trace resources to reduce allocation overhead
|
||||
3. **PubSub synchronous cleanup**: Ensure forward() goroutine exits before Stop() returns
|
||||
|
||||
## Testing Notes
|
||||
|
||||
|
|
@ -98,10 +91,12 @@ These thresholds are intentionally higher than actual growth to:
|
|||
## Related Files
|
||||
|
||||
- `trace/manager.go` - State machine and goroutine management
|
||||
- `trace/state.go` - Channel-based state worker and safeSend
|
||||
- `trace/trace.go` - Release() three-step shutdown
|
||||
- `trace/pubsub/pubsub.go` - PubSub forwarding goroutine
|
||||
- `trace/trace.go` - Release() implementation
|
||||
- `agent/assistant/hook/create_mem_test.go` - Memory leak tests
|
||||
- `trace/trace_lifecycle_test.go` - Boundary condition tests for shutdown races
|
||||
- `trace/BUGFIX.md` - Detailed bug analysis and fix documentation
|
||||
|
||||
---
|
||||
|
||||
_Last updated: December 2025_
|
||||
_Last updated: February 2026_
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ type manager struct {
|
|||
traceID string
|
||||
driver types.Driver
|
||||
stateCmdChan chan stateCommand // Single channel for all state mutations
|
||||
closed int32 // Atomic flag: 1 = closed, safeSend rejects new commands
|
||||
autoArchive bool // Auto-archive on complete/fail
|
||||
pubsub *pubsub.PubSub // Reference to independent pubsub service (for publishing only, doesn't own it)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
package trace
|
||||
|
||||
import (
|
||||
"time"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/yaoapp/kun/log"
|
||||
"github.com/yaoapp/yao/trace/types"
|
||||
|
|
@ -210,9 +211,10 @@ func (c *cmdSpaceKVOp) execute(s *managerState) {
|
|||
c.resp <- err
|
||||
}
|
||||
|
||||
// State worker - processes all commands serially in a single goroutine
|
||||
// State worker - processes all commands serially in a single goroutine.
|
||||
// Exits only when stateCmdChan is closed by Release(). The for-range loop
|
||||
// automatically drains any buffered commands before returning (Go spec guarantee).
|
||||
func (m *manager) startStateWorker() {
|
||||
// Initialize state
|
||||
state := &managerState{
|
||||
rootNode: nil,
|
||||
currentNodes: []*types.TraceNode{},
|
||||
|
|
@ -220,59 +222,29 @@ func (m *manager) startStateWorker() {
|
|||
traceStatus: types.TraceStatusPending,
|
||||
completed: false,
|
||||
updates: make([]*types.TraceUpdate, 0, 100),
|
||||
// subscribers removed - now managed by SubscriptionManager
|
||||
}
|
||||
|
||||
// Process commands until channel is closed (on Release)
|
||||
// Note: We don't exit on context cancellation anymore - state machine should continue
|
||||
// running until Release() is called, which closes the channel
|
||||
for {
|
||||
cmd, ok := <-m.stateCmdChan
|
||||
if !ok {
|
||||
// Channel closed by Release(), exit cleanly
|
||||
return
|
||||
}
|
||||
for cmd := range m.stateCmdChan {
|
||||
cmd.execute(state)
|
||||
|
||||
// Optional: Exit after processing completion (but only after draining)
|
||||
// This is mainly for optimization - the channel will be closed by Release() anyway
|
||||
if state.completed {
|
||||
// Drain remaining commands with timeout
|
||||
drainTimer := time.NewTimer(100 * time.Millisecond)
|
||||
defer drainTimer.Stop()
|
||||
drainLoop:
|
||||
for {
|
||||
select {
|
||||
case cmd, ok := <-m.stateCmdChan:
|
||||
if !ok {
|
||||
// Channel closed during drain
|
||||
break drainLoop
|
||||
}
|
||||
cmd.execute(state)
|
||||
case <-drainTimer.C:
|
||||
break drainLoop
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Helper methods for manager to send commands
|
||||
|
||||
// safeSend checks if context is cancelled before sending to avoid panic on closed channel
|
||||
// safeSend sends a command to the state worker channel. Returns false if the
|
||||
// manager is closed, context is cancelled, or the channel was closed mid-send.
|
||||
// The atomic closed flag provides a fast-path rejection before touching the channel,
|
||||
// which is critical in CGO callback stacks where recover() may not work.
|
||||
func (m *manager) safeSend(cmd stateCommand) (ok bool) {
|
||||
// Use defer/recover to handle the case where channel is closed mid-send
|
||||
if atomic.LoadInt32(&m.closed) == 1 {
|
||||
return false
|
||||
}
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
// Channel was closed, silently return false
|
||||
ok = false
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-m.ctx.Done():
|
||||
// Context cancelled, channel may be closed
|
||||
return false
|
||||
case m.stateCmdChan <- cmd:
|
||||
return true
|
||||
|
|
@ -383,6 +355,8 @@ func (m *manager) stateSetUpdates(updates []*types.TraceUpdate) {
|
|||
// stateExecuteSpaceOp executes a space operation serially through state worker
|
||||
func (m *manager) stateExecuteSpaceOp(spaceID string, fn func() error) error {
|
||||
resp := make(chan error, 1)
|
||||
m.stateCmdChan <- &cmdSpaceKVOp{spaceID: spaceID, fn: fn, resp: resp}
|
||||
if !m.safeSend(&cmdSpaceKVOp{spaceID: spaceID, fn: fn, resp: resp}) {
|
||||
return fmt.Errorf("trace %s: state worker stopped", m.traceID)
|
||||
}
|
||||
return <-resp
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
gonanoid "github.com/matoous/go-nanoid/v2"
|
||||
|
|
@ -487,17 +488,20 @@ func Release(traceID string) error {
|
|||
return fmt.Errorf("trace not found in registry: %s", traceID)
|
||||
}
|
||||
|
||||
// Stop manager
|
||||
// Stop manager with safe three-step shutdown sequence.
|
||||
// Order matters: flag blocks new writes -> cancel unblocks in-flight safeSend ->
|
||||
// close terminates state worker (which drains remaining buffer first).
|
||||
if mgr, ok := info.Manager.(*manager); ok {
|
||||
// Close state machine channel to stop state worker goroutine
|
||||
log.Trace("[TRACE] Release: closing state command channel")
|
||||
close(mgr.stateCmdChan)
|
||||
// Step 1: Set closed flag — new safeSend calls return false immediately
|
||||
atomic.StoreInt32(&mgr.closed, 1)
|
||||
|
||||
// Cancel the manager's context to stop other background operations
|
||||
// Step 2: Cancel context — unblocks any safeSend blocked in select on ctx.Done
|
||||
if mgr.cancel != nil {
|
||||
log.Trace("[TRACE] Release: cancelling manager context")
|
||||
mgr.cancel()
|
||||
}
|
||||
|
||||
// Step 3: Close channel — state worker for-range exits after draining buffer
|
||||
close(mgr.stateCmdChan)
|
||||
}
|
||||
|
||||
// Stop independent PubSub service
|
||||
|
|
|
|||
343
trace/trace_lifecycle_test.go
Normal file
343
trace/trace_lifecycle_test.go
Normal file
|
|
@ -0,0 +1,343 @@
|
|||
package trace_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/yao/trace"
|
||||
"github.com/yaoapp/yao/trace/types"
|
||||
)
|
||||
|
||||
// TestReleaseWhileWriting verifies that calling Release while multiple
|
||||
// goroutines are actively writing to the trace does not panic or deadlock.
|
||||
func TestReleaseWhileWriting(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
|
||||
// Add a root node so logging operations have a target
|
||||
_, err = manager.Add("root", types.TraceNodeOption{Label: "Root"})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Start 20 goroutines continuously writing
|
||||
var wg sync.WaitGroup
|
||||
const numWriters = 20
|
||||
stop := make(chan struct{})
|
||||
|
||||
for i := 0; i < numWriters; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
default:
|
||||
manager.Info("writer %d tick", idx)
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
// Let writers run briefly, then Release while they are still writing
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
err = trace.Release(traceID)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Signal writers to stop and wait
|
||||
close(stop)
|
||||
wg.Wait()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestReleaseDuringSpaceOp verifies that calling Release while space
|
||||
// key-value operations are in flight does not panic.
|
||||
func TestReleaseDuringSpaceOp(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
|
||||
space, err := manager.CreateSpace(types.TraceSpaceOption{Label: "Test Space"})
|
||||
assert.NoError(t, err)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
const numOps = 10
|
||||
stop := make(chan struct{})
|
||||
|
||||
for i := 0; i < numOps; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
default:
|
||||
key := fmt.Sprintf("key_%d", idx)
|
||||
// Errors are expected after Release; we only care about no panic.
|
||||
_ = manager.SetSpaceValue(space.ID, key, idx)
|
||||
_, _ = manager.GetSpaceValue(space.ID, key)
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
err = trace.Release(traceID)
|
||||
assert.NoError(t, err)
|
||||
|
||||
close(stop)
|
||||
wg.Wait()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestReleaseAfterMarkComplete verifies that MarkComplete followed by
|
||||
// immediate Release and further operations does not panic.
|
||||
func TestReleaseAfterMarkComplete(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
|
||||
_, err = manager.Add("root", types.TraceNodeOption{Label: "Root"})
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = manager.MarkComplete()
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = trace.Release(traceID)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Post-release operations should not panic.
|
||||
// They may return errors or be silently dropped.
|
||||
manager.Info("after release")
|
||||
_ = manager.SetOutput("stale output")
|
||||
_, _ = manager.Add("late", types.TraceNodeOption{Label: "Late"})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentReleaseAndMarkCancelled verifies that calling MarkCancelled
|
||||
// and Release concurrently does not panic or deadlock.
|
||||
func TestConcurrentReleaseAndMarkCancelled(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
|
||||
_, err = manager.Add("root", types.TraceNodeOption{Label: "Root"})
|
||||
assert.NoError(t, err)
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = trace.MarkCancelled(traceID, "test cancel")
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = trace.Release(traceID)
|
||||
}()
|
||||
|
||||
wg.Wait()
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
// success — no deadlock
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("deadlock: concurrent MarkCancelled + Release did not finish in 5s")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestSafeSendAfterClosed verifies that operations using safeSend after
|
||||
// Release return gracefully instead of panicking.
|
||||
func TestSafeSendAfterClosed(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
|
||||
_, err = manager.Add("root", types.TraceNodeOption{Label: "Root"})
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = trace.Release(traceID)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// All of these internally use safeSend. After Release they should
|
||||
// return nil/error/zero-value, never panic.
|
||||
manager.Info("post-close info")
|
||||
manager.Debug("post-close debug")
|
||||
manager.Error("post-close error")
|
||||
manager.Warn("post-close warn")
|
||||
|
||||
root, _ := manager.GetRootNode()
|
||||
assert.Nil(t, root)
|
||||
|
||||
nodes, _ := manager.GetCurrentNodes()
|
||||
assert.Nil(t, nodes)
|
||||
|
||||
status := manager.IsComplete()
|
||||
assert.True(t, status)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestRapidCreateReleaseLoop stress-tests the create/release cycle to ensure
|
||||
// no goroutine accumulation or panics over many iterations.
|
||||
func TestRapidCreateReleaseLoop(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
// Warm up and let baseline stabilize
|
||||
runtime.GC()
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
baseGoroutines := runtime.NumGoroutine()
|
||||
|
||||
const iterations = 100
|
||||
for i := 0; i < iterations; i++ {
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
|
||||
_, err = manager.Add("input", types.TraceNodeOption{Label: "Node"})
|
||||
assert.NoError(t, err)
|
||||
manager.Info("iteration %d", i)
|
||||
|
||||
err = manager.MarkComplete()
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = trace.Release(traceID)
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
// Allow goroutines to wind down
|
||||
runtime.GC()
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
|
||||
finalGoroutines := runtime.NumGoroutine()
|
||||
delta := finalGoroutines - baseGoroutines
|
||||
|
||||
// Allow a small margin for runtime goroutines; flag severe leaks
|
||||
assert.LessOrEqual(t, delta, 20,
|
||||
"goroutine leak: base=%d final=%d delta=%d", baseGoroutines, finalGoroutines, delta)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentAllPattern simulates the real All() orchestrator pattern:
|
||||
// parent creates a trace, forks N goroutines that each write to the shared
|
||||
// trace manager, then parent calls MarkComplete + Release while children
|
||||
// may still be writing.
|
||||
func TestConcurrentAllPattern(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
|
||||
// Root node
|
||||
_, err = manager.Add("orchestrator", types.TraceNodeOption{Label: "Orchestrator"})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create a shared space
|
||||
space, err := manager.CreateSpace(types.TraceSpaceOption{Label: "Shared"})
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Fork N "child" goroutines, each doing work on the shared trace
|
||||
const numChildren = 10
|
||||
childStarted := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for i := 0; i < numChildren; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
|
||||
// Signal that this child has started
|
||||
select {
|
||||
case childStarted <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
|
||||
// Simulate work: log, set space values, complete
|
||||
for j := 0; j < 20; j++ {
|
||||
manager.Info("child %d step %d", idx, j)
|
||||
_ = manager.SetSpaceValue(space.ID, fmt.Sprintf("child_%d_%d", idx, j), j)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
// Wait for at least a few children to start, then trigger shutdown
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
|
||||
// Parent completes and releases — some children are still writing
|
||||
_ = manager.MarkComplete()
|
||||
_ = trace.Release(traceID)
|
||||
|
||||
// Wait for all children to finish (they should not panic)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
// success
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("deadlock: child goroutines did not finish in 10s")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue