84 lines
1.4 KiB
Go
84 lines
1.4 KiB
Go
package internal
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type contextKey string
|
|
|
|
var (
|
|
contextKeyApp contextKey = "application-context"
|
|
)
|
|
|
|
func WithAppContext(ctx context.Context, appCtx *AppContext) context.Context {
|
|
appCtx.parent = ctx
|
|
return context.WithValue(ctx, contextKeyApp, appCtx)
|
|
}
|
|
|
|
func GetAppContext(ctx context.Context) *AppContext {
|
|
return ctx.Value(contextKeyApp).(*AppContext)
|
|
}
|
|
|
|
type AppContext struct {
|
|
parent context.Context
|
|
|
|
Logger *zap.Logger
|
|
Data map[string]any
|
|
Cancel context.CancelFunc
|
|
|
|
task []task
|
|
taskCnt atomic.Int32
|
|
taskDone atomic.Int32
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
func (c AppContext) Init() AppContext {
|
|
c.Data = make(map[string]any)
|
|
return c
|
|
}
|
|
|
|
func (c *AppContext) GetContext() context.Context {
|
|
ctx, cancel := context.WithCancel(c.parent)
|
|
c.task = append(c.task, task{
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
})
|
|
return ctx
|
|
}
|
|
|
|
// DoTask spawn task on a goroutine
|
|
func (p *AppContext) DoTask(f func(ctx context.Context)) {
|
|
ctx := p.GetContext()
|
|
p.wg.Add(1)
|
|
p.taskCnt.Add(1)
|
|
go func() {
|
|
defer p.wg.Done()
|
|
defer p.taskDone.Add(1)
|
|
f(ctx)
|
|
}()
|
|
}
|
|
|
|
func (p *AppContext) TaskCount() int {
|
|
return int(p.taskCnt.Load())
|
|
}
|
|
|
|
func (p *AppContext) TaskDone() int {
|
|
return int(p.taskDone.Load())
|
|
}
|
|
|
|
func (p *AppContext) TaskRemaining() int {
|
|
return p.TaskCount() - p.TaskDone()
|
|
}
|
|
|
|
func (p *AppContext) WaitTask() {
|
|
p.wg.Wait()
|
|
}
|
|
|
|
type task struct {
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
}
|