import threading
import logging

try:
    from drmaa import Session, JobControlAction
except OSError as e:
    LOAD_ERROR_MESSAGE = "OSError - problem loading shared library [%s]." % e
    Session = None
except ImportError as e:
    LOAD_ERROR_MESSAGE = "ImportError - problem importing library (`pip install drmaa` may fix this) [%s]." % e
    # Will not be able to use DRMAA
    Session = None

NO_DRMAA_MESSAGE = "Attempt to use DRMAA, but DRMAA Python library cannot be loaded. "

log = logging.getLogger(__name__)


class DrmaaSessionFactory(object):
    """
    Abstraction used to production DrmaaSession wrappers.
    """
    def __init__(self):
        self.session_constructor = Session

    def get(self, **kwds):
        session_constructor = self.session_constructor
        if session_constructor is None:
            raise Exception(NO_DRMAA_MESSAGE + LOAD_ERROR_MESSAGE)
        return DrmaaSession(session_constructor, **kwds)


class DrmaaSession(object):
    """
    Abstraction around `drmaa` module `Session` objects.
    """
    session_lock = threading.Lock()
    session = None

    def __init__(self, session_constructor, **kwds):
        with DrmaaSession.session_lock:
            if DrmaaSession.session is None:
                log.debug("Initializing DRMAA session from thread %s", threading.current_thread().name)
                DrmaaSession.session = session_constructor()
                DrmaaSession.session.initialize()

    def run_job(self, **kwds):
        """
        Create a DRMAA job template, populate with specified properties,
        run the job, and return the external_job_id.
        """
        template = DrmaaSession.session.createJobTemplate()
        try:
            for key in kwds:
                setattr(template, key, kwds[key])
            with DrmaaSession.session_lock:
                return DrmaaSession.session.runJob(template)
        finally:
            DrmaaSession.session.deleteJobTemplate(template)

    def kill(self, external_job_id):
        with DrmaaSession.session_lock:
            return DrmaaSession.session.control(str(external_job_id), JobControlAction.TERMINATE)

    def job_status(self, external_job_id):
        return DrmaaSession.session.jobStatus(str(external_job_id))

    def close(self):
        return DrmaaSession.session.exit()


__all__ = ['DrmaaSessionFactory']

@@ -1,24 +1,25 @@
 import 
-threading
+contextlib
 %0Aimport 
@@ -25,16 +25,33 @@
  logging
+%0Aimport threading
 %0A%0Atry:%0A 
@@ -1094,16 +1094,38 @@
 .Lock()%0A
+    session_count = 0%0A
     sess
@@ -1198,37 +1198,30 @@
        with 
-DrmaaSession.
+self._
 session_lock
@@ -1212,32 +1212,34 @@
 lf._session_lock
+()
 :%0A            if
@@ -1261,32 +1261,199 @@
 ession is None:%0A
+                if DrmaaSession.session_count != 0:%0A                    log.warn(%22DrmaaSession.session is None but session_count is non-zero - logic error occurred.%22)%0A
                 
@@ -1650,16 +1650,180 @@
 ialize()
+%0A            DrmaaSession.session_count += 1%0A%0A%0A    @contextlib.contextmanager%0A    def _session_lock(self):%0A        with DrmaaSession.session_lock:%0A            yield
 %0A%0A    de
@@ -2654,42 +2654,558 @@
     
-return DrmaaSession.session.exit()
+with self._session_lock():%0A            if DrmaaSession.session_count == 0:%0A                log.warn(%22close() called with zero active session counted - logic error.%22)%0A                return%0A%0A            DrmaaSession.session_count -= 1%0A            if DrmaaSession.session_count == 0:%0A                if DrmaaSession.session is None:%0A                    log.warn(%22close() called with a non-zero session count but no session is defined.%22)%0A                    return%0A%0A                DrmaaSession.session.exit()%0A                DrmaaSession.session = None
 %0A%0A%0A_

Revise handling of global session in drmaa factory.