Replace current, not global logger in safe loggers (JuliaGPU#622)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

src/utils.jl

@@ -60,13 +60,16 @@ for level in [:debug, :info, :warn, :error]
             # NOTE: `@placeholder` in order to avoid hard-coding @__LINE__ etc
             macrocall.args[1] = Symbol($"@$level")
             quote
-                old_logger = global_logger()
                 io = IOContext(Core.stderr, :color=>STDERR_HAS_COLOR[])
-                min_level = _invoked_min_enabled_level(old_logger)
-                global_logger(Logging.ConsoleLogger(io, min_level))
-                ret = $(esc(macrocall))
-                global_logger(old_logger)
-                ret
+                # ideally we call Logging.shouldlog() here, but that is likely to yield,
+                # so instead we rely on the min_enabled_level of the logger.
+                # in the case of custom loggers that may be an issue, because,
+                # they may expect Logging.shouldlog() getting called, so we use
+                # the global_logger()'s min level which is more likely to be usable.
+                min_level = _invoked_min_enabled_level(global_logger())
+                with_logger(Logging.ConsoleLogger(io, min_level)) do
+                    $(esc(macrocall))
+                end
             end
         end
     end

test/Project.toml

@@ -2,6 +2,7 @@
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"

test/util_tests.jl

@@ -19,4 +19,56 @@ end
     @test GPUCompiler.mangle_sig(Tuple{typeof(sin), XX{Int64(-10)}}) == "_Z3sin2XXILln10EE"  # "sin(XX<-10l>)"
 end
 
+@testset "safe loggers" begin
+    using Logging: Logging
+
+    struct YieldingLogger <: Logging.AbstractLogger
+        logger::Logging.AbstractLogger
+        YieldingLogger() = new(Logging.current_logger())
+    end
+
+    function Logging.handle_message(logger::YieldingLogger, args...)
+        yield()
+        return Logging.handle_message(logger.logger, args...)
+    end
+
+    Logging.shouldlog(::YieldingLogger, ::Any...) = true
+    Logging.min_enabled_level(::YieldingLogger) = Logging.Debug
+
+    GPUCompiler.@locked function f()
+        GPUCompiler.@safe_debug "safe_debug"
+        GPUCompiler.@safe_info "safe_info"
+        GPUCompiler.@safe_warn "safe_warn"
+        GPUCompiler.@safe_error "safe_error"
+        GPUCompiler.@safe_show "safe_show"
+    end
+
+    @test begin
+        @sync begin
+            Threads.@spawn begin
+                sleep(0.1)
+                @debug "debug"
+                sleep(0.1)
+                @info "info"
+                sleep(0.1)
+                @warn "warn"
+                sleep(0.1)
+                @error "error"
+                sleep(0.1)
+                @show "show"
+                sleep(0.1)
+            end
+            pipe = Pipe()
+            Base.link_pipe!(pipe; reader_supports_async=true, writer_supports_async=true)
+            Threads.@spawn print(stdout, read(pipe, String))
+            Threads.@spawn Logging.with_logger(YieldingLogger()) do
+                sleep(0.1)
+                redirect_stdout(f, pipe)
+                close(pipe)
+            end
+        end
+        true
+    end
+end
+
 end