165 lines
4.1 KiB
V
165 lines
4.1 KiB
V
module onecontext
|
|
|
|
import context
|
|
import sync
|
|
import time
|
|
|
|
// canceled is the error returned when the cancel function is called on a merged context
|
|
pub const canceled = error('canceled context')
|
|
|
|
@[heap]
|
|
struct OneContext {
|
|
mut:
|
|
ctx context.Context
|
|
ctxs []context.Context
|
|
done chan int
|
|
err IError = none
|
|
err_mutex &sync.Mutex = sync.new_mutex()
|
|
cancel_fn context.CancelFn = unsafe { nil }
|
|
cancel_ctx context.Context
|
|
}
|
|
|
|
// merge allows to merge multiple contexts
|
|
// it returns the merged context
|
|
pub fn merge(ctx context.Context, ctxs ...context.Context) (context.Context, context.CancelFn) {
|
|
mut background := context.background()
|
|
cancel_ctx, cancel := context.with_cancel(mut background)
|
|
mut octx := &OneContext{
|
|
done: chan int{cap: 3}
|
|
ctx: ctx
|
|
ctxs: ctxs
|
|
cancel_fn: cancel
|
|
cancel_ctx: cancel_ctx
|
|
}
|
|
spawn octx.run()
|
|
return context.Context(octx), context.CancelFn(cancel)
|
|
}
|
|
|
|
// deadline returns the earliest deadline among all merged contexts,
|
|
// or none if no context has a deadline set.
|
|
pub fn (octx OneContext) deadline() ?time.Time {
|
|
mut min := time.Time{}
|
|
mut found := false
|
|
|
|
if deadline := octx.ctx.deadline() {
|
|
min = deadline
|
|
found = true
|
|
}
|
|
|
|
for ctx in octx.ctxs {
|
|
if deadline := ctx.deadline() {
|
|
if !found || deadline < min {
|
|
min = deadline
|
|
}
|
|
found = true
|
|
}
|
|
}
|
|
|
|
if !found {
|
|
return none
|
|
}
|
|
|
|
return min
|
|
}
|
|
|
|
// done returns the done channel, which is closed when the merged context is canceled.
|
|
pub fn (octx OneContext) done() chan int {
|
|
return octx.done
|
|
}
|
|
|
|
// err returns the error from the merged context, or `none` if not yet canceled.
|
|
pub fn (mut octx OneContext) err() IError {
|
|
octx.err_mutex.lock()
|
|
defer {
|
|
octx.err_mutex.unlock()
|
|
}
|
|
return octx.err
|
|
}
|
|
|
|
// value looks up a value by key across all merged contexts, returning
|
|
// the first match found or none if no context holds the key.
|
|
pub fn (octx OneContext) value(key context.Key) ?context.Any {
|
|
if value := octx.ctx.value(key) {
|
|
return value
|
|
}
|
|
|
|
for ctx in octx.ctxs {
|
|
if value := ctx.value(key) {
|
|
return value
|
|
}
|
|
}
|
|
|
|
return none
|
|
}
|
|
|
|
// run starts listening for cancellation signals from all merged contexts.
|
|
pub fn (mut octx OneContext) run() {
|
|
mut wrapped_ctx := &octx.ctx
|
|
if octx.ctxs.len == 1 {
|
|
mut first_ctx := &octx.ctxs[0]
|
|
octx.run_two_contexts(mut wrapped_ctx, mut first_ctx)
|
|
return
|
|
}
|
|
|
|
octx.run_multiple_contexts(mut wrapped_ctx)
|
|
for mut ctx in octx.ctxs {
|
|
octx.run_multiple_contexts(mut ctx)
|
|
}
|
|
}
|
|
|
|
// str returns a string representation of the OneContext.
|
|
pub fn (octx OneContext) str() string {
|
|
return ''
|
|
}
|
|
|
|
// cancel cancels the merged context with the given error, closing
|
|
// the done channel and propagating cancellation to the underlying context.
|
|
pub fn (mut octx OneContext) cancel(err IError) {
|
|
octx.cancel_fn()
|
|
octx.err_mutex.lock()
|
|
octx.err = err
|
|
octx.err_mutex.unlock()
|
|
if !octx.done.closed {
|
|
octx.done <- 0
|
|
octx.done.close()
|
|
}
|
|
}
|
|
|
|
// run_two_contexts spawns a listener that cancels the merged context
|
|
// when either of the two given contexts is done.
|
|
pub fn (mut octx OneContext) run_two_contexts(mut ctx1 context.Context, mut ctx2 context.Context) {
|
|
octx_cancel_done := octx.cancel_ctx.done()
|
|
c1done := ctx1.done()
|
|
c2done := ctx2.done()
|
|
spawn fn (mut octx OneContext, octx_cancel_done chan int, c1done chan int, c2done chan int, mut ctx1 context.Context, mut ctx2 context.Context) {
|
|
select {
|
|
_ := <-octx_cancel_done {
|
|
octx.cancel(canceled)
|
|
}
|
|
_ := <-c1done {
|
|
octx.cancel(ctx1.err())
|
|
}
|
|
_ := <-c2done {
|
|
octx.cancel(ctx2.err())
|
|
}
|
|
}
|
|
}(mut &octx, octx_cancel_done, c1done, c2done, mut ctx1, mut ctx2)
|
|
}
|
|
|
|
// run_multiple_contexts spawns a listener that cancels the merged context
|
|
// when the given context is done.
|
|
pub fn (mut octx OneContext) run_multiple_contexts(mut ctx context.Context) {
|
|
octx_cancel_done := octx.cancel_ctx.done()
|
|
cdone := ctx.done()
|
|
spawn fn (mut octx OneContext, octx_cancel_done chan int, cdone chan int, mut ctx context.Context) {
|
|
select {
|
|
_ := <-octx_cancel_done {
|
|
octx.cancel(canceled)
|
|
}
|
|
_ := <-cdone {
|
|
octx.cancel(ctx.err())
|
|
}
|
|
}
|
|
}(mut &octx, octx_cancel_done, cdone, mut ctx)
|
|
}
|