Release per-VM connections and cached balloon targets once the VMM is gone

lib/hypervisor/balloon_target_cache.go

@@ -30,10 +30,16 @@ func (c *BalloonTargetCache) Load(socketPath string) (int64, bool) {
 		return value, true
 	}
 
-	if indexedKey, ok := c.keys.Load(socketPath); ok {
-		if keyString, ok := indexedKey.(string); ok {
-			if value, ok := c.loadKey(keyString); ok {
-				return value, true
+	// SocketCacheKey falls back to the bare path when the socket file is
+	// missing, which can happen while its VMM is still running. A socket that
+	// exists under another key belongs to a newer process, which must not
+	// inherit the old target.
+	if key == socketPath {
+		if indexedKey, ok := c.keys.Load(socketPath); ok {
+			if keyString, ok := indexedKey.(string); ok {
+				if value, ok := c.loadKey(keyString); ok {
+					return value, true
+				}
 			}
 		}
 	}

lib/hypervisor/balloon_target_cache_test.go

@@ -47,6 +47,57 @@ func TestBalloonTargetCacheDeleteClearsIndexedKeyAfterSocketRemoval(t *testing.T
 	assert.False(t, ok)
 }
 
+func TestBalloonTargetCacheSurvivesMissingSocket(t *testing.T) {
+	t.Parallel()
+
+	socketPath := testSocketPath(t)
+	listener, err := net.Listen("unix", socketPath)
+	require.NoError(t, err)
+
+	var cache BalloonTargetCache
+	cache.Store(socketPath, 512)
+
+	require.NoError(t, listener.Close())
+	if err := os.Remove(socketPath); err != nil && !os.IsNotExist(err) {
+		require.NoError(t, err)
+	}
+
+	value, ok := cache.Load(socketPath)
+	require.True(t, ok)
+	assert.Equal(t, int64(512), value)
+}
+
+func TestBalloonTargetCacheMissesForRecreatedSocket(t *testing.T) {
+	t.Parallel()
+
+	socketPath := testSocketPath(t)
+	listener, err := net.Listen("unix", socketPath)
+	require.NoError(t, err)
+
+	var cache BalloonTargetCache
+	cache.Store(socketPath, 512)
+	oldKey := SocketCacheKey(socketPath)
+
+	require.NoError(t, listener.Close())
+	if err := os.Remove(socketPath); err != nil && !os.IsNotExist(err) {
+		require.NoError(t, err)
+	}
+	// Some filesystems hand a freed inode straight back, so occupy it.
+	placeholder := socketPath + ".hold"
+	require.NoError(t, os.WriteFile(placeholder, nil, 0o600))
+	t.Cleanup(func() { _ = os.Remove(placeholder) })
+
+	listener, err = net.Listen("unix", socketPath)
+	require.NoError(t, err)
+	defer listener.Close()
+	if SocketCacheKey(socketPath) == oldKey {
+		t.Skip("filesystem reused the socket inode")
+	}
+
+	_, ok := cache.Load(socketPath)
+	assert.False(t, ok)
+}
+
 func testSocketPath(t *testing.T) string {
 	t.Helper()
 

lib/hypervisor/cloudhypervisor/cloudhypervisor_test.go

@@ -3,6 +3,7 @@ package cloudhypervisor
 import (
 	"testing"
 
+	"github.com/kernel/hypeman/lib/hypervisor"
 	"github.com/kernel/hypeman/lib/vmm"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -34,6 +35,21 @@ func TestAssignedGuestMemoryBytes(t *testing.T) {
 	})
 }
 
+func TestReleaseVMMStateForgetsBalloonTarget(t *testing.T) {
+	t.Parallel()
+
+	socketPath := t.TempDir() + "/cloud-hypervisor.sock"
+	balloonTargetCache.Store(socketPath, int64(384))
+	t.Cleanup(func() {
+		clearBalloonTargetCache(socketPath)
+	})
+
+	hypervisor.ReleaseVMMState(hypervisor.TypeCloudHypervisor, socketPath)
+
+	_, cached := balloonTargetCache.Load(socketPath)
+	assert.False(t, cached)
+}
+
 func TestGetTargetGuestMemoryBytesUsesWarmCacheBeforeVMInfo(t *testing.T) {
 	t.Parallel()
 

lib/hypervisor/cloudhypervisor/process.go

@@ -59,6 +59,7 @@ func init() {
 	hypervisor.RegisterClientFactory(hypervisor.TypeCloudHypervisor, func(socketPath string) (hypervisor.Hypervisor, error) {
 		return New(socketPath)
 	})
+	hypervisor.RegisterVMMStateReleaser(hypervisor.TypeCloudHypervisor, clearBalloonTargetCache)
 }
 
 // Starter implements hypervisor.VMStarter for Cloud Hypervisor.
@@ -128,7 +129,9 @@ func (s *Starter) StartVM(ctx context.Context, p *paths.Paths, version string, s
 		return 0, nil, fmt.Errorf("start serial reader: %w", err)
 	}
 
-	// 1. Start the Cloud Hypervisor process
+	// 1. Start the Cloud Hypervisor process. One that exited on its own leaves
+	// its balloon target cached, so drop it before the socket path is reused.
+	clearBalloonTargetCache(socketPath)
 	processCtx, processSpan := hypervisor.StartProcessSpan(ctx, hypervisor.TypeCloudHypervisor)
 	pid, err := vmm.StartProcess(processCtx, p, chVersion, socketPath)
 	hypervisor.FinishTraceSpan(processSpan, err)
@@ -209,7 +212,9 @@ func (s *Starter) RestoreVM(ctx context.Context, p *paths.Paths, version string,
 		return 0, nil, fmt.Errorf("rewrite snapshot serial config: %w", err)
 	}
 
-	// 1. Start the Cloud Hypervisor process
+	// 1. Start the Cloud Hypervisor process. One that exited on its own leaves
+	// its balloon target cached, so drop it before the socket path is reused.
+	clearBalloonTargetCache(socketPath)
 	processStartTime := time.Now()
 	processCtx, processSpan := hypervisor.StartProcessSpan(ctx, hypervisor.TypeCloudHypervisor)
 	pid, err := vmm.StartProcess(processCtx, p, chVersion, socketPath)

lib/hypervisor/hypervisor.go

@@ -468,3 +468,25 @@ func NewClient(hvType Type, socketPath string) (Hypervisor, error) {
 	}
 	return WrapHypervisor(hvType, client), nil
 }
+
+// VMMStateReleaser drops what a hypervisor type caches for a socket whose VMM
+// has exited, such as a pooled client or the last balloon target.
+type VMMStateReleaser func(socketPath string)
+
+// vmmStateReleasers maps hypervisor types that cache per-socket state to their releasers.
+var vmmStateReleasers = make(map[Type]VMMStateReleaser)
+
+// RegisterVMMStateReleaser registers the releaser for a hypervisor type that
+// caches state per socket.
+func RegisterVMMStateReleaser(t Type, releaser VMMStateReleaser) {
+	vmmStateReleasers[t] = releaser
+}
+
+// ReleaseVMMState drops any state cached for socketPath, closing a cached
+// client. Call it only once the VMM behind socketPath is confirmed gone.
+// Hypervisor types that cache nothing ignore it.
+func ReleaseVMMState(hvType Type, socketPath string) {
+	if release, ok := vmmStateReleasers[hvType]; ok {
+		release(socketPath)
+	}
+}

lib/hypervisor/qemu/pool.go

@@ -63,14 +63,6 @@ func poolTypeMismatchError(socketPath string, cached, requested hypervisor.Type)
 	return fmt.Errorf("QEMU client for %s is pooled as hypervisor %s, not %s", socketPath, cached, requested)
 }
 
-// resetClient drops a pooled connection before a new QEMU process reuses the
-// same socket path. Disconnect happens outside the pool lock and stale users
-// can only remove their own client generation.
-func resetClient(socketPath string) {
-	client := takeClient(socketPath, nil)
-	closeClientAsync(client)
-}
-
 // removeClient removes client only if it is still the current generation for
 // its socket path. This prevents a late error from an old QEMU process from
 // evicting the replacement process's client.
@@ -88,6 +80,14 @@ func Remove(socketPath string) {
 	closeClientAsync(client)
 }
 
+// releaseVMMState drops what is cached for a socket whose QEMU has exited.
+// Unlike Remove, it also forgets the balloon target, which has to outlive a
+// reconnect to a QEMU that is still running.
+func releaseVMMState(socketPath string) {
+	Remove(socketPath)
+	clearBalloonTargetCache(socketPath)
+}
+
 // takeClient removes the current client. When expected is non-nil, removal is
 // conditional on pointer identity to protect against socket-path reuse.
 func takeClient(socketPath string, expected *QEMU) *QEMU {

lib/hypervisor/qemu/pool_test.go

@@ -1,7 +1,13 @@
 package qemu
 
 import (
+	"bufio"
+	"io"
+	"net"
+	"os"
+	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/kernel/hypeman/lib/hypervisor"
 	"github.com/stretchr/testify/require"
@@ -47,3 +53,62 @@ func TestGetOrCreateForTypeRejectsCachedBackendMismatch(t *testing.T) {
 	clientPool.RUnlock()
 	require.Equal(t, hypervisor.TypeQEMU, cached.profile.hypervisorType(), "a stale caller must not evict the live client")
 }
+
+func TestReleaseVMMStateClosesPooledConnection(t *testing.T) {
+	for _, hvType := range []hypervisor.Type{hypervisor.TypeQEMU, hypervisor.TypeQEMUMicroVM} {
+		t.Run(string(hvType), func(t *testing.T) {
+			// t.TempDir() overflows the unix socket path limit on macOS.
+			dir, err := os.MkdirTemp("/tmp", "qmp")
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = os.RemoveAll(dir) })
+			socketPath := filepath.Join(dir, "qemu.sock")
+			t.Cleanup(func() { releaseVMMState(socketPath) })
+			closed := serveFakeQMP(t, socketPath)
+
+			_, err = GetOrCreateForType(socketPath, hvType)
+			require.NoError(t, err)
+			balloonTargetCache.Store(socketPath, 512<<20)
+
+			hypervisor.ReleaseVMMState(hvType, socketPath)
+
+			clientPool.RLock()
+			_, pooled := clientPool.clients[socketPath]
+			clientPool.RUnlock()
+			require.False(t, pooled)
+			_, cached := balloonTargetCache.Load(socketPath)
+			require.False(t, cached, "balloon target must not outlive the VM")
+			select {
+			case <-closed:
+			case <-time.After(5 * time.Second):
+				t.Fatal("released client did not close its QMP connection")
+			}
+		})
+	}
+}
+
+// serveFakeQMP accepts one QMP connection on socketPath, completes the
+// capabilities handshake, and closes the returned channel when the client
+// disconnects.
+func serveFakeQMP(t *testing.T, socketPath string) <-chan struct{} {
+	t.Helper()
+	ln, err := net.Listen("unix", socketPath)
+	require.NoError(t, err)
+	t.Cleanup(func() { ln.Close() })
+
+	closed := make(chan struct{})
+	go func() {
+		defer close(closed)
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+		_, _ = io.WriteString(conn, `{"QMP":{"version":{"qemu":{"major":8,"minor":2,"micro":0},"package":""},"capabilities":[]}}`+"\n")
+		if _, err := bufio.NewReader(conn).ReadBytes('\n'); err != nil {
+			return
+		}
+		_, _ = io.WriteString(conn, `{"return":{}}`+"\n")
+		_, _ = io.Copy(io.Discard, conn)
+	}()
+	return closed
+}

lib/hypervisor/qemu/process.go

@@ -55,6 +55,7 @@ func init() {
 		hypervisor.RegisterClientFactory(t, func(socketPath string) (hypervisor.Hypervisor, error) {
 			return NewForType(socketPath, t)
 		})
+		hypervisor.RegisterVMMStateReleaser(t, releaseVMMState)
 	}
 }
 
@@ -309,6 +310,13 @@ func (p *startedProcess) cleanup() {
 	_ = os.Remove(p.socketPath)
 }
 
+// cleanupFailedStart kills a QEMU process whose start failed. The failure can
+// come after the QMP client is pooled, so the client is dropped with it.
+func cleanupFailedStart(proc *startedProcess) {
+	proc.cleanup()
+	releaseVMMState(proc.socketPath)
+}
+
 func vfioTermGraceFor(cfg hypervisor.VMConfig) time.Duration {
 	if cfg.VGPUDevicePath != "" || len(cfg.PCIDevices) > 0 {
 		return hypervisor.VFIOTermGrace
@@ -343,9 +351,9 @@ func (s *Starter) startQEMUProcess(ctx context.Context, p *paths.Paths, version
 		return 0, nil, nil, fmt.Errorf("socket already in use, QEMU may be running at %s", socketPath)
 	}
 
-	// A graceful guest shutdown can leave the old QMP client pooled even though
-	// the process exited. Drop it before this process reuses the socket path.
-	resetClient(socketPath)
+	// A QEMU that exited on its own leaves its QMP client and balloon target
+	// cached. Drop both before this process reuses the socket path.
+	releaseVMMState(socketPath)
 	os.Remove(socketPath)
 
 	// Create command
@@ -394,7 +402,7 @@ func (s *Starter) startQEMUProcess(ctx context.Context, p *paths.Paths, version
 	log.DebugContext(processCtx, "QEMU process started", "pid", pid, "duration_ms", time.Since(processStartTime).Milliseconds())
 
 	// Setup cleanup to kill, reap, and remove the socket if subsequent steps fail.
-	cu := cleanup.Make(proc.cleanup)
+	cu := cleanup.Make(func() { cleanupFailedStart(proc) })
 
 	// Wait for socket to be ready
 	socketWaitStart := time.Now()

lib/hypervisor/qemu/process_test.go

@@ -394,6 +394,57 @@ func TestStartManagedProcessCleanupRemovesSocketAndReapsExitedProcess(t *testing
 	assert.True(t, cmd.ProcessState.Exited())
 }
 
+func TestCleanupFailedStartDropsPooledClient(t *testing.T) {
+	socketPath := filepath.Join(t.TempDir(), "qemu.sock")
+	clientPool.Lock()
+	clientPool.clients[socketPath] = &QEMU{socketPath: socketPath, profile: StandardProfile{}}
+	clientPool.Unlock()
+	t.Cleanup(func() {
+		clientPool.Lock()
+		delete(clientPool.clients, socketPath)
+		clientPool.Unlock()
+	})
+
+	proc, err := startManagedProcess(exec.Command("sleep", "30"), socketPath)
+	require.NoError(t, err)
+
+	cleanupFailedStart(proc)
+
+	clientPool.RLock()
+	_, pooled := clientPool.clients[socketPath]
+	clientPool.RUnlock()
+	require.False(t, pooled)
+}
+
+// A QEMU that exits on its own never goes through stop or delete, so the next
+// start on the same socket path has to drop what the old process left cached.
+func TestStartQEMUProcessDropsStaleBalloonTarget(t *testing.T) {
+	starter := NewStarter()
+	if _, err := starter.GetBinaryPath(nil, ""); err != nil {
+		t.Skipf("QEMU not available: %v", err)
+	}
+
+	// t.TempDir() overflows the unix socket path limit on macOS.
+	dir, err := os.MkdirTemp("/tmp", "qmp")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = os.RemoveAll(dir) })
+	socketPath := filepath.Join(dir, "qemu.sock")
+	balloonTargetCache.Store(socketPath, 512<<20)
+	t.Cleanup(func() { clearBalloonTargetCache(socketPath) })
+
+	args := append(buildQMPArgs(socketPath), "-machine", "none", "-nodefaults", "-display", "none")
+	_, _, cu, err := starter.startQEMUProcess(t.Context(), paths.New(dir), "", socketPath, args, 0)
+	require.NoError(t, err)
+	t.Cleanup(cu.Clean)
+
+	// Whether a leftover target is still reachable through the new socket
+	// depends on the filesystem reusing its inode. With the socket file gone
+	// the lookup is by path alone, so it finds a leftover on any filesystem.
+	require.NoError(t, os.Remove(socketPath))
+	_, cached := balloonTargetCache.Load(socketPath)
+	require.False(t, cached)
+}
+
 func TestWaitForSocketOrExitReturnsEarlyWhenProcessDies(t *testing.T) {
 	socketPath := filepath.Join(t.TempDir(), "qemu.sock")
 

lib/instances/delete.go

@@ -143,6 +143,14 @@ func (m *manager) deleteInstanceWithOptions(
 		}
 	}
 
+	// The VMM is gone. Close its cached control connection, and the guest-agent
+	// connection that the graceful shutdown RPC opened again after step 3. Both
+	// would otherwise stay open for the life of the process.
+	hypervisor.ReleaseVMMState(inst.HypervisorType, inst.SocketPath)
+	if dialer, err := hypervisor.NewVsockDialer(inst.HypervisorType, inst.VsockSocket, inst.VsockCID); err == nil {
+		guest.CloseConn(dialer.Key())
+	}
+
 	// Release before deleting metadata so a failed release can be retried safely.
 	if storedVGPUDevicePath(stored) != "" {
 		log.InfoContext(ctx, "destroying vGPU", "instance_id", id, "uuid", stored.GPUMdevUUID)