diff --git a/pkg/urpc/urpc.go b/pkg/urpc/urpc.go index 6d528b180..719f0e92f 100644 --- a/pkg/urpc/urpc.go +++ b/pkg/urpc/urpc.go @@ -175,13 +175,23 @@ type Server struct { // wg is a wait group for all outstanding clients. wg sync.WaitGroup + + // afterRPCCallback is called after each RPC is successfully completed. + afterRPCCallback func() } // NewServer returns a new server. func NewServer() *Server { + return NewServerWithCallback(nil) +} + +// NewServerWithCallback returns a new server, who upon completion of each RPC +// calls the given function. +func NewServerWithCallback(afterRPCCallback func()) *Server { return &Server{ - methods: make(map[string]registeredMethod), - clients: make(map[*unet.Socket]clientState), + methods: make(map[string]registeredMethod), + clients: make(map[*unet.Socket]clientState), + afterRPCCallback: afterRPCCallback, } } @@ -274,6 +284,11 @@ func (s *Server) handleOne(client *unet.Socket) error { return err } + defer func() { + if s.afterRPCCallback != nil { + s.afterRPCCallback() + } + }() // Explicitly close all these files after the call. // // This is also explicitly a reference to the files after the call,