Add ability to record graphics console to video files

Use old mp4v so that it is quick, tolerating huge files

confluent_client/bin/nodeconsole

@@ -39,6 +39,8 @@ import select
 import signal
 import socket
 import re
+import queue
+import threading
 import tty
 import termios
 import fcntl
@@ -88,6 +90,11 @@ argparser.add_option('-i', '--interval', type='float',
                           'works for one node')
 argparser.add_option('-v', '--video', action='store_true', default=False,
                      help='Attempt to continuously stream video from nodes that support streaming console via confluent')
+argparser.add_option('-o', '--outputfile', type='string', default=None,
+                     help='Record VNC mode output to the specified video file. When '
+                          'multiple nodes are recorded, use a confluent attribute '
+                          'expression (e.g. "{node}.mp4") to make the filename unique '
+                          'for each node.  If multiple nodes specified but not an expression, node name will be inserted.')
 argparser.add_option('-w','--windowed', action='store_true', default=False,
                      help='Open terminal windows for each node.  The '
                           'environment variable NODECONSOLE_WINDOWED_COMMAND '
@@ -827,27 +834,34 @@ cheight = 0
 imagedatabynode = {}
 firstnodename = None
 
+# Terminal drawing is offloaded to a single worker thread so slow image
+# encoding/output doesn't stall the asyncio loop. The lock keeps the worker
+# from interleaving output with redraw() on the main thread.
+draw_lock = threading.Lock()
+draw_queue = queue.Queue()
+
 def redraw():
-    for node in imagedatabynode:
-        imgdata = imagedatabynode[node]
-        if node in nodepositions:
-            prep_node_tile(node)
-            cursor_save()
-        else:
-            if options.interval is not None:
-                if node != firstnodename:
-                    sys.stderr.write('Multiple nodes not supported for interval')
-                    sys.exit(1)
-                sticky_cursor()
-            sys.stdout.write('{}: '.format(node))
-        # one row is used by our own name, so cheight - 1 for that allowance
-        draw_image(imgdata, cwidth, cheight - 1 if cheight else cheight)
-        if node in nodepositions:
-            cursor_restore()
-            reset_cursor(node)
-        else:
-            sys.stdout.write('\n')
-        sys.stdout.flush()
+    with draw_lock:
+        for node in imagedatabynode:
+            imgdata = imagedatabynode[node]
+            if node in nodepositions:
+                prep_node_tile(node)
+                cursor_save()
+            else:
+                if options.interval is not None:
+                    if node != firstnodename:
+                        sys.stderr.write('Multiple nodes not supported for interval')
+                        sys.exit(1)
+                    sticky_cursor()
+                sys.stdout.write('{}: '.format(node))
+            # one row is used by our own name, so cheight - 1 for that allowance
+            draw_image(imgdata, cwidth, cheight - 1 if cheight else cheight)
+            if node in nodepositions:
+                cursor_restore()
+                reset_cursor(node)
+            else:
+                sys.stdout.write('\n')
+            sys.stdout.flush()
 resized = False
 inputwatcher = None
 
@@ -961,9 +975,29 @@ async def grab_vncs(urlbynode):
     try:
         if streaming:
             directed = direct_console()
+        outputbynode = {}
+        if options.outputfile:
+            expr = options.outputfile
+            if '{' in expr:
+                # Let the server evaluate the expression to a unique name per node
+                sess = client.Command()
+                noderange = ','.join(urlbynode)
+                async for res in sess.create(
+                        '/noderange/{}/attributes/expression'.format(noderange),
+                        {'expression': expr}):
+                    for node in res.get('databynode', {}):
+                        outputbynode[node] = res['databynode'][node]['value']
+            else:
+                for node in urlbynode:
+                    if len(urlbynode) > 1:
+                        base, ext = os.path.splitext(options.outputfile)
+                        outputbynode[node] = '{}-{}{}'.format(base, node, ext)
+                    else:
+                        outputbynode[node] = expr
         for node in urlbynode:
             url = urlbynode[node]
-            tasks.append(asyncio.create_task(do_vnc(node, url)))
+            tasks.append(asyncio.create_task(
+                do_vnc(node, url, outputbynode.get(node))))
         await asyncio.gather(*tasks)
     except Exception as e:
         sys.stderr.write(f"Error in grab_vncs: {e}\n")
@@ -1046,13 +1080,13 @@ def toggle_focus_all():
         del focused_nodes[node]
 
 
-async def do_vnc(node, url):
+async def do_vnc(node, url, outputfile=None):
     global streaming
     keeprunning = True
     retries = 5
     while keeprunning:
         try:
-            async with await vnc.VNCClient.create(url) as client:
+            async with await vnc.VNCClient.create(url, outputfile=outputfile) as client:
                 vncclientsbynode[node] = client
                 while True:
                     # Retrieve pixels as a 3D numpy array
@@ -1090,6 +1124,16 @@ async def do_vnc(node, url):
 
 
 def draw_node(node, imgdata, errorstr, firstnodename, cwidth, cheight):
+    # Queue the frame for the drawing thread; frames arriving mid-draw wait here
+    draw_queue.put((node, imgdata, errorstr, firstnodename, cwidth, cheight))
+
+def _draw_worker():
+    while True:
+        node, imgdata, errorstr, firstnodename, cwidth, cheight = draw_queue.get()
+        with draw_lock:
+            _draw_node(node, imgdata, errorstr, firstnodename, cwidth, cheight)
+
+def _draw_node(node, imgdata, errorstr, firstnodename, cwidth, cheight):
     imagedatabynode[node] = imgdata
     if node in nodepositions:
         prep_node_tile(node)
@@ -1113,6 +1157,9 @@ def draw_node(node, imgdata, errorstr, firstnodename, cwidth, cheight):
         sys.stdout.write('\n')
     sys.stdout.flush()
 
+draw_thread = threading.Thread(target=_draw_worker, daemon=True)
+draw_thread.start()
+
 if options.screenshot or options.video:
     if not sys.stdout.isatty():
         sys.stderr.write(

confluent_client/confluent/vnc.py

@@ -2,6 +2,9 @@ import asyncio
 from PIL import Image
 import io
 import numpy as np
+import queue
+import threading
+import time
 import zlib
 
 # This results in an RGBA organization of pixels
@@ -48,8 +51,27 @@ class VNCClient:
         return False
     
     @classmethod
-    async def create(cls, url):
+    async def create(cls, url, outputfile=None, fps=30):
         self = cls()
+        self.outputfile = outputfile
+        self.fps = fps
+        self.video_writer = None
+        self._video_size = None
+        self._last_frame = None
+        self._last_frame_time = None
+        self._video_queue = None
+        self._video_thread = None
+        self._cv2 = None
+        if outputfile:
+            try:
+                import cv2
+            except ImportError:
+                raise ImportError("OpenCV is required for video output but is not installed.")
+            self._cv2 = cv2
+            self._video_queue = queue.Queue()
+            self._video_thread = threading.Thread(
+                target=self._video_worker, daemon=True)
+            self._video_thread.start()
         if url.startswith('unix://'):
             url = url.replace('unix://', '')
         if url.startswith('/'):
@@ -212,8 +234,52 @@ class VNCClient:
         for _ in range(num_rects):
             await self._handle_rectangle()
         self._updating = False
+        self._write_video_frame()
         self._request_screen_update(incremental=True)
 
+    def _write_video_frame(self):
+        if not self._cv2 or self.framebuffer is None:
+            return
+        # Snapshot the framebuffer now and hand it to the writer thread. Frames
+        # captured while a write is in progress simply queue up behind it.
+        frame = np.ascontiguousarray(
+            np.array(self.framebuffer.convert('RGB'))[:, :, ::-1])
+        self._video_queue.put((frame, time.monotonic()))
+
+    def _video_worker(self):
+        cv2 = self._cv2
+        while True:
+            item = self._video_queue.get()
+            if item is None:
+                # Flush the final frame for the time it stayed on screen
+                if self.video_writer is not None and self._last_frame is not None:
+                    nframes = max(1, round(
+                        (time.monotonic() - self._last_frame_time) * self.fps))
+                    for _ in range(nframes):
+                        self.video_writer.write(self._last_frame)
+                if self.video_writer is not None:
+                    self.video_writer.release()
+                    self.video_writer = None
+                return
+            frame, now = item
+            if self.video_writer is None:
+                self._video_size = (frame.shape[1], frame.shape[0])
+                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+                self.video_writer = cv2.VideoWriter(
+                    self.outputfile, fourcc, self.fps, self._video_size)
+            if (frame.shape[1], frame.shape[0]) != self._video_size:
+                frame = cv2.resize(frame, self._video_size)
+            if self._last_frame is None:
+                self._last_frame = frame
+                self._last_frame_time = now
+                continue
+            # Hold the previous frame for the real time it was displayed
+            nframes = max(1, round((now - self._last_frame_time) * self.fps))
+            for _ in range(nframes):
+                self.video_writer.write(self._last_frame)
+            self._last_frame = frame
+            self._last_frame_time = now
+
     async def _handle_rectangle(self):
         if self.framebuffer is None:
             self.framebuffer = Image.new('RGBA', (self.width, self.height))
@@ -272,5 +338,10 @@ class VNCClient:
                 break
         return length
     async def close(self):
+        if self._video_thread is not None:
+            # Signal the writer thread to flush and finalize the file
+            self._video_queue.put(None)
+            await asyncio.to_thread(self._video_thread.join)
+            self._video_thread = None
         self.writer.close()
         await self.writer.wait_closed()