Skip to content

Commit 2dae25b

Browse files
authored
Run the GC when failing to find a handle, but lots are active. (#2583)
This should help use cases where lots of short-lived tasks are used, without the penalty of hitting the GC for every lookup.
1 parent f22c9b4 commit 2dae25b

File tree

1 file changed

+15
-7
lines changed

1 file changed

+15
-7
lines changed

lib/utils/cache.jl

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,17 +28,25 @@ end
2828
# remove a handle from the cache, or create a new one
2929
function Base.pop!(cache::HandleCache{K,V}, key::K) where {K,V}
3030
# check the cache
31-
handle = @lock cache.lock begin
32-
if !haskey(cache.idle_handles, key) || isempty(cache.idle_handles[key])
33-
nothing
31+
handle, num_active_handles = @lock cache.lock begin
32+
if haskey(cache.idle_handles, key) && !isempty(cache.idle_handles[key])
33+
pop!(cache.idle_handles[key]), length(cache.active_handles)
3434
else
35-
pop!(cache.idle_handles[key])
35+
nothing, length(cache.active_handles)
3636
end
3737
end
3838

39-
# if we didn't find anything, create a new handle.
40-
# we could (and used to) run `GC.gc(false)` here to free up old handles,
41-
# but that can be expensive when using lots of short-lived tasks.
39+
# if we didn't find anything, but lots of handles are active, try to free some
40+
if handle === nothing && num_active_handles > cache.max_entries
41+
GC.gc(false)
42+
@lock cache.lock begin
43+
if haskey(cache.idle_handles, key) && isempty(cache.idle_handles[key])
44+
handle = pop!(cache.idle_handles[key])
45+
end
46+
end
47+
end
48+
49+
# if we still didn't find anything, create a new handle
4250
if handle === nothing
4351
CUDA.maybe_collect()
4452
handle = cache.ctor(key)

0 commit comments

Comments
 (0)