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