Merge pull request #4573 from Xopherus/fix-4413

Close multiplexer when context is cancelled
hashicorp · Aug 14, 2018 · a24dd6d · a24dd6d
2 parents dafe580 + 3f24197
commit a24dd6d
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/nomad/rpc.go b/nomad/rpc.go
@@ -208,6 +208,11 @@ func (s *Server) handleMultiplex(ctx context.Context, conn net.Conn, rpcCtx *RPC
 	s.setupRpcServer(rpcServer, rpcCtx)
 
 	for {
+		// stop handling connections if context was cancelled
+		if ctx.Err() != nil {
+			return
+		}
+
 		sub, err := server.Accept()
 		if err != nil {
 			if err != io.EOF {
@@ -311,6 +316,11 @@ func (s *Server) handleMultiplexV2(ctx context.Context, conn net.Conn, rpcCtx *R
 	s.setupRpcServer(rpcServer, rpcCtx)
 
 	for {
+		// stop handling connections if context was cancelled
+		if ctx.Err() != nil {
+			return
+		}
+
 		// Accept a new stream
 		sub, err := server.Accept()
 		if err != nil {