72 lines
1.4 KiB
Go
72 lines
1.4 KiB
Go
package postresponse
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
)
|
|
|
|
type callback func()
|
|
|
|
type callbacksKey struct{}
|
|
|
|
type callbacks struct {
|
|
mu sync.Mutex
|
|
list []callback
|
|
}
|
|
|
|
func WithCallbacks(ctx context.Context) context.Context {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
if _, ok := ctx.Value(callbacksKey{}).(*callbacks); ok {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, callbacksKey{}, &callbacks{})
|
|
}
|
|
|
|
func Register(ctx context.Context, cb func()) bool {
|
|
if cb == nil {
|
|
return false
|
|
}
|
|
cbs, ok := ctx.Value(callbacksKey{}).(*callbacks)
|
|
if !ok || cbs == nil {
|
|
return false
|
|
}
|
|
cbs.mu.Lock()
|
|
cbs.list = append(cbs.list, cb)
|
|
cbs.mu.Unlock()
|
|
return true
|
|
}
|
|
|
|
func Run(ctx context.Context) {
|
|
run := Take(ctx)
|
|
if run != nil {
|
|
run()
|
|
}
|
|
}
|
|
|
|
// Take transfers ownership of every currently registered callback to the caller.
|
|
// The returned function is idempotent and may safely outlive the request context.
|
|
// MTProto uses this to release a business worker after admitting rpc_result while
|
|
// still delaying follow-up updates until the result reaches the reliable stream.
|
|
func Take(ctx context.Context) func() {
|
|
cbs, ok := ctx.Value(callbacksKey{}).(*callbacks)
|
|
if !ok || cbs == nil {
|
|
return nil
|
|
}
|
|
cbs.mu.Lock()
|
|
list := append([]callback(nil), cbs.list...)
|
|
cbs.list = nil
|
|
cbs.mu.Unlock()
|
|
if len(list) == 0 {
|
|
return nil
|
|
}
|
|
var once sync.Once
|
|
return func() {
|
|
once.Do(func() {
|
|
for _, cb := range list {
|
|
cb()
|
|
}
|
|
})
|
|
}
|
|
}
|