diff options
Diffstat (limited to 'rpc/server.go')
| -rw-r--r-- | rpc/server.go | 11 |
1 files changed, 6 insertions, 5 deletions
diff --git a/rpc/server.go b/rpc/server.go index 11373b504..0f29035ed 100644 --- a/rpc/server.go +++ b/rpc/server.go @@ -125,7 +125,7 @@ func (s *Server) RegisterName(name string, rcvr interface{}) error { // If singleShot is true it will process a single request, otherwise it will handle // requests until the codec returns an error when reading a request (in most cases // an EOF). It executes requests in parallel when singleShot is false. -func (s *Server) serveRequest(codec ServerCodec, singleShot bool, options CodecOption) error { +func (s *Server) serveRequest(codec ServerCodec, singleShot bool, options CodecOption, ctx context.Context) error { var pend sync.WaitGroup defer func() { @@ -140,7 +140,8 @@ func (s *Server) serveRequest(codec ServerCodec, singleShot bool, options CodecO s.codecsMu.Unlock() }() - ctx, cancel := context.WithCancel(context.Background()) + // ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(ctx) defer cancel() // if the codec supports notification include a notifier that callbacks can use @@ -215,14 +216,14 @@ func (s *Server) serveRequest(codec ServerCodec, singleShot bool, options CodecO // stopped. In either case the codec is closed. func (s *Server) ServeCodec(codec ServerCodec, options CodecOption) { defer codec.Close() - s.serveRequest(codec, false, options) + s.serveRequest(codec, false, options, context.Background()) } // ServeSingleRequest reads and processes a single RPC request from the given codec. It will not // close the codec unless a non-recoverable error has occurred. Note, this method will return after // a single request has been processed! -func (s *Server) ServeSingleRequest(codec ServerCodec, options CodecOption) { - s.serveRequest(codec, true, options) +func (s *Server) ServeSingleRequest(codec ServerCodec, options CodecOption, ctx context.Context) { + s.serveRequest(codec, true, options, ctx) } // Stop will stop reading new requests, wait for stopPendingRequestTimeout to allow pending requests to finish, |