Add session object and timeouts to API

debian/changelog

@@ -1,3 +1,9 @@
+olimp-control (1.2-1) unstable; urgency=low
+
+  * Use session object with retry and shorter timeout
+
+ -- LMIO TC <[email protected]>  Fri, 03 May 2024 12:43:00 +0000
+
 olimp-control (1.1-1) unstable; urgency=low
 
   * Update systemd unit file

makedeb.sh

@@ -6,7 +6,7 @@
 #   https://www.debian.org/doc/manuals/maint-guide/
 #   https://wiki.debian.org/Packaging
 
-ver='1.1'
+ver='1.2'
 projdir="olimp-control-${ver}"
 projtar="olimp-control_${ver}.orig.tar.gz"
 builddir="build"

olimp-control.py

@@ -5,6 +5,7 @@
 import hmac
 import json
 import requests
+import requests.adapters
 import signal
 import subprocess
 import threading
@@ -15,10 +16,11 @@ class LmioCtrlApi:
     _FAILED_AUTH = 'Response failed authentication validation'
     _NO_TICKETS = 'No tickets available'
 
-    def __init__(self, url, key, mid):
+    def __init__(self, url, key, mid, timeout):
         self._url = url
         self._key = key.encode('utf-8')
         self._mid = mid
+        self._timeout = timeout if timeout > 0.0 else None
 
     def _get_body_hmac(self, body):
         return hmac.HMAC(self._key, body, 'sha1').hexdigest()
@@ -48,15 +50,23 @@ def _validate_response(self, body, timestamp, headers):
         except Exception as e:
             return False
 
-    def do_ping(self):
+    def get_session(self):
+        session = requests.Session()
+        retry = requests.adapters.Retry(connect=5, backoff_factor=0.5)
+        adapter = requests.adapters.HTTPAdapter(max_retries=retry)
+        session.mount('http://', adapter)
+        session.mount('https://', adapter)
+        return session
+
+    def do_ping(self, session):
         payload = self._get_basic_payload()
         payload['uptime'] = subprocess.check_output(["uptime"]).decode('utf-8').strip()
         payload['hasRoot'] = 0
         post_data = json.dumps(payload)
         headers = self._get_auth_headers(post_data)
 
         try:
-            r = requests.post(self._url + '/ping', data=post_data, headers=headers)
+            r = session.post(self._url + '/ping', data=post_data, headers=headers, timeout=self._timeout)
             if self._validate_response(r.content, payload['timestamp'], r.headers):
                 response = r.json()
                 print(response['status'], response['message'])
@@ -66,13 +76,13 @@ def do_ping(self):
         except Exception as e:
             print(e)
 
-    def do_get_ticket(self):
+    def do_get_ticket(self, session):
         payload = self._get_basic_payload()
         get_data = json.dumps(payload)
         headers = self._get_auth_headers(get_data)
 
         try:
-            r = requests.get(self._url + '/ticket', data=get_data, headers=headers)
+            r = session.get(self._url + '/ticket', data=get_data, headers=headers, timeout=self._timeout)
             if self._validate_response(r.content, payload['timestamp'], r.headers):
                 if r.status_code != 404:
                     response = r.json()
@@ -93,7 +103,7 @@ def do_get_ticket(self):
 
         return {}
 
-    def do_post_ticket_results(self, results):
+    def do_post_ticket_results(self, session, results):
         payload = self._get_basic_payload()
         payload['tid'] = results['tid']
         payload['exectime'] = results['exectime']
@@ -104,7 +114,7 @@ def do_post_ticket_results(self, results):
         headers = self._get_auth_headers(post_data)
 
         try:
-            r = requests.post(self._url + '/ticket', data=post_data, headers=headers)
+            r = session.post(self._url + '/ticket', data=post_data, headers=headers, timeout=self._timeout)
             if self._validate_response(r.content, payload['timestamp'], r.headers):
                 response = r.json()
                 print(response['status'], response['message'])
@@ -136,14 +146,15 @@ def get_machine_id():
     all_macs = subprocess.check_output("cat /sys/class/net/*/address | sort", shell=True)
     return hashlib.sha1(all_macs).hexdigest()
 
-def main_loop(url, key, mid, poll_frequency, exit_event):
-    api = LmioCtrlApi(url, key, mid)
+def main_loop(url, key, mid, poll_frequency, timeout, exit_event):
+    api = LmioCtrlApi(url, key, mid, timeout)
     while True:
-        api.do_ping()
-        ticket = api.do_get_ticket()
+        session = api.get_session()
+        api.do_ping(session)
+        ticket = api.do_get_ticket(session)
         if ticket:
             ticket_results = execute_ticket(ticket)
-            api.do_post_ticket_results(ticket_results)
+            api.do_post_ticket_results(session, ticket_results)
 
         wait_result = exit_event.wait(poll_frequency)
         if wait_result:
@@ -155,6 +166,7 @@ def main_loop(url, key, mid, poll_frequency, exit_event):
     parser.add_argument('-f', '--poll-frequency', default=60.0, type=float,
         help='frequency of server check-in, in seconds')
     parser.add_argument('-k', '--key-file', default='/etc/olimp-control/key', help='path to key file')
+    parser.add_argument('-t', '--timeout', default=20.0, type=float, help='timeout for API operations, in seconds')
     parser.add_argument('url', nargs='?', default='https://ctrl.lmio.lt/olimp/api',
         help='url of control api base')
 
@@ -168,4 +180,4 @@ def main_loop(url, key, mid, poll_frequency, exit_event):
     signal.signal(signal.SIGINT, sigterm_handler)
     signal.signal(signal.SIGTERM, sigterm_handler)
 
-    main_loop(args.url, get_key(args.key_file), get_machine_id(), args.poll_frequency, shutdown_event)
+    main_loop(args.url, get_key(args.key_file), get_machine_id(), args.poll_frequency, args.timeout, shutdown_event)