Add retry logic in case of websocket connection closing. (Qiskit#324)

* Add retry logic in case of websocket connection closing.

* Change from allowing multiple retries to only allowing one retry.

* Updated doc string.

* Add test and mock websocket server handler for a successful retry.

* Remove currently unused global vars.

* Remove unused global import, add decorator

* Update doc string.

* Add test and mock websocket server handler for a failed retry.

* Add test and mock websocket server handler for job not found.

* Update doc strings for test cases.

* Update docstrings and documentation.

* moved helper decorator to test/decorators.py

* Fixed lint and style

* Added logger warning for retry when websocket is closed.

* Refactored handle_token_retry_success

* Removed decorator to keep track of retry

* Wrong function was called, fix.

* Updated logger.warning message.

* Improve readibility and consistensy.

Author: Luis Garcia

qiskit/providers/ibmq/api_v2/clients/websocket.py

@@ -135,7 +135,8 @@ def get_job_status(self, job_id, timeout=None):
 
         Reads status messages from the API, which are issued at regular
         intervals (20 seconds). When a final state is reached, the server
-        closes the socket.
+        closes the socket. If the websocket connection is closed without
+        a reason, there is an attempt to retry one time.
 
         Args:
             job_id (str): id of the job.
@@ -154,6 +155,7 @@ def get_job_status(self, job_id, timeout=None):
 
         original_timeout = timeout
         start_time = time.time()
+        attempt_retry = True  # By default, attempt to retry if the websocket connection closes.
         last_status = None
 
         try:
@@ -198,7 +200,17 @@ def get_job_status(self, job_id, timeout=None):
                     elif ex.code == 4002:
                         break
                     elif ex.code == 4003:
+                        attempt_retry = False  # No point in retrying.
                         message = 'Job id not found'
+
+                    if attempt_retry:
+                        logger.warning('Connection with the websocket closed '
+                                       'unexpectedly: %s(status_code=%s). '
+                                       'Retrying get_job_status.', message, ex.code)
+                        attempt_retry = False  # Disallow further retries.
+                        websocket = yield from self._connect(url)
+                        continue
+
                     raise WebsocketError('Connection with websocket closed '
                                          'unexpectedly: {}'.format(message)) from ex
         finally:

test/ibmq/websocket/test_websocket.py

@@ -29,7 +29,9 @@
 
 from .websocket_server import (
     TOKEN_JOB_COMPLETED, TOKEN_JOB_TRANSITION, TOKEN_WRONG_FORMAT,
-    TOKEN_TIMEOUT, websocket_handler)
+    TOKEN_TIMEOUT, TOKEN_WEBSOCKET_RETRY_SUCCESS,
+    TOKEN_WEBSOCKET_RETRY_FAILURE, TOKEN_WEBSOCKET_JOB_NOT_FOUND,
+    websocket_handler)
 
 TEST_IP_ADDRESS = '127.0.0.1'
 INVALID_PORT = 9876
@@ -111,3 +113,29 @@ def test_invalid_response(self):
         with self.assertRaises(WebsocketIBMQProtocolError):
             _ = asyncio.get_event_loop().run_until_complete(
                 client.get_job_status('job_id'))
+
+    def test_websocket_retry_success(self):
+        """Test retrieving a job status during a retry attempt."""
+        client = WebsocketClient('ws://{}:{}'.format(
+            TEST_IP_ADDRESS, VALID_PORT), TOKEN_WEBSOCKET_RETRY_SUCCESS)
+        response = asyncio.get_event_loop().run_until_complete(
+            client.get_job_status('job_id'))
+        self.assertIsInstance(response, dict)
+        self.assertIn('status', response)
+        self.assertEqual(response['status'], 'COMPLETED')
+
+    def test_websocket_retry_failure(self):
+        """Test exceeding the retry limit for retrieving a job status."""
+        client = WebsocketClient('ws://{}:{}'.format(
+            TEST_IP_ADDRESS, VALID_PORT), TOKEN_WEBSOCKET_RETRY_FAILURE)
+        with self.assertRaises(WebsocketError):
+            _ = asyncio.get_event_loop().run_until_complete(
+                client.get_job_status('job_id'))
+
+    def test_websocket_job_not_found(self):
+        """Test retrieving a job status for an non existent id."""
+        client = WebsocketClient('ws://{}:{}'.format(
+            TEST_IP_ADDRESS, VALID_PORT), TOKEN_WEBSOCKET_JOB_NOT_FOUND)
+        with self.assertRaises(WebsocketError):
+            _ = asyncio.get_event_loop().run_until_complete(
+                client.get_job_status('job_id'))

test/ibmq/websocket/websocket_server.py

@@ -24,6 +24,9 @@
 TOKEN_JOB_TRANSITION = 'token_job_transition'
 TOKEN_TIMEOUT = 'token_timeout'
 TOKEN_WRONG_FORMAT = 'token_wrong_format'
+TOKEN_WEBSOCKET_RETRY_SUCCESS = 'token_websocket_retry_success'
+TOKEN_WEBSOCKET_RETRY_FAILURE = 'token_websocket_retry_failure'
+TOKEN_WEBSOCKET_JOB_NOT_FOUND = 'token_websocket_job_not_found'
 
 
 @asyncio.coroutine
@@ -39,7 +42,10 @@ def websocket_handler(websocket, path):
     if token in (TOKEN_JOB_COMPLETED,
                  TOKEN_JOB_TRANSITION,
                  TOKEN_TIMEOUT,
-                 TOKEN_WRONG_FORMAT):
+                 TOKEN_WRONG_FORMAT,
+                 TOKEN_WEBSOCKET_RETRY_SUCCESS,
+                 TOKEN_WEBSOCKET_RETRY_FAILURE,
+                 TOKEN_WEBSOCKET_JOB_NOT_FOUND):
         msg_out = json.dumps({'type': 'authenticated'})
         yield from websocket.send(msg_out.encode('utf8'))
     else:
@@ -55,6 +61,12 @@ def websocket_handler(websocket, path):
         yield from handle_token_timeout(websocket)
     elif token == TOKEN_WRONG_FORMAT:
         yield from handle_token_wrong_format(websocket)
+    elif token == TOKEN_WEBSOCKET_RETRY_SUCCESS:
+        yield from handle_token_retry_success(websocket)
+    elif token == TOKEN_WEBSOCKET_RETRY_FAILURE:
+        yield from handle_token_retry_failure(websocket)
+    elif token == TOKEN_WEBSOCKET_JOB_NOT_FOUND:
+        yield from handle_token_job_not_found(websocket)
 
 
 @asyncio.coroutine
@@ -94,3 +106,25 @@ def handle_token_wrong_format(websocket):
     """Return a status in an invalid format."""
     yield from websocket.send('INVALID'.encode('utf8'))
     yield from websocket.close()
+
+
+@asyncio.coroutine
+def handle_token_retry_success(websocket):
+    """Close the socket once and force a retry."""
+    if not hasattr(handle_token_retry_success, 'retry_attempt'):
+        setattr(handle_token_retry_success, 'retry_attempt', True)
+        yield from handle_token_retry_failure(websocket)
+    else:
+        yield from handle_token_job_completed(websocket)
+
+
+@asyncio.coroutine
+def handle_token_retry_failure(websocket):
+    """Continually close the socket, until both the first attempt and retry fail."""
+    yield from websocket.close()
+
+
+@asyncio.coroutine
+def handle_token_job_not_found(websocket):
+    """Close the socket, specifying code for job not found."""
+    yield from websocket.close(code=4003)