Thoroughly nospecialize all functions; add no-alloc, no-specialize test.

Make sure that _all_ functions in ExceptionUnwrapping are marked
nospecialize. We don't want to pay the wasted compilation time at
runtime, since these are all going to be called in _exceptional_ cases,
certainly not in a hot loop, and because we've seen crashes caused by a
stackoverflow in type inference while attempting to specialize the code
to handle a StackOverflowException! 😅

This time, we add a unit test to ensure that that these functions do not
allocate and do not incur new compilation when called with novel
arguments.

src/ExceptionUnwrapping.jl

@@ -105,50 +105,46 @@ end
 UnwrappedExceptionNotFound{R}(e::E) where {R,E} = UnwrappedExceptionNotFound{R,E}(e)
 
 
+# We have confirmed via Cthulhu and the Allocations profiler that these seem to correctly
+# not be specializing, and not allocating.
+# ... For some reason, it seems that we also need to put this attribute on the arguments
+# in the function definitions as well. Without that, it is still specializing. Not sure why.
+@nospecialize
+
 # Base case is that e -> e
-unwrap_exception(e) = e
+unwrap_exception(@nospecialize(e)) = e
 # Add overloads for wrapped exception types to unwrap the exception.
 # TaskFailedExceptions wrap a failed task, which contains the exception that caused it
 # to fail. You can unwrap the exception to discover the root cause of the failure.
 unwrap_exception(e::Base.TaskFailedException) = e.task.exception
 unwrap_exception(e::Base.CapturedException) = e.ex
 
-has_wrapped_exception(::T, ::Type{T}) where T = true
-
-# We have confirmed via Cthulhu and the Allocations profiler that these seem to correctly
-# not be specializing, and not allocating.
-@nospecialize
-
-# Types don't match, do the unrolling, but prevent inference since this happens at runtime
-# and only during exception catch blocks, and might have arbitrarily nested types. And in
-# practice, we've seen julia's inference really struggles here.
+# If types don't match, do the unrolling, but prevent inference since this happens at
+# runtime and only during exception catch blocks, and might have arbitrarily nested types.
+# And in practice, we've seen julia's inference really struggles here.
 # The inferencebarrier blocks the callee from being inferred until it's actually called at
 # runtime, so that we don't pay for expensive inference if the exception path isn't
 # triggered.
-function has_wrapped_exception(e, ::Type{T}) where T
+function has_wrapped_exception(@nospecialize(e), @nospecialize(T::Type))
+    e isa T && return true
     Base.inferencebarrier(_has_wrapped_exception)(e, T)
 end
-function _has_wrapped_exception(e, ::Type{T}) where T
+function _has_wrapped_exception(@nospecialize(e), @nospecialize(T::Type))
     while !(e isa T) && is_wrapped_exception(e)
         e::Any = unwrap_exception(e)
     end
     return e isa T
 end
 
-function is_wrapped_exception(e)
+function is_wrapped_exception(@nospecialize(e))
     return e !== unwrap_exception(e)
 end
 
-@specialize
-
-unwrap_exception_until(e::T, ::Type{T}) where T = e
-
-@nospecialize
-
-function unwrap_exception_until(e, ::Type{T}) where T
+function unwrap_exception_until(@nospecialize(e), @nospecialize(T::Type))
+    e isa T && return e
     Base.inferencebarrier(_unwrap_exception_until)(e, T)
 end
-function _unwrap_exception_until(e, ::Type{T}) where T
+function _unwrap_exception_until(@nospecialize(e), @nospecialize(T::Type))
     while !(e isa T) && is_wrapped_exception(e)
         e::Any = unwrap_exception(e)
     end
@@ -159,10 +155,10 @@ function _unwrap_exception_until(e, ::Type{T}) where T
     end
 end
 
-function unwrap_exception_to_root(e)
+function unwrap_exception_to_root(@nospecialize(e))
     Base.inferencebarrier(_unwrap_exception_to_root)(e)
 end
-function _unwrap_exception_to_root(e)
+function _unwrap_exception_to_root(@nospecialize(e))
     while is_wrapped_exception(e)
         e::Any = unwrap_exception(e)
     end
@@ -171,4 +167,18 @@ end
 
 @specialize
 
+function __init__()
+    # Can't use `(Any,)` for unwrap_exception because it has a more-specific subtype variant
+    @assert precompile(unwrap_exception, (ErrorException,))  # nospecialized variant
+    @assert precompile(unwrap_exception, (Base.TaskFailedException,))
+
+    @assert precompile(is_wrapped_exception, (Any,))  # nospecialized
+    @assert precompile(unwrap_exception_to_root, (Any,))  # nospecialized
+    @static if VERSION >= v"1.7.0-"
+        @assert precompile(summarize_current_exceptions, (IO, Task))  # nospecialized
+    end
+
+    @assert precompile(has_wrapped_exception, (Any, Type))  # nospecialized
+end
+
 end # module

src/exception_summary.jl

@@ -6,6 +6,8 @@
 - Seen set, for deduplication
 =#
 
+@nospecialize
+
 # Consider adding a _summarize_exception() overload for DistributedException
 #     Pros: less noise
 #     Cons: possibly hiding intermediate exceptions that might have been helpful to see.
@@ -110,7 +112,10 @@ function _summarize_exception(io::IO, e::CompositeException, stack; prefix = not
     end
 end
 # This is the overload that prints the actual exception that occurred.
-function _summarize_exception(io::IO, exc, stack; prefix = nothing)
+function _summarize_exception(io::IO, @nospecialize(exc), stack; prefix = nothing)
+    @show exc
+    @show is_wrapped_exception(exc)
+    global EXC = exc
     # First, check that this exception isn't some other kind of user-defined
     # wrapped exception. We want to unwrap this layer as well, so that we are
     # printing just the true exceptions in the summary, not any exception
@@ -156,3 +161,5 @@ function _summarize_exception(io::IO, exc, stack; prefix = nothing)
         println(io)
     end
 end
+
+@specialize

test/ExceptionUnwrapping.jl

@@ -24,7 +24,7 @@ if VERSION >= v"1.3.0-"
             @test_throws UnwrappedExceptionNotFound{ArgumentError} unwrap_exception_until(e, ArgumentError) isa ErrorException
         end
     end
-    
+
     @testset "Wrapped CapturedException" begin
         e = CapturedException(ErrorException("oh no"), backtrace())
         @test unwrap_exception(e) == ErrorException("oh no")
@@ -88,5 +88,27 @@ ExceptionUnwrapping.unwrap_exception(e::MyWrappedException2) = e.exc
     end
 end
 
+@testset "allocations" begin
+    t = @async throw(ArgumentError("foo"))
+    try wait(t) catch end
+    TE = TaskFailedException(t)
+
+    # Precompile it once
+    @test ExceptionUnwrapping.has_wrapped_exception(TE, ArgumentError) == true
+    @test ExceptionUnwrapping.unwrap_exception(TE) isa ArgumentError
+
+    # Test no allocations
+    @test @allocated(ExceptionUnwrapping.has_wrapped_exception(TE, ArgumentError)) == 0
+    @test @allocated(ExceptionUnwrapping.unwrap_exception(TE)) == 0
+
+    # Test that there's nothing being compiled, even for novel types
+    @eval struct Foo <: Exception end
+    e = Foo()
+    @test @allocated(ExceptionUnwrapping.has_wrapped_exception(e, ArgumentError)) == 0
+    @test @allocated(ExceptionUnwrapping.has_wrapped_exception(e, Foo)) == 0
+    @test @allocated(ExceptionUnwrapping.unwrap_exception(e)) == 0
+end
+
+
 
 end # module