1  @doc """ Module for managing time-to-live cache.
  2  
  3  """
  4  module TimeToLive
  5  using ConcurrentCollections: modify!, Delete, ConcurrentDict
  6  
  7  export TTL, safettl
  8  
  9  using Base.Iterators: peel
 10  using Dates: DateTime, Period, now
 11  using ..DocStringExtensions
 12  
 13  struct Node{T}
 14      value::T
 15      expiry::DateTime
 16  end
 17  
 18  isexpired(v::Node) = now() > v.expiry
 19  isexpired(time::DateTime) = v::Node -> time > v.expiry
 20  
 21  @doc """
 22      TTL(ttl::Period; refresh_on_access::Bool=false)
 23      TTL{K, V}(ttl::Period; refresh_on_access::Bool=false)
 24  
 25  $(TYPEDFIELDS)
 26  
 27  An associative [TTL](https://en.wikipedia.org/wiki/Time_to_live) cache.
 28  If `refresh_on_access` is set, expiries are reset whenever they are accessed.
 29  """
 30  struct TTL{K,V,D<:AbstractDict,P<:Period} <: AbstractDict{K,V}
 31      dict::D where {D<:AbstractDict{K,Node{V}}}
 32      ttl::P
 33      refresh::Bool
 34  
 35      function TTL{K,V}(
 36          ttl::P; refresh_on_access::Bool=false, dict_type=Dict
 37      ) where {K,V,P<:Period}
 38          new{K,V,dict_type,P}(dict_type{K,Node{V}}(), ttl, refresh_on_access)
 39      end
 40      function TTL(ttl::Period; refresh_on_access::Bool=false)
 41          TTL{Any,Any}(ttl; refresh_on_access=refresh_on_access)
 42      end
 43  end
 44  
 45  @doc """ Safely instantiate a TTL dictionary.
 46  
 47  $(TYPEDSIGNATURES)
 48  
 49  This function safely creates a Time-to-Live (TTL) dictionary with specified key and value types, along with an optional ttl parameter.
 50  """
 51  function safettl(K::Type, V::Type, ttl; kwargs...)
 52      TTL{K,V}(ttl; dict_type=ConcurrentDict, kwargs...)
 53  end
 54  
 55  Base.delete!(t::TTL, key) = (delete!(t.dict, key); t)
 56  Base.empty!(t::ConcurrentDict{K,V}) where {K,V} =
 57      for k in keys(t)
 58          modify!(t, k) do value
 59              Delete(value)
 60          end
 61      end
 62  Base.delete!(t::ConcurrentDict{K,V}, k) where {K,V} =
 63      modify!(t, k) do value
 64          Delete(value)
 65      end
 66  
 67  Base.empty!(t::TTL) = (empty!(t.dict); t)
 68  # Specifying ::Function fixes some method invalidations
 69  Base.get(f::Function, t::TTL, key) = haskey(t, key) ? t[key] : f()
 70  Base.get!(t::TTL, key, default) = haskey(t, key) ? t[key] : (t[key] = default)
 71  Base.length(t::TTL) = count(!isexpired(now()), values(t.dict))
 72  Base.push!(t::TTL, p::Pair) = (t[p.first] = p.second; t)
 73  Base.setindex!(t::TTL{K,V}, v, k) where {K,V} = t.dict[k] = Node{V}(v, now() + t.ttl)
 74  Base.sizehint!(t::TTL, newsz) = (sizehint!(t.dict, newsz); t)
 75  
 76  function Base.pop!(t::TTL)
 77      p = pop!(t.dict)
 78      return isexpired(p.second) ? pop!(t) : p.first => p.second.value
 79  end
 80  
 81  function Base.pop!(t::TTL, key)
 82      v = pop!(t.dict, key)
 83      isexpired(v) && throw(KeyError(key))
 84      return v.value
 85  end
 86  
 87  function Base.get(t::TTL, key, default)
 88      haskey(t.dict, key) || return default
 89      v = t.dict[key]
 90      if isexpired(v)
 91          delete!(t, key)
 92          return default
 93      end
 94      t.refresh && (t[key] = v.value)
 95      return v.value
 96  end
 97  
 98  function Base.getkey(t::TTL, key, default)
 99      return if haskey(t, key)
100          if isexpired(t.dict[key])
101              delete!(t, key)
102              default
103          else
104              key
105          end
106      else
107          default
108      end
109  end
110  
111  function Base.iterate(t::TTL, ks=keys(t.dict))
112      isempty(ks) && return nothing
113      k, rest = peel(ks)
114      v = t.dict[k]
115      return if isexpired(v)
116          delete!(t, k)
117          iterate(t, rest)
118      else
119          k => v.value, rest
120      end
121  end
122  
123  function Base.getindex(t::TTL, key)
124      v = t.dict[key]
125      if isexpired(v)
126          delete!(t, key)
127          throw(KeyError(key))
128      elseif t.refresh
129          t[key] = v.value
130      end
131      return v.value
132  end
133  
134  end