cubefs

1
package clock
2

3
import (
4
	"context"
5
	"fmt"
6
	"sync"
7
	"time"
8
)
9

10
func (m *Mock) WithTimeout(parent context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
11
	return m.WithDeadline(parent, m.Now().Add(timeout))
12
}
13

14
func (m *Mock) WithDeadline(parent context.Context, deadline time.Time) (context.Context, context.CancelFunc) {
15
	if cur, ok := parent.Deadline(); ok && cur.Before(deadline) {
16
		// The current deadline is already sooner than the new one.
17
		return context.WithCancel(parent)
18
	}
19
	ctx := &timerCtx{clock: m, parent: parent, deadline: deadline, done: make(chan struct{})}
20
	propagateCancel(parent, ctx)
21
	dur := m.Until(deadline)
22
	if dur <= 0 {
23
		ctx.cancel(context.DeadlineExceeded) // deadline has already passed
24
		return ctx, func() {}
25
	}
26
	ctx.Lock()
27
	defer ctx.Unlock()
28
	if ctx.err == nil {
29
		ctx.timer = m.AfterFunc(dur, func() {
30
			ctx.cancel(context.DeadlineExceeded)
31
		})
32
	}
33
	return ctx, func() { ctx.cancel(context.Canceled) }
34
}
35

36
// propagateCancel arranges for child to be canceled when parent is.
37
func propagateCancel(parent context.Context, child *timerCtx) {
38
	if parent.Done() == nil {
39
		return // parent is never canceled
40
	}
41
	go func() {
42
		select {
43
		case <-parent.Done():
44
			child.cancel(parent.Err())
45
		case <-child.Done():
46
		}
47
	}()
48
}
49

50
type timerCtx struct {
51
	sync.Mutex
52

53
	clock    Clock
54
	parent   context.Context
55
	deadline time.Time
56
	done     chan struct{}
57

58
	err   error
59
	timer *Timer
60
}
61

62
func (c *timerCtx) cancel(err error) {
63
	c.Lock()
64
	defer c.Unlock()
65
	if c.err != nil {
66
		return // already canceled
67
	}
68
	c.err = err
69
	close(c.done)
70
	if c.timer != nil {
71
		c.timer.Stop()
72
		c.timer = nil
73
	}
74
}
75

76
func (c *timerCtx) Deadline() (deadline time.Time, ok bool) { return c.deadline, true }
77

78
func (c *timerCtx) Done() <-chan struct{} { return c.done }
79

80
func (c *timerCtx) Err() error { return c.err }
81

82
func (c *timerCtx) Value(key interface{}) interface{} { return c.parent.Value(key) }
83

84
func (c *timerCtx) String() string {
85
	return fmt.Sprintf("clock.WithDeadline(%s [%s])", c.deadline, c.deadline.Sub(c.clock.Now()))
86
}
87