sql: finish root span of the query

Fixes #853

Signed-off-by: Miguel Molina <miguel@erizocosmi.co>

Author: erizocosmico

engine_test.go

@@ -9,6 +9,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/opentracing/opentracing-go"
+
 	sqle "github.com/src-d/go-mysql-server"
 	"github.com/src-d/go-mysql-server/auth"
 	"github.com/src-d/go-mysql-server/memory"
@@ -3420,6 +3422,32 @@ func TestDeleteFromErrors(t *testing.T) {
 	}
 }
 
+type mockSpan struct {
+	opentracing.Span
+	finished bool
+}
+
+func (m *mockSpan) Finish() {
+	m.finished = true
+}
+
+func TestRootSpanFinish(t *testing.T) {
+	e := newEngine(t)
+	fakeSpan := &mockSpan{Span: opentracing.NoopTracer{}.StartSpan("")}
+	ctx := sql.NewContext(
+		context.Background(),
+		sql.WithRootSpan(fakeSpan),
+	)
+
+	_, iter, err := e.Query(ctx, "SELECT 1")
+	require.NoError(t, err)
+
+	_, err = sql.RowIterToRows(iter)
+	require.NoError(t, err)
+
+	require.True(t, fakeSpan.finished)
+}
+
 var generatorQueries = []struct {
 	query    string
 	expected []sql.Row

server/context.go

@@ -92,15 +92,13 @@ func (s *SessionManager) NewContextWithQuery(
 	s.mu.Unlock()
 
 	context := sql.NewContext(
-		opentracing.ContextWithSpan(
-			context.Background(),
-			s.tracer.StartSpan("query"),
-		),
+		context.Background(),
 		sql.WithSession(sess),
 		sql.WithTracer(s.tracer),
 		sql.WithPid(s.nextPid()),
 		sql.WithQuery(query),
 		sql.WithMemoryManager(s.memory),
+		sql.WithRootSpan(s.tracer.StartSpan("query")),
 	)
 
 	return context

sql/analyzer/process.go

@@ -85,5 +85,10 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
 		return nil, err
 	}
 
-	return plan.NewQueryProcess(node, func() { processList.Done(ctx.Pid()) }), nil
+	return plan.NewQueryProcess(node, func() {
+		processList.Done(ctx.Pid())
+		if span := ctx.RootSpan(); span != nil {
+			span.Finish()
+		}
+	}), nil
 }

sql/session.go

@@ -211,10 +211,11 @@ func NewBaseSession() Session {
 type Context struct {
 	context.Context
 	Session
-	Memory *MemoryManager
-	pid    uint64
-	query  string
-	tracer opentracing.Tracer
+	Memory   *MemoryManager
+	pid      uint64
+	query    string
+	tracer   opentracing.Tracer
+	rootSpan opentracing.Span
 }
 
 // ContextOption is a function to configure the context.
@@ -255,6 +256,13 @@ func WithMemoryManager(m *MemoryManager) ContextOption {
 	}
 }
 
+// WithRootSpan sets the root span of the context.
+func WithRootSpan(s opentracing.Span) ContextOption {
+	return func(ctx *Context) {
+		ctx.rootSpan = s
+	}
+}
+
 // NewContext creates a new query context. Options can be passed to configure
 // the context. If some aspect of the context is not configure, the default
 // value will be used.
@@ -264,7 +272,7 @@ func NewContext(
 	ctx context.Context,
 	opts ...ContextOption,
 ) *Context {
-	c := &Context{ctx, NewBaseSession(), nil, 0, "", opentracing.NoopTracer{}}
+	c := &Context{ctx, NewBaseSession(), nil, 0, "", opentracing.NoopTracer{}, nil}
 	for _, opt := range opts {
 		opt(c)
 	}
@@ -298,12 +306,17 @@ func (c *Context) Span(
 	span := c.tracer.StartSpan(opName, opts...)
 	ctx := opentracing.ContextWithSpan(c.Context, span)
 
-	return span, &Context{ctx, c.Session, c.Memory, c.Pid(), c.Query(), c.tracer}
+	return span, &Context{ctx, c.Session, c.Memory, c.Pid(), c.Query(), c.tracer, c.rootSpan}
 }
 
 // WithContext returns a new context with the given underlying context.
 func (c *Context) WithContext(ctx context.Context) *Context {
-	return &Context{ctx, c.Session, c.Memory, c.Pid(), c.Query(), c.tracer}
+	return &Context{ctx, c.Session, c.Memory, c.Pid(), c.Query(), c.tracer, c.rootSpan}
+}
+
+// RootSpan returns the root span, if any.
+func (c *Context) RootSpan() opentracing.Span {
+	return c.rootSpan
 }
 
 // Error adds an error as warning to the session.