diff options
Diffstat (limited to 'pkg/urpc')
| -rw-r--r-- | pkg/urpc/urpc.go | 19 |
1 files changed, 19 insertions, 0 deletions
diff --git a/pkg/urpc/urpc.go b/pkg/urpc/urpc.go index dfd23032c..0e9a829f6 100644 --- a/pkg/urpc/urpc.go +++ b/pkg/urpc/urpc.go @@ -170,6 +170,9 @@ type Server struct { // methods is the set of server methods. methods map[string]registeredMethod + // stoppers are all registered stoppers. + stoppers []Stopper + // clients is a map of clients. clients map[*unet.Socket]clientState @@ -195,6 +198,12 @@ func NewServerWithCallback(afterRPCCallback func()) *Server { } } +// Stopper is an optional interface, that when implemented, allows an object +// to have a callback executed when the server is shutting down. +type Stopper interface { + Stop() +} + // Register registers the given object as an RPC receiver. // // This functions is the same way as the built-in RPC package, but it does not @@ -206,6 +215,7 @@ func (s *Server) Register(obj interface{}) { defer s.mu.Unlock() typ := reflect.TypeOf(obj) + stopper, hasStop := obj.(Stopper) // If we got a pointer, deref it to the underlying object. We need this to // obtain the name of the underlying type. @@ -221,6 +231,10 @@ func (s *Server) Register(obj interface{}) { // Can't be anonymous. panic("type not named.") } + if hasStop && method.Name == "Stop" { + s.stoppers = append(s.stoppers, stopper) + continue // Legal stop method. + } prettyName := typDeref.Name() + "." + method.Name if _, ok := s.methods[prettyName]; ok { @@ -448,6 +462,11 @@ func (s *Server) Stop() { // Wait for all outstanding requests. defer s.wg.Wait() + // Call any Stop callbacks. + for _, stopper := range s.stoppers { + stopper.Stop() + } + // Close all known clients. s.mu.Lock() defer s.mu.Unlock() |