Last active
August 18, 2019 09:18
-
-
Save Tomotoes/8ff92cf557584cc70e8fe352619121bd to your computer and use it in GitHub Desktop.
Officially provided Context package
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package context | |
import ( | |
"errors" | |
"reflect" | |
"sync" | |
"time" | |
) | |
// Context 四个方法 | |
type Context interface { | |
Deadline() (deadline time.Time, ok bool) | |
Done() <-chan struct{} | |
Err() error | |
Value(key interface{}) interface{} | |
} | |
// 网络等其他领域的定义的 Error 对象 | |
type deadlineExceededError struct{} | |
func (deadlineExceededError) Error() string { return "context deadline exceeded" } | |
func (deadlineExceededError) Timeout() bool { return true } | |
func (deadlineExceededError) Temporary() bool { return true } | |
var ( | |
Canceled error = errors.New("context canceled") | |
DeadlineExceeded error = deadlineExceededError{} | |
) | |
// 定义一个空的实现 Context 类型, 返回的Background TODO 基于此类型 | |
type emptyCtx int | |
func (*emptyCtx) Deadline() (deadline time.Time, ok bool) { | |
return | |
} | |
func (*emptyCtx) Done() <-chan struct{} { | |
return nil | |
} | |
func (*emptyCtx) Err() error { | |
return nil | |
} | |
func (*emptyCtx) Value(key interface{}) interface{} { | |
return nil | |
} | |
var ( | |
background = new(emptyCtx) | |
todo = new(emptyCtx) | |
) | |
func Background() Context { | |
return background | |
} | |
func TODO() Context { | |
return todo | |
} | |
type CancelFunc func() | |
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) { | |
c := newCancelCtx(parent) | |
propagateCancel(parent, &c) | |
return &c, func() { c.cancel(true, Canceled) } | |
} | |
func newCancelCtx(parent Context) cancelCtx { | |
return cancelCtx{Context: parent} | |
} | |
func propagateCancel(parent Context, child canceler) { | |
if parent.Done() == nil { | |
// TODO or Background | |
return // parent is never canceled | |
} | |
if p, ok := parentCancelCtx(parent); ok { | |
p.mu.Lock() | |
if p.err != nil { | |
// parent has already been canceled | |
child.cancel(false, p.err) | |
} else { | |
if p.children == nil { | |
p.children = make(map[canceler]struct{}) | |
} | |
p.children[child] = struct{}{} | |
} | |
p.mu.Unlock() | |
} else { | |
go func() { | |
select { | |
case <-parent.Done(): | |
child.cancel(false, parent.Err()) | |
case <-child.Done(): | |
} | |
}() | |
} | |
} | |
func parentCancelCtx(parent Context) (*cancelCtx, bool) { | |
for { | |
switch c := parent.(type) { | |
case *cancelCtx: | |
return c, true | |
case *timerCtx: | |
return &c.cancelCtx, true | |
case *valueCtx: | |
parent = c.Context | |
default: | |
return nil, false | |
} | |
} | |
} | |
func removeChild(parent Context, child canceler) { | |
p, ok := parentCancelCtx(parent) | |
if !ok { | |
return | |
} | |
p.mu.Lock() | |
if p.children != nil { | |
delete(p.children, child) | |
} | |
p.mu.Unlock() | |
} | |
type canceler interface { | |
cancel(removeFromParent bool, err error) | |
Done() <-chan struct{} | |
} | |
// closedchan is a reusable closed channel. | |
var closedchan = make(chan struct{}) | |
func init() { | |
close(closedchan) | |
} | |
type cancelCtx struct { | |
Context | |
mu sync.Mutex // protects following fields | |
done chan struct{} // created lazily, closed by first cancel call | |
children map[canceler]struct{} // set to nil by the first cancel call | |
err error // set to non-nil by the first cancel call | |
} | |
func (c *cancelCtx) Done() <-chan struct{} { | |
c.mu.Lock() | |
if c.done == nil { | |
c.done = make(chan struct{}) | |
} | |
d := c.done | |
c.mu.Unlock() | |
return d | |
} | |
func (c *cancelCtx) Err() error { | |
c.mu.Lock() | |
err := c.err | |
c.mu.Unlock() | |
return err | |
} | |
func (c *cancelCtx) cancel(removeFromParent bool, err error) { | |
if err == nil { | |
panic("context: internal error: missing cancel error") | |
} | |
c.mu.Lock() | |
if c.err != nil { | |
c.mu.Unlock() | |
return // already canceled | |
} | |
c.err = err | |
if c.done == nil { | |
c.done = closedchan | |
} else { | |
close(c.done) | |
} | |
for child := range c.children { | |
// NOTE: acquiring the child's lock while holding parent's lock. | |
child.cancel(false, err) | |
} | |
c.children = nil | |
c.mu.Unlock() | |
if removeFromParent { | |
removeChild(c.Context, c) | |
} | |
} | |
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) { | |
if cur, ok := parent.Deadline(); ok && cur.Before(d) { | |
// The current deadline is already sooner than the new one. | |
return WithCancel(parent) | |
} | |
c := &timerCtx{ | |
cancelCtx: newCancelCtx(parent), | |
deadline: d, | |
} | |
propagateCancel(parent, c) | |
dur := time.Until(d) | |
if dur <= 0 { | |
c.cancel(true, DeadlineExceeded) // deadline has already passed | |
return c, func() { c.cancel(false, Canceled) } | |
} | |
c.mu.Lock() | |
defer c.mu.Unlock() | |
if c.err == nil { | |
c.timer = time.AfterFunc(dur, func() { | |
c.cancel(true, DeadlineExceeded) | |
}) | |
} | |
return c, func() { c.cancel(true, Canceled) } | |
} | |
// A timerCtx carries a timer and a deadline. It embeds a cancelCtx to | |
// implement Done and Err. It implements cancel by stopping its timer then | |
// delegating to cancelCtx.cancel. | |
type timerCtx struct { | |
cancelCtx | |
timer *time.Timer // Under cancelCtx.mu. | |
deadline time.Time | |
} | |
func (c *timerCtx) Deadline() (deadline time.Time, ok bool) { | |
return c.deadline, true | |
} | |
func (c *timerCtx) cancel(removeFromParent bool, err error) { | |
c.cancelCtx.cancel(false, err) | |
if removeFromParent { | |
// Remove this timerCtx from its parent cancelCtx's children. | |
removeChild(c.cancelCtx.Context, c) | |
} | |
c.mu.Lock() | |
if c.timer != nil { | |
c.timer.Stop() | |
c.timer = nil | |
} | |
c.mu.Unlock() | |
} | |
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) { | |
return WithDeadline(parent, time.Now().Add(timeout)) | |
} | |
func WithValue(parent Context, key, val interface{}) Context { | |
if key == nil { | |
panic("nil key") | |
} | |
if !reflect.TypeOf(key).Comparable() { | |
panic("key is not comparable") | |
} | |
return &valueCtx{parent, key, val} | |
} | |
type valueCtx struct { | |
Context | |
key, val interface{} | |
} | |
func (c *valueCtx) Value(key interface{}) interface{} { | |
if c.key == key { | |
return c.val | |
} | |
return c.Context.Value(key) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment