Merge pull request #6222 from avagin:stop

PiperOrigin-RevId: 381561785

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

pkg/urpc/urpc.go

@@ -27,6 +27,7 @@ import (
 	"os"
 	"reflect"
 	"runtime"
+	"time"
 
 	"gvisor.dev/gvisor/pkg/fd"
 	"gvisor.dev/gvisor/pkg/log"
@@ -458,29 +459,45 @@ 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()
-
+//
+// timeout is the time for clients to complete ongoing RPCs.
+func (s *Server) Stop(timeout time.Duration) {
 	// Call any Stop callbacks.
 	for _, stopper := range s.stoppers {
 		stopper.Stop()
 	}
 
-	// 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
+	done := make(chan bool, 1)
+	go func() {
+		if timeout != 0 {
+			timer := time.NewTicker(timeout)
+			defer timer.Stop()
+			select {
+			case <-done:
+				return
+			case <-timer.C:
+			}
 		}
-	}
+
+		// 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.

runsc/boot/controller.go

@@ -18,6 +18,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	gtime "time"
 
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"golang.org/x/sys/unix"
@@ -165,8 +166,11 @@ func newController(fd int, l *Loader) (*controller, error) {
 	return ctrl, nil
 }
 
+// stopRPCTimeout is the time for clients to complete ongoing RPCs.
+const stopRPCTimeout = 15 * gtime.Second
+
 func (c *controller) stop() {
-	c.srv.Stop()
+	c.srv.Stop(stopRPCTimeout)
 }
 
 // containerManager manages sandbox containers.

runsc/boot/loader_test.go

@@ -214,7 +214,7 @@ func doStartSignal(t *testing.T, vfsEnabled bool) {
 
 	// We aren't going to wait on this application, so the control server
 	// needs to be shut down manually.
-	defer l.ctrl.srv.Stop()
+	defer l.ctrl.srv.Stop(time.Hour)
 
 	// Start a goroutine that calls WaitForStartSignal and writes to a
 	// channel when it returns.

runsc/sandbox/sandbox.go

@@ -797,7 +797,13 @@ func (s *Sandbox) Wait(cid string) (unix.WaitStatus, error) {
 		// Try the Wait RPC to the sandbox.
 		var ws unix.WaitStatus
 		err = conn.Call(boot.ContainerWait, &cid, &ws)
+		conn.Close()
 		if err == nil {
+			if s.IsRootContainer(cid) {
+				if err := s.waitForStopped(); err != nil {
+					return unix.WaitStatus(0), err
+				}
+			}
 			// It worked!
 			return ws, nil
 		}