Repository: dwavesystems/dwave-cloud-client (python)

Issue:
Robust `Client.close()`, no client/solver/future reuse after close
After `Client.close()`:
- all workers have to be killed, immediately if user doesn't want to wait for results (which they are now too, but in a blocking manner)
- all solvers have to be invalidated, so no future submissions are possible (channel is closed),
- all futures have to be resolved (with timeout/shutdown error if `Client.close(wait=False)`)

Pull request description:
Robust client close
Close #217.

Patch:
diff --git a/dwave/cloud/client/base.py b/dwave/cloud/client/base.py
index b000c9a..7eea948 100644
--- a/dwave/cloud/client/base.py
+++ b/dwave/cloud/client/base.py
@@ -65,6 +65,7 @@ from plucky import pluck
 from dwave.cloud import api
 from dwave.cloud.package_info import __packagename__, __version__
 from dwave.cloud.exceptions import *    # TODO: fix
+from dwave.cloud.exceptions import UseAfterCloseError
 from dwave.cloud.computation import Future
 from dwave.cloud.config import load_config, update_config, validate_config_v1
 from dwave.cloud.config import constants as config_constants
@@ -85,6 +86,15 @@ __all__ = ['Client']
 logger = logging.getLogger(__name__)
 
 
+def _require_initialized(method):
+    def _method(obj, *args, **kwargs):
+        if not getattr(obj, '_initialized', False):
+            raise UseAfterCloseError(
+                f"{method.__name__} cannot be called after client has been closed")
+        return method(obj, *args, **kwargs)
+    return _method
+
+
 class Client(object):
     """Base client class for all D-Wave API clients. Used by QPU, software and
     hybrid :term:`sampler` classes.
@@ -439,6 +449,8 @@ class Client(object):
     def __init__(self, **kwargs):
         logger.debug("Client init called with: %r", kwargs)
 
+        self._initialized = False
+
         # derive instance-level defaults from class defaults and init defaults
         self.defaults = copy.deepcopy(self.DEFAULTS)
         user_defaults = kwargs.pop('defaults', None)
@@ -519,6 +531,8 @@ class Client(object):
         self._download_answer_executor = \
             ThreadPoolExecutor(self._DOWNLOAD_ANSWER_THREAD_COUNT)
 
+        self._initialized = True
+
     def create_session(self):
         """Create a new requests session based on client's (self) params.
 
@@ -579,6 +593,9 @@ class Client(object):
             >>> client.close()    # doctest: +SKIP
 
         """
+        if not self._initialized:
+            return
+
         # Finish all the work that requires the connection
         logger.debug("Joining submission queue")
         self._submission_queue.join()
@@ -590,14 +607,13 @@ class Client(object):
         self._load_queue.join()
 
         logger.debug("Shutting down problem upload executor")
-        self._upload_problem_executor.shutdown()
+        self._upload_problem_executor.shutdown(wait=True)
         logger.debug("Shutting down problem part upload executor")
-        self._upload_part_executor.shutdown()
+        self._upload_part_executor.shutdown(wait=True)
         logger.debug("Shutting down problem encoder executor")
-        self._encode_problem_executor.shutdown()
-
+        self._encode_problem_executor.shutdown(wait=True)
         logger.debug("Shutting down answer download executor")
-        self._download_answer_executor.shutdown()
+        self._download_answer_executor.shutdown(wait=True)
 
         # Send kill-task to all worker threads
         # Note: threads can't be 'killed' in Python, they have to die by
@@ -623,6 +639,9 @@ class Client(object):
         if solvers_session:
             solvers_session.close()
 
+        # this client can't be used anymore
+        self._initialized = False
+
     def __enter__(self):
         """Let connections be used in with blocks."""
         return self
@@ -1200,12 +1219,14 @@ class Client(object):
         except IndexError:
             raise SolverNotFoundError("Solver with the requested features not available")
 
+    @_require_initialized
     def _submit(self, body, future):
         """Enqueue a problem for submission to the server.
 
         This method is thread safe.
         """
         self._submission_queue.put(self._submit.Message(body, future))
+
     _submit.Message = namedtuple('Message', ['body', 'future'])
 
     def _do_submit_problems(self):
@@ -1407,6 +1428,7 @@ class Client(object):
             # lock in the future, otherwise deadlock occurs.
             future._set_exception(exc)
 
+    @_require_initialized
     def _cancel(self, id_, future):
         """Enqueue a problem to be canceled.
 
@@ -1458,6 +1480,7 @@ class Client(object):
         finally:
             session.close()
 
+    @_require_initialized
     def _poll(self, future: Future) -> None:
         """Enqueue a problem to poll the server for status."""
 
