pkg: Add activate_with_early_release for early lock release (#348)

Add a variant of activate_do that allows early release of ACTIVATE_LOCK
before the callback completes. The caller can notify the passed event to
signal that the activated environment is no longer needed, enabling the
environment to be restored and the lock released while the callback
continues executing.

Use this in full-analysis to release the activation lock after the code
loading phase finishes, allowing other analysis requests to proceed
while type inference continues in the background.

Author: aviatesk

src/analysis/Interpreter.jl

@@ -23,24 +23,38 @@ struct LSInterpreter{S<:Server} <: JET.ConcreteInterpreter
     request::AnalysisRequest
     analyzer::LSAnalyzer
     counter::Counter
+    activation_done::Union{Nothing,Base.Event}
     state::JET.InterpretationState
-    function LSInterpreter(server::S, request::AnalysisRequest, analyzer::LSAnalyzer, counter::Counter) where S<:Server
-        return new{S}(server, request, analyzer, counter)
+    function LSInterpreter(
+            server::S, request::AnalysisRequest, analyzer::LSAnalyzer, counter::Counter,
+            activation_done::Union{Nothing,Base.Event}
+        ) where S<:Server
+        return new{S}(server, request, analyzer, counter, activation_done)
     end
-    function LSInterpreter(server::S, request::AnalysisRequest, analyzer::LSAnalyzer, counter::Counter, state::JET.InterpretationState) where S<:Server
-        return new{S}(server, request, analyzer, counter, state)
+    function LSInterpreter(
+            server::S, request::AnalysisRequest, analyzer::LSAnalyzer, counter::Counter,
+            activation_done::Union{Nothing,Base.Event}, state::JET.InterpretationState
+        ) where S<:Server
+        return new{S}(server, request, analyzer, counter, activation_done, state)
     end
 end
 
 # The main constructor
-LSInterpreter(server::Server, request::AnalysisRequest) = LSInterpreter(server, request, LSAnalyzer(request.entry), Counter())
+function LSInterpreter(
+        server::Server, request::AnalysisRequest;
+        activation_done::Union{Nothing,Base.Event} = nothing
+    )
+    return LSInterpreter(server, request, LSAnalyzer(request.entry), Counter(), activation_done)
+end
 
 # `JET.ConcreteInterpreter` interface
 JET.InterpretationState(interp::LSInterpreter) = interp.state
 function JET.ConcreteInterpreter(interp::LSInterpreter, state::JET.InterpretationState)
     # add `state` to `interp`, and update `interp.analyzer.cache`
     initialize_cache!(interp.analyzer, state.res.analyzed_files)
-    return LSInterpreter(interp.server, interp.request, interp.analyzer, interp.counter, state)
+    return LSInterpreter(
+        interp.server, interp.request, interp.analyzer, interp.counter,
+        interp.activation_done, state)
 end
 JET.ToplevelAbstractAnalyzer(interp::LSInterpreter) = interp.analyzer
 
@@ -58,6 +72,13 @@ function cache_intermediate_analysis_result!(interp::LSInterpreter)
 end
 
 function JET.analyze_from_definitions!(interp::LSInterpreter, config::JET.ToplevelConfig)
+    activation_done = interp.activation_done
+    if activation_done !== nothing
+        # The phase that requires code loading has finished, so this environment is no
+        # longer needed, so let's release the `ACTIVATION_LOCK`
+        notify(activation_done)
+    end
+
     # Cache intermediate analysis results after file analysis completes
     # This makes module context information available immediately for LS features
     cache_intermediate_analysis_result!(interp)

src/analysis/full-analysis.jl

@@ -386,18 +386,18 @@ function execute_analysis_request(server::Server, request::AnalysisRequest)
         result = analyze_parsed_if_exist(server, request)
 
     elseif entry isa ScriptInEnvAnalysisEntry
-        result = activate_do(entry.env_path) do
-            analyze_parsed_if_exist(server, request)
+        result = activate_with_early_release(entry.env_path) do activation_done::Base.Event
+            analyze_parsed_if_exist(server, request; activation_done)
         end
 
     elseif entry isa PackageSourceAnalysisEntry
-        result = activate_do(entry.env_path) do
-            analyze_parsed_if_exist(server, request, entry.pkgid)
+        result = activate_with_early_release(entry.env_path) do activation_done::Base.Event
+            analyze_parsed_if_exist(server, request, entry.pkgid; activation_done)
         end
 
     elseif entry isa PackageTestAnalysisEntry
-        result = activate_do(entry.env_path) do
-            analyze_parsed_if_exist(server, request)
+        result = activate_with_early_release(entry.env_path) do activation_done::Base.Event
+            analyze_parsed_if_exist(server, request; activation_done)
         end
 
     else error("Unsupported analysis entry $entry") end
@@ -428,16 +428,21 @@ function end_full_analysis_progress(server::Server, cancellable_token::Cancellab
         WorkDoneProgressEnd(; message = "Analysis completed"))
 end
 
-function analyze_parsed_if_exist(server::Server, request::AnalysisRequest, args...)
+function analyze_parsed_if_exist(
+        server::Server, request::AnalysisRequest, args...;
+        activation_done::Union{Nothing,Base.Event} = nothing
+    )
     uri = entryuri(request.entry)
     jetconfigs = getjetconfigs(request.entry)
     fi = get_saved_file_info(server.state, uri)
     if !isnothing(fi)
         filename = @something uri2filename(uri) error(lazy"Unsupported URI: $uri")
-        return JET.analyze_and_report_expr!(LSInterpreter(server, request), fi.syntax_node, filename, args...; jetconfigs...)
+        interp = LSInterpreter(server, request; activation_done)
+        return JET.analyze_and_report_expr!(interp, fi.syntax_node, filename, args...; jetconfigs...)
     else
         filepath = @something uri2filepath(uri) error(lazy"Unsupported URI: $uri")
-        return JET.analyze_and_report_file!(LSInterpreter(server, request), filepath, args...; jetconfigs...)
+        interp = LSInterpreter(server, request; activation_done)
+        return JET.analyze_and_report_file!(interp, filepath, args...; jetconfigs...)
     end
 end
 
@@ -675,23 +680,25 @@ end
 
 function ensure_instantiated!(server::Server, env_path::String)
     if get_config(server.state.config_manager, :full_analysis, :auto_instantiate)
+        manifest_name = "Manifest-v$(VERSION.major).$(VERSION.minor).toml"
+        manifest_path = joinpath(dirname(env_path), manifest_name)
+        io = IOBuffer()
         try
-            manifest_name = "Manifest-v$(VERSION.major).$(VERSION.minor).toml"
-            manifest_path = joinpath(dirname(env_path), manifest_name)
             if !isfile(manifest_path)
                 JETLS_DEV_MODE && @info "Touching versioned manifest file" env_path
                 touch(manifest_path)
             end
             JETLS_DEV_MODE && @info "Resolving package environment" env_path
-            Pkg.resolve()
+            Pkg.resolve(; io)
             JETLS_DEV_MODE && @info "Instantiating package environment" env_path
-            Pkg.instantiate()
+            Pkg.instantiate(; io)
         catch e
             @error """Failed to instantiate package environment;
             Unable to instantiate the environment of the target package for analysis,
             so this package will be analyzed as a script instead.
             This may cause various features such as diagnostics to not function properly.
             It is recommended to fix the problem by referring to the following error""" env_path
+            print(stderr, String(take!(io)))
             Base.showerror(stderr, e, catch_backtrace())
             show_warning_message(server, """
                 Failed to instantiate package environment at $env_path.

src/utils/pkg.jl

@@ -42,7 +42,7 @@ function find_pkg_name(env_path::AbstractString)
     return pkg_name isa String ? pkg_name : nothing
 end
 
-const ACTIVATE_LOCK = ReentrantLock()
+const PKG_ACTIVATION_LOCK = ReentrantLock()
 
 """
     activate_do(func, env_path::String)
@@ -51,15 +51,41 @@ Temporarily activate the environment at `env_path`, execute `func`, and restore
 previous environment. Uses a global lock to prevent concurrent environment switching.
 """
 function activate_do(func, env_path::String)
-    @lock ACTIVATE_LOCK begin
-        old_env = Pkg.project().path
-        try
-            Pkg.activate(env_path; io=devnull)
-            func()
-        finally
-            Pkg.activate(old_env; io=devnull)
-        end
+    lock(PKG_ACTIVATION_LOCK)
+    old_env = Pkg.project().path
+    try
+        Pkg.activate(env_path; io=devnull)
+        return func()
+    finally
+        Pkg.activate(old_env; io=devnull)
+        unlock(PKG_ACTIVATION_LOCK)
+    end
+end
+
+"""
+    activate_with_early_release(func, env_path::String)
+
+Temporarily activate the environment at `env_path` and execute `func(activation_done)`.
+Unlike [`activate_do`](@ref), this allows early release of `PKG_ACTIVATION_LOCK` before `func`
+completes: the caller can `notify(activation_done)` to signal that the activated environment
+is no longer needed, allowing the environment to be restored and the lock released while
+`func` continues executing. `func` is allowed to return without notifying, in which case
+the event is automatically notified in the `finally` block.
+"""
+function activate_with_early_release(func, env_path::String)
+    activation_done = Base.Event()
+    lock(PKG_ACTIVATION_LOCK)
+    old_env = Pkg.project().path
+    Pkg.activate(env_path; io=devnull)
+    t = Threads.@spawn try
+        func(activation_done)
+    finally
+        notify(activation_done)
     end
+    wait(activation_done)
+    Pkg.activate(old_env; io=devnull)
+    unlock(PKG_ACTIVATION_LOCK)
+    return fetch(t)
 end
 
 function find_package_directory(path::String, env_path::String)