summaryrefslogtreecommitdiff
path: root/Util/MemoryCache.cs
blob: bcbd4d2039d86b5e25d7bc399ffa9064eaa3b3a2 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
using System.Collections;

namespace HyperBooru.Util;

public class MemoryCache<TKey, TValue> : IEnumerable<KeyValuePair<TKey, TValue>> where TKey : struct {
    /// <summary>Maximum number of items this cache may hold (unlimited if null)</summary>
    public int?                 MaxItems   { get; init; }
    /// <summary>Maximum amount of time an item may stay in cache before it is removed (unlimited if null)</summary>
    public TimeSpan?            MaxAge     { get; init; }
    /// <summary>Function that will be called to populate the cache in the event of a cache-miss</summary>
    public Func<TKey, TValue?>? DataSource { get; init; }

    private Dictionary<TKey, CacheItem> cache = new();

    private struct CacheItem {
        public DateTime CreateTime     { get; set; }
        public DateTime LastAccessTime { get; set; }
        public TValue   Value          { get; set; }
    }

    public TValue this[TKey key] {
        get => GetValue(key);
        set {
            Prune();
            cache[key] = new() {
                CreateTime     = DateTime.Now,
                LastAccessTime = DateTime.Now,
                Value          = value
            };
        }
    }

    public TValue GetValue(TKey key) {
        bool success = cache.TryGetValue(key, out var result);
        if(success) {
            result.LastAccessTime = DateTime.Now;
            if(MaxAge is null)
                return result.Value;
            if(result.CreateTime > DateTime.Now - MaxAge)
                return result.Value;
        }

        if(DataSource is null)
            throw new KeyNotFoundException();

        TValue? value = DataSource(key);

        if(value is null)
            throw new KeyNotFoundException();

        Prune();
        cache[key] = new() {
            CreateTime     = DateTime.Now,
            LastAccessTime = DateTime.Now,
            Value          = value
        };
        return value;
    }

    public bool TryGetValue(TKey key, out TValue? value) {
        try {
            value = GetValue(key);
            return true;
        } catch(KeyNotFoundException) {
            value = default;
            return false;
        }
    }

    public bool Remove(TKey key) => cache.Remove(key);

    public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator() {
        DateTime? expiry = MaxAge is null ? null : DateTime.Now - MaxAge;

        foreach(var kv in cache) {
            // Don't return expired cache items
            if(expiry is not null)
                if(kv.Value.CreateTime < expiry)
                    continue;

            yield return new KeyValuePair<TKey, TValue>(kv.Key, kv.Value.Value);
        }
    }

    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

    private void Prune() {
        DateTime? expiry = MaxAge is null ? null : DateTime.Now - MaxAge;

        // If an expiry time for cache items was
        // specified, remove expired cache items.
        if(expiry is not null) {
            foreach(var kv in cache) {
                if(kv.Value.CreateTime < expiry)
                    cache.Remove(kv.Key);
            }
        }

        // If this cache was created with a maximum size,
        // remove elements until that size is reached.
        if(MaxItems is null || cache.Count() < MaxItems)
            return;

        var toRemove = cache
            .OrderBy(kv => kv.Value.LastAccessTime)
            .Take(cache.Count() - (int) MaxItems + 1)
            .Select(kv => kv.Key)
            .ToArray();

        foreach(var key in toRemove)
            cache.Remove(key);
    }
}