Vihren Atlas

Atlas › Test

TestStreamBatcher_AddTimeout

go.temporal.io/server/common/stream_batcher/TestStreamBatcher_AddTimeout

Package
go.temporal.io/server/common/stream_batcher
Suite / test hierarchy
TestStreamBatcher_AddTimeout
Test
TestStreamBatcher_AddTimeout
Introduced at
batcher.go ×1 Joint frontier
Covered ranges
44
Covered lines
182
Covered files
4

Covered source

Expand a file to inspect source; the > gutter marks covered lines.

go.temporal.io/server/common/clock/event_time_source.go 78 covered LOC · 23 ranges

Open complete file

39
40 // NewEventTimeSource returns a EventTimeSource with the current time set to Unix zero: 1970-01-01 00:00:00 +0000 UTC.
41 > func NewEventTimeSource() *EventTimeSource { event_time_source.go
42 > return &EventTimeSource{
43 > now: time.Unix(0, 0),
44 > }
45 > }
46
47 // Some clients depend on the fact that the runtime's timers do _not_ run synchronously.
55
56 // Now return the current time.
57 > func (ts *EventTimeSource) Now() time.Time { event_time_source.go
58 > ts.mu.RLock()
59 > defer ts.mu.RUnlock()
60 >
61 > return ts.now
62 > }
63
64 func (ts *EventTimeSource) Since(t time.Time) time.Duration {
71 // wrap all such calls in a goroutine. If the duration is non-positive, the callback will fire immediately before
72 // AfterFunc returns.
73 > func (ts *EventTimeSource) AfterFunc(d time.Duration, f func()) Timer { event_time_source.go
74 > if d < 0 {
75 d = 0
76 }
77 > timer := &fakeTimer{timeSource: ts, deadline: ts.Now().Add(d), callback: f} event_time_source.go
78 > ts.addTimer(timer)
79 > return timer
80 }
81
82 // NewTimer creates a Timer that will send the current time on a channel after at least
83 // duration d. It returns the channel and the Timer.
84 > func (ts *EventTimeSource) NewTimer(d time.Duration) (<-chan time.Time, Timer) { event_time_source.go
85 > c := make(chan time.Time, 1)
86 > // we can't call ts.Now() from the callback so just calculate what it should be
87 > target := ts.Now().Add(d)
88 > timer := &fakeTimer{
89 > timeSource: ts,
90 > deadline: target,
91 > callback: func() { c <- target },
92 c: c,
93 }
94 > ts.addTimer(timer) event_time_source.go
95 > return c, timer
96 }
97
98 > func (ts *EventTimeSource) addTimer(t *fakeTimer) { event_time_source.go
99 > ts.mu.Lock()
100 > defer ts.mu.Unlock()
101 > t.index = len(ts.timers)
102 > ts.timers = append(ts.timers, t)
103 > ts.fireTimers()
104 > }
105
106 // Update the fake current time. It returns the timeSource so that you can chain calls like this:
125
126 // AdvanceNext advances to the next timer.
127 > func (ts *EventTimeSource) AdvanceNext() { event_time_source.go
128 > ts.mu.Lock()
129 > defer ts.mu.Unlock()
130 >
131 > if len(ts.timers) == 0 {
132 return
133 }
134 // just do linear search, this is efficient enough for now
135 > tmin := ts.timers[0].deadline event_time_source.go
136 > for _, t := range ts.timers[1:] {
137 > tmin = util.MinTime(tmin, t.deadline) event_time_source.go
138 > }
139 > ts.now = tmin event_time_source.go
140 > ts.fireTimers()
141 }
142
150
151 // Sleep is a convenience function for waiting on a new timer.
152 > func (ts *EventTimeSource) Sleep(d time.Duration) { event_time_source.go
153 > t, _ := ts.NewTimer(d)
154 > <-t
155 > }
156
157 // fireTimers fires all timers that are ready.
158 > func (ts *EventTimeSource) fireTimers() { event_time_source.go
159 > n := 0
160 > for _, t := range ts.timers {
161 > if t.deadline.After(ts.now) { event_time_source.go
162 > ts.timers[n] = t event_time_source.go
163 > t.index = n
164 > n++
165 > } else { event_time_source.go
166 > if ts.async { event_time_source.go
167 go t.callback()
168 > } else { event_time_source.go
169 > t.callback() event_time_source.go
170 > }
171 > t.done = true event_time_source.go
172 }
173 }
174 > ts.timers = ts.timers[:n] event_time_source.go
175 }
176
200
201 // Stop the timer. Returns true if the timer was active.
202 > func (t *fakeTimer) Stop() bool { event_time_source.go
203 > t.timeSource.mu.Lock()
204 > defer t.timeSource.mu.Unlock()
205 >
206 > if t.done {
207 return false
208 }
209
210 > i := t.index event_time_source.go
211 > timers := t.timeSource.timers
212 >
213 > timers[i] = timers[len(timers)-1] // swap with last timer
214 > timers[i].index = i // update index of swapped timer
215 > timers = timers[:len(timers)-1] // shrink list
216 >
217 > t.timeSource.timers = timers
218 > t.done = true // ensure that the timer is not reused
219 >
220 > return true
221 }
go.temporal.io/server/common/stream_batcher/batcher.go 65 covered LOC · 11 ranges

Open complete file

45 // NewBatcher creates a Batcher. `fn` is the processing function, `opts` are the timing options.
46 // `clock` is usually clock.NewRealTimeSource but can be a fake time source for testing.
47 > func NewBatcher[T, R any](fn func([]T) R, opts BatcherOptions, timeSource clock.TimeSource) *Batcher[T, R] { batcher.go
48 > return &Batcher[T, R]{
49 > fn: fn,
50 > opts: opts,
51 > timeSource: timeSource,
52 > submitC: make(chan batchPair[T, R]),
53 > }
54 > }
55
56 // Add adds an item to the stream and returns when it has been processed, or if the context is
58 // for the whole batch that the item ended up in, and a context error. Even if Add returns a
59 // context error, the item may still be processed in the future!
60 > func (b *Batcher[T, R]) Add(ctx context.Context, t T) (R, error) { batcher.go
61 > resp := make(chan R, 1)
62 > pair := batchPair[T, R]{resp: resp, item: t}
63 >
64 > for {
65 > runningC := b.running.Load()
66 > for runningC == nil {
67 > // goroutine is not running, try to start it
68 > newRunningC := make(chan struct{})
69 > if b.running.CompareAndSwap(nil, &newRunningC) {
70 > // we were the first one to notice the nil, start it now
71 > go b.loop(&newRunningC)
72 > }
73 // if CompareAndSwap failed, someone else was calling Add at the same time and
74 // started the goroutine already. reload to get the new running channel.
75 > runningC = b.running.Load() batcher.go
76 }
77
78 > select { batcher.go
79 case <-(*runningC):
80 // we loaded a non-nil running channel, but it closed while we're waiting to
81 // submit. the goroutine must have just exited. try again.
82 continue
83 > case b.submitC <- pair: batcher.go
84 > select {
85 > case r := <-resp:
86 > return r, nil
87 > case <-ctx.Done(): batcher.go
88 > var zeroR R
89 > return zeroR, ctx.Err()
90 }
91 case <-ctx.Done():
96 }
97
98 > func (b *Batcher[T, R]) loop(runningC *chan struct{}) { batcher.go
99 > defer func() {
100 // store nil so that Add knows it should start a goroutine
101 b.running.Store(nil)
106 }()
107
108 > var items []T batcher.go
109 > var resps []chan R
110 > for {
111 > clear(items)
112 > clear(resps)
113 > items, resps = items[:0], resps[:0]
114 >
115 > // wait for first item. if no item after a while, exit the goroutine
116 > idleC, idleT := b.timeSource.NewTimer(b.opts.IdleTime)
117 > select {
118 > case pair := <-b.submitC:
119 > items = append(items, pair.item)
120 > resps = append(resps, pair.resp)
121 case <-idleC:
122 return
123 }
124 > idleT.Stop() batcher.go
125 >
126 > // try to add more items. stop after a gap of MinDelay, total time of MaxDelay,
127 > // or MaxItems items.
128 > maxWaitC, maxWaitT := b.timeSource.NewTimer(b.opts.MaxDelay)
129 > loop:
130 > for len(items) < b.opts.MaxItems {
131 > gapC, gapT := b.timeSource.NewTimer(b.opts.MinDelay)
132 > select {
133 case pair := <-b.submitC:
134 items = append(items, pair.item)
135 resps = append(resps, pair.resp)
136 > case <-gapC: batcher.go
137 > break loop
138 case <-maxWaitC:
139 gapT.Stop()
142 gapT.Stop()
143 }
144 > maxWaitT.Stop() batcher.go
145 >
146 > // process batch
147 > r := b.fn(items)
148 >
149 > // send the single response to all items in the batch
150 > for _, resp := range resps {
151 > resp <- r
152 > }
153 }
154 }
go.temporal.io/server/common/clock/context.go 34 covered LOC · 7 ranges

Open complete file

20 }
21
22 > func (ctx *ctxWithDeadline) Done() <-chan struct{} { context.go
23 > return ctx.done
24 > }
25
26 > func (ctx *ctxWithDeadline) Err() error { context.go
27 > select {
28 > case <-ctx.done:
29 > return ctx.err
30 default:
31 return nil
33 }
34
35 > func (ctx *ctxWithDeadline) deadlineExceeded() { context.go
36 > ctx.once.Do(func() {
37 > ctx.err = context.DeadlineExceeded
38 > close(ctx.done)
39 > })
40 }
41
42 > func (ctx *ctxWithDeadline) cancel() { context.go
43 > ctx.once.Do(func() {
44 > // We'd like to call ctx.timer.Stop() here, but we can't: the time source may call context.go
45 > // deadlineExceeded while holding its lock, which acquires the once mutex. Here we have
46 > // the once mutex and want to cancel a timer, which would create a potential lock
47 > // cycle. So just leave the timer as a no-op.
48 > ctx.err = context.Canceled
49 > close(ctx.done)
50 > })
51 }
52
55 deadline time.Time,
56 timeSource TimeSource,
57 > ) (context.Context, context.CancelFunc) { context.go
58 > ctxd := &ctxWithDeadline{
59 > Context: ctx,
60 > deadline: deadline,
61 > done: make(chan struct{}),
62 > }
63 > timer := timeSource.AfterFunc(deadline.Sub(timeSource.Now()), ctxd.deadlineExceeded)
64 > ctxd.timer = timer
65 > return ctxd, ctxd.cancel
66 > }
67
68 func ContextWithTimeout(
70 timeout time.Duration,
71 timeSource TimeSource,
72 > ) (context.Context, context.CancelFunc) { context.go
73 > return ContextWithDeadline(ctx, timeSource.Now().Add(timeout), timeSource)
74 > }
go.temporal.io/server/common/util/util.go 5 covered LOC · 3 ranges

Open complete file

12
13 // MinTime returns the earlier of two given time.Time
14 > func MinTime(a, b time.Time) time.Time { util.go
15 > if a.Before(b) {
16 > return a util.go
17 > }
18 > return b util.go
19 }
20