proposal: context: enable first class citizenship of third party implementations

[gobwas]

Hello there,

This proposal concerns a way to implement custom context.Context type, that is clever enough to cancel its children contexts during cancelation.

Suppose I have few goroutines which executes some tasks periodically. I need to provide two things to the tasks – first is the goroutine ID and second is a context-similar cancelation mechanism:

type WorkerContext struct {
    ID uint
    done chan struct{}
}

// WorkerContext implements context.Context.

func worker(id uint, tasks <-chan func(*WorkerContext)) {
    done := make(chan struct{})
    defer close(done)

    ctx := WorkerContext{
        ID: id,
        done: make(chan struct{}),
    }

    for task := range w.tasks {
        task(&ctx)
    }
}

go worker(1, tasks)
go worker(N, tasks)

Then, on the caller's side, the use of goroutines looks like this:

tasks <- func(wctx *WorkerContext) {
    // Do some worker related things with wctx.ID.

    ctx, cancel := context.WithTimeout(wctx, time.Second)
    defer cancel()

    doSomeWork(ctx)
}

Looking at the context.go code, if the given parent context is not of *cancelCtx type, it starts a separate goroutine to stick on parent.Done() channel and then propagates the cancelation.

The point is that for the performance sensitive applications starting of a separate goroutine could be an issue. And it could be avoided if WorkerContext could implement some interface to handle cancelation and track children properly. I mean something like this:

package context

type Canceler interface {
    Cancel(removeFromParent bool, err error)
}

type Tracker interface {
    AddChild(Canceler)
    RemoveChild(Canceler)
}

func propagateCancel(parent Context, child Canceler) {
    if parent.Done() == nil {
        return // parent is never canceled
    }
    if p, ok := parentCancelCtx(parent); ok { // p is now Tracker.
        if err := p.Err(); err != nil {
            // parent has already been canceled
            child.Cancel(false, err)
        } else {
            p.AddChild(child)
        }
    } else {
        go func() {
            select {
            case <-parent.Done():
                child.Cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

func parentCancelCtx(parent Context) (Tracker, bool) {
    for {
        switch c := parent.(type) {
                case Tracker:
            return c, true                
        case *timerCtx:
            return &c.cancelCtx, true
        case *valueCtx:
            parent = c.Context
        default:
            return nil, false
        }
    }
}

Also, with that changes we could even use custom contexts as children of created with context.WithCancel() and others:

type myCustonContext struct {} // Implements context.Canceler.

parent, cancel := context.WithCancel()
child := new(myCustomContext)
context.Bind(parent, child)

cancel() // Cancels child too.

Sergey.