feat: add post-response callbacks
This commit is contained in:
parent
2f399d5df2
commit
a246fd5417
2 changed files with 60 additions and 1 deletions
|
|
@ -22,6 +22,7 @@ import (
|
|||
|
||||
"telesrv/internal/compat/layerwire"
|
||||
"telesrv/internal/observability/dbtrace"
|
||||
"telesrv/internal/postresponse"
|
||||
"telesrv/internal/store"
|
||||
)
|
||||
|
||||
|
|
@ -495,6 +496,7 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buf
|
|||
return nil
|
||||
}
|
||||
|
||||
ctx = postresponse.WithCallbacks(ctx)
|
||||
ctx, dbStats := dbtrace.WithStats(ctx)
|
||||
start := s.clock.Now()
|
||||
result, err := s.rpc.Dispatch(ctx, c.authKeyID, c.sessionID, b)
|
||||
|
|
@ -547,7 +549,11 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buf
|
|||
}
|
||||
|
||||
s.log.Info("RPC handled", fields...)
|
||||
return s.sendResult(ctx, c, msgID, result)
|
||||
if err := s.sendResult(ctx, c, msgID, result); err != nil {
|
||||
return err
|
||||
}
|
||||
postresponse.Run(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
// sendResult 把 RPC 结果包成 rpc_result 并加密回发。
|
||||
|
|
|
|||
53
internal/postresponse/callbacks.go
Normal file
53
internal/postresponse/callbacks.go
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
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) {
|
||||
cbs, ok := ctx.Value(callbacksKey{}).(*callbacks)
|
||||
if !ok || cbs == nil {
|
||||
return
|
||||
}
|
||||
cbs.mu.Lock()
|
||||
list := append([]callback(nil), cbs.list...)
|
||||
cbs.list = nil
|
||||
cbs.mu.Unlock()
|
||||
for _, cb := range list {
|
||||
cb()
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue