add extended span to traceutils (#28)

util/traceutil/trace/span.go

@@ -4,15 +4,19 @@ import (
 	"context"
 	"encoding/json"
 	"sync"
+	"time"
 
-	"github.com/eluv-io/common-go/format/duration"
 	"github.com/eluv-io/utc-go"
+
+	"github.com/eluv-io/common-go/format/duration"
 )
 
 // Span represents an operation that is named and timed. It may contain sub-spans, representing nested sub-operations. A
 // span can also record an arbitrary number of attributes (key-value pairs) and events (bundled attributes with a name
 // and timestamp)
 type Span interface {
+	ExtendedSpan
+
 	// End completes the span. No updates are allowed to span after it
 	// ends. The only exception is setting status of the span.
 	End()
@@ -43,6 +47,23 @@ type Span interface {
 	FindByName(name string) Span
 }
 
+type ExtendedSpan interface {
+	// StartTime returns the start time of the span.
+	StartTime() utc.UTC
+
+	// EndTime returns the end time of the span.
+	EndTime() utc.UTC
+
+	// Duration returns the duration of the span.
+	Duration() time.Duration
+
+	// MarshalExtended returns true if the span is set to use an extended JSON representation during marshaling.
+	MarshalExtended() bool
+
+	// SetMarshalExtended sets the span to use an extended JSON representation during marshaling.
+	SetMarshalExtended()
+}
+
 type Event struct {
 	Name string                 `json:"name"`
 	Time utc.UTC                `json:"-"`
@@ -65,33 +86,50 @@ func (n NoopSpan) Json() string                                         { return
 func (n NoopSpan) Attributes() map[string]interface{}                   { return nil }
 func (n NoopSpan) Events() []*Event                                     { return nil }
 func (n NoopSpan) FindByName(string) Span                               { return nil }
+func (n NoopSpan) StartTime() utc.UTC                                   { return utc.Zero }
+func (n NoopSpan) EndTime() utc.UTC                                     { return utc.Zero }
+func (n NoopSpan) Duration() time.Duration                              { return 0 }
+func (n NoopSpan) MarshalExtended() bool                                { return false }
+func (n NoopSpan) SetMarshalExtended()                                  {}
 
 // ---------------------------------------------------------------------------------------------------------------------
 
 func newSpan(name string) *RecordingSpan {
 	s := &RecordingSpan{
-		StartTime: utc.Now(),
+		startTime: utc.Now(),
 	}
 	s.Data.Name = name
+	s.Data.Start = s.startTime.String()
 	return s
 }
 
 type RecordingSpan struct {
 	mutex     sync.Mutex
 	Parent    Span
-	StartTime utc.UTC
-	EndTime   utc.UTC
-	Data      struct {
-		Name     string                 `json:"name"`
-		Duration duration.Spec          `json:"time"`
-		Attr     map[string]interface{} `json:"attr,omitempty"`
-		Events   []*Event               `json:"evnt,omitempty"`
-		Subs     []Span                 `json:"subs,omitempty"`
-	}
+	Data      recordingExtendedData
+	startTime utc.UTC
+	endTime   utc.UTC
+	duration  time.Duration
+	extended  bool
+}
+
+type recordingData struct {
+	Name     string                 `json:"name"`
+	Duration duration.Spec          `json:"time"`
+	Attr     map[string]interface{} `json:"attr,omitempty"`
+	Events   []*Event               `json:"evnt,omitempty"`
+	Subs     []Span                 `json:"subs,omitempty"`
+}
+
+type recordingExtendedData struct {
+	recordingData
+	Start string `json:"start"`
+	End   string `json:"end"`
 }
 
 func (s *RecordingSpan) Start(ctx context.Context, name string) (context.Context, Span) {
 	sub := newSpan(name)
+	sub.Parent = s
 
 	s.mutex.Lock()
 	s.Data.Subs = append(s.Data.Subs, sub)
@@ -104,11 +142,13 @@ func (s *RecordingSpan) End() {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
-	if s.EndTime != utc.Zero {
+	if s.endTime != utc.Zero {
 		return
 	}
-	s.EndTime = utc.Now()
-	s.Data.Duration = duration.Spec(s.EndTime.Sub(s.StartTime)).RoundTo(1)
+	s.endTime = utc.Now()
+	s.duration = s.endTime.Sub(s.startTime)
+	s.Data.Duration = duration.Spec(s.duration).RoundTo(1)
+	s.Data.End = s.endTime.String()
 }
 
 func (s *RecordingSpan) Attribute(name string, val interface{}) {
@@ -174,6 +214,45 @@ func (s *RecordingSpan) FindByName(name string) Span {
 	return nil
 }
 
+func (s *RecordingSpan) StartTime() utc.UTC {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	return s.startTime
+}
+
+func (s *RecordingSpan) EndTime() utc.UTC {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	return s.endTime
+}
+
+func (s *RecordingSpan) Duration() time.Duration {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	return s.duration
+}
+
+func (s *RecordingSpan) MarshalExtended() bool {
+	if s.extended {
+		return true
+	} else if s.Parent != nil {
+		// Check parent for extended flag
+		return s.Parent.MarshalExtended()
+	}
+	return false
+}
+
+func (s *RecordingSpan) SetMarshalExtended() {
+	s.extended = true
+}
+
 func (s *RecordingSpan) MarshalJSON() ([]byte, error) {
-	return json.Marshal(s.Data)
+	if s.MarshalExtended() {
+		return json.Marshal(s.Data)
+	} else {
+		return json.Marshal(s.Data.recordingData)
+	}
 }

util/traceutil/traceutil_test.go

@@ -3,14 +3,22 @@ package traceutil_test
 import (
 	"context"
 	"io"
+	"regexp"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/eluv-io/utc-go"
+
 	"github.com/eluv-io/common-go/util/traceutil"
 	"github.com/eluv-io/common-go/util/traceutil/trace"
 )
 
+const spanJson = `{"name":"root-span","time":"1s","subs":[{"name":"sub-span","time":"1s"}]}`
+
+const spanExtendedJson = `{"name":"root-span","time":"1s","subs":[{"name":"sub-span","time":"1s","start":"2020-02-02T00:00:00","end":"2020-02-02T00:00:01"}],"start":"2020-02-02T00:00:00","end":"2020-02-02T00:00:01"}`
+
 func TestStartSubSpan(t *testing.T) {
 	ctx, span := trace.StartRootSpan(context.Background(), "root-span")
 	require.NotNil(t, ctx)
@@ -58,3 +66,37 @@ func TestWithSubSpan(t *testing.T) {
 	require.Equal(t, "sub-a", subA.Data.Name)
 	require.Equal(t, "sub-b", subB.Data.Name)
 }
+
+func TestExtendedSpan(t *testing.T) {
+	var now time.Time
+	defer utc.MockNowFn(func() utc.UTC {
+		return utc.MustParse("2020-02-02").Add(time.Since(now))
+	})()
+
+	now = time.Now()
+	ctx, span := trace.StartRootSpan(context.Background(), "root-span")
+	require.NotNil(t, ctx)
+	require.NotNil(t, span)
+
+	ctx, sub := traceutil.StartSubSpan(ctx, "sub-span")
+	require.NotNil(t, ctx)
+	require.NotNil(t, sub)
+
+	time.Sleep(time.Second)
+
+	sub.End()
+	span.End()
+
+	require.False(t, span.MarshalExtended())
+	require.False(t, sub.MarshalExtended())
+	require.Equal(t, spanJson, span.Json())
+
+	span.SetMarshalExtended()
+	require.True(t, span.MarshalExtended())
+	require.True(t, sub.MarshalExtended())
+	require.Equal(t, spanExtendedJson, removeMs(span.Json()))
+}
+
+func removeMs(s string) string {
+	return regexp.MustCompile(`\.00\dZ`).ReplaceAllString(s, "")
+}