Minimize allocations when unpacking TimeZones from cache

src/TimeZones.jl

@@ -66,10 +66,10 @@ 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", "zoneddatetime.jl"))
+include(joinpath("types", "timezone.jl"))
 include(joinpath("tzfile", "TZFile.jl"))
 include(joinpath("tzjfile", "TZJFile.jl"))
 include("exceptions.jl")

src/types/timezone.jl

@@ -1,27 +1,43 @@
-# Thread-local TimeZone cache, which caches time zones _per thread_, allowing thread-safe
+# Thread-local TimeZone caches, which caches time zones _per thread_, allowing thread-safe
 # caching. Note that this means the cache will grow in size, and may store redundant objects
 # accross multiple threads, but this extra space usage allows for fast, lock-free access
 # to the cache, while still being thread-safe.
-const THREAD_TZ_CACHES = Vector{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 THREAD_FTZ_CACHES = Vector{Dict{String,Tuple{FixedTimeZone,Class}}}()
+const THREAD_VTZ_CACHES = Vector{Dict{String,Tuple{VariableTimeZone,Class}}}()
 
 # Based upon the thread-safe Global RNG implementation in the Random stdlib:
 # https://github.com/JuliaLang/julia/blob/e4fcdf5b04fd9751ce48b0afc700330475b42443/stdlib/Random/src/RNGs.jl#L369-L385
-@inline _tz_cache() = _tz_cache(Threads.threadid())
-@noinline function _tz_cache(tid::Int)
-    0 < tid <= length(THREAD_TZ_CACHES) || _tz_cache_length_assert()
-    if @inbounds isassigned(THREAD_TZ_CACHES, tid)
-        @inbounds cache = THREAD_TZ_CACHES[tid]
+@inline _ftz_cache() = _ftz_cache(Threads.threadid())
+@inline _vtz_cache() = _vtz_cache(Threads.threadid())
+@noinline function _ftz_cache(tid::Int)
+    0 < tid <= length(THREAD_FTZ_CACHES) || _ftz_cache_length_assert()
+    if @inbounds isassigned(THREAD_FTZ_CACHES, tid)
+        @inbounds cache = THREAD_FTZ_CACHES[tid]
     else
-        cache = eltype(THREAD_TZ_CACHES)()
-        @inbounds THREAD_TZ_CACHES[tid] = cache
+        cache = eltype(THREAD_FTZ_CACHES)()
+        @inbounds THREAD_FTZ_CACHES[tid] = cache
     end
     return cache
 end
-@noinline _tz_cache_length_assert() = @assert false "0 < tid <= length(THREAD_TZ_CACHES)"
+@noinline function _vtz_cache(tid::Int)
+    0 < tid <= length(THREAD_VTZ_CACHES) || _vtz_cache_length_assert()
+    if @inbounds isassigned(THREAD_VTZ_CACHES, tid)
+        @inbounds cache = THREAD_VTZ_CACHES[tid]
+    else
+        cache = eltype(THREAD_VTZ_CACHES)()
+        @inbounds THREAD_VTZ_CACHES[tid] = cache
+    end
+    return cache
+end
+@noinline _ftz_cache_length_assert() = @assert false "0 < tid <= length(THREAD_FTZ_CACHES)"
+@noinline _vtz_cache_length_assert() = @assert false "0 < tid <= length(THREAD_VTZ_CACHES)"
 
 function _reset_tz_cache()
     # ensures that we didn't save a bad object
-    resize!(empty!(THREAD_TZ_CACHES), Threads.nthreads())
+    resize!(empty!(THREAD_FTZ_CACHES), Threads.nthreads())
+    resize!(empty!(THREAD_VTZ_CACHES), Threads.nthreads())
 end
 
 """
@@ -67,32 +83,56 @@ TimeZone(::AbstractString, ::Class)
 function TimeZone(str::AbstractString, mask::Class=Class(:DEFAULT))
     # Note: If the class `mask` does not match the time zone we'll still load the
     # information into the cache to ensure the result is consistent.
-    tz, class = get!(_tz_cache(), str) do
-        tz_path = joinpath(_COMPILED_DIR[], split(str, "/")...)
+    ftz, class = get(_ftz_cache(), str, (nothing, Class(:NONE)))
+    if ftz !== nothing
+        _check_class(mask, class, str)
+        return ftz::FixedTimeZone
+    end
 
-        if isfile(tz_path)
-            open(TZJFile.read, tz_path, "r")(str)
-        elseif occursin(FIXED_TIME_ZONE_REGEX, str)
-            FixedTimeZone(str), Class(:FIXED)
-        elseif !isdir(_COMPILED_DIR[]) || isempty(readdir(_COMPILED_DIR[]))
-            # Note: Julia 1.0 supresses the build logs which can hide issues in time zone
-            # compliation which result in no tzdata time zones being available.
-            throw(ArgumentError(
-                "Unable to find time zone \"$str\". Try running `TimeZones.build()`."
-            ))
-        else
-            throw(ArgumentError("Unknown time zone \"$str\""))
+    vtz, class = get(_vtz_cache(), str, (nothing, Class(:NONE)))
+    if vtz !== nothing
+        _check_class(mask, class, str)
+        return vtz::VariableTimeZone
+    end
+
+    # We need to compute the timezone
+    tz_path = joinpath(_COMPILED_DIR[], split(str, "/")...)
+    if isfile(tz_path)
+        tz, class = open(TZJFile.read, tz_path, "r")(str)::Tuple{TimeZone,Class}
+        if tz isa FixedTimeZone
+            _ftz_cache()[str] = (tz, class)
+            _check_class(mask, class, str)
+            return tz::FixedTimeZone
+        elseif tz isa VariableTimeZone
+            _vtz_cache()[str] = (tz, class)
+            _check_class(mask, class, str)
+            return tz::VariableTimeZone
         end
+    elseif occursin(FIXED_TIME_ZONE_REGEX, str)
+        ftz = FixedTimeZone(str)
+        class = Class(:FIXED)
+        _ftz_cache()[str] = (ftz, class)
+        _check_class(mask, class, str)
+        return ftz
+    elseif !isdir(_COMPILED_DIR[]) || isempty(readdir(_COMPILED_DIR[]))
+        # Note: Julia 1.0 supresses the build logs which can hide issues in time zone
+        # compliation which result in no tzdata time zones being available.
+        throw(ArgumentError(
+            "Unable to find time zone \"$str\". Try running `TimeZones.build()`."
+        ))
+    else
+        throw(ArgumentError("Unknown time zone \"$str\""))
     end
+end
 
+function _check_class(mask::Class, class::Class, str)
     if mask & class == Class(:NONE)
         throw(ArgumentError(
             "The time zone \"$str\" is of class `$(repr(class))` which is " *
             "currently not allowed by the mask: `$(repr(mask))`"
         ))
     end
-
-    return tz
+    return nothing
 end
 
 """
@@ -110,6 +150,13 @@ macro tz_str(str)
     TimeZone(str)
 end
 
+function _get_from_cache(str)
+    (tz, class) = get(_ftz_cache(), str, (nothing, Class(:NONE)))
+    tz !== nothing && return (tz, class)
+    (tz, class) = get(_vtz_cache(), str, (nothing, Class(:NONE)))
+    return (tz, class)
+end
+
 """
     istimezone(str::AbstractString, mask::Class=Class(:DEFAULT)) -> Bool
 
@@ -122,16 +169,18 @@ function istimezone(str::AbstractString, mask::Class=Class(:DEFAULT))
     end
 
     # Perform more expensive checks against pre-compiled time zones
-    tz, class = get(_tz_cache(), str) do
+    tz, class = _get_from_cache(str)
+    if tz === nothing
         tz_path = joinpath(_COMPILED_DIR[], split(str, "/")...)
-
         if isfile(tz_path)
             # Cache the data since we're already performing the deserialization
-            _tz_cache()[str] = open(TZJFile.read, tz_path, "r")(str)
-        else
-            nothing, Class(:NONE)
+            tz, class = open(TZJFile.read, tz_path, "r")(str)
+            if tz isa FixedTimeZone
+                _ftz_cache()[str] = (tz, class)
+            elseif tz isa VariableTimeZone
+                _vtz_cache()[str] = (tz, class)
+            end
         end
     end
-
     return tz !== nothing && mask & class != Class(:NONE)
 end

src/tzdata/compile.jl

@@ -1,7 +1,7 @@
 using Dates
 using Dates: parse_components
 
-using ...TimeZones: _tz_cache, _tz_source_dir, _compiled_dir
+using ...TimeZones: _ftz_cache, _vtz_cache, _tz_source_dir, _compiled_dir
 using ...TimeZones: TimeZones, TimeZone, FixedTimeZone, VariableTimeZone, Transition, Class
 using ...TimeZones: rename
 using ..TZData: TimeOffset, ZERO, MIN_GMT_OFFSET, MAX_GMT_OFFSET, MIN_SAVE, MAX_SAVE,
@@ -699,7 +699,8 @@ function compile(tz_source::TZSource, dest_dir::AbstractString; kwargs...)
     # thread's local cache.
     Threads.@threads :static for i in 1:Threads.nthreads()
         @assert Threads.threadid() === i "TimeZones.TZData.compile() must be called from the main, top-level Task."
-        empty!(_tz_cache())
+        empty!(_ftz_cache())
+        empty!(_vtz_cache())
     end
 
     for (tz, class) in results

test/helpers.jl

@@ -1,5 +1,17 @@
 # Utility functions for testing
 
+if VERSION < v"1.9.0-"  # 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
@@ -36,6 +48,11 @@ show_compact = (io, args...) -> show(IOContext(io, :compact => true), args...)
 # not be used and only should be required if the test tzdata version and built tzdata
 # version do not match.
 function cache_tz((tz, class)::Tuple{TimeZone, TimeZones.Class})
-    TimeZones._tz_cache()[TimeZones.name(tz)] = (tz, class)
+    tz_cache = if tz isa FixedTimeZone
+        TimeZones._ftz_cache()
+    elseif tz isa VariableTimeZone
+        TimeZones._vtz_cache()
+    end
+    tz_cache[TimeZones.name(tz)] = (tz, class)
     return tz
 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 0 == @allocations(TimeZone("UTC"))
+
+    tz = TimeZone("America/Winnipeg")  # populate cache
+    @assert tz isa VariableTimeZone
+    @test 1 == @allocations(TimeZone("America/Winnipeg"))
+end
+
 @testset "istimezone" begin
     # Invalidate the cache to ensure that `istimezone` works for non-loaded time zones.
     TimeZones._reset_tz_cache()