Reset control server on restore with the new kernel.

This is necessary to ensure that the control server is aware of the new kernel
that was installed on restore.

PiperOrigin-RevId: 644792458

pkg/control/server/server.go

@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"sync/atomic"
 	"time"
 
 	"golang.org/x/sys/unix"
@@ -42,7 +43,7 @@ type Server struct {
 	socket *unet.ServerSocket
 
 	// server is our rpc server.
-	server *urpc.Server
+	server atomic.Pointer[urpc.Server]
 
 	// wg waits for the accept loop to terminate.
 	wg sync.WaitGroup
@@ -50,9 +51,18 @@ type Server struct {
 
 // New returns a new bound control server.
 func New(socket *unet.ServerSocket) *Server {
-	return &Server{
+	s := &Server{
 		socket: socket,
-		server: urpc.NewServer(),
+	}
+	s.server.Store(urpc.NewServer())
+	return s
+}
+
+// ResetServer resets the server, clearing all registered objects. It stops the
+// old server asynchronously.
+func (s *Server) ResetServer() {
+	if old := s.server.Swap(urpc.NewServer()); old != nil {
+		go old.Stop(0)
 	}
 }
 
@@ -75,7 +85,7 @@ func (s *Server) Stop(timeout time.Duration) {
 
 	// This will cause existing clients to be terminated safely. If the
 	// registered handlers have a Stop callback, it will be called.
-	s.server.Stop(timeout)
+	s.server.Load().Stop(timeout)
 }
 
 // StartServing starts listening for connect and spawns the main service
@@ -107,13 +117,13 @@ func (s *Server) serve() {
 		}
 
 		// Handle the connection non-blockingly.
-		s.server.StartHandling(conn)
+		s.server.Load().StartHandling(conn)
 	}
 }
 
 // Register registers a specific control interface with the server.
 func (s *Server) Register(obj any) {
-	s.server.Register(obj)
+	s.server.Load().Register(obj)
 }
 
 // CreateFromFD creates a new control bound to the given 'fd'. It has no

pkg/urpc/urpc.go

@@ -460,7 +460,10 @@ func (s *Server) StartHandling(client *unet.Socket) {
 // will be closed after completing any pending RPCs. This method will block
 // until all clients have disconnected.
 //
-// timeout is the time for clients to complete ongoing RPCs.
+// timeout is the time for clients to complete pending RPCs. After timeout
+// expires, all clients are drained (i.e. their ongoing RPC is allowed to
+// complete) and closed. Any new RPCs will not be processed. Note that ongoing
+// RPCs are *not* interrupted or cancelled.
 func (s *Server) Stop(timeout time.Duration) {
 	// Call any Stop callbacks.
 	for _, stopper := range s.stoppers {
@@ -569,7 +572,7 @@ func marshal(s *unet.Socket, v any, fs []*os.File) error {
 	return nil
 }
 
-// unmarhsal receives an FD (optional) and unmarshals the given struct.
+// unmarshal receives an FD (optional) and unmarshals the given struct.
 func unmarshal(s *unet.Socket, v any) ([]*os.File, error) {
 	// Receive a single byte.
 	r := s.Reader(true)

runsc/boot/controller.go

@@ -180,29 +180,42 @@ func newController(fd int, l *Loader) (*controller, error) {
 		},
 		srv: srv,
 	}
-	ctrl.srv.Register(ctrl.manager)
-	ctrl.srv.Register(&control.Cgroups{Kernel: l.k})
-	ctrl.srv.Register(&control.Lifecycle{Kernel: l.k})
-	ctrl.srv.Register(&control.Logging{})
-	ctrl.srv.Register(&control.Proc{Kernel: l.k})
-	ctrl.srv.Register(&control.State{Kernel: l.k})
-	ctrl.srv.Register(&control.Usage{Kernel: l.k})
-	ctrl.srv.Register(&control.Metrics{})
-	ctrl.srv.Register(&debug{})
+	ctrl.registerHandlers()
+	return ctrl, nil
+}
+
+func (c *controller) registerHandlers() {
+	l := c.manager.l
+	c.srv.Register(c.manager)
+	c.srv.Register(&control.Cgroups{Kernel: l.k})
+	c.srv.Register(&control.Lifecycle{Kernel: l.k})
+	c.srv.Register(&control.Logging{})
+	c.srv.Register(&control.Proc{Kernel: l.k})
+	c.srv.Register(&control.State{Kernel: l.k})
+	c.srv.Register(&control.Usage{Kernel: l.k})
+	c.srv.Register(&control.Metrics{})
+	c.srv.Register(&debug{})
 
 	if eps, ok := l.k.RootNetworkNamespace().Stack().(*netstack.Stack); ok {
-		ctrl.srv.Register(&Network{
+		c.srv.Register(&Network{
 			Stack:  eps.Stack,
 			Kernel: l.k,
 		})
 	}
 	if l.root.conf.ProfileEnable {
-		ctrl.srv.Register(control.NewProfile(l.k))
+		c.srv.Register(control.NewProfile(l.k))
 	}
-	return ctrl, nil
 }
 
-// stopRPCTimeout is the time for clients to complete ongoing RPCs.
+// refreshHandlers resets the server and re-registers all handlers using l.
+// Useful when l.k has been replaced (e.g. during a restore).
+func (c *controller) refreshHandlers() {
+	c.srv.ResetServer()
+	c.registerHandlers()
+}
+
+// stopRPCTimeout is the time for clients to finish making any RPCs. Note that
+// ongoing RPCs after this timeout still run to completion.
 const stopRPCTimeout = 15 * gtime.Second
 
 func (c *controller) stop() {

runsc/boot/restore.go

@@ -246,7 +246,6 @@ func (r *restorer) restore(l *Loader) error {
 	l.watchdog = dog
 	l.root.procArgs = kernel.CreateProcessArgs{}
 	l.restore = true
-
 	l.sandboxID = l.root.cid
 
 	l.mu.Lock()
@@ -292,6 +291,9 @@ func (r *restorer) restore(l *Loader) error {
 
 	l.k.RestoreContainerMapping(l.containerIDs)
 
+	// Refresh the control server with the newly created kernel.
+	l.ctrl.refreshHandlers()
+
 	// Release `l.mu` before calling into callbacks.
 	cu.Clean()
 

runsc/container/container_test.go

@@ -2787,6 +2787,39 @@ func TestUsageFD(t *testing.T) {
 	if total == 0 {
 		t.Errorf("UsageFD total got zero")
 	}
+
+	// Set the image path, which is where the checkpoint image will be saved.
+	dir, err := ioutil.TempDir(testutil.TmpDir(), "checkpoint")
+	if err != nil {
+		t.Fatalf("ioutil.TempDir failed: %v", err)
+	}
+	defer os.RemoveAll(dir)
+	if err := os.Chmod(dir, 0777); err != nil {
+		t.Fatalf("error chmoding file: %q, %v", dir, err)
+	}
+
+	// Checkpoint running container.
+	if err := cont.Checkpoint(dir, false /* direct */, statefile.Options{Compression: statefile.CompressionLevelDefault}, pgalloc.SaveOpts{}); err != nil {
+		t.Fatalf("error checkpointing container: %v", err)
+	}
+	cont.Destroy()
+	cont = nil
+
+	cont2, err := New(conf, args)
+	if err != nil {
+		t.Fatalf("error creating container: %v", err)
+	}
+	defer cont2.Destroy()
+
+	if err := cont2.Restore(conf, dir, false /* direct */); err != nil {
+		t.Fatalf("error restoring container: %v", err)
+	}
+
+	// Ensure UsageFD() is still working after restore.
+	_, err = cont2.Sandbox.UsageFD()
+	if err != nil {
+		t.Fatalf("error usageFD from restored container: %v", err)
+	}
 }
 
 // TestProfile checks that profiling options generate profiles.