Robustify python3 version + fixes

http-disk-server.py → http_disk_server.py

@@ -216,6 +216,7 @@ def __init__(self, *args, **kwargs):
         def get_content_size(self):
             return get_device_size(self.disk_fd)
 
+        # Ignore broken pipe for python2 version.
         # See: https://stackoverflow.com/questions/6063416/python-basehttpserver-how-do-i-catch-trap-broken-pipe-errors#answer-14355079
         def finish(self):
             try:
@@ -348,7 +349,22 @@ def do_PUT(self):
             self.end_headers()
             response = BytesIO()
             response.write(b'Ok.')
-            self.wfile.write(response.getvalue())
+
+            # Write response for all python versions.
+            # Handle broken pipe error for python3 version.
+            try:
+                self.wfile.write(response.getvalue())
+            except Exception:
+                if sys.version_info > (3,):
+                    from builtins import BrokenPipeError
+                    try:
+                        raise
+                    except BrokenPipeError:
+                        # The client closed the connection too early,
+                        # so we should ignore the error.
+                        pass
+                else:
+                    raise
 
     return RequestHandler
 
@@ -409,6 +425,10 @@ def emit_server_ready():
                 disk_fd = open_device(disk, retry=False)
             else:
                 disk_fd = open_device(disk)
+
+            if SIGTERM_RECEIVED:
+                break
+
             is_block_device = stat.S_ISBLK(os.fstat(disk_fd).st_mode)
 
             HandlerClass = MakeRequestHandler(disk_fd, is_block_device)

nbd-http-server.py → nbd_http_server.py

@@ -25,8 +25,9 @@
 import signal
 import subprocess
 import sys
+import threading
 
-WORKING_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), '')
+WORKING_DIR = os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), '')
 NBDKIT_PLUGIN = WORKING_DIR + '../lib64/nbdkit/plugins/nbdkit-multi-http-plugin.so'
 
 # ==============================================================================
@@ -65,8 +66,14 @@ def run_or_ignore(fun):
 
 # -----------------------------------------------------------------------------
 
+THREAD_PRINT_LOCK = threading.Lock()
+
+def thread_print(str):
+    with THREAD_PRINT_LOCK:
+        print(str)
+
 def eprint(str):
-    print(OUTPUT_PREFIX + str, file=sys.stderr)
+    thread_print(OUTPUT_PREFIX + str)
 
 # -----------------------------------------------------------------------------
 
@@ -242,15 +249,43 @@ def clean_paths():
     ]
     if device_size is not None:
         arguments.append('device-size=' + str(device_size))
-    server = subprocess.Popen(arguments)
+
+    # Start nbdkit process.
+    server = subprocess.Popen(
+        arguments,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        universal_newlines=True,
+        env=dict(os.environ, PYTHONUNBUFFERED='1')
+    )
+
+    # Wait for init.
+    try:
+        with timeout(10):
+            while server.poll() is None:
+                line = server.stdout.readline().rstrip('\n')
+                if not line:
+                    continue
+                print(line)
+                if 'written pidfile' in line:
+                    break
+    except Exception as e:
+        raise Exception('Failed to start nbdkit server: {}.'.format(e))
+
+    # Continue to log server messages in stdout.
+    def log_server_messages():
+        while server.poll() is None:
+            line = server.stdout.readline().rstrip('\n')
+            if line:
+                thread_print(line)
+
+    server_stdout_thread = threading.Thread(target=log_server_messages)
+    server_stdout_thread.start()
+
     nbd = None
 
     try:
         nbd = attach_nbd(socket_path, nbd_name, pid_path)
-        # Flush to ensure we can read output (like NBD device used)
-        # in another process without waiting.
-        sys.stdout.flush()
-        sys.stderr.flush()
         while True:
             try:
                 if SIGTERM_RECEIVED:
@@ -265,6 +300,7 @@ def clean_paths():
             nbd.disconnect()
         server.send_signal(signal.SIGQUIT)
         server.wait()
+        server_stdout_thread.join()
         clean_paths()
 
 # ==============================================================================

setup.py

@@ -20,12 +20,12 @@
 
 setup(
     name="http-nbd-transfer",
-    version="1.4.0",
+    version="1.5.0",
     description="Set of tools to transfer NBD requests to an HTTP server",
     author="Ronan Abhamon <ronan.abhamon@vates.tech>",
     author_email="ronan.abhamon@vates.tech",
     url="https://vates.tech",
     license="GPLv3",
-    py_modules=["http-disk-server", "nbd-http-server"],
+    py_modules=["http_disk_server", "nbd_http_server"],
     scripts=["scripts/http-disk-server", "scripts/nbd-http-server"]
 )

tests/bin/http-disk-server

@@ -1 +1 @@
-../../http-disk-server.py
+../../http_disk_server.py

tests/bin/nbd-http-server

@@ -1 +1 @@
-../../nbd-http-server.py
+../../nbd_http_server.py

tests/conftest.py

@@ -137,7 +137,8 @@ def start_http_server(backing_path):
         arguments,
         stdout=sys.stdout,
         stderr=subprocess.STDOUT,
-        preexec_fn=os.setsid
+        preexec_fn=os.setsid,
+        env=dict(os.environ, PYTHONUNBUFFERED='1')
     )
 
 # ------------------------------------------------------------------------------
@@ -160,7 +161,8 @@ def start_nbd_server(volume_name, device_size):
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
         preexec_fn=os.setsid,
-        universal_newlines=True
+        universal_newlines=True,
+        env=dict(os.environ, PYTHONUNBUFFERED='1')
     )
 
     try: