Add device server (#4)

* Add create_device

* Fix the resource

* Add device server

* Add Device server and destroy function

* Fix credo

* Format the code

* Implement reviewers suggestions

---------

Co-authored-by: root <root@gpu2+sandbox>
Co-authored-by: Jerzy Wilczek <jerzy.wilczek@swmansion.com>

Author: 3 people

lib/decoder.ex

@@ -5,7 +5,7 @@ defmodule Membrane.VKVideo.Decoder do
   """
   use Membrane.Filter
 
-  alias Membrane.VKVideo.Native
+  alias Membrane.VKVideo.{DeviceServer, Native}
 
   def_input_pad :input, accepted_format: %Membrane.H264{stream_structure: :annexb, alignment: :au}
   def_output_pad :output, accepted_format: %Membrane.RawVideo{pixel_format: :NV12}
@@ -18,7 +18,8 @@ defmodule Membrane.VKVideo.Decoder do
 
   @impl true
   def handle_setup(_ctx, state) do
-    {:ok, decoder} = Native.new_decoder()
+    {:ok, device} = DeviceServer.get_device()
+    {:ok, decoder} = Native.new_decoder(device)
     state = %{state | decoder: decoder}
     {[], state}
   end
@@ -71,7 +72,8 @@ defmodule Membrane.VKVideo.Decoder do
   @impl true
   def handle_end_of_stream(:input, _ctx, state) do
     {:ok, flushed_frames} = Native.flush_decoder(state.decoder)
+    :ok = Native.destroy(state.decoder)
     {actions, state} = Enum.flat_map_reduce(flushed_frames, state, &prepare_actions(&1, &2))
-    {actions ++ [end_of_stream: :output], state}
+    {actions ++ [end_of_stream: :output], %{state | decoder: nil}}
   end
 end

lib/device_server.ex

@@ -0,0 +1,28 @@
+defmodule Membrane.VKVideo.DeviceServer do
+  @moduledoc false
+  use GenServer
+  alias Membrane.VKVideo.Native
+
+  @impl true
+  def init(_opts) do
+    {:ok, %{device: nil}}
+  end
+
+  @impl true
+  def handle_call(:get_device, _from, state) do
+    state = maybe_create_device(state)
+    {:reply, state.device, state}
+  end
+
+  defp maybe_create_device(%{device: nil} = state) do
+    device = Native.create_device()
+    %{state | device: device}
+  end
+
+  defp maybe_create_device(state), do: state
+
+  @spec get_device() :: {:ok, Native.t()} | no_return()
+  def get_device() do
+    GenServer.call(__MODULE__, :get_device)
+  end
+end

lib/encoder.ex

@@ -7,7 +7,7 @@ defmodule Membrane.VKVideo.Encoder do
 
   require Membrane.Logger
 
-  alias Membrane.VKVideo.Native
+  alias Membrane.VKVideo.{DeviceServer, Native}
 
   def_input_pad :input, accepted_format: %Membrane.RawVideo{pixel_format: :NV12}
 
@@ -84,8 +84,11 @@ defmodule Membrane.VKVideo.Encoder do
   end
 
   defp spawn_encoder(state) do
+    {:ok, device} = DeviceServer.get_device()
+
     {:ok, encoder} =
       Native.new_encoder(
+        device,
         state.width,
         state.height,
         state.framerate,
@@ -130,6 +133,7 @@ defmodule Membrane.VKVideo.Encoder do
 
   @impl true
   def handle_end_of_stream(:input, _ctx, state) do
+    :ok = Native.destroy(state.encoder)
     state = %{state | encoder: nil}
     {[end_of_stream: :output], state}
   end

lib/native.ex

@@ -15,8 +15,11 @@ defmodule Membrane.VKVideo.Native do
           pts: non_neg_integer() | nil
         }
 
-  @spec new_decoder() :: {:ok, t()} | no_return()
-  def new_decoder(), do: :erlang.nif_error(:nif_not_loaded)
+  @spec create_device() :: {:ok, t()} | no_return()
+  def create_device(), do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec new_decoder(t()) :: {:ok, t()} | no_return()
+  def new_decoder(_device), do: :erlang.nif_error(:nif_not_loaded)
 
   @spec decode(t(), binary(), pts_ns :: non_neg_integer() | nil) ::
           {:ok, raw_frame()} | no_return()
@@ -27,6 +30,7 @@ defmodule Membrane.VKVideo.Native do
   def flush_decoder(_decoder), do: :erlang.nif_error(:nif_not_loaded)
 
   @spec new_encoder(
+          t(),
           non_neg_integer(),
           non_neg_integer(),
           {non_neg_integer(), non_neg_integer()},
@@ -36,10 +40,13 @@ defmodule Membrane.VKVideo.Native do
           | {:variable_bitrate, Membrane.VKVideo.Encoder.VariableBitrate.t()}
           | {:constant_bitrate, Membrane.VKVideo.Encoder.ConstantBitrate.t()}
         ) :: {:ok, t()} | no_return()
-  def new_encoder(_width, _height, _framerate, _tune, _rate_control),
+  def new_encoder(_device, _width, _height, _framerate, _tune, _rate_control),
     do: :erlang.nif_error(:nif_not_loaded)
 
   @spec encode(t(), binary(), pts_ns :: non_neg_integer() | nil) ::
           {:ok, encoded_frame()} | no_return()
   def encode(_encoder, _raw_frame, _pts \\ nil), do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec destroy(t()) :: :ok
+  def destroy(_resource), do: :erlang.nif_error(:nif_not_loaded)
 end

lib/vk_video.ex

@@ -0,0 +1,11 @@
+defmodule Membrane.VKVideo do
+  @moduledoc false
+  use Application
+
+  alias Membrane.VKVideo.DeviceServer
+
+  @impl true
+  def start(_start_type, _start_args) do
+    GenServer.start_link(DeviceServer, nil, name: DeviceServer)
+  end
+end

mix.exs

@@ -28,7 +28,8 @@ defmodule Membrane.VKVideo.Mixfile do
 
   def application do
     [
-      extra_applications: []
+      extra_applications: [],
+      mod: {Membrane.VKVideo, []}
     ]
   end
 

native/vkvideo/src/decoder.rs

@@ -1,11 +1,9 @@
+use crate::ok;
 use crate::Resource;
 use rustler::{Atom, Binary, Env, Error, NifStruct, OwnedBinary, ResourceArc};
 use std::sync::Mutex;
 use vk_video::{parameters::DecoderParameters, BytesDecoder, EncodedInputChunk};
 
-rustler::atoms! {
-  ok,
-}
 pub struct DecoderResource {
     pub decoder_mutex: Mutex<BytesDecoder>,
 }
@@ -19,16 +17,14 @@ pub struct RawFrame<'a> {
     pub height: u32,
 }
 
-pub fn new() -> Result<(Atom, ResourceArc<Resource>), Error> {
-    let instance = vk_video::VulkanInstance::new()
-        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
-    let adapter = instance
-        .create_adapter(None)
-        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
-    let device = adapter
-        .create_device(wgpu::Features::empty(), wgpu::Limits::default())
-        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
-    let decoder = device
+pub fn new(
+    _env: Env,
+    resource: ResourceArc<Resource>,
+) -> Result<(Atom, ResourceArc<Resource>), Error> {
+    let decoder = resource
+        .device()
+        .ok_or_else(|| Error::BadArg)?
+        .device
         .create_bytes_decoder(DecoderParameters::default())
         .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
     let decoder_mutex = Mutex::new(decoder);

native/vkvideo/src/encoder.rs

@@ -1,16 +1,13 @@
+use crate::ok;
 use crate::Resource;
 use rustler::{Atom, Error, NifTaggedEnum, ResourceArc};
 use rustler::{Binary, Env, NifStruct, NifUnitEnum, OwnedBinary};
 use std::sync::Mutex;
 use vk_video::parameters::{RateControl, Rational, VideoParameters};
 use vk_video::{BytesEncoder, Frame, RawFrameData};
 
-rustler::atoms! {
-  ok
-}
-
 pub struct EncoderResource {
-    pub encoder_mutex: Mutex<BytesEncoder>,
+    pub encoder_mutex: Mutex<Option<BytesEncoder>>,
     pub width: u32,
     pub height: u32,
 }
@@ -75,22 +72,17 @@ impl Into<RateControl> for EncoderRateControl {
 }
 
 pub fn new(
+    _env: Env,
+    resource: ResourceArc<Resource>,
     width: u32,
     height: u32,
     frame_rate: (u32, u32),
     tune: EncoderTune,
     rate_control: EncoderRateControl,
 ) -> Result<(Atom, ResourceArc<Resource>), Error> {
+    let device_resource = &resource.device().ok_or_else(|| Error::BadArg)?.device;
     let non_zero_width = std::num::NonZero::new(width).ok_or(Error::BadArg)?;
     let non_zero_height = std::num::NonZero::new(height).ok_or(Error::BadArg)?;
-    let instance = vk_video::VulkanInstance::new()
-        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
-    let adapter = instance
-        .create_adapter(None)
-        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
-    let device = adapter
-        .create_device(wgpu::Features::empty(), wgpu::Limits::default())
-        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
 
     let video_parameters = VideoParameters {
         width: non_zero_width,
@@ -102,18 +94,18 @@ pub fn new(
     };
 
     let parameters = match tune {
-        EncoderTune::LowLatency => device
+        EncoderTune::LowLatency => device_resource
             .encoder_parameters_low_latency(video_parameters, rate_control.into())
             .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?,
-        EncoderTune::HighQuality => device
+        EncoderTune::HighQuality => device_resource
             .encoder_parameters_high_quality(video_parameters, rate_control.into())
             .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?,
     };
 
-    let encoder = device
+    let encoder = device_resource
         .create_bytes_encoder(parameters)
         .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
-    let encoder_mutex = Mutex::new(encoder);
+    let encoder_mutex = Mutex::new(Some(encoder));
     let encoder_resource = EncoderResource {
         encoder_mutex,
         width,
@@ -140,11 +132,13 @@ pub fn encode<'a>(
         pts: pts_ns,
     };
 
-    let mut encoder = encoder_resource
+    let mut guard = encoder_resource
         .encoder_mutex
         .lock()
         .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
 
+    let encoder = guard.as_mut().ok_or(Error::BadArg)?;
+
     let encoded_frame = encoder
         .encode(&frame, false)
         .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;

native/vkvideo/src/lib.rs

@@ -1,38 +1,77 @@
 use decoder::{DecoderResource, RawFrame};
 use encoder::{EncodedFrame, EncoderRateControl, EncoderResource, EncoderTune};
 use rustler::{Atom, Binary, Env, Error, ResourceArc, Term};
+use std::sync::Arc;
+use vk_video::VulkanDevice;
+
+rustler::atoms! {
+  ok,
+}
 
 pub mod decoder;
 pub mod encoder;
 
 pub enum Resource {
     Encoder(EncoderResource),
     Decoder(DecoderResource),
+    Device(DeviceResource),
+}
+
+pub struct DeviceResource {
+    pub device: Arc<VulkanDevice>,
 }
 
+impl std::panic::RefUnwindSafe for DeviceResource {}
+
 impl Resource {
     pub fn encoder(&self) -> Option<&EncoderResource> {
         match self {
             Self::Encoder(encoder_resource) => Some(encoder_resource),
-            Self::Decoder(_) => None,
+            _ => None,
         }
     }
 
     pub fn decoder(&self) -> Option<&DecoderResource> {
         match self {
             Self::Decoder(decoder_resource) => Some(decoder_resource),
-            Self::Encoder(_) => None,
+            _ => None,
+        }
+    }
+
+    pub fn device(&self) -> Option<&DeviceResource> {
+        match self {
+            Self::Device(device_resource) => Some(device_resource),
+            _ => None,
         }
     }
 }
 
+#[allow(non_local_definitions)]
 fn load(env: Env, _: Term) -> bool {
     rustler::resource!(Resource, env)
 }
 
 #[rustler::nif(schedule = "DirtyIo")]
-fn new_decoder() -> Result<(Atom, ResourceArc<Resource>), Error> {
-    decoder::new()
+fn create_device() -> Result<(Atom, ResourceArc<Resource>), Error> {
+    let instance = vk_video::VulkanInstance::new()
+        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
+    let adapter = instance
+        .create_adapter(None)
+        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
+    let device = adapter
+        .create_device(wgpu::Features::empty(), wgpu::Limits::default())
+        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
+
+    let device_resource = ResourceArc::new(Resource::Device(DeviceResource { device }));
+    Ok((ok(), device_resource))
+}
+
+#[rustler::nif(schedule = "DirtyIo")]
+fn new_decoder(
+    env: Env,
+    resource: ResourceArc<Resource>,
+) -> Result<(Atom, ResourceArc<Resource>), Error> {
+    decoder::new(env, resource)
 }
 
 #[rustler::nif(schedule = "DirtyIo")]
@@ -55,13 +94,15 @@ pub fn flush_decoder(
 
 #[rustler::nif(schedule = "DirtyIo")]
 fn new_encoder(
+    env: Env,
+    resource: ResourceArc<Resource>,
     width: u32,
     height: u32,
     frame_rate: (u32, u32),
     tune: EncoderTune,
     rate_control: EncoderRateControl,
 ) -> Result<(Atom, ResourceArc<Resource>), Error> {
-    encoder::new(width, height, frame_rate, tune, rate_control)
+    encoder::new(env, resource, width, height, frame_rate, tune, rate_control)
 }
 
 #[rustler::nif(schedule = "DirtyIo")]
@@ -74,4 +115,17 @@ fn encode<'a>(
     encoder::encode(env, resource, bytes, pts_ns)
 }
 
+#[rustler::nif(schedule = "DirtyIo")]
+fn destroy<'a>(env: Env<'a>, resource: ResourceArc<Resource>) -> Result<Atom, Error> {
+    if let Resource::Encoder(encoder) = &*resource {
+        let mut encoder = encoder
+            .encoder_mutex
+            .lock()
+            .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
+        *encoder = None;
+    }
+
+    Ok(ok())
+}
+
 rustler::init!("Elixir.Membrane.VKVideo.Native", load = load);