virsh: Register and unregister the close callback also in cmdConnect

This patch improves the error message after disconnecting from the
hypervisor and adds the close callback operations required not to leak
the callback reference.

diff --git a/tools/virsh.c b/tools/virsh.c
index 297591d13f6161dbf89af8ca9631c8e6bbca721a..b7a5cc1f27c75cc8012680169be68739ce919d23 100644 (file)
--- a/tools/virsh.c
+++ b/tools/virsh.c
@@ -391,10 +391,14 @@ cmdConnect(vshControl *ctl, const vshCmd *cmd)
 
     if (ctl->conn) {
         int ret;
-        if ((ret = virConnectClose(ctl->conn)) != 0) {
-            vshError(ctl, _("Failed to disconnect from the hypervisor, %d leaked reference(s)"), ret);
-            return false;
-        }
+
+        virConnectUnregisterCloseCallback(ctl->conn, vshCatchDisconnect);
+        ret = virConnectClose(ctl->conn);
+        if (ret < 0)
+            vshError(ctl, "%s", _("Failed to disconnect from the hypervisor"));
+        else if (ret > 0)
+            vshError(ctl, "%s", _("One or more references were leaked after "
+                                  "disconnect from the hypervisor"));
         ctl->conn = NULL;
     }
 
@@ -411,10 +415,16 @@ cmdConnect(vshControl *ctl, const vshCmd *cmd)
     ctl->conn = virConnectOpenAuth(ctl->name, virConnectAuthPtrDefault,
                                    ctl->readonly ? VIR_CONNECT_RO : 0);
 
-    if (!ctl->conn)
+    if (!ctl->conn) {
         vshError(ctl, "%s", _("Failed to connect to the hypervisor"));
+        return false;
+    }
+
+    if (virConnectRegisterCloseCallback(ctl->conn, vshCatchDisconnect,
+                                        NULL, NULL) < 0)
+        vshError(ctl, "%s", _("Unable to register disconnect callback"));
 
-    return !!ctl->conn;
+    return true;
 }