@@ -1566,8 +1589,8 @@ class Client(object):
                         if not add(future):
                             return
                     else:
-                        task_done()
                         self._poll_queue.put(task)
+                        task_done()
                         break
 
                 # build a query string with ids of all futures in this frame
@@ -1634,6 +1657,7 @@ class Client(object):
         finally:
             session.close()
 
+    @_require_initialized
     def _load(self, future):
         """Enqueue a problem to download results from the server.
 
@@ -1685,6 +1709,7 @@ class Client(object):
         finally:
             session.close()
 
+    @_require_initialized
     def _download_answer_binary_ref(self, *, auth_method: str, url: str,
                                     output: Optional[io.IOBase] = None) -> concurrent.futures.Future:
         """Initiate binary-ref answer download, returning the binary data in
@@ -1724,6 +1749,7 @@ class Client(object):
 
         return output
 
+    @_require_initialized
     def upload_problem_encoded(self, problem, problem_id=None, **kwargs):
         """Initiate multipart problem upload, returning the Problem ID in a
         :class:`~concurrent.futures.Future`.
diff --git a/dwave/cloud/exceptions.py b/dwave/cloud/exceptions.py
index 47ce4d2..01bfe74 100644
--- a/dwave/cloud/exceptions.py
+++ b/dwave/cloud/exceptions.py
@@ -77,6 +77,10 @@ class ProblemUploadError(Exception):
     """Problem multipart upload failed."""
 
 
+class UseAfterCloseError(Exception):
+    """Use attempted after client/solver/connection was closed."""
+
+
 # for backward compatibility
 from dwave.cloud.api.exceptions import RequestTimeout
 
diff --git a/dwave/cloud/solver.py b/dwave/cloud/solver.py
index d621d07..c9fdc2a 100644
--- a/dwave/cloud/solver.py
+++ b/dwave/cloud/solver.py
@@ -33,6 +33,7 @@ import concurrent.futures
 import copy
 import logging
 import orjson
+import weakref
 from collections import abc
 from functools import partial
 from tempfile import SpooledTemporaryFile
@@ -113,9 +114,25 @@ class BaseSolver(object):
     _handled_problem_types = {}
     _handled_encoding_formats = {}
 
-    def __init__(self, client, data):
-        # client handles async api requests (via local thread pool)
-        self.client = client
+    _client_ref = None
+
+    @property
+    def client(self):
+        """Returns a reference to client if it still exists, or None otherwise."""
+
+        client = None
+        if self._client_ref is not None:
+            client = self._client_ref()
+            if client is None:
+                raise RuntimeError('Client unavailable / closed')
+
+        return client
+
+    def __init__(self, client: Optional['dwave.cloud.Client'], data: dict):
+        # note: we use a weakref so that client can be closed and gc'ed
+        #       while keeping solver instances alive
+        # dev note: allow None value for testing
+        self._client_ref = weakref.ref(client) if client is not None else None
 
         # data for each solver includes at least: id, description, properties,
         # status and avg_load
diff --git a/releasenotes/notes/robust-client-close-00028d8fc087c441.yaml b/releasenotes/notes/robust-client-close-00028d8fc087c441.yaml
new file mode 100644
index 0000000..99617d7
--- /dev/null
+++ b/releasenotes/notes/robust-client-close-00028d8fc087c441.yaml
@@ -0,0 +1,13 @@
+---
+features:
+  - |
+    More robust ``Client.close()``.
+
+    Client use after close is now disabled. An attempt to submit a problem, poll
+    for a status, or download an answer will now result in an ``UseAfterCloseError``
+    exception.
+
+    Also, solvers and computations now only weakly reference the client, so they
+    don't block client resources cleanup.
+
+    See `#217 <https://github.com/dwavesystems/dwave-cloud-client/issues/217>`_.


Test patch:
diff --git a/tests/test_client.py b/tests/test_client.py
index d39eacc..88296a7 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -20,8 +20,11 @@ test_mock_solver_loading.py duplicates some of these tests against a mock server
 
 import importlib
 import json
+import sys
+import threading
 import time
 import unittest
+import weakref
 from collections import defaultdict
 from contextlib import contextmanager
 from unittest import mock
@@ -661,6 +664,25 @@ class VerifyLazyClientImport(unittest.TestCase):
             self.assertTrue(issubclass(getattr(mod, 'Client'), Client))
 
 
+class VerifyClientCleanup(unittest.TestCase):
+
+    def test_client_close(self):
+        thread_cnt_base = threading.active_count()
+
+        client = Client(endpoint='endpoint', token='token')
+        ref = weakref.ref(client)
+
+        self.assertGreater(threading.active_count(), thread_cnt_base)
+        self.assertGreater(sys.getrefcount(client), 1)
+
+        client.close()
+        del client
+
+        # verify all threads closed and no client referrers
+        self.assertEqual(threading.active_count(), thread_cnt_base)
+        self.assertIsNone(ref())
+
+
 @mock.patch("dwave.cloud.regions.get_regions", get_default_regions)
 class ClientConfigIntegration(unittest.TestCase):
 
diff --git a/tests/test_mock_solver_loading.py b/tests/test_mock_solver_loading.py
index cc80a78..bb6fc01 100644
--- a/tests/test_mock_solver_loading.py
+++ b/tests/test_mock_solver_loading.py
@@ -171,6 +171,16 @@ class MockSolverLoading(unittest.TestCase):
                 self.assertEqual(solver.data['status'], ref.data['status'])
                 self.assertEqual(solver.data['avg_load'], ref.data['avg_load'])
 
+    def test_client_ref(self):
+        with Client(endpoint=self.endpoint, token=self.token) as client:
+            solver = client.get_solver(self.solver1_name)
+            self.assertEqual(solver.client, client)
+
+        del client
+
+        with self.assertRaises(RuntimeError):
+            solver.client
+
     def test_load_all_solvers(self):
         refs = [solver_object(self.solver1_name), solver_object(self.solver2_name)]
 
diff --git a/tests/test_mock_submission.py b/tests/test_mock_submission.py
index 88fb92d..c1b247b 100644
--- a/tests/test_mock_submission.py
+++ b/tests/test_mock_submission.py
@@ -17,6 +17,7 @@
 import threading
 import time
 import unittest
+import weakref
 import zlib
 from concurrent.futures import TimeoutError, ThreadPoolExecutor
 from typing import Any
@@ -36,7 +37,7 @@ from dwave.cloud.client import Client
 from dwave.cloud.computation import Future
 from dwave.cloud.exceptions import (
     SolverFailureError, CanceledFutureError, SolverError,
-    InvalidAPIResponseError)
+    InvalidAPIResponseError, UseAfterCloseError)
 from dwave.cloud.solver import Solver
 from dwave.cloud.utils.qubo import evaluate_ising
 from dwave.cloud.utils.time import utcrel
@@ -873,6 +874,7 @@ class TestComputationID(unittest.TestCase):
         self.assertEqual(f.sampleset.wait_id(), submission_id)
         self.assertEqual(f.sampleset.wait_id(timeout=1), submission_id)
 
+
 @mock.patch('time.sleep', lambda *x: None)
 class TestOffsetHandling(_QueryTest):
 
@@ -1116,6 +1118,55 @@ class TestComputationSamplesetCaching(unittest.TestCase):
                 self.assertIsNone(response._sampleset())
 
 
+@mock.patch.object(Client, 'create_session', lambda _: mock.Mock())
+class TestClientClose(MockSubmissionBase, unittest.TestCase):
+    # exception is raised when client is used after it's been closed
+
+    def test_submit_fails(self):
+        with Client(**self.config) as client:
+            solver = Solver(client, self.sapi.solver.data)
+
+        with self.subTest('problem submit fails'):
+            with self.assertRaises(UseAfterCloseError):
+                solver.sample_ising(*self.sapi.problem)
+
+        with self.subTest('problem data multipart upload fails'):
+            with self.assertRaises(UseAfterCloseError):
+                client.upload_problem_encoded('mock')
+
+    def test_retrieve_fails(self):
+        with Client(**self.config) as client:
+            solver = Solver(client, self.sapi.solver.data)
+
+        with self.subTest('problem status poll fails'):
+            with self.assertRaises(UseAfterCloseError):
+                solver._retrieve_problem('mock-id')
+
+        with self.subTest('qpu answer load fails'):
+            with self.assertRaises(UseAfterCloseError):
+                client.retrieve_answer('problem-id')
+
+        with self.subTest('binary answer download fails'):
+            with self.assertRaises(UseAfterCloseError):
+                solver._download_binary_ref(auth_method='mock', url='mock')
+
+    def test_cancel_fails(self):
+        with Client(**self.config) as client:
+            ...
+
+        with self.assertRaises(UseAfterCloseError):
+            client._cancel('future-id', 'future')
+
+    def test_ref_cleanup(self):
+        with Client(**self.config) as client:
+            ref = weakref.ref(client)
+            solver = Solver(client, self.sapi.solver.data)
+
+        del client
+
+        self.assertIsNone(ref())
+
+
 class TestThreadSafety(unittest.TestCase):
 
     def test_result_loading(self):