Make Time Zone Cache thread-safe with lazily-allocated thread-local caches (#344)

Co-authored-by: Curtis Vogt <[email protected]>

Author: NHDaly

Diff for: src/TimeZones.jl

@@ -38,6 +38,9 @@ const DEPS_DIR = joinpath(PKG_DIR, "deps")
 abstract type Local <: TimeZone end
 
 function __init__()
+    # Initialize the thread-local TimeZone cache (issue #342)
+    _reset_tz_cache()
+
     # Base extension needs to happen everytime the module is loaded (issue #24)
     Dates.CONVERSION_SPECIFIERS['z'] = TimeZone
     Dates.CONVERSION_SPECIFIERS['Z'] = TimeZone

Diff for: src/build.jl

@@ -15,7 +15,7 @@ function build(version::AbstractString=tzdata_version(); force::Bool=false)
     end
 
     # Reset cached information
-    empty!(TIME_ZONE_CACHE)
+    _reset_tz_cache()
 
     @info "Successfully built TimeZones"
 end

Diff for: src/types/timezone.jl

@@ -1,4 +1,28 @@
-const TIME_ZONE_CACHE = Dict{String,Tuple{TimeZone,Class}}()
+# Thread-local TimeZone cache, 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}}}()
+
+# 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]
+    else
+        cache = eltype(THREAD_TZ_CACHES)()
+        @inbounds THREAD_TZ_CACHES[tid] = cache
+    end
+    return cache
+end
+@noinline _tz_cache_length_assert() = @assert false "0 < tid <= length(THREAD_TZ_CACHES)"
+
+function _reset_tz_cache()
+    # ensures that we didn't save a bad object
+    resize!(empty!(THREAD_TZ_CACHES), Threads.nthreads())
+end
 
 """
     TimeZone(str::AbstractString) -> TimeZone
@@ -43,7 +67,7 @@ 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!(TIME_ZONE_CACHE, str) do
+    tz, class = get!(_tz_cache(), str) do
         tz_path = joinpath(TZData.COMPILED_DIR, split(str, "/")...)
 
         if isfile(tz_path)
@@ -98,12 +122,12 @@ function istimezone(str::AbstractString, mask::Class=Class(:DEFAULT))
     end
 
     # Perform more expensive checks against pre-compiled time zones
-    tz, class = get(TIME_ZONE_CACHE, str) do
+    tz, class = get(_tz_cache(), str) do
         tz_path = joinpath(TZData.COMPILED_DIR, split(str, "/")...)
 
         if isfile(tz_path)
             # Cache the data since we're already performing the deserialization
-            TIME_ZONE_CACHE[str] = open(deserialize, tz_path, "r")
+            _tz_cache()[str] = open(deserialize, tz_path, "r")
         else
             nothing, Class(:NONE)
         end

Diff for: src/tzdata/compile.jl

@@ -2,7 +2,7 @@ using Dates
 using Serialization
 using Dates: parse_components
 
-using ...TimeZones: TIME_ZONE_CACHE
+using ...TimeZones: _tz_cache
 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,
@@ -694,7 +694,14 @@ function compile(tz_source::TZSource, dest_dir::AbstractString; kwargs...)
     results = compile(tz_source; kwargs...)
 
     isdir(dest_dir) || error("Destination directory doesn't exist")
-    empty!(TIME_ZONE_CACHE)
+    # When we recompile the TimeZones from a new source, we clear all the existing cached
+    # TimeZone objects, so that newly constructed objects pick up the newly compiled rules.
+    # Since we use thread-local caches, we spawn a task on _each thread_ to clear that
+    # thread's local cache.
+    Threads.@threads 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())
+    end
 
     for (tz, class) in results
         parts = split(TimeZones.name(tz), '/')

Diff for: test/helpers.jl

@@ -36,6 +36,6 @@ 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.TIME_ZONE_CACHE[TimeZones.name(tz)] = (tz, class)
+    TimeZones._tz_cache()[TimeZones.name(tz)] = (tz, class)
     return tz
 end

Diff for: test/runtests.jl

@@ -72,6 +72,7 @@ include("helpers.jl")
     include("rounding.jl")
     include("parse.jl")
     include("plotting.jl")
+    VERSION >= v"1.3" && include("thread-safety.jl")
 
     # Note: Run the build tests last to ensure that re-compiling the time zones files
     # doesn't interfere with other tests.

Diff for: test/thread-safety.jl

@@ -0,0 +1,67 @@
+# Test that TimeZones.jl can be safely used in a multithreaded environment.
+# Note that the number of threads being used cannot be changed dynamically, so
+# this test file spawns a new julia process running with multiple threads.
+
+using Test
+
+const program = """
+using TimeZones
+using Test
+
+@assert Threads.nthreads() > 1 "This system does not support multiple threads, so the thread-safety tests cannot be run."
+
+@testset "Multithreaded TimeZone brute force test" begin
+    function create_zdt(year, month, day, tz_name)
+        ZonedDateTime(DateTime(year, month, day), TimeZone(tz_name))
+    end
+    function cycle_zdts()
+        return [
+            try
+                create_zdt(year, month, day, tz_name)
+            catch e
+                # Ignore ZonedDateTimes that aren't valid
+                e isa Union{ArgumentError,AmbiguousTimeError,NonExistentTimeError} || rethrow()
+                nothing
+            end
+            for year in 2000:2020
+            for month in 1:5
+            for day in 10:15
+            for tz_name in timezone_names()
+        ]
+    end
+
+    outputs = Channel(Inf)
+    @sync begin
+        for _ in 1:15
+            Threads.@spawn begin
+                put!(outputs, cycle_zdts())
+            end
+        end
+    end
+    close(outputs)
+
+    tzs = collect(outputs)
+
+    # Test that every Task produced the same result
+    allsame(x) = all(y -> y == first(x), x)
+    @test allsame(tzs)
+end
+
+#----------------------------------------------------
+
+@testset "Interleaved compile() and TimeZone construction" begin
+    @sync for i in 1:20
+        if (i % 5 == 0)
+            TimeZones.TZData.compile()
+        end
+        Threads.@spawn begin
+            TimeZone("US/Eastern", TimeZones.Class(:LEGACY))
+        end
+    end
+end
+"""
+
+@info "Running Thread Safety tests"
+@testset "Multithreaded TimeZone construction" begin
+    run(`$(Base.julia_cmd()) -t8 --proj -E $(program)`)
+end

Diff for: test/types/timezone.jl

@@ -2,7 +2,7 @@ using TimeZones: Class
 
 @testset "istimezone" begin
     # Invalidate the cache to ensure that `istimezone` works for non-loaded time zones.
-    empty!(TimeZones.TIME_ZONE_CACHE)
+    TimeZones._reset_tz_cache()
 
     @test istimezone("Europe/Warsaw")
     @test istimezone("UTC+02")