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/compat/layerwire"
|
||||||
"telesrv/internal/observability/dbtrace"
|
"telesrv/internal/observability/dbtrace"
|
||||||
|
"telesrv/internal/postresponse"
|
||||||
"telesrv/internal/store"
|
"telesrv/internal/store"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -495,6 +496,7 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buf
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
ctx = postresponse.WithCallbacks(ctx)
|
||||||
ctx, dbStats := dbtrace.WithStats(ctx)
|
ctx, dbStats := dbtrace.WithStats(ctx)
|
||||||
start := s.clock.Now()
|
start := s.clock.Now()
|
||||||
result, err := s.rpc.Dispatch(ctx, c.authKeyID, c.sessionID, b)
|
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...)
|
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 并加密回发。
|
// 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