From 65e668dbc2914a3ddb8161934c14e46f0fbb5409 Mon Sep 17 00:00:00 2001
From: Jeff Wendling <leterip@gmail.com>
Date: Tue, 15 Dec 2020 12:28:09 -0500
Subject: [PATCH] reparent orphaned spans to grandparents

---
 ctx.go     | 12 ++++++++++--
 span.go    | 17 +++++++++++++++--
 spanbag.go |  5 +++++
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/ctx.go b/ctx.go
index 0bdcbae..449d940 100644
--- a/ctx.go
+++ b/ctx.go
@@ -121,9 +121,17 @@ func newSpan(ctx context.Context, f *Func, args []interface{},
 		s.children.Iterate(func(child *Span) {
 			children = append(children, child)
 		})
+		s.children.Clear()
 		s.mtx.Unlock()
-		for _, child := range children {
-			child.orphan()
+
+		if s.parent != nil && !orphaned {
+			for _, child := range children {
+				s.parent.addChild(child)
+			}
+		} else {
+			for _, child := range children {
+				child.orphan()
+			}
 		}
 
 		if s.parent != nil {
diff --git a/span.go b/span.go
index 6326c6e..00706a9 100644
--- a/span.go
+++ b/span.go
@@ -34,18 +34,31 @@ type Annotation struct {
 
 func (s *Span) addChild(child *Span) {
 	s.mtx.Lock()
-	s.children.Add(child)
 	done := s.done
+	parent := s.parent
+	if !done {
+		s.children.Add(child)
+	}
 	s.mtx.Unlock()
+
 	if done {
-		child.orphan()
+		if parent == nil {
+			child.orphan()
+		} else {
+			parent.addChild(child)
+		}
 	}
 }
 
 func (s *Span) removeChild(child *Span) {
 	s.mtx.Lock()
+	parent := s.parent
 	s.children.Remove(child)
 	s.mtx.Unlock()
+
+	if parent != nil {
+		parent.removeChild(child)
+	}
 }
 
 func (s *Span) orphan() {
diff --git a/spanbag.go b/spanbag.go
index 4e9a6d2..a36c4da 100644
--- a/spanbag.go
+++ b/spanbag.go
@@ -23,6 +23,11 @@ type spanBag struct {
 	rest  map[*Span]int32
 }
 
+func (b *spanBag) Clear() {
+	b.first = nil
+	b.rest = nil
+}
+
 func (b *spanBag) Add(s *Span) {
 	if b.first == nil {
 		b.first = s