Fix qemu pool eviction locking and stale removal races

Author: cursoragent

lib/hypervisor/qemu/pool.go

@@ -43,7 +43,8 @@ func GetOrCreateForType(socketPath string, hypervisorType hypervisor.Type) (*QEM
 		}
 		delete(clientPool.clients, socketPath)
 		if client.client != nil {
-			_ = client.client.Close()
+			// Match Remove's behavior so a stuck QMP close doesn't block the pool lock.
+			go client.client.Close()
 		}
 	}
 
@@ -57,29 +58,33 @@ func GetOrCreateForType(socketPath string, hypervisorType hypervisor.Type) (*QEM
 	return client, nil
 }
 
-// resetClient synchronously drops a pooled connection before a new QEMU
+// resetClient drops a pooled connection before a new QEMU
 // process reuses the same socket path.
 func resetClient(socketPath string) {
-	clientPool.Lock()
-	defer clientPool.Unlock()
-	if client, ok := clientPool.clients[socketPath]; ok {
-		delete(clientPool.clients, socketPath)
-		if client.client != nil {
-			_ = client.client.Close()
-		}
-	}
+	removeIfCurrent(socketPath, nil)
 }
 
 // Remove closes and removes a client from the pool.
 // Called automatically on errors to allow fresh reconnection.
 // Close is done asynchronously to avoid blocking if the connection is in a bad state.
 func Remove(socketPath string) {
+	removeIfCurrent(socketPath, nil)
+}
+
+// removeIfCurrent removes and closes the pooled client if the current entry
+// matches expected. Passing nil expected removes whatever is currently pooled.
+func removeIfCurrent(socketPath string, expected *QEMU) {
 	clientPool.Lock()
 	defer clientPool.Unlock()
 
 	if client, ok := clientPool.clients[socketPath]; ok {
+		if expected != nil && client != expected {
+			return
+		}
 		delete(clientPool.clients, socketPath)
-		// Close asynchronously to avoid blocking on stuck connections
-		go client.client.Close()
+		// Close asynchronously to avoid blocking on stuck connections.
+		if client.client != nil {
+			go client.client.Close()
+		}
 	}
 }

lib/hypervisor/qemu/pool_test.go

@@ -26,3 +26,42 @@ func TestGetOrCreateForTypeReplacesCachedBackendMismatch(t *testing.T) {
 	clientPool.RUnlock()
 	require.False(t, stillCached, "stale backend client must be removed before reconnect")
 }
+
+func TestRemoveIfCurrentSkipsReplacementClient(t *testing.T) {
+	socketPath := t.TempDir() + "/qemu.sock"
+	stale := &QEMU{socketPath: socketPath, hypervisorType: hypervisor.TypeQEMU}
+	replacement := &QEMU{socketPath: socketPath, hypervisorType: hypervisor.TypeQEMUMicroVM}
+
+	clientPool.Lock()
+	clientPool.clients[socketPath] = replacement
+	clientPool.Unlock()
+	t.Cleanup(func() {
+		clientPool.Lock()
+		delete(clientPool.clients, socketPath)
+		clientPool.Unlock()
+	})
+
+	removeIfCurrent(socketPath, stale)
+
+	clientPool.RLock()
+	cached, ok := clientPool.clients[socketPath]
+	clientPool.RUnlock()
+	require.True(t, ok, "replacement client should remain pooled")
+	require.Same(t, replacement, cached, "stale remove must not evict replacement client")
+}
+
+func TestRemoveIfCurrentRemovesMatchingClient(t *testing.T) {
+	socketPath := t.TempDir() + "/qemu.sock"
+	client := &QEMU{socketPath: socketPath, hypervisorType: hypervisor.TypeQEMU}
+
+	clientPool.Lock()
+	clientPool.clients[socketPath] = client
+	clientPool.Unlock()
+
+	removeIfCurrent(socketPath, client)
+
+	clientPool.RLock()
+	_, ok := clientPool.clients[socketPath]
+	clientPool.RUnlock()
+	require.False(t, ok, "matching pooled client should be removed")
+}

lib/hypervisor/qemu/qemu.go

@@ -71,7 +71,7 @@ func capabilities(hypervisorType hypervisor.Type) hypervisor.Capabilities {
 // This sends a graceful shutdown signal to the guest.
 func (q *QEMU) DeleteVM(ctx context.Context) error {
 	if err := q.client.SystemPowerdown(); err != nil {
-		Remove(q.socketPath)
+		removeIfCurrent(q.socketPath, q)
 		return err
 	}
 	clearBalloonTargetCache(q.socketPath)
@@ -81,11 +81,11 @@ func (q *QEMU) DeleteVM(ctx context.Context) error {
 // Shutdown stops the QEMU process.
 func (q *QEMU) Shutdown(ctx context.Context) error {
 	if err := q.client.Quit(); err != nil {
-		Remove(q.socketPath)
+		removeIfCurrent(q.socketPath, q)
 		return err
 	}
 	// Connection is gone after quit, remove from pool
-	Remove(q.socketPath)
+	removeIfCurrent(q.socketPath, q)
 	clearBalloonTargetCache(q.socketPath)
 	return nil
 }
@@ -94,7 +94,7 @@ func (q *QEMU) Shutdown(ctx context.Context) error {
 func (q *QEMU) GetVMInfo(ctx context.Context) (*hypervisor.VMInfo, error) {
 	status, err := q.client.Status()
 	if err != nil {
-		Remove(q.socketPath)
+		removeIfCurrent(q.socketPath, q)
 		return nil, fmt.Errorf("query status: %w", err)
 	}
 
@@ -129,7 +129,7 @@ func (q *QEMU) GetVMInfo(ctx context.Context) (*hypervisor.VMInfo, error) {
 // Pause suspends VM execution.
 func (q *QEMU) Pause(ctx context.Context) error {
 	if err := q.client.Stop(); err != nil {
-		Remove(q.socketPath)
+		removeIfCurrent(q.socketPath, q)
 		return err
 	}
 	return nil
@@ -138,7 +138,7 @@ func (q *QEMU) Pause(ctx context.Context) error {
 // Resume continues VM execution.
 func (q *QEMU) Resume(ctx context.Context) error {
 	if err := q.client.Continue(); err != nil {
-		Remove(q.socketPath)
+		removeIfCurrent(q.socketPath, q)
 		return err
 	}
 	return nil
@@ -153,13 +153,13 @@ func (q *QEMU) Snapshot(ctx context.Context, destPath string) error {
 	memoryFile := destPath + "/memory"
 	uri := "exec:cat > " + memoryFile
 	if err := q.client.Migrate(uri); err != nil {
-		Remove(q.socketPath)
+		removeIfCurrent(q.socketPath, q)
 		return fmt.Errorf("migrate: %w", err)
 	}
 
 	// Wait for migration to complete
 	if err := q.client.WaitMigration(ctx, migrationTimeout); err != nil {
-		Remove(q.socketPath)
+		removeIfCurrent(q.socketPath, q)
 		return fmt.Errorf("wait migration: %w", err)
 	}
 
@@ -197,7 +197,7 @@ func (q *QEMU) SetTargetGuestMemoryBytes(ctx context.Context, bytes int64) error
 		return fmt.Errorf("target guest memory %d must be non-negative", bytes)
 	}
 	if err := q.client.Balloon(bytes); err != nil {
-		Remove(q.socketPath)
+		removeIfCurrent(q.socketPath, q)
 		return fmt.Errorf("set balloon target: %w", err)
 	}
 	balloonTargetCache.Store(q.socketPath, bytes)