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)
 ifclient.client!=nil {
-_=client.client.Close()
+// Match Remove's behavior so a stuck QMP close doesn't block the pool lock.
+goclient.client.Close()
  }
 	}
 
@@ -57,29 +58,33 @@ func GetOrCreateForType(socketPath string, hypervisorType hypervisor.Type) (*QEM
 returnclient, 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.
 funcresetClient(socketPathstring) {
-clientPool.Lock()
-deferclientPool.Unlock()
-ifclient, ok:=clientPool.clients[socketPath]; ok {
-delete(clientPool.clients, socketPath)
-ifclient.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.
 funcRemove(socketPathstring) {
+removeIfCurrent(socketPath, nil)
+}
+
+// removeIfCurrent removes and closes the pooled client if the current entry
+// matches expected. Passing nil expected removes whatever is currently pooled.
+funcremoveIfCurrent(socketPathstring, expected*QEMU) {
 clientPool.Lock()
 deferclientPool.Unlock()
 
 ifclient, ok:=clientPool.clients[socketPath]; ok {
+ifexpected!=nil&&client!=expected {
+return
+ }
 delete(clientPool.clients, socketPath)
-// Close asynchronously to avoid blocking on stuck connections
-goclient.client.Close()
+// Close asynchronously to avoid blocking on stuck connections.
+ifclient.client!=nil {
+goclient.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")
 }
+
+funcTestRemoveIfCurrentSkipsReplacementClient(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")
+}
+
+funcTestRemoveIfCurrentRemovesMatchingClient(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 {
 iferr:=q.client.SystemPowerdown(); err!=nil {
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 returnerr
 	}
 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 {
 iferr:=q.client.Quit(); err!=nil {
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 returnerr
 	}
 // Connection is gone after quit, remove from pool
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 clearBalloonTargetCache(q.socketPath)
 returnnil
 }
@@ -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()
 iferr!=nil {
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 returnnil, 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 {
 iferr:=q.client.Stop(); err!=nil {
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 returnerr
 	}
 returnnil
@@ -138,7 +138,7 @@ func (q *QEMU) Pause(ctx context.Context) error {
 // Resume continues VM execution.
 func (q*QEMU) Resume(ctx context.Context) error {
 iferr:=q.client.Continue(); err!=nil {
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 returnerr
 	}
 returnnil
@@ -153,13 +153,13 @@ func (q *QEMU) Snapshot(ctx context.Context, destPath string) error {
 memoryFile:=destPath+"/memory"
 uri:="exec:cat > "+memoryFile
 iferr:=q.client.Migrate(uri); err!=nil {
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 returnfmt.Errorf("migrate: %w", err)
 	}
 
 // Wait for migration to complete
 iferr:=q.client.WaitMigration(ctx, migrationTimeout); err!=nil {
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 returnfmt.Errorf("wait migration: %w", err)
 	}
 
@@ -197,7 +197,7 @@ func (q *QEMU) SetTargetGuestMemoryBytes(ctx context.Context, bytes int64) error
 returnfmt.Errorf("target guest memory %d must be non-negative", bytes)
 	}
 iferr:=q.client.Balloon(bytes); err!=nil {
-Remove(q.socketPath)
+removeIfCurrent(q.socketPath, q)
 returnfmt.Errorf("set balloon target: %w", err)
 	}
 balloonTargetCache.Store(q.socketPath, bytes)