2 files changed, 30 insertions, 18 deletions
diff --git a/pkg/control/server/server.go b/pkg/control/server/server.go
index 629dae8f4..eca06791c 100644
--- a/pkg/control/server/server.go
+++ b/pkg/control/server/server.go
@@ -21,6 +21,7 @@ implementations of the control interface.
 package server
 
 import (
+	"context"
 	"os"
 
 	"gvisor.dev/gvisor/pkg/log"
@@ -65,13 +66,13 @@ func (s *Server) Wait() {
 
 // Stop stops the server. Note that this function should only be called once
 // and the server should not be used afterwards.
-func (s *Server) Stop() {
+func (s *Server) Stop(ctx context.Context) {
 	s.socket.Close()
 	s.Wait()
 
 	// This will cause existing clients to be terminated safely. If the
 	// registered handlers have a Stop callback, it will be called.
-	s.server.Stop()
+	s.server.Stop(ctx)
 }
 
 // StartServing starts listening for connect and spawns the main service
diff --git a/pkg/urpc/urpc.go b/pkg/urpc/urpc.go
index 0e9a829f6..7872d6fa1 100644
--- a/pkg/urpc/urpc.go
+++ b/pkg/urpc/urpc.go
@@ -20,6 +20,7 @@ package urpc
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -458,29 +459,39 @@ func (s *Server) StartHandling(client *unet.Socket) {
 // No new requests should be initiated after calling Stop. Existing clients
 // will be closed after completing any pending RPCs. This method will block
 // until all clients have disconnected.
-func (s *Server) Stop() {
-	// Wait for all outstanding requests.
-	defer s.wg.Wait()
+func (s *Server) Stop(ctx context.Context) {
+	done := make(chan bool)
 
 	// Call any Stop callbacks.
 	for _, stopper := range s.stoppers {
 		stopper.Stop()
 	}
+	go func() {
+		select {
+		case <-done:
+			return
+		case <-ctx.Done():
+		}
 
-	// Close all known clients.
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	for client, state := range s.clients {
-		switch state {
-		case idle:
-			// Close connection now.
-			client.Close()
-			s.clients[client] = closed
-		case processing:
-			// Request close when done.
-			s.clients[client] = closeRequested
+		// Close all known clients.
+		s.mu.Lock()
+		defer s.mu.Unlock()
+		for client, state := range s.clients {
+			switch state {
+			case idle:
+				// Close connection now.
+				client.Close()
+				s.clients[client] = closed
+			case processing:
+				// Request close when done.
+				s.clients[client] = closeRequested
+			}
 		}
-	}
+	}()
+	// Wait for all outstanding requests.
+	s.wg.Wait()
+	done <- true
+
 }
 
 // Client is a urpc client.