Add IVFReader.close/1 and IVFWriter.close/1

Author: mickel8

lib/ex_webrtc/media/ivf_reader.ex

@@ -84,6 +84,9 @@ defmodule ExWebRTC.Media.IVFReader do
 
   @opaque t() :: File.io_device()
 
+  @doc """
+  Opens an IVF file and reads its header.
+  """
   @spec open(Path.t()) :: {:ok, IVFHeader.t(), t()} | {:error, term()}
   def open(path) do
     with {:ok, file} <- File.open(path),
@@ -111,6 +114,9 @@ defmodule ExWebRTC.Media.IVFReader do
     end
   end
 
+  @doc """
+  Reads the next video frame from an IVF file.
+  """
   @spec next_frame(t()) :: {:ok, IVFFrame.t()} | {:error, term()} | :eof
   def next_frame(reader) do
     with <<len_frame::little-integer-size(32), timestamp::little-integer-size(64)>> <-
@@ -124,4 +130,14 @@ defmodule ExWebRTC.Media.IVFReader do
       _other -> {:error, :invalid_file}
     end
   end
+
+  @doc """
+  Closes an IVF reader.
+
+  When a process owning the IVF reader exits, IVF reader is closed automatically. 
+  """
+  @spec close(t()) :: :ok | {:error, File.posix() | term()}
+  def close(reader) do
+    File.close(reader)
+  end
 end

lib/ex_webrtc/media/ivf_writer.ex

@@ -13,15 +13,18 @@ defmodule ExWebRTC.Media.IVFWriter do
   @enforce_keys [:file]
   defstruct @enforce_keys ++ [update_header_after: 0, frames_cnt: 0]
 
-  defguard update_header?(writer)
-           when writer.frames_cnt >= writer.update_header_after and
-                  rem(writer.frames_cnt, writer.update_header_after) == 0
+  defguardp update_header?(writer)
+            when writer.frames_cnt >= writer.update_header_after and
+                   rem(writer.frames_cnt, writer.update_header_after) == 0
 
   @doc """
   Creates a new IVF writer.
 
   Initially, IVF header is written with `num_frames` and 
   is updated every `num_frames` by `num_frames`.
+  To have a precise number of frames in the header,
+  either write exactly `num_frames` or call `close/1` 
+  at the end of writing.
   """
   @spec open(Path.t(),
           fourcc: non_neg_integer(),
@@ -53,20 +56,35 @@ defmodule ExWebRTC.Media.IVFWriter do
   end
 
   @doc """
-  Writes IVF frame into a file.
+  Writes an IVF frame into a file.
   """
   @spec write_frame(t(), IVFFrame.t()) :: {:ok, t()} | {:error, term()}
   def write_frame(writer, frame) when update_header?(writer) and frame.data != <<>> do
-    case update_header(writer) do
+    case update_header(writer, writer.frames_cnt + writer.update_header_after) do
       :ok -> do_write_frame(writer, frame)
       {:error, _reason} = error -> error
     end
   end
 
   def write_frame(writer, frame) when frame.data != <<>>, do: do_write_frame(writer, frame)
 
-  defp update_header(writer) do
-    num_frames = <<writer.frames_cnt + writer.update_header_after::little-32>>
+  @doc """
+  Updates a number of frames in the header and closes the writer.
+
+  If a process owning an IVF writer exits, a file open by the IVF writer
+  will be closed automatically but header will not be updated.
+  See also `open/2` for more information on automatic header updates.
+  """
+  @spec close(t()) :: :ok | {:error, File.posix() | term()}
+  def close(writer) do
+    case update_header(writer, writer.frames_cnt) do
+      :ok -> File.close(writer.file)
+      {:error, _reason} = error -> error
+    end
+  end
+
+  defp update_header(writer, num_frames) do
+    num_frames = <<num_frames::little-32>>
     ret = :file.pwrite(writer.file, 24, num_frames)
     {:ok, _position} = :file.position(writer.file, :eof)
     ret

test/ex_webrtc/media/ivf_reader_test.exs

@@ -26,6 +26,7 @@ defmodule ExWebRTC.Media.IVFReaderTest do
     end
 
     assert :eof == IVFReader.next_frame(reader)
+    assert :ok == IVFReader.close(reader)
   end
 
   test "empty file" do

test/ex_webrtc/media/ivf_writer_test.exs

@@ -7,7 +7,7 @@ defmodule ExWebRTC.Media.IVFWritertTest do
   test "IVF writer", %{tmp_dir: tmp_dir} do
     path = Path.join([tmp_dir, "output.ivf"])
     <<fourcc::little-32>> = "VP80"
-    num_frames = 2
+    num_frames = 5
 
     {:ok, writer} =
       IVFWriter.open(path,
@@ -42,7 +42,7 @@ defmodule ExWebRTC.Media.IVFWritertTest do
     # `num_frames + 1` frame
     expected_num_frames = 2 * num_frames
     frame = %IVFFrame{timestamp: num_frames, data: <<0, 1, 2, 3, 4>>}
-    assert {:ok, _writer} = IVFWriter.write_frame(writer, frame)
+    assert {:ok, writer} = IVFWriter.write_frame(writer, frame)
 
     assert {:ok,
             %IVFHeader{
@@ -65,5 +65,22 @@ defmodule ExWebRTC.Media.IVFWritertTest do
       assert frame.timestamp == i
       assert frame.data == <<0, 1, 2, 3, 4>>
     end
+
+    # assert that after calling close/1, number of frames
+    # in the header is equal to the 
+    # exact number of frames that were written
+    assert :ok = IVFWriter.close(writer)
+    exact_num_frames = num_frames + 1
+
+    assert {:ok,
+            %IVFHeader{
+              fourcc: ^fourcc,
+              height: 640,
+              width: 480,
+              num_frames: ^exact_num_frames,
+              timebase_denum: 30,
+              timebase_num: 1,
+              unused: 0
+            }, _reader} = IVFReader.open(path)
   end
 end