Add listen_and_write thread to handle io

manim/scene/scene_file_writer.py

@@ -5,10 +5,11 @@
 __all__ = ["SceneFileWriter"]
 
 import json
-import os
 import shutil
 import subprocess
 from pathlib import Path
+from queue import Queue
+from threading import Thread
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -17,7 +18,7 @@
 from pydub import AudioSegment
 
 from manim import __version__
-
+from .section import DefaultSectionType, Section
 from .. import config, logger
 from .._config.logger_utils import set_file_logger
 from ..constants import RendererType
@@ -33,7 +34,6 @@
     write_to_movie,
 )
 from ..utils.sounds import get_full_sound_file_path
-from .section import DefaultSectionType, Section
 
 if TYPE_CHECKING:
     from manim.renderer.opengl_renderer import OpenGLRenderer
@@ -78,6 +78,7 @@ def __init__(self, renderer, scene_name, **kwargs):
         self.partial_movie_files: list[str] = []
         self.subcaptions: list[srt.Subtitle] = []
         self.sections: list[Section] = []
+        self.queue = Queue()
         # first section gets automatically created for convenience
         # if you need the first section to be skipped, add a first section by hand, it will replace this one
         self.next_section(
@@ -388,15 +389,14 @@ def write_frame(self, frame_or_renderer: np.ndarray | OpenGLRenderer):
         elif config.renderer == RendererType.CAIRO:
             frame = frame_or_renderer
             if write_to_movie():
-                self.writing_process.stdin.write(frame.tobytes())
+                self.queue.put(frame.tobytes())
             if is_png_format() and not config["dry_run"]:
                 self.output_image_from_array(frame)
 
     def write_opengl_frame(self, renderer: OpenGLRenderer):
         if write_to_movie():
-            self.writing_process.stdin.write(
-                renderer.get_raw_frame_buffer_object_data(),
-            )
+            self.queue.put(renderer.get_raw_frame_buffer_object_data())
+
         elif is_png_format() and not config["dry_run"]:
             target_dir = self.image_file_path.parent / self.image_file_path.stem
             extension = self.image_file_path.suffix
@@ -453,6 +453,8 @@ def finish(self):
         """
         if write_to_movie():
             if hasattr(self, "writing_process"):
+                self.queue.put(-1)
+                self.writer_thread.join()
                 self.writing_process.terminate()
             self.combine_to_movie()
             if config.save_sections:
@@ -515,12 +517,26 @@ def open_movie_pipe(self, file_path=None):
         else:
             command += ["-vcodec", "libx264", "-pix_fmt", "yuv420p"]
         command += [file_path]
+
         self.writing_process = subprocess.Popen(command, stdin=subprocess.PIPE)
+        self.writer_thread = Thread(target=self.listen_and_write, args=())
+        self.writer_thread.start()
+
+    def listen_and_write(self) -> None:
+        while True:
+            buf = self.queue.get()
+            if buf == -1:
+                return
+
+            self.writing_process.stdin.write(buf)
 
+    # TODO(saveliy): this method needs to close the thread
     def close_movie_pipe(self):
         """
         Used internally by Manim to gracefully stop writing to FFMPEG's input buffer
         """
+        self.queue.put(-1)
+        self.writer_thread.join()
         self.writing_process.stdin.close()
         self.writing_process.wait()
 
@@ -667,7 +683,8 @@ def combine_to_movie(self):
         self.print_file_ready_message(str(movie_file_path))
         if write_to_movie():
             for file_path in partial_movie_files:
-                # We have to modify the accessed time so if we have to clean the cache we remove the one used the longest.
+                # We have to modify the accessed time so if we have to clean the cache we remove the one used the
+                # longest.
                 modify_atime(file_path)
 
     def combine_to_section_videos(self) -> None: