Minimize allocations when unpacking TimeZones from cache (updated)

src/TimeZones.jl

@@ -70,9 +70,9 @@ include("indexable_generator.jl")
 
 include("class.jl")
 include("utcoffset.jl")
-include(joinpath("types", "timezone.jl"))
 include(joinpath("types", "fixedtimezone.jl"))
 include(joinpath("types", "variabletimezone.jl"))
+include(joinpath("types", "timezone.jl"))
 include(joinpath("types", "zoneddatetime.jl"))
 include(joinpath("tzfile", "TZFile.jl"))
 include(joinpath("tzjfile", "TZJFile.jl"))

src/types/timezone.jl

@@ -1,15 +1,20 @@
 # Retains the compiled tzdata in memory. Read-only access is thread-safe and any changes
 # to this structure can result in inconsistent behaviour.
-const _TZ_CACHE = Dict{String,Tuple{TimeZone,Class}}()
+#
+# Use a separate cache for FixedTimeZone (which is `isbits`) so the container is concretely
+# typed and we avoid allocating a FixedTimeZone every time we get one from the cache.
+const _FTZ_CACHE = Dict{String,Tuple{FixedTimeZone,Class}}()
+const _VTZ_CACHE = Dict{String,Tuple{VariableTimeZone,Class}}()
 
 function _reload_cache(compiled_dir::AbstractString)
- _reload_cache!(_TZ_CACHE, compiled_dir)
- !isempty(_TZ_CACHE) || error("Cache remains empty after loading")
+ _reload_cache!(_FTZ_CACHE, _VTZ_CACHE, compiled_dir)
+ !isempty(_FTZ_CACHE) && !isempty(_VTZ_CACHE) || error("Cache remains empty after loading")
  return nothing
 end
 
-function _reload_cache!(cache::AbstractDict, compiled_dir::AbstractString)
- empty!(cache)
+function _reload_cache!(ftz_cache::AbstractDict, vtz_cache::AbstractDict, compiled_dir::AbstractString)
+ empty!(ftz_cache)
+ empty!(vtz_cache)
  check = Tuple{String,String}[(compiled_dir, "")]
 
  for (dir, partial) in check
@@ -22,12 +27,26 @@ function _reload_cache!(cache::AbstractDict, compiled_dir::AbstractString)
  if isdir(path)
  push!(check, (path, name))
  else
- cache[name] = open(TZJFile.read, path, "r")(name)
+ tz, class = open(TZJFile.read, path, "r")(name)
+ if isa(tz, FixedTimeZone)
+ ftz_cache[name] = (tz, class)
+ elseif isa(tz, VariableTimeZone)
+ vtz_cache[name] = (tz, class)
+ else
+ error("Unhandled TimeZone class encountered: $(typeof(tz))")
+ end
  end
  end
  end
+ return nothing
+end
 
- return cache
+function _get_from_cache(f_default::Function, name::AbstractString)
+ return get(_FTZ_CACHE, name) do
+ get(_VTZ_CACHE, name) do
+ return f_default()
+ end
+ end
 end
 
 """
@@ -71,7 +90,7 @@ US/Pacific (UTC-8/UTC-7)
 TimeZone(::AbstractString, ::Class)
 
 function TimeZone(str::AbstractString, mask::Class=Class(:DEFAULT))
- tz, class = get(_TZ_CACHE, str) do
+ tz, class = _get_from_cache(str) do
  if occursin(FIXED_TIME_ZONE_REGEX, str)
  FixedTimeZone(str), Class(:FIXED)
  else
@@ -116,7 +135,7 @@ function istimezone(str::AbstractString, mask::Class=Class(:DEFAULT))
  end
 
  # Checks against pre-compiled time zones
- tz, class = get(_TZ_CACHE, str) do
+ tz, class = _get_from_cache(str) do
  nothing, Class(:NONE)
  end
 

test/helpers.jl

@@ -1,5 +1,17 @@
 # Utility functions for testing
 
+if VERSION < v"1.9.0-DEV.1744" # https://github.com/JuliaLang/julia/pull/47367
+ macro allocations(ex)
+ quote
+ while false; end # want to force compilation, but v1.6 doesn't have `@force_compile`
+ local stats = Base.gc_num()
+ $(esc(ex))
+ local diff = Base.GC_Diff(Base.gc_num(), stats)
+ Base.gc_alloc_count(diff)
+ end
+ end
+end
+
 function ignore_output(body::Function; stdout::Bool=true, stderr::Bool=true)
  out_old = Base.stdout
  err_old = Base.stderr
@@ -55,12 +67,23 @@ function add!(cache::Dict, tz::FixedTimeZone)
 end
 
 function with_tz_cache(f, cache::Dict{String,Tuple{TimeZone,TimeZones.Class}})
- old_cache = deepcopy(TimeZones._TZ_CACHE)
- copy!(TimeZones._TZ_CACHE, cache)
+ old_ftz_cache = deepcopy(TimeZones._FTZ_CACHE)
+ old_vtz_cache = deepcopy(TimeZones._VTZ_CACHE)
+
+ # Split the contents of `cache` between the fixed and variable caches
+ # as appropriate.
+ empty!(TimeZones._FTZ_CACHE)
+ empty!(TimeZones._VTZ_CACHE)
+ foreach(cache) do (k, v)
+ tz = first(v)
+ cache = tz isa FixedTimeZone ? TimeZones._FTZ_CACHE : TimeZones._VTZ_CACHE
+ setindex!(cache, v, k)
+ end
 
  try
  return f()
  finally
- copy!(TimeZones._TZ_CACHE, old_cache)
+ copy!(TimeZones._FTZ_CACHE, old_ftz_cache)
+ copy!(TimeZones._VTZ_CACHE, old_vtz_cache)
  end
 end

test/types/timezone.jl

@@ -1,5 +1,15 @@
 using TimeZones: Class
 
+@testset "TimeZone allocations" begin
+ tz = TimeZone("UTC") # run once for compilation and to populate cache
+ @assert tz isa FixedTimeZone
+ @test @allocations(TimeZone("UTC")) == 0
+
+ tz = TimeZone("America/Winnipeg") # populate cache
+ @assert tz isa VariableTimeZone
+ @test @allocations(TimeZone("America/Winnipeg")) == 2
+end
+
 @testset "istimezone" begin
  @test istimezone("Europe/Warsaw")
  @test istimezone("UTC+02")