Minimize allocations when unpacking TimeZones from cache (updated)

src/TimeZones.jl

@@ -53,9 +53,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,32 +1,29 @@
-# Retains the compiled tzdata in memory. Read-only access is thread-safe and any changes
-# to this structure can result in inconsistent behaviour.
-# Do not access this object directly, instead use `get_tz_cache()` to access the cache.
-const _TZ_CACHE = Dict{String,Tuple{TimeZone,Class}}()
-const _TZ_CACHE_LOCK = ReentrantLock()
-const _TZ_CACHE_INITIALIZED = Threads.Atomic{Bool}(false)
-
-function _init_tz_cache()
-    # Write out our compiled tzdata representations into a scratchspace
-    desired_version = TZData.tzdata_version()
+# 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.
+struct TimeZoneCache
+    ftz::Dict{String,Tuple{FixedTimeZone,Class}}
+    vtz::Dict{String,Tuple{VariableTimeZone,Class}}
+    lock::ReentrantLock
+    initialized::Threads.Atomic{Bool}
+end
 
-    _COMPILED_DIR[] = if desired_version == TZJData.TZDATA_VERSION
-        TZJData.ARTIFACT_DIR
-    else
-        TZData.build(desired_version, _scratch_dir())
-    end
+TimeZoneCache() = TimeZoneCache(Dict(), Dict(), ReentrantLock(), Threads.Atomic{Bool}(false))
 
-    # Load the pre-computed TZData into memory.
-    return _reload_tz_cache(_COMPILED_DIR[])
-end
+# Retains the compiled tzdata in memory. Read-only access to the cache is thread-safe and
+# any changes to this structure can result in inconsistent behaviour. Do not access this
+# object directly, instead use `get` to access the cache content.
+const _TZ_CACHE = TimeZoneCache()
 
-function _reload_tz_cache(compiled_dir::AbstractString)
-    _reload_tz_cache!(_TZ_CACHE, compiled_dir)
-    !isempty(_TZ_CACHE) || error("Cache remains empty after loading")
-    return _TZ_CACHE
+function Base.copy!(dst::TimeZoneCache, src::TimeZoneCache)
+    copy!(dst.ftz, src.ftz)
+    copy!(dst.vtz, src.vtz)
+    dst.initialized[] = src.initialized[]
+    return dst
 end
 
-function _reload_tz_cache!(cache::AbstractDict, compiled_dir::AbstractString)
-    empty!(cache)
+function reload!(cache::TimeZoneCache, compiled_dir::AbstractString=_COMPILED_DIR[])
+    empty!(cache.ftz)
+    empty!(cache.vtz)
     check = Tuple{String,String}[(compiled_dir, "")]
 
     for (dir, partial) in check
@@ -39,26 +36,57 @@ function _reload_tz_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 tz isa FixedTimeZone
+                    cache.ftz[name] = (tz, class)
+                elseif tz isa VariableTimeZone
+                    cache.vtz[name] = (tz, class)
+                else
+                    error("Unhandled TimeZone class encountered: $(typeof(tz))")
+                end
             end
         end
     end
 
+    !isempty(cache.ftz) && !isempty(cache.vtz) || error("Cache remains empty after loading")
+
     return cache
 end
 
-function _get_tz_cache()
-    if !_TZ_CACHE_INITIALIZED[]
-        lock(_TZ_CACHE_LOCK) do
-            if !_TZ_CACHE_INITIALIZED[]
-                _init_tz_cache()
-                _TZ_CACHE_INITIALIZED[] = true
+function Base.get(body::Function, cache::TimeZoneCache, name::AbstractString)
+    if !cache.initialized[]
+        lock(cache.lock) do
+            if !cache.initialized[]
+                _initialize()
+                reload!(cache)
+                cache.initialized[] = true
             end
         end
     end
-    return _TZ_CACHE
+
+    return get(cache.ftz, name) do
+        get(cache.vtz, name) do
+            body()
+        end
+    end
 end
 
+function _initialize()
+    # Write out our compiled tzdata representations into a scratchspace
+    desired_version = TZData.tzdata_version()
+
+    _COMPILED_DIR[] = if desired_version == TZJData.TZDATA_VERSION
+        TZJData.ARTIFACT_DIR
+    else
+        TZData.build(desired_version, _scratch_dir())
+    end
+
+    return nothing
+end
+
+_reload_tz_cache(compiled_dir::AbstractString) = reload!(_TZ_CACHE, compiled_dir)
+
 """
     TimeZone(str::AbstractString) -> TimeZone
 
@@ -100,7 +128,7 @@ US/Pacific (UTC-8/UTC-7)
 TimeZone(::AbstractString, ::Class)
 
 function TimeZone(str::AbstractString, mask::Class=Class(:DEFAULT))
-    tz, class = get(_get_tz_cache(), str) do
+    tz, class = get(_TZ_CACHE, str) do
         if occursin(FIXED_TIME_ZONE_REGEX, str)
             FixedTimeZone(str), Class(:FIXED)
         else
@@ -145,6 +173,6 @@ function istimezone(str::AbstractString, mask::Class=Class(:DEFAULT))
     end
 
     # Checks against pre-compiled time zones
-    class = get(_get_tz_cache(), str, (UTC_ZERO, Class(:NONE)))[2]
+    class = get(() -> (UTC_ZERO, Class(:NONE)), _TZ_CACHE, str)[2]
     return mask & class != 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
@@ -35,26 +47,31 @@ show_compact = (io, args...) -> show(IOContext(io, :compact => true), args...)
 # Modified the internal TimeZones cache. Should only be used as part of testing and only is
 # needed when the data between the test tzdata version and the built tzdata versions differ.
 
-function add!(cache::Dict, t::Tuple{TimeZone,TimeZones.Class})
+function add!(dict::Dict, t::Tuple{TimeZone,TimeZones.Class})
     tz, class = t
     name = TimeZones.name(tz)
-    push!(cache, name => t)
+    push!(dict, name => t)
     return tz
 end
 
-function add!(cache::Dict, tz::VariableTimeZone)
+function add!(cache::TimeZones.TimeZoneCache, t::Tuple{T,TimeZones.Class}) where {T<:TimeZone}
+    dict = T == FixedTimeZone ? cache.ftz : cache.vtz
+    return add!(dict, t)
+end
+
+function add!(cache::TimeZones.TimeZoneCache, tz::VariableTimeZone)
     # Not all `VariableTimeZone`s are the STANDARD class. However, for testing purposes
     # the class doesn't need to be precise.
     class = TimeZones.Class(:STANDARD)
-    return add!(cache, (tz, class))
+    return add!(cache.vtz, (tz, class))
 end
 
-function add!(cache::Dict, tz::FixedTimeZone)
+function add!(cache::TimeZones.TimeZoneCache, tz::FixedTimeZone)
     class = TimeZones.Class(:FIXED)
-    return add!(cache, (tz, class))
+    return add!(cache.ftz, (tz, class))
 end
 
-function with_tz_cache(f, cache::Dict{String,Tuple{TimeZone,TimeZones.Class}})
+function with_tz_cache(f, cache::TimeZones.TimeZoneCache)
     old_cache = deepcopy(TimeZones._TZ_CACHE)
     copy!(TimeZones._TZ_CACHE, cache)
 

test/io.jl

@@ -1,7 +1,8 @@
 using TimeZones.TZData: parse_components
 using TimeZones: Transition
 
-cache = Dict{String,Tuple{TimeZone,TimeZones.Class}}()
+cache = TimeZones.TimeZoneCache()
+cache.initialized[] = true
 
 dt = DateTime(1942,12,25,1,23,45)
 custom_dt = DateTime(1800,1,1)

test/types/timezone.jl

@@ -25,3 +25,15 @@ end
     @test TimeZone("Etc/GMT+12", Class(:LEGACY)) == FixedTimeZone("Etc/GMT+12", -12 * 3600)
     @test TimeZone("Etc/GMT-14", Class(:LEGACY)) == FixedTimeZone("Etc/GMT-14", 14 * 3600)
 end
+
+@testset "allocations" begin
+    tz = TimeZone("UTC")  # Trigger compilation and ensure the cache is populated
+    @test tz isa FixedTimeZone
+    @test @allocations(TimeZone("UTC")) == 0
+    @test @allocations(istimezone("UTC")) == 0
+
+    tz = TimeZone("America/Winnipeg")  # Trigger compilation and ensure the cache is populated
+    @test tz isa VariableTimeZone
+    @test @allocations(TimeZone("America/Winnipeg")) == 2
+    @test @allocations(istimezone("America/Winnipeg")) == 1
+end