Internal change.

PiperOrigin-RevId: 233124342
Change-Id: Id4b4857af89815ffb9254cc30df4244b2768d9f2

1 files 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,