mirror of
https://github.com/Luzifer/ansible-role-version.git
synced 2024-12-23 19:11:20 +00:00
120 lines
2.5 KiB
Go
120 lines
2.5 KiB
Go
|
package ctxext
|
||
|
|
||
|
import (
|
||
|
"sync"
|
||
|
"time"
|
||
|
|
||
|
context "golang.org/x/net/context"
|
||
|
)
|
||
|
|
||
|
// WithParents returns a Context that listens to all given
|
||
|
// parents. It effectively transforms the Context Tree into
|
||
|
// a Directed Acyclic Graph. This is useful when a context
|
||
|
// may be cancelled for more than one reason. For example,
|
||
|
// consider a database with the following Get function:
|
||
|
//
|
||
|
// func (db *DB) Get(ctx context.Context, ...) {}
|
||
|
//
|
||
|
// DB.Get may have to stop for two different contexts:
|
||
|
// * the caller's context (caller might cancel)
|
||
|
// * the database's context (might be shut down mid-request)
|
||
|
//
|
||
|
// WithParents saves the day by allowing us to "merge" contexts
|
||
|
// and continue on our merry contextual way:
|
||
|
//
|
||
|
// ctx = ctxext.WithParents(ctx, db.ctx)
|
||
|
//
|
||
|
// Passing related (mutually derived) contexts to WithParents is
|
||
|
// actually ok. The child is cancelled when any of its parents is
|
||
|
// cancelled, so if any of its parents are also related, the cancel
|
||
|
// propagation will reach the child via the shortest path.
|
||
|
func WithParents(ctxts ...context.Context) context.Context {
|
||
|
if len(ctxts) < 1 {
|
||
|
panic("no contexts provided")
|
||
|
}
|
||
|
|
||
|
ctx := &errCtx{
|
||
|
done: make(chan struct{}),
|
||
|
dead: earliestDeadline(ctxts),
|
||
|
}
|
||
|
|
||
|
// listen to all contexts and use the first.
|
||
|
for _, c2 := range ctxts {
|
||
|
go func(pctx context.Context) {
|
||
|
select {
|
||
|
case <-ctx.Done(): // cancelled by another parent
|
||
|
return
|
||
|
case <-pctx.Done(): // this parent cancelled
|
||
|
// race: two parents may have cancelled at the same time.
|
||
|
// break tie with mutex (inside c.cancel)
|
||
|
ctx.cancel(pctx.Err())
|
||
|
}
|
||
|
}(c2)
|
||
|
}
|
||
|
|
||
|
return ctx
|
||
|
}
|
||
|
|
||
|
func earliestDeadline(ctxts []context.Context) *time.Time {
|
||
|
var d1 *time.Time
|
||
|
for _, c := range ctxts {
|
||
|
if c == nil {
|
||
|
panic("given nil context.Context")
|
||
|
}
|
||
|
|
||
|
// use earliest deadline.
|
||
|
d2, ok := c.Deadline()
|
||
|
if !ok {
|
||
|
continue
|
||
|
}
|
||
|
|
||
|
if d1 == nil || (*d1).After(d2) {
|
||
|
d1 = &d2
|
||
|
}
|
||
|
}
|
||
|
return d1
|
||
|
}
|
||
|
|
||
|
type errCtx struct {
|
||
|
dead *time.Time
|
||
|
done chan struct{}
|
||
|
err error
|
||
|
mu sync.RWMutex
|
||
|
}
|
||
|
|
||
|
func (c *errCtx) cancel(err error) {
|
||
|
c.mu.Lock()
|
||
|
defer c.mu.Unlock()
|
||
|
|
||
|
select {
|
||
|
case <-c.Done():
|
||
|
return
|
||
|
default:
|
||
|
}
|
||
|
|
||
|
c.err = err
|
||
|
close(c.done) // signal done to all
|
||
|
}
|
||
|
|
||
|
func (c *errCtx) Done() <-chan struct{} {
|
||
|
return c.done
|
||
|
}
|
||
|
|
||
|
func (c *errCtx) Err() error {
|
||
|
c.mu.Lock()
|
||
|
defer c.mu.Unlock()
|
||
|
return c.err
|
||
|
}
|
||
|
|
||
|
func (c *errCtx) Value(key interface{}) interface{} {
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
func (c *errCtx) Deadline() (deadline time.Time, ok bool) {
|
||
|
if c.dead == nil {
|
||
|
return
|
||
|
}
|
||
|
|
||
|
return *c.dead, true
|
||
|
}
|