[dbnode] Avoid crashing on stack overflow (#3401)

* Added max depth protection to x/context

* Adjusted to work with code from master #2

* Add debug to understand why script fails in CI

* Fix linter errors

* Remove debug from ci lint

src/x/context/context.go

@@ -22,19 +22,27 @@ package context
 
 import (
 	stdctx "context"
+	"fmt"
 	"sync"
 
 	xopentracing "github.com/m3db/m3/src/x/opentracing"
 	xresource "github.com/m3db/m3/src/x/resource"
 
 	lightstep "github.com/lightstep/lightstep-tracer-go"
 	"github.com/opentracing/opentracing-go"
+	"github.com/opentracing/opentracing-go/ext"
 	"github.com/opentracing/opentracing-go/mocktracer"
 	"github.com/uber/jaeger-client-go"
 )
 
+const (
+	maxDistanceFromRootContext = 100
+)
+
 var (
 	noopTracer opentracing.NoopTracer
+
+	errSpanTooDeep = fmt.Errorf("span created exceeds maximum depth allowed (%d)", maxDistanceFromRootContext)
 )
 
 // NB(r): using golang.org/x/net/context is too GC expensive.
@@ -48,6 +56,7 @@ type ctx struct {
 	wg                   sync.WaitGroup
 	finalizeables        *finalizeableList
 	parent               Context
+	distanceFromRoot     uint16
 	checkedAndNotSampled bool
 }
 
@@ -283,6 +292,7 @@ func (c *ctx) Reset() {
 
 	c.Lock()
 	c.done, c.finalizeables, c.goCtx, c.checkedAndNotSampled = false, nil, stdctx.Background(), false
+	c.distanceFromRoot = 0
 	c.Unlock()
 }
 
@@ -315,6 +325,7 @@ func (c *ctx) newChildContext() Context {
 func (c *ctx) setParentCtx(parentCtx Context) {
 	c.Lock()
 	c.parent = parentCtx
+	c.distanceFromRoot = parentCtx.DistanceFromRootContext() + 1
 	c.Unlock()
 }
 
@@ -326,8 +337,16 @@ func (c *ctx) parentCtx() Context {
 	return parent
 }
 
+func (c *ctx) DistanceFromRootContext() uint16 {
+	c.RLock()
+	distanceFromRootContext := c.distanceFromRoot
+	c.RUnlock()
+
+	return distanceFromRootContext
+}
+
 func (c *ctx) StartSampledTraceSpan(name string) (Context, opentracing.Span, bool) {
-	if c.checkedAndNotSampled {
+	if c.checkedAndNotSampled || c.DistanceFromRootContext() >= maxDistanceFromRootContext {
 		return c, noopTracer.StartSpan(name), false
 	}
 	goCtx := c.GoContext()
@@ -340,6 +359,9 @@ func (c *ctx) StartSampledTraceSpan(name string) (Context, opentracing.Span, boo
 
 	child := c.newChildContext()
 	child.SetGoContext(childGoCtx)
+	if child.DistanceFromRootContext() == maxDistanceFromRootContext {
+		ext.LogError(span, errSpanTooDeep)
+	}
 	return child, span, true
 }
 

src/x/context/context_test.go

@@ -22,6 +22,7 @@ package context
 
 import (
 	stdctx "context"
+	"fmt"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -290,3 +291,47 @@ func TestGoContext(t *testing.T) {
 	returnCtx = xCtx.GoContext()
 	assert.Equal(t, stdctx.Background(), returnCtx)
 }
+
+// Test that too deep span tree created is prevented
+// Span is marked as error in that case that which is well defined
+func TestTooDeepSpanTreeIsPreventedAndMarked(t *testing.T) {
+	// use a mock tracer to ensure sampling rate is set to 1.
+	tracer := mocktracer.New()
+
+	span := tracer.StartSpan("root-span")
+	defer span.Finish()
+
+	context := NewWithGoContext(opentracing.ContextWithSpan(stdctx.Background(), span))
+
+	var (
+		lastChildSpanCreated    opentracing.Span
+		lastChildContextCreated = context
+	)
+	for i := 1; i <= maxDistanceFromRootContext; i++ {
+		lastChildContextCreated, lastChildSpanCreated, _ =
+			lastChildContextCreated.StartSampledTraceSpan(fmt.Sprintf("test-action-depth-%d", i))
+	}
+
+	mockSpan := lastChildSpanCreated.(*mocktracer.MockSpan)
+
+	errorTagValue := mockSpan.Tag("error")
+	assert.NotNil(t, errorTagValue)
+	assert.Equal(t, true, errorTagValue)
+	spanLogs := mockSpan.Logs()
+	assert.Len(t, spanLogs, 1)
+	spanLog := spanLogs[0]
+	assert.True(t, fieldsContains(spanLog.Fields, "event", "error") &&
+		fieldsContains(spanLog.Fields, "error.object", errSpanTooDeep.Error()))
+
+	childContext, _, _ := lastChildContextCreated.StartSampledTraceSpan("another-span-beyond-max-depth")
+	assert.Equal(t, lastChildContextCreated, childContext)
+}
+
+func fieldsContains(fields []mocktracer.MockKeyValue, key string, value string) bool {
+	for _, field := range fields {
+		if field.Key == key && field.ValueString == value {
+			return true
+		}
+	}
+	return false
+}

src/x/context/types.go

@@ -89,6 +89,9 @@ type Context interface {
 	// and a bool if the span is being sampled. This is used over StartTraceSpan()
 	// for hot paths where performance is crucial.
 	StartSampledTraceSpan(string) (Context, opentracing.Span, bool)
+
+	// DistanceFromRootContext returns the distance from root context (root context tree)
+	DistanceFromRootContext() uint16
 }
 
 // Pool provides a pool for contexts.