From 66b86de18c0669e63a3a6ba4671267ad98123eea Mon Sep 17 00:00:00 2001 From: root Date: Wed, 17 Dec 2025 11:28:02 +0100 Subject: [PATCH 1/4] Puts encoder and decoder into a single NIF --- Cargo.lock | 11 +--- Cargo.toml | 2 +- lib/decoder.ex | 6 +- lib/decoder/native.ex | 23 ------- lib/encoder.ex | 6 +- lib/{encoder => }/native.ex | 25 ++++++-- .../{vkvideo_decoder => vkvideo}/Cargo.toml | 2 +- .../src/lib.rs => vkvideo/src/decoder.rs} | 51 ++++++++-------- .../src/lib.rs => vkvideo/src/encoder.rs} | 54 ++++++++-------- native/vkvideo/src/lib.rs | 61 +++++++++++++++++++ native/vkvideo_encoder/Cargo.toml | 13 ---- test/decoder_native_test.exs | 6 +- test/encoder_native_test.exs | 4 +- 13 files changed, 149 insertions(+), 115 deletions(-) delete mode 100644 lib/decoder/native.ex rename lib/{encoder => }/native.ex (53%) rename native/{vkvideo_decoder => vkvideo}/Cargo.toml (90%) rename native/{vkvideo_decoder/src/lib.rs => vkvideo/src/decoder.rs} (81%) rename native/{vkvideo_encoder/src/lib.rs => vkvideo/src/encoder.rs} (85%) create mode 100644 native/vkvideo/src/lib.rs delete mode 100644 native/vkvideo_encoder/Cargo.toml diff --git a/Cargo.lock b/Cargo.lock index 9c7f891..d91bb2d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1009,16 +1009,7 @@ dependencies = [ ] [[package]] -name = "vkvideo_decoder" -version = "0.1.0" -dependencies = [ - "rustler", - "vk-video", - "wgpu", -] - -[[package]] -name = "vkvideo_encoder" +name = "vkvideo" version = "0.1.0" dependencies = [ "rustler", diff --git a/Cargo.toml b/Cargo.toml index a992cab..39025aa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] resolver = "2" -members = ["native/vkvideo_decoder", "native/vkvideo_encoder"] +members = ["native/vkvideo"] diff --git a/lib/decoder.ex b/lib/decoder.ex index 2d91a51..f4190da 100644 --- a/lib/decoder.ex +++ b/lib/decoder.ex @@ -5,7 +5,7 @@ defmodule Membrane.VKVideo.Decoder do """ use Membrane.Filter - alias __MODULE__.Native + alias Membrane.VKVideo.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,7 @@ defmodule Membrane.VKVideo.Decoder do @impl true def handle_setup(_ctx, state) do - {:ok, decoder} = Native.new() + {:ok, decoder} = Native.new_decoder() state = %{state | decoder: decoder} {[], state} end @@ -70,7 +70,7 @@ 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) {actions, state} = Enum.flat_map_reduce(flushed_frames, state, &prepare_actions(&1, &2)) {actions ++ [end_of_stream: :output], state} end diff --git a/lib/decoder/native.ex b/lib/decoder/native.ex deleted file mode 100644 index f4c7930..0000000 --- a/lib/decoder/native.ex +++ /dev/null @@ -1,23 +0,0 @@ -defmodule Membrane.VKVideo.Decoder.Native do - @moduledoc false - use Rustler, otp_app: :membrane_vk_video_plugin, crate: :vkvideo_decoder - - @type t :: reference() - @type raw_frame :: %{ - payload: binary(), - pts: non_neg_integer() | nil, - width: non_neg_integer(), - height: non_neg_integer() - } - - @spec new() :: {:ok, t()} | no_return() - def new(), 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(t()) :: - {:ok, raw_frame()} | no_return() - def flush(_decoder), do: :erlang.nif_error(:nif_not_loaded) -end diff --git a/lib/encoder.ex b/lib/encoder.ex index 11e41db..010fff2 100644 --- a/lib/encoder.ex +++ b/lib/encoder.ex @@ -7,7 +7,7 @@ defmodule Membrane.VKVideo.Encoder do require Membrane.Logger - alias __MODULE__.Native + alias Membrane.VKVideo.Native def_input_pad :input, accepted_format: %Membrane.RawVideo{pixel_format: :NV12} @@ -68,7 +68,7 @@ defmodule Membrane.VKVideo.Encoder do end {:ok, encoder} = - Native.new( + Native.new_encoder( stream_format.width, stream_format.height, state.framerate, @@ -93,7 +93,7 @@ defmodule Membrane.VKVideo.Encoder do if stream_format.width != state.width or stream_format.height != state.height or stream_format.framerate != state.framerate do {:ok, encoder} = - Native.new( + Native.new_encoder( stream_format.width, stream_format.height, stream_format.framerate, diff --git a/lib/encoder/native.ex b/lib/native.ex similarity index 53% rename from lib/encoder/native.ex rename to lib/native.ex index 663c58a..bba5327 100644 --- a/lib/encoder/native.ex +++ b/lib/native.ex @@ -1,15 +1,32 @@ -defmodule Membrane.VKVideo.Encoder.Native do +defmodule Membrane.VKVideo.Native do @moduledoc false - use Rustler, otp_app: :membrane_vk_video_plugin, crate: :vkvideo_encoder + 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 new( + @spec new_decoder() :: {:ok, t()} | no_return() + def new_decoder(), 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( non_neg_integer(), non_neg_integer(), {non_neg_integer(), non_neg_integer()}, @@ -19,7 +36,7 @@ defmodule Membrane.VKVideo.Encoder.Native do | {:variable_bitrate, Membrane.VKVideo.Encoder.VariableBitrate.t()} | {:constant_bitrate, Membrane.VKVideo.Encoder.ConstantBitrate.t()} ) :: {:ok, t()} | no_return() - def new(_width, _height, _framerate, _tune, _rate_control), + def new_encoder(_width, _height, _framerate, _tune, _rate_control), do: :erlang.nif_error(:nif_not_loaded) @spec encode(t(), binary(), pts_ns :: non_neg_integer() | nil) :: diff --git a/native/vkvideo_decoder/Cargo.toml b/native/vkvideo/Cargo.toml similarity index 90% rename from native/vkvideo_decoder/Cargo.toml rename to native/vkvideo/Cargo.toml index f9f93cd..7372c72 100644 --- a/native/vkvideo_decoder/Cargo.toml +++ b/native/vkvideo/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "vkvideo_decoder" +name = "vkvideo" version = "0.1.0" authors = [] edition = "2021" diff --git a/native/vkvideo_decoder/src/lib.rs b/native/vkvideo/src/decoder.rs similarity index 81% rename from native/vkvideo_decoder/src/lib.rs rename to native/vkvideo/src/decoder.rs index 26f0502..81be4a3 100644 --- a/native/vkvideo_decoder/src/lib.rs +++ b/native/vkvideo/src/decoder.rs @@ -1,22 +1,25 @@ +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, } -fn load(env: Env, _: Term) -> bool { - rustler::resource!(DecoderResource, env) +#[derive(NifStruct)] +#[module = "Membrane.VKVideo.RawFrame"] +pub struct RawFrame<'a> { + pub payload: Binary<'a>, + pub pts_ns: Option, + pub width: u32, + pub height: u32, } -#[rustler::nif(schedule = "DirtyIo")] -fn new() -> Result<(Atom, ResourceArc), Error> { +pub fn new() -> Result<(Atom, ResourceArc), Error> { let instance = vk_video::VulkanInstance::new() .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?; let adapter = instance @@ -29,27 +32,24 @@ fn new() -> Result<(Atom, ResourceArc), Error> { .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 }); + let decoder_resource = DecoderResource { decoder_mutex }; + let resource = ResourceArc::new(Resource { + decoder: Some(decoder_resource), + encoder: None, + }); Ok((ok(), resource)) } -#[derive(NifStruct)] -#[module = "Membrane.VKVideo.RawFrame"] -struct RawFrame<'a> { - pub payload: Binary<'a>, - pub pts_ns: Option, - pub width: u32, - pub height: u32, -} - -#[rustler::nif(schedule = "DirtyIo")] -fn decode<'a>( +pub fn decode<'a>( env: Env<'a>, - resource: ResourceArc, + resource: ResourceArc, bytes: Binary, pts_ns: Option, ) -> Result<(Atom, Vec>), Error> { let mut decoder = resource + .decoder + .as_ref() + .expect("Decoder not initialized") .decoder_mutex .lock() .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?; @@ -77,9 +77,12 @@ fn decode<'a>( } Ok((ok(), results)) } -#[rustler::nif(schedule = "DirtyIo")] -fn flush(env: Env, resource: ResourceArc) -> Result<(Atom, Vec), Error> { + +pub fn flush(env: Env, resource: ResourceArc) -> Result<(Atom, Vec), Error> { let mut decoder = resource + .decoder + .as_ref() + .expect("Decoder not initialized") .decoder_mutex .lock() .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?; @@ -104,5 +107,3 @@ fn flush(env: Env, resource: ResourceArc) -> Result<(Atom, Vec< } Ok((ok(), results)) } - -rustler::init!("Elixir.Membrane.VKVideo.Decoder.Native", load = load); diff --git a/native/vkvideo_encoder/src/lib.rs b/native/vkvideo/src/encoder.rs similarity index 85% rename from native/vkvideo_encoder/src/lib.rs rename to native/vkvideo/src/encoder.rs index e110f0b..c086758 100644 --- a/native/vkvideo_encoder/src/lib.rs +++ b/native/vkvideo/src/encoder.rs @@ -1,38 +1,36 @@ +use crate::Resource; use rustler::{Atom, Error, NifTaggedEnum, ResourceArc}; -use rustler::{Binary, Env, NifStruct, NifUnitEnum, OwnedBinary, Term}; +use rustler::{Binary, Env, NifStruct, NifUnitEnum, OwnedBinary}; use std::sync::Mutex; use vk_video::parameters::{RateControl, Rational, VideoParameters}; use vk_video::{BytesEncoder, Frame, RawFrameData}; -struct EncoderResource { - pub encoder_mutex: Mutex, - pub width: u32, - pub height: u32, -} - -fn load(env: Env, _: Term) -> bool { - rustler::resource!(EncoderResource, env) -} rustler::atoms! { ok } +pub struct EncoderResource { + pub encoder_mutex: Mutex, + pub width: u32, + pub height: u32, +} + #[derive(NifStruct)] #[module = "Membrane.VKVideo.EncodedFrame"] -struct EncodedFrame<'a> { +pub struct EncodedFrame<'a> { pub payload: Binary<'a>, pub pts_ns: Option, } #[derive(NifUnitEnum)] -enum EncoderTune { +pub enum EncoderTune { LowLatency, HighQuality, } #[derive(NifStruct)] #[module = "Membrane.VKVideo.Encoder.VariableBitrate"] -struct VariableBitrate { +pub struct VariableBitrate { pub average_bitrate: u64, pub max_bitrate: u64, pub virtual_buffer_size_ms: u64, @@ -40,13 +38,13 @@ struct VariableBitrate { #[derive(NifStruct)] #[module = "Membrane.VKVideo.Encoder.ConstantBitrate"] -struct ConstantBitrate { +pub struct ConstantBitrate { pub bitrate: u64, pub virtual_buffer_size_ms: u64, } #[derive(NifTaggedEnum)] -enum EncoderRateControl { +pub enum EncoderRateControl { EncoderDefault, VariableBitrate(VariableBitrate), ConstantBitrate(ConstantBitrate), @@ -76,14 +74,13 @@ impl Into for EncoderRateControl { } } -#[rustler::nif(schedule = "DirtyIo")] -fn new( +pub fn new( width: u32, height: u32, frame_rate: (u32, u32), tune: EncoderTune, rate_control: EncoderRateControl, -) -> Result<(Atom, ResourceArc), Error> { +) -> Result<(Atom, ResourceArc), Error> { 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() @@ -117,31 +114,36 @@ fn new( .create_bytes_encoder(parameters) .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?; let encoder_mutex = Mutex::new(encoder); - let resource = ResourceArc::new(EncoderResource { + let encoder_resource = EncoderResource { encoder_mutex, width, height, + }; + + let resource = ResourceArc::new(Resource { + encoder: Some(encoder_resource), + decoder: None, }); Ok((ok(), resource)) } -#[rustler::nif(schedule = "DirtyIo")] -fn encode<'a>( +pub fn encode<'a>( env: Env<'a>, - resource: ResourceArc, + resource: ResourceArc, bytes: Binary, pts_ns: Option, ) -> Result<(Atom, EncodedFrame<'a>), Error> { + let encoder_resource = resource.encoder.as_ref().expect("Encoder not initialized"); let frame = Frame { data: RawFrameData { frame: bytes.to_vec(), - width: resource.width, - height: resource.height, + width: encoder_resource.width, + height: encoder_resource.height, }, pts: pts_ns, }; - let mut encoder = resource + let mut encoder = encoder_resource .encoder_mutex .lock() .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?; @@ -163,5 +165,3 @@ fn encode<'a>( }, )) } - -rustler::init!("Elixir.Membrane.VKVideo.Encoder.Native", load = load); diff --git a/native/vkvideo/src/lib.rs b/native/vkvideo/src/lib.rs new file mode 100644 index 0000000..68ada05 --- /dev/null +++ b/native/vkvideo/src/lib.rs @@ -0,0 +1,61 @@ +use decoder::{DecoderResource, RawFrame}; +use encoder::{EncodedFrame, EncoderRateControl, EncoderResource, EncoderTune}; +use rustler::{Atom, Binary, Env, Error, ResourceArc, Term}; + +pub mod decoder; +pub mod encoder; + +pub struct Resource { + pub decoder: Option, + pub encoder: Option, +} + +fn load(env: Env, _: Term) -> bool { + rustler::resource!(Resource, env) +} + +#[rustler::nif(schedule = "DirtyIo")] +fn new_decoder() -> Result<(Atom, ResourceArc), Error> { + decoder::new() +} + +#[rustler::nif(schedule = "DirtyIo")] +pub fn decode<'a>( + env: Env<'a>, + resource: ResourceArc, + bytes: Binary, + pts_ns: Option, +) -> Result<(Atom, Vec>), Error> { + decoder::decode(env, resource, bytes, pts_ns) +} + +#[rustler::nif(schedule = "DirtyIo")] +pub fn flush_decoder( + env: Env, + resource: ResourceArc, +) -> Result<(Atom, Vec), Error> { + decoder::flush(env, resource) +} + +#[rustler::nif(schedule = "DirtyIo")] +fn new_encoder( + width: u32, + height: u32, + frame_rate: (u32, u32), + tune: EncoderTune, + rate_control: EncoderRateControl, +) -> Result<(Atom, ResourceArc), Error> { + encoder::new(width, height, frame_rate, tune, rate_control) +} + +#[rustler::nif(schedule = "DirtyIo")] +fn encode<'a>( + env: Env<'a>, + resource: ResourceArc, + bytes: Binary, + pts_ns: Option, +) -> Result<(Atom, EncodedFrame<'a>), Error> { + encoder::encode(env, resource, bytes, pts_ns) +} + +rustler::init!("Elixir.Membrane.VKVideo.Native", load = load); diff --git a/native/vkvideo_encoder/Cargo.toml b/native/vkvideo_encoder/Cargo.toml deleted file mode 100644 index f0f42af..0000000 --- a/native/vkvideo_encoder/Cargo.toml +++ /dev/null @@ -1,13 +0,0 @@ -[package] -name = "vkvideo_encoder" -version = "0.1.0" -authors = [] -edition = "2021" - -[lib] -crate-type = ["cdylib"] - -[dependencies] -rustler = "0.37.2" -vk-video = { git = "https://github.com/software-mansion/smelter.git", features=["expose_parsers"] } -wgpu = "25.0.2" diff --git a/test/decoder_native_test.exs b/test/decoder_native_test.exs index aac6383..b02c0bb 100644 --- a/test/decoder_native_test.exs +++ b/test/decoder_native_test.exs @@ -1,15 +1,15 @@ defmodule Decoder.NativeTest do use ExUnit.Case, async: false - alias Membrane.VKVideo.Decoder.Native + alias Membrane.VKVideo.Native @tag :requires_gpu test "Decoder decodes H.264 stream" do in_path = "./fixtures/input-100.h264" |> Path.expand(__DIR__) assert {:ok, file} = File.read(in_path) - {:ok, decoder_ref} = Native.new() + {:ok, decoder_ref} = Native.new_decoder() {:ok, decoded_frames} = Native.decode(decoder_ref, file, 0) - {:ok, flushed_frames} = Native.flush(decoder_ref) + {:ok, flushed_frames} = Native.flush_decoder(decoder_ref) all_frames = decoded_frames ++ flushed_frames assert length(all_frames) == 100 [first_frame | _rest_of_frames] = all_frames diff --git a/test/encoder_native_test.exs b/test/encoder_native_test.exs index 8bf8313..ce94065 100644 --- a/test/encoder_native_test.exs +++ b/test/encoder_native_test.exs @@ -1,6 +1,6 @@ defmodule Encoder.NativeTest do use ExUnit.Case, async: false - alias Membrane.VKVideo.Encoder.Native + alias Membrane.VKVideo.Native @width 1280 @height 720 @@ -40,7 +40,7 @@ defmodule Encoder.NativeTest do assert {:ok, file} = File.read(in_path) {:ok, encoder_ref} = - Native.new(@width, @height, @framerate, unquote(tune), unquote(rate_control_ast)) + Native.new_encoder(@width, @height, @framerate, unquote(tune), unquote(rate_control_ast)) raw_frames = for <>, do: chunk From 004c9f0104394218809d40fadc6c0d9861ada6aa Mon Sep 17 00:00:00 2001 From: root Date: Wed, 17 Dec 2025 11:47:25 +0100 Subject: [PATCH 2/4] Make Resource an Enum --- native/vkvideo/src/decoder.rs | 15 +++++---------- native/vkvideo/src/encoder.rs | 7 ++----- native/vkvideo/src/lib.rs | 22 +++++++++++++++++++--- 3 files changed, 26 insertions(+), 18 deletions(-) diff --git a/native/vkvideo/src/decoder.rs b/native/vkvideo/src/decoder.rs index 81be4a3..0e5e4fb 100644 --- a/native/vkvideo/src/decoder.rs +++ b/native/vkvideo/src/decoder.rs @@ -33,10 +33,7 @@ pub fn new() -> Result<(Atom, ResourceArc), Error> { .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: Some(decoder_resource), - encoder: None, - }); + let resource = ResourceArc::new(Resource::Decoder(decoder_resource)); Ok((ok(), resource)) } @@ -47,9 +44,8 @@ pub fn decode<'a>( pts_ns: Option, ) -> Result<(Atom, Vec>), Error> { let mut decoder = resource - .decoder - .as_ref() - .expect("Decoder not initialized") + .decoder() + .ok_or_else(|| Error::BadArg)? .decoder_mutex .lock() .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?; @@ -80,9 +76,8 @@ pub fn decode<'a>( pub fn flush(env: Env, resource: ResourceArc) -> Result<(Atom, Vec), Error> { let mut decoder = resource - .decoder - .as_ref() - .expect("Decoder not initialized") + .decoder() + .ok_or_else(|| Error::BadArg)? .decoder_mutex .lock() .map_err(|err| Error::RaiseTerm(Box::new(err.to_string())))?; diff --git a/native/vkvideo/src/encoder.rs b/native/vkvideo/src/encoder.rs index c086758..ba2419f 100644 --- a/native/vkvideo/src/encoder.rs +++ b/native/vkvideo/src/encoder.rs @@ -120,10 +120,7 @@ pub fn new( height, }; - let resource = ResourceArc::new(Resource { - encoder: Some(encoder_resource), - decoder: None, - }); + let resource = ResourceArc::new(Resource::Encoder(encoder_resource)); Ok((ok(), resource)) } @@ -133,7 +130,7 @@ pub fn encode<'a>( bytes: Binary, pts_ns: Option, ) -> Result<(Atom, EncodedFrame<'a>), Error> { - let encoder_resource = resource.encoder.as_ref().expect("Encoder not initialized"); + let encoder_resource = resource.encoder().ok_or_else(|| Error::BadArg)?; let frame = Frame { data: RawFrameData { frame: bytes.to_vec(), diff --git a/native/vkvideo/src/lib.rs b/native/vkvideo/src/lib.rs index 68ada05..2ffe1b1 100644 --- a/native/vkvideo/src/lib.rs +++ b/native/vkvideo/src/lib.rs @@ -5,9 +5,25 @@ use rustler::{Atom, Binary, Env, Error, ResourceArc, Term}; pub mod decoder; pub mod encoder; -pub struct Resource { - pub decoder: Option, - pub encoder: Option, +pub enum Resource { + Encoder(EncoderResource), + Decoder(DecoderResource), +} + +impl Resource { + pub fn encoder(&self) -> Option<&EncoderResource> { + match self { + Self::Encoder(encoder_resource) => Some(encoder_resource), + Self::Decoder(_) => None, + } + } + + pub fn decoder(&self) -> Option<&DecoderResource> { + match self { + Self::Decoder(decoder_resource) => Some(decoder_resource), + Self::Encoder(_) => None, + } + } } fn load(env: Env, _: Term) -> bool { From 02fffd1304a5f2fcfe6ef51a6c9bb1457934bbdc Mon Sep 17 00:00:00 2001 From: root Date: Mon, 22 Dec 2025 15:32:30 +0100 Subject: [PATCH 3/4] Fix function names --- lib/encoder.ex | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/encoder.ex b/lib/encoder.ex index 5b60985..097996e 100644 --- a/lib/encoder.ex +++ b/lib/encoder.ex @@ -85,7 +85,7 @@ defmodule Membrane.VKVideo.Encoder do defp spawn_encoder(state) do {:ok, encoder} = - Native.new( + Native.new_encoder( state.width, state.height, state.framerate, From 8c15bfa881d7e600e6a54f26bbdc704ea75f7dd1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20Kita?= Date: Tue, 23 Dec 2025 11:57:54 +0100 Subject: [PATCH 4/4] 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 Co-authored-by: Jerzy Wilczek --- lib/decoder.ex | 8 +++-- lib/device_server.ex | 28 +++++++++++++++ lib/encoder.ex | 6 +++- lib/native.ex | 13 +++++-- lib/vk_video.ex | 11 ++++++ mix.exs | 3 +- native/vkvideo/src/decoder.rs | 22 +++++------- native/vkvideo/src/encoder.rs | 30 +++++++--------- native/vkvideo/src/lib.rs | 64 ++++++++++++++++++++++++++++++++--- test/decoder_native_test.exs | 6 ++-- test/encoder_native_test.exs | 15 ++++++-- 11 files changed, 158 insertions(+), 48 deletions(-) create mode 100644 lib/device_server.ex create mode 100644 lib/vk_video.ex diff --git a/lib/decoder.ex b/lib/decoder.ex index f4190da..ef895c7 100644 --- a/lib/decoder.ex +++ b/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 diff --git a/lib/device_server.ex b/lib/device_server.ex new file mode 100644 index 0000000..c2a0494 --- /dev/null +++ b/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 diff --git a/lib/encoder.ex b/lib/encoder.ex index 097996e..16516eb 100644 --- a/lib/encoder.ex +++ b/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 diff --git a/lib/native.ex b/lib/native.ex index bba5327..cd1f49a 100644 --- a/lib/native.ex +++ b/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 diff --git a/lib/vk_video.ex b/lib/vk_video.ex new file mode 100644 index 0000000..9fcf078 --- /dev/null +++ b/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 diff --git a/mix.exs b/mix.exs index e59330b..4a3962b 100644 --- a/mix.exs +++ b/mix.exs @@ -28,7 +28,8 @@ defmodule Membrane.VKVideo.Mixfile do def application do [ - extra_applications: [] + extra_applications: [], + mod: {Membrane.VKVideo, []} ] end diff --git a/native/vkvideo/src/decoder.rs b/native/vkvideo/src/decoder.rs index 0e5e4fb..344ef15 100644 --- a/native/vkvideo/src/decoder.rs +++ b/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, } @@ -19,16 +17,14 @@ pub struct RawFrame<'a> { pub height: u32, } -pub fn new() -> Result<(Atom, ResourceArc), 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, +) -> Result<(Atom, ResourceArc), 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); diff --git a/native/vkvideo/src/encoder.rs b/native/vkvideo/src/encoder.rs index ba2419f..906a2b0 100644 --- a/native/vkvideo/src/encoder.rs +++ b/native/vkvideo/src/encoder.rs @@ -1,3 +1,4 @@ +use crate::ok; use crate::Resource; use rustler::{Atom, Error, NifTaggedEnum, ResourceArc}; use rustler::{Binary, Env, NifStruct, NifUnitEnum, OwnedBinary}; @@ -5,12 +6,8 @@ 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, + pub encoder_mutex: Mutex>, pub width: u32, pub height: u32, } @@ -75,22 +72,17 @@ impl Into for EncoderRateControl { } pub fn new( + _env: Env, + resource: ResourceArc, width: u32, height: u32, frame_rate: (u32, u32), tune: EncoderTune, rate_control: EncoderRateControl, ) -> Result<(Atom, ResourceArc), 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())))?; diff --git a/native/vkvideo/src/lib.rs b/native/vkvideo/src/lib.rs index 2ffe1b1..d0e1e19 100644 --- a/native/vkvideo/src/lib.rs +++ b/native/vkvideo/src/lib.rs @@ -1,6 +1,12 @@ 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; @@ -8,31 +14,64 @@ pub mod encoder; pub enum Resource { Encoder(EncoderResource), Decoder(DecoderResource), + Device(DeviceResource), +} + +pub struct DeviceResource { + pub device: Arc, } +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), Error> { - decoder::new() +fn create_device() -> Result<(Atom, ResourceArc), 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, +) -> Result<(Atom, ResourceArc), 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, width: u32, height: u32, frame_rate: (u32, u32), tune: EncoderTune, rate_control: EncoderRateControl, ) -> Result<(Atom, ResourceArc), 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) -> Result { + 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); diff --git a/test/decoder_native_test.exs b/test/decoder_native_test.exs index b02c0bb..e590faa 100644 --- a/test/decoder_native_test.exs +++ b/test/decoder_native_test.exs @@ -1,13 +1,14 @@ defmodule Decoder.NativeTest do use ExUnit.Case, async: false - alias Membrane.VKVideo.Native + alias Membrane.VKVideo.{DeviceServer, Native} @tag :requires_gpu test "Decoder decodes H.264 stream" do in_path = "./fixtures/input-100.h264" |> Path.expand(__DIR__) assert {:ok, file} = File.read(in_path) - {:ok, decoder_ref} = Native.new_decoder() + {:ok, device} = DeviceServer.get_device() + {:ok, decoder_ref} = Native.new_decoder(device) {:ok, decoded_frames} = Native.decode(decoder_ref, file, 0) {:ok, flushed_frames} = Native.flush_decoder(decoder_ref) all_frames = decoded_frames ++ flushed_frames @@ -17,5 +18,6 @@ defmodule Decoder.NativeTest do assert first_frame.width == 1280 assert first_frame.height == 720 assert <<213, 213, 213, 213, _rest::binary>> = first_frame.payload + Native.destroy(decoder_ref) end end diff --git a/test/encoder_native_test.exs b/test/encoder_native_test.exs index ce94065..51adeb6 100644 --- a/test/encoder_native_test.exs +++ b/test/encoder_native_test.exs @@ -1,6 +1,6 @@ defmodule Encoder.NativeTest do use ExUnit.Case, async: false - alias Membrane.VKVideo.Native + alias Membrane.VKVideo.{DeviceServer, Native} @width 1280 @height 720 @@ -39,8 +39,17 @@ defmodule Encoder.NativeTest do assert {:ok, file} = File.read(in_path) + {:ok, device} = DeviceServer.get_device() + {:ok, encoder_ref} = - Native.new_encoder(@width, @height, @framerate, unquote(tune), unquote(rate_control_ast)) + Native.new_encoder( + device, + @width, + @height, + @framerate, + unquote(tune), + unquote(rate_control_ast) + ) raw_frames = for <>, do: chunk @@ -53,6 +62,8 @@ defmodule Encoder.NativeTest do end) assert File.read!(ref_path) == Enum.join(encoded_frames) + + Native.destroy(encoder_ref) end end) end