lima-guestagent: Handle SIGTERM, support save/restore eventState on stop/start

`eventState` is saved at `/run/lima-guestagent/event-state.json`.
The saved `eventState` is expected to be removed on OS restart.

Signed-off-by: Norio Nomura <norio.nomura@gmail.com>

Author: norio-nomura

cmd/lima-guestagent/daemon_linux.go

@@ -7,6 +7,8 @@ import (
 	"errors"
 	"net"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/mdlayher/vsock"
@@ -26,6 +28,7 @@ func newDaemonCommand() *cobra.Command {
 		Short: "Run the daemon",
 		RunE:  daemonAction,
 	}
+	daemonCommand.Flags().String("runtime-dir", "/run/lima-guestagent", "Directory to store runtime state")
 	daemonCommand.Flags().Duration("tick", 3*time.Second, "Tick for polling events")
 	daemonCommand.Flags().Int("vsock-port", 0, "Use vsock server instead a UNIX socket")
 	daemonCommand.Flags().String("virtio-port", "", "Use virtio server instead a UNIX socket")
@@ -34,6 +37,13 @@ func newDaemonCommand() *cobra.Command {
 
 func daemonAction(cmd *cobra.Command, _ []string) error {
 	ctx := cmd.Context()
+	runtimeDir, err := cmd.Flags().GetString("runtime-dir")
+	if err != nil {
+		return err
+	}
+	if err := os.MkdirAll(runtimeDir, 0o755); err != nil {
+		return err
+	}
 	socket := "/run/lima-guestagent.sock"
 	tick, err := cmd.Flags().GetDuration("tick")
 	if err != nil {
@@ -66,11 +76,16 @@ func daemonAction(cmd *cobra.Command, _ []string) error {
 		tickerInst = ticker.NewCompoundTicker(simpleTicker, ebpfTicker)
 	}
 
-	agent, err := guestagent.New(ctx, tickerInst)
+	ctx, stop := signal.NotifyContext(ctx, syscall.SIGTERM)
+	defer stop()
+	go func() {
+		<-ctx.Done()
+		logrus.Debug("Received SIGTERM, shutting down the guest agent")
+	}()
+	agent, err := guestagent.New(ctx, tickerInst, runtimeDir)
 	if err != nil {
 		return err
 	}
-	defer agent.Close()
 
 	err = os.RemoveAll(socket)
 	if err != nil {
@@ -104,5 +119,6 @@ func daemonAction(cmd *cobra.Command, _ []string) error {
 		l = socketL
 		logrus.Infof("serving the guest agent on %q", socket)
 	}
-	return server.StartServer(l, &server.GuestServer{Agent: agent, TunnelS: portfwdserver.NewTunnelServer()})
+	defer logrus.Debug("exiting lima-guestagent daemon")
+	return server.StartServer(ctx, l, &server.GuestServer{Agent: agent, TunnelS: portfwdserver.NewTunnelServer()})
 }

cmd/lima-guestagent/lima-guestagent.TEMPLATE.service

@@ -2,7 +2,7 @@
 Description=lima-guestagent
 
 [Service]
-ExecStart={{.Binary}} daemon {{.Args}}
+ExecStart={{.Binary}} daemon {{.Args}} --runtime-dir="%t/%N"
 Type=simple
 Restart=on-failure
 OOMPolicy=continue

cmd/lima-guestagent/main_linux.go

@@ -11,12 +11,14 @@ import (
 
 	"github.com/lima-vm/lima/v2/cmd/yq"
 	"github.com/lima-vm/lima/v2/pkg/debugutil"
+	"github.com/lima-vm/lima/v2/pkg/osutil"
 	"github.com/lima-vm/lima/v2/pkg/version"
 )
 
 func main() {
 	yq.MaybeRunYQ()
 	if err := newApp().Execute(); err != nil {
+		osutil.HandleExitError(err)
 		logrus.Fatal(err)
 	}
 }

pkg/guestagent/api/server/server.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"net"
 
+	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/types/known/emptypb"
 
@@ -15,10 +16,22 @@ import (
 	"github.com/lima-vm/lima/v2/pkg/portfwdserver"
 )
 
-func StartServer(lis net.Listener, guest *GuestServer) error {
+func StartServer(ctx context.Context, lis net.Listener, guest *GuestServer) error {
 	server := grpc.NewServer()
 	api.RegisterGuestServiceServer(server, guest)
-	return server.Serve(lis)
+	go func() {
+		<-ctx.Done()
+		logrus.Debug("Stopping the gRPC server")
+		server.GracefulStop()
+		logrus.Debug("Closing the listener used by the gRPC server")
+		lis.Close()
+	}()
+	err := server.Serve(lis)
+	// grpc.Server.Serve() expects to return a non-nil error caused by lis.Accept()
+	if opErr, ok := err.(*net.OpError); ok && opErr.Err.Error() == "use of closed network connection" {
+		return nil
+	}
+	return err
 }
 
 type GuestServer struct {
@@ -33,6 +46,7 @@ func (s *GuestServer) GetInfo(ctx context.Context, _ *emptypb.Empty) (*api.Info,
 
 func (s *GuestServer) GetEvents(_ *emptypb.Empty, stream api.GuestService_GetEventsServer) error {
 	responses := make(chan *api.Event)
+	// expects Events() to close the channel when stream.Context() is done or ticker stops
 	go s.Agent.Events(stream.Context(), responses)
 	for response := range responses {
 		err := stream.Send(response)

pkg/guestagent/guestagent_linux.go

@@ -5,7 +5,10 @@ package guestagent
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
 	"os"
+	"path/filepath"
 	"reflect"
 	"time"
 
@@ -19,7 +22,7 @@ import (
 	"github.com/lima-vm/lima/v2/pkg/guestagent/timesync"
 )
 
-func New(ctx context.Context, ticker ticker.Ticker) (Agent, error) {
+func New(ctx context.Context, ticker ticker.Ticker, runtimeDir string) (Agent, error) {
 	socketsLister, err := sockets.NewLister()
 	if err != nil {
 		return nil, err
@@ -28,25 +31,37 @@ func New(ctx context.Context, ticker ticker.Ticker) (Agent, error) {
 		ticker:                   ticker,
 		socketLister:             socketsLister,
 		kubernetesServiceWatcher: kubernetesservice.NewServiceWatcher(),
+		runtimeDir:               runtimeDir,
 	}
 
 	go a.kubernetesServiceWatcher.Start(ctx)
-	go a.fixSystemTimeSkew()
+	go a.fixSystemTimeSkew(ctx)
+
+	go func() {
+		<-ctx.Done()
+		logrus.Debug("Closing the agent")
+		if err := a.Close(); err != nil {
+			logrus.Errorf("error on agent.Close(): %v", err)
+		}
+	}()
 
 	return a, nil
 }
 
+var _ Agent = (*agent)(nil)
+
 type agent struct {
 	// Ticker is like time.Ticker.
 	// We can't use inotify for /proc/net/tcp, so we need this ticker to
 	// reload /proc/net/tcp.
 	ticker                   ticker.Ticker
 	socketLister             *sockets.Lister
 	kubernetesServiceWatcher *kubernetesservice.ServiceWatcher
+	runtimeDir               string
 }
 
 type eventState struct {
-	ports []*api.IPPort
+	Ports []*api.IPPort `json:"ports,omitempty"`
 }
 
 func comparePorts(old, neww []*api.IPPort) (added, removed []*api.IPPort) {
@@ -82,13 +97,13 @@ func (a *agent) collectEvent(ctx context.Context, st eventState) (*api.Event, ev
 		err error
 	)
 	newSt := st
-	newSt.ports, err = a.LocalPorts(ctx)
+	newSt.Ports, err = a.LocalPorts(ctx)
 	if err != nil {
 		ev.Errors = append(ev.Errors, err.Error())
 		ev.Time = timestamppb.Now()
 		return ev, newSt
 	}
-	ev.AddedLocalPorts, ev.RemovedLocalPorts = comparePorts(st.ports, newSt.ports)
+	ev.AddedLocalPorts, ev.RemovedLocalPorts = comparePorts(st.Ports, newSt.Ports)
 	ev.Time = timestamppb.Now()
 	return ev, newSt
 }
@@ -102,8 +117,16 @@ func isEventEmpty(ev *api.Event) bool {
 func (a *agent) Events(ctx context.Context, ch chan *api.Event) {
 	defer close(ch)
 	tickerCh := a.ticker.Chan()
-	defer a.ticker.Stop()
-	var st eventState
+
+	st, err := a.LoadEventState()
+	if err != nil {
+		logrus.Errorf("failed to load state: %v", err)
+	}
+	defer func() {
+		if err := a.SaveEventState(st); err != nil {
+			logrus.Errorf("failed to save state: %v", err)
+		}
+	}()
 	for {
 		var ev *api.Event
 		ev, st = a.collectEvent(ctx, st)
@@ -115,6 +138,7 @@ func (a *agent) Events(ctx context.Context, ch chan *api.Event) {
 			return
 		case _, ok := <-tickerCh:
 			if !ok {
+				logrus.Debug("ticker channel closed")
 				return
 			}
 			logrus.Debug("tick!")
@@ -190,7 +214,7 @@ func (a *agent) Info(ctx context.Context) (*api.Info, error) {
 
 const deltaLimit = 2 * time.Second
 
-func (a *agent) fixSystemTimeSkew() {
+func (a *agent) fixSystemTimeSkew(ctx context.Context) {
 	logrus.Info("fixSystemTimeSkew(): monitoring system time skew")
 	for {
 		ok, err := timesync.HasRTC()
@@ -217,6 +241,13 @@ func (a *agent) fixSystemTimeSkew() {
 				logrus.Infof("fixSystemTimeSkew: system time synchronized with rtc")
 				break
 			}
+			select {
+			case <-ctx.Done():
+				logrus.Debug("fixSystemTimeSkew: context done, exiting")
+				ticker.Stop()
+				return
+			default:
+			}
 		}
 		ticker.Stop()
 	}
@@ -239,5 +270,42 @@ func (a *agent) Close() error {
 			return err
 		}
 	}
+	a.ticker.Stop()
 	return nil
 }
+
+const eventStateFileName = "event-state.json"
+
+// LoadEventState loads the event state from a file in JSON format.
+// If the file does not exist, it returns an empty eventState with no error.
+// The saved eventState is expected to be removed on OS restart.
+func (a *agent) LoadEventState() (eventState, error) {
+	logrus.Debug("Loading event state")
+	path := filepath.Join(a.runtimeDir, eventStateFileName)
+	data, err := os.ReadFile(path)
+	if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			return eventState{}, nil
+		}
+		return eventState{}, err
+	}
+	var st eventState
+	if err := json.Unmarshal(data, &st); err != nil {
+		return eventState{}, err
+	}
+	// We don't remove the file after loading for debugging purposes.
+	return st, nil
+}
+
+// SaveEventState saves the event state to a file in JSON format.
+// It overwrites the file if it already exists.
+// The saved eventState is expected to be removed on OS restart.
+func (a *agent) SaveEventState(st eventState) error {
+	logrus.Debug("Saving event state")
+	data, err := json.Marshal(st)
+	if err != nil {
+		return err
+	}
+	path := filepath.Join(a.runtimeDir, eventStateFileName)
+	return os.WriteFile(path, data, 0o644)
+}

pkg/guestagent/ticker/compound.go

@@ -5,6 +5,8 @@ package ticker
 
 import (
 	"time"
+
+	"github.com/sirupsen/logrus"
 )
 
 func NewCompoundTicker(t1, t2 Ticker) Ticker {
@@ -15,10 +17,13 @@ type compoundTicker struct {
 	t1, t2 Ticker
 }
 
+var _ Ticker = (*compoundTicker)(nil)
+
 func (ticker *compoundTicker) Chan() <-chan time.Time {
 	ch := make(chan time.Time)
 	go func() {
-		defer ticker.Stop()
+		defer close(ch)
+		defer logrus.Debug("compoundTicker: exiting")
 		for {
 			select {
 			case v, ok := <-ticker.t1.Chan():

pkg/guestagent/ticker/ebpf_linux.go

@@ -4,6 +4,7 @@
 package ticker
 
 import (
+	"errors"
 	"os"
 	"strings"
 	"time"
@@ -20,45 +21,49 @@ func NewEbpfTicker(tracepoints []string) (Ticker, error) {
 		ticker ebpfTicker
 		err    error
 	)
+	defer func() {
+		if err != nil {
+			ticker.Stop()
+		}
+	}()
 	ticker.events, err = ebpf.NewMap(&ebpf.MapSpec{
 		Name:       "lima_ticker_events",
 		Type:       ebpf.RingBuf,
 		MaxEntries: 1 << 20,
 	})
 	if err != nil {
-		ticker.Stop()
 		return nil, err
 	}
 
 	ticker.prog, err = buildEbpfProg(ticker.events)
 	if err != nil {
-		ticker.Stop()
 		return nil, err
 	}
 
 	for _, tp := range tracepoints {
 		tpPair := strings.SplitN(tp, ":", 2)
 		tpLink, err := link.Tracepoint(tpPair[0], tpPair[1], ticker.prog, nil)
 		if err != nil {
-			ticker.Stop()
 			return nil, err
 		}
 		ticker.links = append(ticker.links, tpLink)
 	}
 
 	ticker.reader, err = ringbuf.NewReader(ticker.events)
 	if err != nil {
-		ticker.Stop()
 		return nil, err
 	}
 
 	ticker.ch = make(chan time.Time)
 	go func() {
+		defer close(ticker.ch)
 		for {
 			_, rdErr := ticker.reader.Read()
 			if rdErr != nil {
-				logrus.WithError(rdErr).Warn("ebpfTicker: failed to read ringbuf")
-				ticker.Stop()
+				if !errors.Is(rdErr, ringbuf.ErrClosed) {
+					logrus.WithError(rdErr).Warn("ebpfTicker: failed to read ringbuf")
+				}
+				logrus.Debug("ebpfTicker: exiting")
 				return
 			}
 			ticker.ch <- time.Now()
@@ -68,6 +73,8 @@ func NewEbpfTicker(tracepoints []string) (Ticker, error) {
 	return &ticker, nil
 }
 
+var _ Ticker = (*ebpfTicker)(nil)
+
 type ebpfTicker struct {
 	events *ebpf.Map
 	prog   *ebpf.Program
@@ -93,9 +100,7 @@ func (ticker *ebpfTicker) Stop() {
 	if ticker.reader != nil {
 		_ = ticker.reader.Close()
 	}
-	if ticker.ch != nil {
-		close(ticker.ch)
-	}
+	// ticker.ch will be closed in go routine in NewEbpfTicker() to avoid sending on closed channel
 }
 
 func buildEbpfProg(events *ebpf.Map) (*ebpf.Program, error) {