control: forward traces in a non-blocking goroutine

Traces are now forwarded in a non-blocking goroutine when sent through
the traces exporter. This prevents traces forwarded from the client from
being stalled while waiting for an upstream uploader to appear.

In addition, adds a shutdown context to `appcontext` that will only
cancel when an interrupt has been received twice. One interrupt will
signal the program should clean up and shut down, the second indicates
we should skip shutdown procedures (more forceful), and the third will
indicate that we should immediately terminate the program.

This gives a bit more of a degree of control to shutdown procedures like
the traces and metrics exporter so there's a difference between forcibly
calling exit and just waiting a long time for the shutdown to happen.

Includes a more aggressive shutdown timeout for `buildctl` that is
similar to the export timeout on `docker-buildx` for the tracing
shutdown as another preventative measure to ensure the CLI hangs up at
an appropriate time interval.

Signed-off-by: Jonathan A. Sternberg <jonathan.sternberg@docker.com>

cmd/buildctl/common/trace.go

@@ -3,10 +3,12 @@ package common
 import (
 	"context"
 	"os"
+	"time"
 
 	"github.com/moby/buildkit/util/appcontext"
 	"github.com/moby/buildkit/util/tracing/delegated"
 	"github.com/moby/buildkit/util/tracing/detect"
+	"github.com/pkg/errors"
 	"github.com/urfave/cli"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
@@ -14,6 +16,8 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
+const exportTimeout = 50 * time.Millisecond
+
 func AttachAppContext(app *cli.App) error {
 	ctx := appcontext.Context()
 
@@ -69,7 +73,14 @@ func AttachAppContext(app *cli.App) error {
 		if span != nil {
 			span.End()
 		}
-		return tp.Shutdown(context.TODO())
+
+		// Set a rather aggressive timeout for shutting down the tracer provider
+		// to ensure we don't stall on a non-responsive tracing endpoint for too long
+		// on shutdown.
+		ctx, cancel := context.WithTimeoutCause(appcontext.Shutdown(), exportTimeout, errors.WithStack(context.DeadlineExceeded))
+		defer cancel()
+
+		return tp.Shutdown(ctx)
 	}
 	return nil
 }

control/control.go

@@ -46,6 +46,7 @@ import (
 	"github.com/moby/buildkit/util/imageutil"
 	"github.com/moby/buildkit/util/leaseutil"
 	"github.com/moby/buildkit/util/throttle"
+	"github.com/moby/buildkit/util/tracing/forwarder"
 	"github.com/moby/buildkit/util/tracing/transform"
 	"github.com/moby/buildkit/version"
 	"github.com/moby/buildkit/worker"
@@ -89,6 +90,7 @@ type Controller struct { // TODO: ControlService
 	history                      *history.Queue
 	cache                        solver.CacheManager
 	gatewayForwarder             *controlgateway.GatewayForwarder
+	traceForwarder               *forwarder.Exporter
 	throttledGC                  func()
 	throttledReleaseUnreferenced func()
 	gcmu                         sync.Mutex
@@ -137,6 +139,14 @@ func NewController(opt Opt) (*Controller, error) {
 	// use longer interval for releaseUnreferencedCache deleting links quickly is less important
 	c.throttledReleaseUnreferenced = throttle.After(5*time.Minute, func() { c.releaseUnreferencedCache(context.TODO()) })
 
+	if opt.TraceCollector != nil {
+		fwd, err := forwarder.New(context.Background(), opt.TraceCollector)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to start trace forwarder")
+		}
+		c.traceForwarder = fwd
+	}
+
 	defer func() {
 		time.AfterFunc(time.Second, c.throttledGC)
 	}()
@@ -149,6 +159,11 @@ func (c *Controller) Close() error {
 	if err := c.opt.HistoryDB.Close(); err != nil {
 		errs = append(errs, err)
 	}
+	if c.traceForwarder != nil {
+		if err := c.traceForwarder.Shutdown(context.TODO()); err != nil {
+			errs = append(errs, err)
+		}
+	}
 	if err := c.opt.WorkerController.Close(); err != nil {
 		errs = append(errs, err)
 	}
@@ -299,11 +314,11 @@ func (c *Controller) Prune(req *controlapi.PruneRequest, stream controlapi.Contr
 }
 
 func (c *Controller) Export(ctx context.Context, req *tracev1.ExportTraceServiceRequest) (*tracev1.ExportTraceServiceResponse, error) {
-	if c.opt.TraceCollector == nil {
+	if c.traceForwarder == nil {
 		return nil, status.Errorf(codes.Unavailable, "trace collector not configured")
 	}
-	err := c.opt.TraceCollector.ExportSpans(ctx, transform.Spans(req.GetResourceSpans()))
-	if err != nil {
+
+	if err := c.traceForwarder.ExportSpans(ctx, transform.Spans(req.GetResourceSpans())); err != nil {
 		return nil, err
 	}
 	return &tracev1.ExportTraceServiceResponse{}, nil

util/appcontext/appcontext.go

@@ -10,39 +10,56 @@ import (
 	"github.com/pkg/errors"
 )
 
-var appContextCache context.Context
-var appContextOnce sync.Once
-
 // Context returns a static context that reacts to termination signals of the
 // running process. Useful in CLI tools.
 func Context() context.Context {
-	appContextOnce.Do(func() {
+	initContexts()
+	return appContext
+}
+
+// Shutdown returns a static context that closes when multiple interrupt signals
+// have been received to indicate a faster shutdown. Useful in CLI tools.
+func Shutdown() context.Context {
+	initContexts()
+	return shutdownContext
+}
+
+var (
+	appContext       context.Context
+	shutdownContext  context.Context
+	initContextsOnce sync.Once
+)
+
+func initContexts() {
+	initContextsOnce.Do(func() {
 		signals := make(chan os.Signal, 2048)
 		signal.Notify(signals, terminationSignals...)
 
-		const exitLimit = 3
-		retries := 0
-
 		ctx := context.Background()
 		for _, f := range inits {
 			ctx = f(ctx) //nolint:fatcontext
 		}
 
 		ctx, cancel := context.WithCancelCause(ctx)
-		appContextCache = ctx //nolint:fatcontext
+		appContext = ctx //nolint:fatcontext
 
+		shutdownCtx, shutdownCancel := context.WithCancelCause(context.Background())
+		shutdownContext = shutdownCtx
+
+		// We just allow this goroutine to be orphaned since program termination
+		// will clean it up.
 		go func() {
-			for {
-				<-signals
-				retries++
-				err := errors.Errorf("got %d SIGTERM/SIGINTs, forcing shutdown", retries)
-				cancel(err)
-				if retries >= exitLimit {
-					bklog.G(ctx).Error(err.Error())
-					os.Exit(1)
-				}
-			}
+			<-signals
+			err := errors.New("got SIGTERM/SIGINT, forcing shutdown")
+			cancel(err)
+
+			<-signals
+			err = errors.New("got 2 SIGTERM/SIGINTs, skipping shutdown")
+			shutdownCancel(err)
+
+			<-signals
+			err = errors.New("got 3 SIGTERM/SIGINTs, forcibly terminating")
+			bklog.G(ctx).Fatal(err.Error())
 		}()
 	})
-	return appContextCache
 }

util/tracing/forwarder/forwarder.go

@@ -0,0 +1,183 @@
+package forwarder
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/moby/buildkit/util/bklog"
+	"github.com/pkg/errors"
+	"go.opentelemetry.io/otel"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+)
+
+const (
+	pendingBufferSize = 1000
+	exportTimeout     = 30 * time.Second
+)
+
+var (
+	errUnstarted      = errors.New("not started")
+	errAlreadyStarted = errors.New("already started")
+	errShutdown       = errors.New("shutdown")
+)
+
+type exportRequest struct {
+	deadline time.Time
+	spans    []sdktrace.ReadOnlySpan
+}
+
+type Exporter struct {
+	exp sdktrace.SpanExporter
+
+	ctx      context.Context
+	cancel   context.CancelCauseFunc
+	wg       sync.WaitGroup
+	ch       chan<- exportRequest
+	shutdown bool
+	mu       sync.RWMutex
+
+	startOnce sync.Once
+	stopOnce  sync.Once
+}
+
+// New constructs a new Exporter and starts it.
+func New(ctx context.Context, exp sdktrace.SpanExporter) (*Exporter, error) {
+	e := NewUnstarted(exp)
+	if err := e.Start(ctx); err != nil {
+		return nil, err
+	}
+	return e, nil
+}
+
+// NewUnstarted constructs a new Exporter and does not start it.
+func NewUnstarted(exp sdktrace.SpanExporter) *Exporter {
+	return &Exporter{exp: exp}
+}
+
+// Start marks the Exporter as started.
+func (e *Exporter) Start(ctx context.Context) error {
+	err := errAlreadyStarted
+	e.startOnce.Do(func() {
+		e.mu.Lock()
+		defer e.mu.Unlock()
+
+		e.ctx, e.cancel = context.WithCancelCause(context.Background())
+		ch := make(chan exportRequest, pendingBufferSize)
+		e.ch = ch
+		e.wg.Go(func() {
+			e.exportLoop(e.ctx, ch)
+		})
+		err = nil
+	})
+	return err
+}
+
+func (e *Exporter) ExportSpans(ctx context.Context, spans []sdktrace.ReadOnlySpan) error {
+	if len(spans) == 0 {
+		return nil
+	}
+
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+
+	if e.ctx == nil {
+		return errUnstarted
+	} else if e.shutdown {
+		return errShutdown
+	}
+
+	select {
+	case e.ch <- exportRequest{
+		deadline: time.Now().Add(exportTimeout),
+		spans:    spans,
+	}:
+	default:
+		bklog.G(ctx).Warnf("dropped %d spans: exporter buffer full", len(spans))
+	}
+	return nil
+}
+
+// exportLoop reads spans from the buffered channel and exports them. It exits
+// when ctx is canceled.
+func (e *Exporter) exportLoop(ctx context.Context, ch <-chan exportRequest) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case req, ok := <-ch:
+			if !ok {
+				return
+			}
+			e.exportSpans(ctx, req)
+		}
+	}
+}
+
+func (e *Exporter) exportSpans(ctx context.Context, req exportRequest) {
+	if time.Since(req.deadline) >= 0 {
+		bklog.G(ctx).Warnf("dropped %d spans: export deadline missed", len(req.spans))
+	}
+
+	ctx, cancel := context.WithDeadlineCause(ctx, req.deadline, errors.WithStack(context.DeadlineExceeded))
+	defer cancel()
+
+	if err := e.exp.ExportSpans(ctx, req.spans); err != nil {
+		otel.Handle(err)
+	}
+}
+
+func (e *Exporter) Shutdown(ctx context.Context) error {
+	e.mu.Lock()
+	e.shutdown = true
+	cancel := e.cancel
+	e.mu.Unlock()
+
+	if cancel == nil {
+		return nil
+	}
+
+	var err error
+	e.stopOnce.Do(func() {
+		// Close the channel to signal that there are no more spans to export.
+		close(e.ch)
+
+		// Wait for the spans to drain or the shutdown context to have
+		// been canceled. Whichever occurs first.
+		var wg sync.WaitGroup
+		done := make(chan struct{})
+		wg.Go(func() {
+			select {
+			// All pending exports have finished.
+			case <-done:
+				// The shutdown deadline has been reached.
+			case <-ctx.Done():
+			}
+
+			// Either all exports have completed or we've hit
+			// our limit for the shutdown context. Cancel the
+			// context either way.
+			cause := context.Cause(ctx)
+			if cause == nil {
+				cause = context.Canceled
+			}
+			cancel(cause)
+		})
+
+		// Wait for the run goroutine (and any in-flight export) to
+		// finish and then mark the exporter as done when it is.
+		e.wg.Wait()
+		close(done)
+
+		// Ensure we don't leave orphaned goroutines.
+		wg.Wait()
+
+		// Finally finish by calling shutdown on the underlying exporter.
+		// If the context already got canceled, this will probably
+		// do nothing but just going with it anyway.
+		err = e.exp.Shutdown(ctx)
+	})
+	return err
+}
+
+var _ sdktrace.SpanExporter = (*Exporter)(nil)