Enable reserved_stack argument to Threads.@spawn

NEWS.md

@@ -103,6 +103,7 @@ New library features
 * `RegexMatch` objects can now be used to construct `NamedTuple`s and `Dict`s ([#50988])
 * `Lockable` is now exported ([#54595])
 * New `ltruncate`, `rtruncate` and `ctruncate` functions for truncating strings to text width, accounting for char widths ([#55351])
+* `Threads.@spawn` can now take a `reserved_stack` argument to set a stack size limit ([#55201])
 
 Standard library changes
 ------------------------

base/threadingconstructs.jl

@@ -438,14 +438,17 @@ function _spawn_set_thrpool(t::Task, tp::Symbol)
 end
 
 """
-    Threads.@spawn [:default|:interactive] expr
+    Threads.@spawn [:default|:interactive] [reserved_stack=0] expr
 
 Create a [`Task`](@ref) and [`schedule`](@ref) it to run on any available
 thread in the specified threadpool (`:default` if unspecified). The task is
 allocated to a thread once one becomes available. To wait for the task to
 finish, call [`wait`](@ref) on the result of this macro, or call
 [`fetch`](@ref) to wait and then obtain its return value.
 
+The optional argument `reserved_stack` requests a stack size in bytes
+for the task. See [`Task`](@ref) for details.
+
 Values can be interpolated into `@spawn` via `\$`, which copies the value
 directly into the constructed underlying closure. This allows you to insert
 the _value_ of a variable, isolating the asynchronous code from changes to
@@ -466,6 +469,9 @@ the variable's value in the current task.
 !!! compat "Julia 1.9"
     A threadpool may be specified as of Julia 1.9.
 
+!!! compat "Julia 1.12"
+    The `reserved_stack` argument is available as of Julia 1.12.
+
 # Examples
 ```julia-repl
 julia> t() = println("Hello from ", Threads.threadid());
@@ -478,23 +484,37 @@ Hello from 4
 ```
 """
 macro spawn(args...)
-    tp = QuoteNode(:default)
-    na = length(args)
-    if na == 2
-        ttype, ex = args
-        if ttype isa QuoteNode
-            ttype = ttype.value
+    default_tp = default_stack = nothing  # guard against setting threadpool multiple times
+    tp = default_tp
+    reserved_stack = default_stack
+    1 <= length(args) <= 3 || throw(ArgumentError("wrong number of arguments in @spawn"))
+    ex = last(args)
+    for arg in args[1:end-1]
+        if arg isa QuoteNode
+            @assert tp == default_tp "cannot set threadpool multiple times"
+            ttype = arg.value
             if ttype !== :interactive && ttype !== :default
                 throw(ArgumentError(LazyString("unsupported threadpool in @spawn: ", ttype)))
             end
             tp = QuoteNode(ttype)
+        elseif arg isa Expr && arg.head == :(=)
+            if arg.args[1] == :reserved_stack
+                @assert reserved_stack == default_stack "cannot set reserved_stack multiple times"
+                reserved_stack = arg.args[2]
+            else
+                throw(ArgumentError(LazyString("unknown argument in @spawn: ", arg.args[1])))
+            end
         else
-            tp = ttype
+            @assert tp == default_tp "cannot set threadpool multiple times"
+            tp = arg
         end
-    elseif na == 1
-        ex = args[1]
-    else
-        throw(ArgumentError("wrong number of arguments in @spawn"))
+    end
+
+    if tp == default_tp
+        tp = QuoteNode(:default)
+    end
+    if reserved_stack == default_stack
+        reserved_stack = 0
     end
 
     letargs = Base._lift_one_interp!(ex)
@@ -503,7 +523,7 @@ macro spawn(args...)
     var = esc(Base.sync_varname)
     quote
         let $(letargs...)
-            local task = Task($thunk)
+            local task = Task($thunk, $(esc(reserved_stack)))
             task.sticky = false
             _spawn_set_thrpool(task, $(esc(tp)))
             if $(Expr(:islocal, var))

test/threadpool_use.jl

@@ -17,3 +17,20 @@ tp = :foo
 @test_throws ArgumentError Threads.@spawn tp Threads.threadpool()
 @test Threads.threadpooltids(:interactive) == [1]
 @test Threads.threadpooltids(:default) == [2]
+
+# Test with custom stack size:
+@test fetch(Threads.@spawn reserved_stack=0 Threads.threadpool()) === :default
+
+# With both threadpool and stack size
+@test fetch(Threads.@spawn :default reserved_stack=0 Threads.threadpool()) === :default
+
+# Order should not matter:
+@test fetch(Threads.@spawn reserved_stack=0 :default Threads.threadpool()) === :default
+
+# Can set an expression for reserved_stack
+tp = :interactive
+@test fetch(Threads.@spawn reserved_stack=(4 * 1024^2) tp Threads.threadpool()) === :interactive
+
+# Can also pass variables:
+reserved_stack = 4 * 1024^2
+@test fetch(Threads.@spawn reserved_stack=reserved_stack :interactive Threads.threadpool()) === :interactive