use zmq for storage write

docs/config.rst

@@ -98,6 +98,10 @@ General settings
       Remote port to connect to in order to get console messages (defaults to
       ``5557``).
 
+  * ``data``: integer [optional]
+      Remote port for data transfers between the client and agent (defaults to
+      ``0`` for a dynamic port assignment).
+
   * ``host``: string [optional]
       Remote host name or ip to connect to as a client to interact with the
       MTDA agent (defaults to ``localhost``).

mtda/client.py

@@ -17,6 +17,7 @@
 import subprocess
 import tempfile
 import time
+import zmq
 import zstandard as zstd
 
 from mtda.main import MultiTenantDeviceAccess
@@ -44,6 +45,7 @@ def __init__(self, host=None, session=None, config_files=None,
         else:
             self._impl = agent
         self._agent = agent
+        self._data = None
 
         if session is None:
             HOST = socket.gethostname()
@@ -115,7 +117,12 @@ def storage_open(self):
         while tries > 0:
             tries = tries - 1
             try:
-                self._impl.storage_open(self._session)
+                host = self.remote()
+                port = self._impl.storage_open(self._session)
+                context = zmq.Context()
+                socket = context.socket(zmq.PUSH)
+                socket.connect(f'tcp://{host}:{port}')
+                self._data = socket
                 return
             except Exception:
                 if tries > 0:
@@ -201,7 +208,7 @@ def storage_write_image(self, path, callback=None):
 
         try:
             # Prepare for download/copy
-            file.prepare(image_size)
+            file.prepare(self._data, image_size)
 
             # Copy image to shared storage
             file.copy()
@@ -298,6 +305,7 @@ def __init__(self, path, agent, session, blksz, callback=None):
         self._path = path
         self._session = session
         self._totalread = 0
+        self._totalsent = 0
 
     def bmap(self, path):
         return None
@@ -324,21 +332,25 @@ def flush(self):
         inputsize = self._inputsize
         totalread = self._totalread
         outputsize = self._outputsize
+
+        agent.storage_flush(self._totalsent)
         while True:
-            status, writing, written = agent.storage_status(self._session)
+            status, writing, written = agent.storage_status()
             if callback is not None:
                 callback(imgname, totalread, inputsize, written, outputsize)
             if writing is False:
                 break
             time.sleep(0.5)
+        self._socket.close()
 
     def path(self):
         return self._path
 
-    def prepare(self, output_size=None, compression=None):
+    def prepare(self, socket, output_size=None, compression=None):
         compr = self.compression() if compression is None else compression
         self._inputsize = self.size()
         self._outputsize = output_size
+        self._socket = socket
         # if image is uncompressed, we compress on the fly
         if compr == CONSTS.IMAGE.RAW.value:
             compr = CONSTS.IMAGE.ZST.value
@@ -362,22 +374,8 @@ def size(self):
         return None
 
     def _write_to_storage(self, data):
-        max_tries = int(CONSTS.STORAGE.TIMEOUT / CONSTS.STORAGE.RETRY_INTERVAL)
-
-        for _ in range(max_tries):
-            result = self._agent.storage_write(data, self._session)
-            if result != 0:
-                break
-            time.sleep(CONSTS.STORAGE.RETRY_INTERVAL)
-
-        if result > 0:
-            return result
-        elif result < 0:
-            exc = 'write or decompression error from shared storage'
-            raise IOError(exc)
-        else:
-            exc = 'timeout from shared storage'
-            raise IOError(exc)
+        self._socket.send(data)
+        self._totalsent += len(data)
 
 
 class ImageLocal(ImageFile):

mtda/constants.py

@@ -77,7 +77,6 @@ class STORAGE:
 
 
 class WRITER:
-    QUEUE_SLOTS = 16
-    QUEUE_TIMEOUT = 5
+    RECV_TIMEOUT = 5
     READ_SIZE = 1*1024*1024
     WRITE_SIZE = 1*1024*1024