cancellation: fix memory leak in request cancellation handling (#260)

Implement `FixedSizeFIFOQueue` to keep track of recently handled request
IDs and prevent dead IDs from accumulating in `currently_handled` when
clients send `CancelRequestNotification` for already processed requests.

The `handled_history` queue maintains the last `128` processed IDs with
O(1) membership testing, allowing the server to efficiently ignore
duplicate cancellation requests while preventing unbounded memory
growth.

Co-authored-by: Claude <noreply@anthropic.com>

Author: aviatesk

src/FixedSizeFIFOQueue/FixedSizeFIFOQueue.jl

@@ -0,0 +1,59 @@
+mutable struct FixedSizeFIFOQueue{T}
+    const data::Vector{T}
+    head::Int
+    tail::Int
+    size::Int
+    const capacity::Int
+    const items::Set{T}
+
+    function FixedSizeFIFOQueue{T}(capacity::Int) where T
+        capacity > 0 || throw(ArgumentError("Capacity must be positive"))
+        data = Vector{T}(undef, capacity)
+        new{T}(data, 1, 1, 0, capacity, Set{T}())
+    end
+    FixedSizeFIFOQueue(capacity::Int) = FixedSizeFIFOQueue{Any}(capacity)
+end
+
+function Base.push!(queue::FixedSizeFIFOQueue{T}, item::T) where T
+    if queue.size == queue.capacity
+        # Queue is full, overwrite oldest element
+        old_item = queue.data[queue.head]
+        delete!(queue.items, old_item)
+        queue.data[queue.tail] = item
+        push!(queue.items, item)
+        queue.head = mod1(queue.head + 1, queue.capacity)
+        queue.tail = mod1(queue.tail + 1, queue.capacity)
+    else
+        # Queue has space
+        queue.data[queue.tail] = item
+        push!(queue.items, item)
+        queue.tail = mod1(queue.tail + 1, queue.capacity)
+        queue.size += 1
+    end
+    return queue
+end
+
+Base.in(item::T, queue::FixedSizeFIFOQueue{T}) where T = item in queue.items
+
+Base.isempty(queue::FixedSizeFIFOQueue) = queue.size == 0
+
+isfull(queue::FixedSizeFIFOQueue) = queue.size == queue.capacity
+
+Base.length(queue::FixedSizeFIFOQueue) = queue.size
+
+capacity(queue::FixedSizeFIFOQueue) = queue.capacity
+
+function Base.collect(queue::FixedSizeFIFOQueue{T}) where T
+    result = Vector{T}(undef, queue.size)
+    idx = queue.head
+    for i in 1:queue.size
+        result[i] = queue.data[idx]
+        idx = mod1(idx + 1, queue.capacity)
+    end
+    return result
+end
+
+function Base.show(io::IO, queue::FixedSizeFIFOQueue{T}) where T
+    items = collect(queue)
+    print(io, "FixedSizeFIFOQueue{$T}(capacity=$(queue.capacity), items=$items)")
+end

src/JETLS.jl

@@ -43,6 +43,7 @@ Analyzer.LSAnalyzer(args...; kwargs...) = LSAnalyzer(ScriptAnalysisEntry(filepat
 include("analysis/resolver.jl")
 
 include("AtomicContainers/AtomicContainers.jl")
+include("FixedSizeFIFOQueue/FixedSizeFIFOQueue.jl")
 using .AtomicContainers
 const SWStats  = JETLS_DEV_MODE ? AtomicContainers.SWStats : Nothing
 const LWStats  = JETLS_DEV_MODE ? AtomicContainers.LWStats : Nothing
@@ -212,11 +213,6 @@ function handle_sequential_message(server::Server, @nospecialize msg)
     end
 end
 
-# TODO This cancellation handling still has the following cleanup issues not yet implemented:
-# - When a client mistakenly sends a `CancelRequestNotification` for an already processed request,
-#   it will register a dead ID in `currently_handled`. A fixed size FIFO queue to record
-#   already handled message IDs may solve this problem in many cases.
-
 struct ConcurrentMessageHandler
     queue::Channel{Any}
 end
@@ -226,12 +222,20 @@ end
 function (dispatcher::ConcurrentMessageHandler)(server::Server, @nospecialize msg)
     # Handle `currently_handled` processing serially within the concurrent message worker thread
     if msg isa CancelRequestNotification
+        if msg.params.id in server.state.handled_history
+            return # Request was already handled, ignore cancellation
+        end
         cancel!(get!(()->CancelFlag(true), server.state.currently_handled, msg.params.id))
     elseif msg isa WorkDoneProgressCancelNotification
+        if msg.params.token in server.state.handled_history
+            return # Token was already handled, ignore cancellation
+        end
         cancel!(get!(()->CancelFlag(true), server.state.currently_handled, msg.params.token))
     elseif msg isa HandledToken
         delete!(server.state.currently_handled, msg.id)
+        push!(server.state.handled_history, msg.id) # Add to handled history to prevent dead IDs from accumulating
         # @info "Remaining requests" length(server.state.currently_handled) Base.summarysize(server.state.currently_handled)
+        # @info "Handled history" length(server.state.handled_history) Base.summarysize(server.state.handled_history)
     # Handle regular messages concurrently
     elseif msg isa Dict{Symbol,Any} # ResponseMessage or untyped message
         request_caller = let id = get(msg, :id, nothing)

src/types.jl

@@ -371,13 +371,16 @@ const CompletionResolverInfo = CASContainer{Union{Nothing,Tuple{Module,LSPostPro
 # Type aliases for concurrent updates using LWContainer (non-retriable operations)
 const ConfigManager = LWContainer{ConfigManagerData, LWStats}
 
+const HandledHistory = FixedSizeFIFOQueue{Union{Int,String}}
+
 mutable struct ServerState
     const file_cache::FileCache # syntactic analysis cache (synced with `textDocument/didChange`)
     const saved_file_cache::SavedFileCache # syntactic analysis cache (synced with `textDocument/didSave`)
     const testsetinfos_cache::TestsetInfosCache
     const analysis_manager::AnalysisManager
     const extra_diagnostics::ExtraDiagnostics
     const currently_handled::CurrentlyHandled
+    const handled_history::HandledHistory
     const currently_requested::CurrentlyRequested
     const currently_registered::CurrentlyRegistered
     const config_manager::ConfigManager
@@ -397,6 +400,7 @@ mutable struct ServerState
             #=analysis_manager=# AnalysisManager(#=n_workers=# 1), # TODO multiple workers
             #=extra_diagnostics=# ExtraDiagnostics(ExtraDiagnosticsData()),
             #=currently_handled=# CurrentlyHandled(),
+            #=handled_history=# HandledHistory(128),
             #=currently_requested=# CurrentlyRequested(Base.PersistentDict{String,RequestCaller}()),
             #=currently_registered=# CurrentlyRegistered(Set{Registered}()),
             #=config_manager=# ConfigManager(ConfigManagerData()),

test/FixedSizeFIFOQueue/test_FixedSizeFIFOQueue.jl

@@ -0,0 +1,161 @@
+module test_FixedSizeFIFOQueue
+
+using Test
+using JETLS: FixedSizeFIFOQueue, capacity, isfull
+
+struct TestType
+    x::Int
+end
+
+mutable struct MutableTestType
+    x::Int
+    data::Vector{Int}
+end
+
+@testset "FixedSizeFIFOQueue basic operations" begin
+    let q = FixedSizeFIFOQueue{Int}(3)
+        @test isempty(q)
+        @test !isfull(q)
+        @test length(q) == 0
+        @test capacity(q) == 3
+
+        push!(q, 1)
+        @test !isempty(q)
+        @test length(q) == 1
+        @test 1 in q
+        @test !(2 in q)
+
+        push!(q, 2)
+        push!(q, 3)
+        @test isfull(q)
+        @test length(q) == 3
+        @test 1 in q && 2 in q && 3 in q
+    end
+
+    let q = FixedSizeFIFOQueue{Int}(3)
+        push!(q, 1)
+        push!(q, 2)
+        push!(q, 3)
+        push!(q, 4)  # Should overwrite 1
+        @test isfull(q)
+        @test length(q) == 3
+        @test !(1 in q)
+        @test 2 in q && 3 in q && 4 in q
+
+        push!(q, 5)  # Should overwrite 2
+        @test !(2 in q)
+        @test 3 in q && 4 in q && 5 in q
+    end
+end
+
+@testset "FixedSizeFIFOQueue type flexibility" begin
+    let q = FixedSizeFIFOQueue(2)
+        push!(q, "hello")
+        push!(q, 42)
+        @test "hello" in q
+        @test 42 in q
+
+        push!(q, :symbol)
+        @test !("hello" in q)
+        @test 42 in q
+        @test :symbol in q
+    end
+
+    let q = FixedSizeFIFOQueue{TestType}(2)
+        t1 = TestType(1)
+        t2 = TestType(2)
+        t3 = TestType(3)
+
+        push!(q, t1)
+        push!(q, t2)
+        @test t1 in q
+        @test t2 in q
+
+        push!(q, t3)
+        @test !(t1 in q)
+        @test t2 in q
+        @test t3 in q
+    end
+end
+
+@testset "FixedSizeFIFOQueue edge cases" begin
+    let q = FixedSizeFIFOQueue{String}(1)
+        push!(q, "a")
+        @test "a" in q
+        push!(q, "b")
+        @test !("a" in q)
+        @test "b" in q
+    end
+
+    @test_throws ArgumentError FixedSizeFIFOQueue{Int}(0)
+    @test_throws ArgumentError FixedSizeFIFOQueue{Int}(-1)
+end
+
+@testset "FixedSizeFIFOQueue collect and display" begin
+    let q = FixedSizeFIFOQueue{Int}(3)
+        push!(q, 1)
+        push!(q, 2)
+        push!(q, 3)
+
+        items = collect(q)
+        @test items == [1, 2, 3]
+
+        push!(q, 4)  # Overwrites 1
+        items = collect(q)
+        @test items == [2, 3, 4]
+
+        str = string(q)
+        @test occursin("FixedSizeFIFOQueue{Int", str)
+        @test occursin("capacity=3", str)
+        @test occursin("[2, 3, 4]", str)
+    end
+end
+
+@testset "FixedSizeFIFOQueue large capacity" begin
+    let q = FixedSizeFIFOQueue{Union{Int,String}}(1000)
+        for i in 1:500
+            push!(q, i)
+        end
+        for i in 501:1000
+            push!(q, "id-$i")
+        end
+
+        @test isfull(q)
+        @test 1 in q
+        @test 500 in q
+        @test "id-501" in q
+        @test "id-1000" in q
+
+        push!(q, "extra-1")
+        @test !(1 in q)  # First element should be evicted
+        @test 2 in q      # Second element should still be there
+        @test "extra-1" in q
+
+        for i in 1:100
+            push!(q, "new-$i")
+        end
+
+        @test !(2 in q)
+        @test !(100 in q)
+        @test "new-100" in q
+    end
+end
+
+@testset "FixedSizeFIFOQueue garbage collection" begin
+    let q = FixedSizeFIFOQueue{MutableTestType}(2)
+        obj1 = MutableTestType(1, [1,2,3])
+        obj2 = MutableTestType(2, [4,5,6])
+        obj3 = MutableTestType(3, [7,8,9])
+
+        push!(q, obj1)
+        push!(q, obj2)
+        push!(q, obj3)
+
+        @test obj2 in q
+        @test obj3 in q
+        @test !(obj1 in q.data)
+        @test !(obj1 in q.items)
+    end
+end
+
+end # module test_FixedSizeFIFOQueue

test/runtests.jl

@@ -7,6 +7,7 @@ include("setup.jl")
 end
 
 @testset "AtomicContainers" include("AtomicContainers/test_AtomicContainers.jl")
+@testset "FixedSizeFIFOQueue" include("FixedSizeFIFOQueue/test_FixedSizeFIFOQueue.jl")
 
 @testset "JETLS" begin
     @testset "utils" begin