From d953781f2a8ef6303a56c35b3274ac08f0ccfb04 Mon Sep 17 00:00:00 2001
From: Ian Butterworth <i.r.butterworth@gmail.com>
Date: Thu, 12 Aug 2021 17:44:44 -0400
Subject: [PATCH] Enable adjustable codec multithreading (#332)

* enable adjustable multithreaded decoding

* enable 2 threads during CI

* also enable multithreaded encoding

* default threading for encode, Sys.CPU_THREADS for decode

* Revert "enable 2 threads during CI"

This reverts commit ffd91cc64cc6efb92ca32c66ccd34bfc959b9a91.
---
 src/avio.jl     | 6 +++++-
 src/encoding.jl | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/avio.jl b/src/avio.jl
index c26243b7..51013f27 100644
--- a/src/avio.jl
+++ b/src/avio.jl
@@ -248,7 +248,8 @@ function VideoReader(avin::AVInput{I}, video_stream = 1;
                      target_colorspace_details = nothing,
                      allow_vio_gray_transform = true,
                      swscale_options::OptionsT = (;),
-                     sws_color_options::OptionsT = (;)) where I
+                     sws_color_options::OptionsT = (;),
+                     thread_count::Union{Nothing, Int} = Sys.CPU_THREADS) where I
     bad_px_type = transcode && target_format !== nothing &&
         !is_pixel_type_supported(target_format)
     bad_px_type && error("Unsupported pixel format $target_format")
@@ -268,6 +269,7 @@ function VideoReader(avin::AVInput{I}, video_stream = 1;
 
     # Create decoder context
     codec_context = AVCodecContextPtr(codec) # Allocates
+    codec_context.thread_count = thread_count === nothing ? codec_context.thread_count : thread_count
     # Transfer parameters to decoder context
     ret = avcodec_parameters_to_context(codec_context, stream.codecpar)
     ret < 0 && error("Could not copy the codec parameters to the decoder")
@@ -584,6 +586,8 @@ arguments listed below.
 - `sws_color_options::OptionsT = (;)`: Additional keyword arguments passed to
     [sws_setColorspaceDetails]
     (http://ffmpeg.org/doxygen/2.5/group__libsws.html#ga541bdffa8149f5f9203664f955faa040).
+- `thread_count::Union{Nothing, Int} = Sys.CPU_THREADS`: The number of threads the codec is
+    allowed to use or `nothing` for default codec behavior. Defaults to `Sys.CPU_THREADS`.
 """
 openvideo(s::Union{IO, AbstractString, AVInput}, args...; kwargs...) =
     VideoReader(s, args...; kwargs...)
diff --git a/src/encoding.jl b/src/encoding.jl
index d1b0bafa..cff52f7a 100644
--- a/src/encoding.jl
+++ b/src/encoding.jl
@@ -240,7 +240,8 @@ function VideoWriter(filename::AbstractString, ::Type{T},
                      pix_fmt_loss_flags = 0,
                      input_colorspace_details = nothing,
                      allow_vio_gray_transform = true,
-                     sws_color_options::OptionsT = (;)) where T
+                     sws_color_options::OptionsT = (;),
+                     thread_count::Union{Nothing, Int} = nothing) where T
     framerate > 0 || error("Framerate must be strictly positive")
 
     if haskey(encoder_options, :priv_data)
@@ -300,6 +301,7 @@ function VideoWriter(filename::AbstractString, ::Type{T},
     codec_context.time_base = target_timebase
     codec_context.framerate = framerate_rat
     codec_context.pix_fmt = encoding_pix_fmt
+    codec_context.thread_count = thread_count === nothing ? codec_context.thread_count : thread_count
 
     if format_context.oformat.flags & AVFMT_GLOBALHEADER != 0
         codec_context.flags |= AV_CODEC_FLAG_GLOBAL_HEADER
@@ -453,6 +455,8 @@ occurred.
 - `sws_color_options::OptionsT = (;)`: Additional keyword arguments passed to
     [sws_setColorspaceDetails]
     (http://ffmpeg.org/doxygen/2.5/group__libsws.html#ga541bdffa8149f5f9203664f955faa040).
+- `thread_count::Union{Nothing, Int} = nothing`: The number of threads the codec is
+    allowed to use, or `nothing` for default codec behavior. Defaults to `nothing`.
 
 See also: [`write`](@ref), [`close_video_out!`](@ref)
 """