Put encoder and decoder into a single NIF.

Cargo.toml

@@ -1,3 +1,3 @@
 [workspace]
 resolver = "2"
-members = ["native/vkvideo_decoder", "native/vkvideo_encoder"]
+members = ["native/vkvideo"]

lib/decoder.ex

@@ -5,7 +5,7 @@ defmodule Membrane.VKVideo.Decoder do
   """
   use Membrane.Filter
 
-  alias __MODULE__.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()
+    {:ok, device} = DeviceServer.get_device()
+    {:ok, decoder} = Native.new_decoder(device)
     state = %{state | decoder: decoder}
     {[], state}
   end
@@ -70,8 +71,9 @@ defmodule Membrane.VKVideo.Decoder do
 
   @impl true
   def handle_end_of_stream(:input, _ctx, state) do
-    {:ok, flushed_frames} = Native.flush(state.decoder)
+    {: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 __MODULE__.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(
+      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

@@ -0,0 +1,52 @@
+defmodule Membrane.VKVideo.Native do
+  @moduledoc false
+  use Rustler, otp_app: :membrane_vk_video_plugin, crate: :vkvideo
+
+  @type t :: reference()
+  @type raw_frame :: %{
+          payload: binary(),
+          pts: non_neg_integer() | nil,
+          width: non_neg_integer(),
+          height: non_neg_integer()
+        }
+
+  @type encoded_frame :: %{
+          payload: binary(),
+          pts: non_neg_integer() | nil
+        }
+
+  @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()
+  def decode(_decoder, _frame, _pts \\ nil), do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec flush_decoder(t()) ::
+          {:ok, raw_frame()} | no_return()
+  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()},
+          :low_latency | :high_quality,
+          :encoder_default
+          | :disabled
+          | {:variable_bitrate, Membrane.VKVideo.Encoder.VariableBitrate.t()}
+          | {:constant_bitrate, Membrane.VKVideo.Encoder.ConstantBitrate.t()}
+        ) :: {:ok, t()} | no_return()
+  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_decoder/Cargo.toml → native/vkvideo/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "vkvideo_decoder"
+name = "vkvideo"
 version = "0.1.0"
 authors = []
 edition = "2021"

native/vkvideo_decoder/src/lib.rs → native/vkvideo/src/decoder.rs

@@ -1,55 +1,47 @@
+use crate::ok;
+use crate::Resource;
+use rustler::{Atom, Binary, Env, Error, NifStruct, OwnedBinary, ResourceArc};
 use std::sync::Mutex;
-
-use rustler::{Atom, Binary, Env, Error, NifStruct, OwnedBinary, ResourceArc, Term};
 use vk_video::{parameters::DecoderParameters, BytesDecoder, EncodedInputChunk};
 
-rustler::atoms! {
-  ok,
-}
-
-struct DecoderResource {
+pub struct DecoderResource {
     pub decoder_mutex: Mutex<BytesDecoder>,
 }
 
-fn load(env: Env, _: Term) -> bool {
-    rustler::resource!(DecoderResource, env)
-}
-
-#[rustler::nif(schedule = "DirtyIo")]
-fn new() -> Result<(Atom, ResourceArc<DecoderResource>), 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
-        .create_bytes_decoder(DecoderParameters::default())
-        .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
-    let decoder_mutex = Mutex::new(decoder);
-    let resource = ResourceArc::new(DecoderResource { decoder_mutex });
-    Ok((ok(), resource))
-}
-
 #[derive(NifStruct)]
 #[module = "Membrane.VKVideo.RawFrame"]
-struct RawFrame<'a> {
+pub struct RawFrame<'a> {
     pub payload: Binary<'a>,
     pub pts_ns: Option<u64>,
     pub width: u32,
     pub height: u32,
 }
 
-#[rustler::nif(schedule = "DirtyIo")]
-fn decode<'a>(
+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);
+    let decoder_resource = DecoderResource { decoder_mutex };
+    let resource = ResourceArc::new(Resource::Decoder(decoder_resource));
+    Ok((ok(), resource))
+}
+
+pub fn decode<'a>(
     env: Env<'a>,
-    resource: ResourceArc<DecoderResource>,
+    resource: ResourceArc<Resource>,
     bytes: Binary,
     pts_ns: Option<u64>,
 ) -> Result<(Atom, Vec<RawFrame<'a>>), Error> {
     let mut decoder = resource
+        .decoder()
+        .ok_or_else(|| Error::BadArg)?
         .decoder_mutex
         .lock()
         .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
@@ -77,9 +69,11 @@ fn decode<'a>(
     }
     Ok((ok(), results))
 }
-#[rustler::nif(schedule = "DirtyIo")]
-fn flush(env: Env, resource: ResourceArc<DecoderResource>) -> Result<(Atom, Vec<RawFrame>), Error> {
+
+pub fn flush(env: Env, resource: ResourceArc<Resource>) -> Result<(Atom, Vec<RawFrame>), Error> {
     let mut decoder = resource
+        .decoder()
+        .ok_or_else(|| Error::BadArg)?
         .decoder_mutex
         .lock()
         .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?;
@@ -104,5 +98,3 @@ fn flush(env: Env, resource: ResourceArc<DecoderResource>) -> Result<(Atom, Vec<
     }
     Ok((ok(), results))
 }
-
-rustler::init!("Elixir.Membrane.VKVideo.Decoder.Native", load = load);