diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index 2b771a49..207f0d0c 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/internal/mtprotoedge/encrypted.go @@ -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 并加密回发。 diff --git a/internal/postresponse/callbacks.go b/internal/postresponse/callbacks.go new file mode 100644 index 00000000..95015b18 --- /dev/null +++ b/internal/postresponse/callbacks.go @@ -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() + } +}