From a1df4d0727b067ba49935747822d863654e9c152 Mon Sep 17 00:00:00 2001
From: Mark Syms <mark.syms@citrix.com>
Date: Thu, 1 Feb 2024 10:30:00 +0000
Subject: [PATCH] CA-388451: ensure that xapi sessions are logged out

Signed-off-by: Mark Syms <mark.syms@citrix.com>
---
 drivers/cleanup.py         |  2 +-
 drivers/mpathcount.py      |  2 +-
 drivers/sr_health_check.py | 43 ++++++++++++++++++++------------------
 drivers/util.py            | 35 ++++++++++++++++++-------------
 4 files changed, 45 insertions(+), 37 deletions(-)

diff --git a/drivers/cleanup.py b/drivers/cleanup.py
index b23ed3dee..9c49ae9e5 100755
--- a/drivers/cleanup.py
+++ b/drivers/cleanup.py
@@ -3033,7 +3033,7 @@ def _ensure_xapi_initialised(session):
             time.sleep(15)
     finally:
         if local_session is not None:
-            local_session.logout()
+            local_session.xenapi.session.logout()
 
 def _gc(session, srUuid, dryRun=False, immediate=False):
     init(srUuid)
diff --git a/drivers/mpathcount.py b/drivers/mpathcount.py
index d91052860..a31ec1b75 100755
--- a/drivers/mpathcount.py
+++ b/drivers/mpathcount.py
@@ -52,7 +52,7 @@ def get_dm_major():
 def mpc_exit(session, code):
     if session is not None:
         try:
-            session.xenapi.logout()
+            session.xenapi.session.logout()
         except:
             pass
     sys.exit(code)
diff --git a/drivers/sr_health_check.py b/drivers/sr_health_check.py
index 8c4197f17..135e639e9 100755
--- a/drivers/sr_health_check.py
+++ b/drivers/sr_health_check.py
@@ -34,26 +34,29 @@ def main():
         util.SMlog("Unable to open local XAPI session", priority=util.LOG_ERR)
         return
 
-    localhost = util.get_localhost_ref(session)
-
-    sm_types = [x['type'] for x in session.xenapi.SM.get_all_records_where(
-        'field "required_api_version" = "1.0"').values()]
-    for sm_type in sm_types:
-        srs = session.xenapi.SR.get_all_records_where(
-            f'field "type" = "{sm_type}"')
-        for sr in srs:
-            pbds = session.xenapi.PBD.get_all_records_where(
-                f'field "SR" = "{sr}" and field "host" = "{localhost}"')
-            if not pbds:
-                continue
-
-            pbd_ref, pbd = pbds.popitem()
-            if not pbd['currently_attached']:
-                continue
-
-            sr_uuid = srs[sr]['uuid']
-            sr_obj = SR.SR.from_uuid(session, sr_uuid)
-            sr_obj.check_sr(sr_uuid)
+    try:
+        localhost = util.get_localhost_ref(session)
+
+        sm_types = [x['type'] for x in session.xenapi.SM.get_all_records_where(
+            'field "required_api_version" = "1.0"').values()]
+        for sm_type in sm_types:
+            srs = session.xenapi.SR.get_all_records_where(
+                f'field "type" = "{sm_type}"')
+            for sr in srs:
+                pbds = session.xenapi.PBD.get_all_records_where(
+                    f'field "SR" = "{sr}" and field "host" = "{localhost}"')
+                if not pbds:
+                    continue
+
+                pbd_ref, pbd = pbds.popitem()
+                if not pbd['currently_attached']:
+                    continue
+
+                sr_uuid = srs[sr]['uuid']
+                sr_obj = SR.SR.from_uuid(session, sr_uuid)
+                sr_obj.check_sr(sr_uuid)
+    finally:
+        session.xenapi.session.logout()
 
 
 if __name__ == "__main__":
diff --git a/drivers/util.py b/drivers/util.py
index d90cebc52..f6b9ba8d8 100755
--- a/drivers/util.py
+++ b/drivers/util.py
@@ -378,18 +378,19 @@ def ioretry_stat(path, maxretry=IORETRY_MAX):
 def sr_get_capability(sr_uuid):
     result = []
     session = get_localAPI_session()
-    sr_ref = session.xenapi.SR.get_by_uuid(sr_uuid)
-    sm_type = session.xenapi.SR.get_record(sr_ref)['type']
-    sm_rec = session.xenapi.SM.get_all_records_where(
-        "field \"type\" = \"%s\"" % sm_type)
-
-    # SM expects at least one entry of any SR type
-    if len(sm_rec) > 0:
-        result = list(sm_rec.values())[0]['capabilities']
+    try:
+        sr_ref = session.xenapi.SR.get_by_uuid(sr_uuid)
+        sm_type = session.xenapi.SR.get_record(sr_ref)['type']
+        sm_rec = session.xenapi.SM.get_all_records_where(
+            "field \"type\" = \"%s\"" % sm_type)
 
-    session.xenapi.logout()
-    return result
+        # SM expects at least one entry of any SR type
+        if len(sm_rec) > 0:
+            result = list(sm_rec.values())[0]['capabilities']
 
+        return result
+    finally:
+        session.xenapi.session.logout()
 
 def sr_get_driver_info(driver_info):
     results = {}
@@ -1274,11 +1275,15 @@ def is_active(self, name):
 
     def mark_sr(self, name, sruuid, started):
         session = get_localAPI_session()
-        sr = session.xenapi.SR.get_by_uuid(sruuid)
-        if started:
-            session.xenapi.SR.add_to_other_config(sr, name, "active")
-        else:
-            session.xenapi.SR.remove_from_other_config(sr, name)
+        try:
+            sr = session.xenapi.SR.get_by_uuid(sruuid)
+
+            if started:
+                session.xenapi.SR.add_to_other_config(sr, name, "active")
+            else:
+                session.xenapi.SR.remove_from_other_config(sr, name)
+        finally:
+            session.xenapi.session.logout()
 
     def activate(self, name, sruuid):
         if name in self.points: