From 8121b6f67573f63688d5c5ec02fd096de75b2d1e Mon Sep 17 00:00:00 2001
From: Nathan Daly <nathan.daly@relational.ai>
Date: Wed, 22 May 2024 08:53:20 -0600
Subject: [PATCH] Make TestLogger thread-safe (introduce a lock)

Fixes https://github.com/JuliaLang/julia/issues/54439.

- Lock around concurrent accesses to .logs and .message_limits
- Copy the vector out of the logger at the end, to shield against dangling Tasks.

Before:
```julia
julia> Threads.nthreads()
8

julia> function foo(n)
           @info "Doing foo with n=$n"
           @sync for i=1:n
               Threads.@spawn @info "Iteration $i"
           end
           42
       end
foo (generic function with 1 method)

julia> for _ in 1:1000
           @test_logs (:info,"Doing foo with n=10000") match_mode=:any foo(10_000)
       end
julia+RAI(56155,0x1f5157ac0) malloc: double free for ptr 0x128248000
julia+RAI(56155,0x1f5157ac0) malloc: *** set a breakpoint in malloc_error_break to debug

[56155] signal (6): Abort trap: 6
in expression starting at REPL[8]:1

signal (6) thread (1) __pthread_kill at /usr/lib/system/libsystem_kernel.dylib (unknown line)
Allocations: 54370881 (Pool: 54363911; Big: 6970); GC: 119
[1]    56154 abort      julia -tauto
```
After:
```julia
julia> Threads.nthreads()
8

julia> function foo(n)
           @info "Doing foo with n=$n"
           @sync for i=1:n
               Threads.@spawn @info "Iteration $i"
           end
           42
       end
foo (generic function with 1 method)

julia> for _ in 1:1000
           @test_logs (:info,"Doing foo with n=10000") match_mode=:any foo(10_000)
       end
```
(no crash) :)
---
 stdlib/Test/src/logging.jl | 42 +++++++++++++++++++++++++-------------
 1 file changed, 28 insertions(+), 14 deletions(-)

diff --git a/stdlib/Test/src/logging.jl b/stdlib/Test/src/logging.jl
index 4e444874d0fb8..ee647da10551f 100644
--- a/stdlib/Test/src/logging.jl
+++ b/stdlib/Test/src/logging.jl
@@ -2,6 +2,7 @@
 
 using Logging: Logging, AbstractLogger, LogLevel, Info, with_logger
 import Base: occursin
+using Base: @lock
 
 #-------------------------------------------------------------------------------
 """
@@ -35,11 +36,15 @@ struct Ignored ; end
 #-------------------------------------------------------------------------------
 # Logger with extra test-related state
 mutable struct TestLogger <: AbstractLogger
-    logs::Vector{LogRecord}
+    lock::ReentrantLock
+    logs::Vector{LogRecord}  # Guarded by lock.
     min_level::LogLevel
     catch_exceptions::Bool
-    shouldlog_args
-    message_limits::Dict{Any,Int}
+    # Note: shouldlog_args only maintains the info for the most recent log message, which
+    # may not be meaningful in a multithreaded program. See:
+    # https://github.com/JuliaLang/julia/pull/54497#discussion_r1603691606
+    shouldlog_args  # Guarded by lock.
+    message_limits::Dict{Any,Int}  # Guarded by lock.
     respect_maxlog::Bool
 end
 
@@ -80,15 +85,17 @@ Test Passed
 ```
 """
 TestLogger(; min_level=Info, catch_exceptions=false, respect_maxlog=true) =
-    TestLogger(LogRecord[], min_level, catch_exceptions, nothing, Dict{Any, Int}(), respect_maxlog)
+    TestLogger(ReentrantLock(), LogRecord[], min_level, catch_exceptions, nothing, Dict{Any, Int}(), respect_maxlog)
 Logging.min_enabled_level(logger::TestLogger) = logger.min_level
 
 function Logging.shouldlog(logger::TestLogger, level, _module, group, id)
-    if get(logger.message_limits, id, 1) > 0
-        logger.shouldlog_args = (level, _module, group, id)
-        true
-    else
-        false
+    @lock logger.lock begin
+        if get(logger.message_limits, id, 1) > 0
+            logger.shouldlog_args = (level, _module, group, id)
+            return true
+        else
+            return false
+        end
     end
 end
 
@@ -98,12 +105,17 @@ function Logging.handle_message(logger::TestLogger, level, msg, _module,
     if logger.respect_maxlog
         maxlog = get(kwargs, :maxlog, nothing)
         if maxlog isa Core.BuiltinInts
-            remaining = get!(logger.message_limits, id, Int(maxlog)::Int)
-            logger.message_limits[id] = remaining - 1
-            remaining > 0 || return
+            @lock logger.lock begin
+                remaining = get!(logger.message_limits, id, Int(maxlog)::Int)
+                logger.message_limits[id] = remaining - 1
+                remaining > 0 || return
+            end
         end
     end
-    push!(logger.logs, LogRecord(level, msg, _module, group, id, file, line, kwargs))
+    r = LogRecord(level, msg, _module, group, id, file, line, kwargs)
+    @lock logger.lock begin
+        push!(logger.logs, r)
+    end
 end
 
 # Catch exceptions for the test logger only if specified
@@ -112,7 +124,9 @@ Logging.catch_exceptions(logger::TestLogger) = logger.catch_exceptions
 function collect_test_logs(f; kwargs...)
     logger = TestLogger(; kwargs...)
     value = with_logger(f, logger)
-    logger.logs, value
+    @lock logger.lock begin
+        return copy(logger.logs), value
+    end
 end