From b2aa213dd2e19d88b4edf839d20e5bf51534be62 Mon Sep 17 00:00:00 2001 From: Christopher Koch Date: Fri, 8 Feb 2019 14:13:17 -0800 Subject: [PATCH] Internal change. PiperOrigin-RevId: 233124342 Change-Id: Id4b4857af89815ffb9254cc30df4244b2768d9f2 --- pkg/urpc/urpc.go | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) 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,