CA-388451: ensure that xapi sessions are logged out

From: Mark Syms <mark.syms@citrix.com>

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

diff --git a/drivers/cleanup.py b/drivers/cleanup.py
index 97c332c..dce460e 100755
--- a/drivers/cleanup.py
+++ b/drivers/cleanup.py
@@ -2893,7 +2893,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):
     init(srUuid)
diff --git a/drivers/mpathcount.py b/drivers/mpathcount.py
index 0910954..6b8d622 100755
--- a/drivers/mpathcount.py
+++ b/drivers/mpathcount.py
@@ -55,7 +55,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 9c164fd..6ff4bd7 100755
--- a/drivers/sr_health_check.py
+++ b/drivers/sr_health_check.py
@@ -34,28 +34,30 @@ def main():
         util.SMlog("Unable to open local XAPI session", priority=util.LOG_ERR)
         return
 
-    localhost = util.get_localhost_uuid(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(
-            'field "type" = "{sm_type}"'.format(sm_type=sm_type))
-        for sr in srs:
-            pbds = session.xenapi.PBD.get_all_records_where(
-                'field "SR" = "{sr}" and field "host" = "{localhost}"'.format(
-                    sr=sr, localhost=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_uuid(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(
+                'field "type" = "{sm_type}"'.format(sm_type=sm_type))
+            for sr in srs:
+                pbds = session.xenapi.PBD.get_all_records_where(
+                    'field "SR" = "{sr}" and field "host" = "{localhost}"'.format(
+                        sr=sr, localhost=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__":
     main()
diff --git a/drivers/util.py b/drivers/util.py
index 2353f4b..4f9fa9d 100755
--- a/drivers/util.py
+++ b/drivers/util.py
@@ -360,16 +360,18 @@ def ioretry_stat(f, 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 atleast one entry of any SR type
-    if len(sm_rec) > 0:
-        result = sm_rec.values()[0]['capabilities']
-    
-    session.xenapi.logout()
+    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)
+
+        # SM expects atleast one entry of any SR type
+        if len(sm_rec) > 0:
+            result = sm_rec.values()[0]['capabilities']
+    finally:
+        session.xenapi.logout()
+
     return result
 
 def sr_get_driver_info(driver_info):
@@ -1192,11 +1194,14 @@ class FistPoint:
 
     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:
