Merge pull request #6222 from avagin:stop

PiperOrigin-RevId: 381561785

1 files changed, 3 insertions, 2 deletions

diff --git a/pkg/control/server/server.go b/pkg/control/server/server.go
index 629dae8f4..889568177 100644
--- a/pkg/control/server/server.go
+++ b/pkg/control/server/server.go
@@ -22,6 +22,7 @@ package server
 
 import (
 	"os"
+	"time"
 
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sync"
@@ -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(timeout time.Duration) {
 	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(timeout)
 }
 
 // StartServing starts listening for connect and spawns the main service