@@ -22,7 +22,7 @@ type Dispatcher interface {
 
 // ScheduleFunc is a function that represents work to do.
 // The throughput is the maximum number of messages to process for this scheduling.
-type ScheduleFunc func(throughput int)
+type ScheduleFunc func(ctx context.Context, throughput int)
 
 // poolDispatcher implements Dispatcher using a pool of goroutines.
 type poolDispatcher struct {
package execute

import (
	"context"
	"fmt"
	"runtime/debug"
	"sync"

	"github.com/influxdata/flux/codes"
	"github.com/influxdata/flux/internal/errors"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

// Dispatcher schedules work for a query.
// Each transformation submits work to be done to the dispatcher.
// Then the dispatcher schedules to work based on the available resources.
type Dispatcher interface {
	// Schedule fn to be executed
	Schedule(fn ScheduleFunc)
}

// ScheduleFunc is a function that represents work to do.
// The throughput is the maximum number of messages to process for this scheduling.
type ScheduleFunc func(throughput int)

// poolDispatcher implements Dispatcher using a pool of goroutines.
type poolDispatcher struct {
	work chan ScheduleFunc

	throughput int

	mu      sync.Mutex
	closed  bool
	closing chan struct{}
	wg      sync.WaitGroup
	err     error
	errC    chan error

	logger *zap.Logger
}

func newPoolDispatcher(throughput int, logger *zap.Logger) *poolDispatcher {
	return &poolDispatcher{
		throughput: throughput,
		work:       make(chan ScheduleFunc, 100),
		closing:    make(chan struct{}),
		errC:       make(chan error, 1),
		logger:     logger.With(zap.String("component", "dispatcher")),
	}
}

func (d *poolDispatcher) Schedule(fn ScheduleFunc) {
	select {
	case d.work <- fn:
	case <-d.closing:
	}
}

func (d *poolDispatcher) Start(n int, ctx context.Context) {
	d.wg.Add(n)
	for i := 0; i < n; i++ {
		go func() {
			defer d.wg.Done()
			// Setup panic handling on the worker goroutines
			defer func() {
				if e := recover(); e != nil {
					err, ok := e.(error)
					if !ok {
						err = fmt.Errorf("%v", e)
					}

					if errors.Code(err) == codes.ResourceExhausted {
						d.setErr(err)
						return
					}

					err = errors.Wrap(err, codes.Internal, "panic")
					d.setErr(err)
					if entry := d.logger.Check(zapcore.InfoLevel, "Dispatcher panic"); entry != nil {
						entry.Stack = string(debug.Stack())
						entry.Write(zap.Error(err))
					}
				}
			}()
			d.run(ctx)
		}()
	}
}

// Err returns a channel with will produce an error if encountered.
func (d *poolDispatcher) Err() <-chan error {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.errC
}

func (d *poolDispatcher) setErr(err error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	// TODO(nathanielc): Collect all error information.
	if d.err == nil {
		d.err = err
		d.errC <- err
	}
}

//Stop the dispatcher.
func (d *poolDispatcher) Stop() error {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.closed {
		return d.err
	}
	d.closed = true
	close(d.closing)
	d.wg.Wait()
	return d.err
}

// run is the logic executed by each worker goroutine in the pool.
func (d *poolDispatcher) run(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			// Immediately return, do not process any more work
			return
		case <-d.closing:
			// We are done, nothing left to do.
			return
		case fn := <-d.work:
			fn(d.throughput)
		}
	}
}

For transformations, we carry out the work with this dispatcher which in my understanding is pretty much similar to the executor service we have in Java. When we start to work on a transformation, a span needs to be created to trace this operation. This needs access to the `context.Context` which helps to track the hierarchical structure of the spans.