1e4e1224e2
- Update the Subscribe and SubscribeFrom methods to return a cancel function, ensuring proper resource cleanup when subscriptions are no longer needed. - Modify unsubscribe logic to close channels safely, preventing potential panics from sending on closed channels. - Enhance test cases to utilize the new cancel functionality, ensuring robust handling of subscriptions in various scenarios.
297 lines
7.9 KiB
Go
297 lines
7.9 KiB
Go
package trace_test
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/yaoapp/yao/trace"
|
|
"github.com/yaoapp/yao/trace/types"
|
|
)
|
|
|
|
func TestSubscription(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.Release(traceID)
|
|
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
|
|
|
// Subscribe to updates
|
|
updates, cancel, err := manager.Subscribe()
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, updates)
|
|
defer cancel()
|
|
|
|
// Collect updates in background
|
|
var receivedUpdates []*types.TraceUpdate
|
|
var updatesMu sync.Mutex
|
|
done := make(chan bool)
|
|
|
|
go func() {
|
|
timeout := time.After(2 * time.Second)
|
|
for {
|
|
select {
|
|
case update, ok := <-updates:
|
|
if !ok {
|
|
// Channel closed
|
|
done <- true
|
|
return
|
|
}
|
|
|
|
updatesMu.Lock()
|
|
receivedUpdates = append(receivedUpdates, update)
|
|
updatesMu.Unlock()
|
|
|
|
// Check for trace completion
|
|
if update.Type == types.UpdateTypeComplete {
|
|
done <- true
|
|
return
|
|
}
|
|
case <-timeout:
|
|
done <- true
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
// Perform operations
|
|
manager.Info("Test operation")
|
|
_, err = manager.Add("test node", types.TraceNodeOption{Label: "Test"})
|
|
assert.NoError(t, err)
|
|
|
|
err = manager.Complete(map[string]any{"test": "data"})
|
|
assert.NoError(t, err)
|
|
|
|
// Create space and set value
|
|
space, err := manager.CreateSpace(types.TraceSpaceOption{Label: "Test Space"})
|
|
assert.NoError(t, err)
|
|
|
|
err = manager.SetSpaceValue(space.ID, "key", "value")
|
|
assert.NoError(t, err)
|
|
|
|
// Mark trace complete
|
|
err = manager.MarkComplete()
|
|
assert.NoError(t, err)
|
|
|
|
// Wait for completion or timeout
|
|
<-done
|
|
|
|
// Verify we received updates
|
|
updatesMu.Lock()
|
|
defer updatesMu.Unlock()
|
|
|
|
assert.NotEmpty(t, receivedUpdates)
|
|
|
|
// Check for specific event types
|
|
eventTypes := make(map[string]bool)
|
|
for _, update := range receivedUpdates {
|
|
eventTypes[update.Type] = true
|
|
}
|
|
|
|
assert.True(t, eventTypes[types.UpdateTypeInit], "Should receive init event")
|
|
assert.True(t, eventTypes[types.UpdateTypeNodeStart], "Should receive node_start event")
|
|
assert.True(t, eventTypes[types.UpdateTypeComplete], "Should receive complete event")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSubscribeFrom(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.Release(traceID)
|
|
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
|
|
|
// Real scenario: User starts a trace, performs some operations
|
|
_, err = manager.Add("Step 1", types.TraceNodeOption{Label: "Processing"})
|
|
assert.NoError(t, err)
|
|
manager.Info("Processing step 1")
|
|
err = manager.Complete("step1 result")
|
|
assert.NoError(t, err)
|
|
|
|
// Wait to ensure different timestamp (simulate time passing)
|
|
// Use a longer sleep to account for CI environment variability
|
|
time.Sleep(1100 * time.Millisecond)
|
|
|
|
// Record timestamp (simulate user noting current time before refresh)
|
|
// Subtract 1ms to ensure we capture events that happen "now"
|
|
// This accounts for millisecond precision and timing variability in CI
|
|
resumeTimestamp := time.Now().UnixMilli() - 1
|
|
|
|
// Wait to ensure next operations have a clearly different timestamp
|
|
// Use longer sleep for CI reliability
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Continue with more operations
|
|
_, err = manager.Add("Step 2", types.TraceNodeOption{Label: "Finalizing"})
|
|
assert.NoError(t, err)
|
|
manager.Info("Processing step 2")
|
|
err = manager.Complete("step2 result")
|
|
assert.NoError(t, err)
|
|
|
|
// Mark trace complete
|
|
err = manager.MarkComplete()
|
|
assert.NoError(t, err)
|
|
|
|
// Real scenario: User refreshes page and resumes from last known timestamp
|
|
// This should replay events from resumeTimestamp onwards
|
|
updates, cancelSub, err := manager.SubscribeFrom(resumeTimestamp)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, updates)
|
|
defer cancelSub()
|
|
|
|
// Collect updates
|
|
var receivedUpdates []*types.TraceUpdate
|
|
timeout := time.After(1 * time.Second)
|
|
foundStep2 := false
|
|
|
|
collectLoop:
|
|
for {
|
|
select {
|
|
case update, ok := <-updates:
|
|
if !ok {
|
|
// Channel closed
|
|
break collectLoop
|
|
}
|
|
receivedUpdates = append(receivedUpdates, update)
|
|
// Check if we received step 2 events
|
|
if update.Type == types.UpdateTypeNodeStart {
|
|
if data, ok := update.Data.(*types.NodeStartData); ok {
|
|
if data.Node != nil && data.Node.Label == "Finalizing" {
|
|
foundStep2 = true
|
|
}
|
|
}
|
|
}
|
|
// Stop after receiving trace_complete
|
|
if update.Type == types.UpdateTypeComplete {
|
|
break collectLoop
|
|
}
|
|
case <-timeout:
|
|
break collectLoop
|
|
}
|
|
}
|
|
|
|
// Verify we received events from step 2 onwards
|
|
assert.NotEmpty(t, receivedUpdates, "Should receive events from resume point")
|
|
assert.True(t, foundStep2, "Should receive Step 2 events")
|
|
|
|
// All events should be at or after the resume timestamp
|
|
for _, update := range receivedUpdates {
|
|
assert.GreaterOrEqual(t, update.Timestamp, resumeTimestamp,
|
|
"Event timestamp %d should be >= resume timestamp %d (event type: %s)",
|
|
update.Timestamp, resumeTimestamp, update.Type)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIsComplete(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.Release(traceID)
|
|
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
|
|
|
// Initially not complete
|
|
assert.False(t, manager.IsComplete())
|
|
|
|
// Mark complete
|
|
err = manager.MarkComplete()
|
|
assert.NoError(t, err)
|
|
|
|
// Now should be complete
|
|
assert.True(t, manager.IsComplete())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMultipleSubscribers(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.Release(traceID)
|
|
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
|
|
|
// Create multiple subscribers
|
|
sub1, cancel1, err := manager.Subscribe()
|
|
assert.NoError(t, err)
|
|
defer cancel1()
|
|
|
|
sub2, cancel2, err := manager.Subscribe()
|
|
assert.NoError(t, err)
|
|
defer cancel2()
|
|
|
|
sub3, cancel3, err := manager.Subscribe()
|
|
assert.NoError(t, err)
|
|
defer cancel3()
|
|
|
|
// Collect updates from all subscribers
|
|
var wg sync.WaitGroup
|
|
counts := make([]int, 3)
|
|
var mu sync.Mutex
|
|
|
|
for i, sub := range []<-chan *types.TraceUpdate{sub1, sub2, sub3} {
|
|
wg.Add(1)
|
|
go func(idx int, ch <-chan *types.TraceUpdate) {
|
|
defer wg.Done()
|
|
timeout := time.After(1 * time.Second)
|
|
for {
|
|
select {
|
|
case update := <-ch:
|
|
if update != nil {
|
|
mu.Lock()
|
|
counts[idx]++
|
|
mu.Unlock()
|
|
if update.Type == types.UpdateTypeComplete {
|
|
return
|
|
}
|
|
}
|
|
case <-timeout:
|
|
return
|
|
}
|
|
}
|
|
}(i, sub)
|
|
}
|
|
|
|
// Perform operations
|
|
_, err = manager.Add("test", types.TraceNodeOption{Label: "Test"})
|
|
assert.NoError(t, err)
|
|
err = manager.Complete(nil)
|
|
assert.NoError(t, err)
|
|
err = manager.MarkComplete()
|
|
assert.NoError(t, err)
|
|
|
|
// Wait for all subscribers
|
|
wg.Wait()
|
|
|
|
// All subscribers should receive updates
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
for i, count := range counts {
|
|
assert.Greater(t, count, 0, "Subscriber %d should receive updates", i+1)
|
|
}
|
|
})
|
|
}
|
|
}
|