From adffb1e2180f55075457585b397cc89c2300007f Mon Sep 17 00:00:00 2001
From: Mohsen Rajabi <m.kabir8895@gmail.com>
Date: Mon, 23 Mar 2020 17:09:36 +0430
Subject: [PATCH] add asp memory cache

---
 src/Ocelot/Cache/AspMemoryCache.cs            | 73 ++++++++++++++++
 .../DependencyInjection/OcelotBuilder.cs      |  8 +-
 .../Cache/AspMemoryCacheTests.cs              | 83 +++++++++++++++++++
 3 files changed, 162 insertions(+), 2 deletions(-)
 create mode 100644 src/Ocelot/Cache/AspMemoryCache.cs
 create mode 100644 test/Ocelot.UnitTests/Cache/AspMemoryCacheTests.cs

diff --git a/src/Ocelot/Cache/AspMemoryCache.cs b/src/Ocelot/Cache/AspMemoryCache.cs
new file mode 100644
index 000000000..0bd1fe343
--- /dev/null
+++ b/src/Ocelot/Cache/AspMemoryCache.cs
@@ -0,0 +1,73 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+using Microsoft.Extensions.Caching.Distributed;
+using Microsoft.Extensions.Caching.Memory;
+
+namespace Ocelot.Cache
+{
+    public class AspMemoryCache<T> : IOcelotCache<T>
+    {
+        private readonly IMemoryCache _memoryCache;
+        private readonly Dictionary<string, List<string>> _regions;
+
+        public AspMemoryCache(IMemoryCache memoryCache)
+        {
+            _memoryCache = memoryCache;
+            _regions = new Dictionary<string, List<string>>();
+        }
+
+        public void Add(string key, T value, TimeSpan ttl, string region)
+        {
+            if (ttl.TotalMilliseconds <= 0)
+                return;
+
+            _memoryCache.Set(key, value, ttl);
+
+            SetRegion(region, key);
+        }
+
+        public T Get(string key, string region)
+        {
+            _memoryCache.TryGetValue<T>(key, out T value);
+
+            return value;
+        }
+
+        public void ClearRegion(string region)
+        {
+            if (_regions.ContainsKey(region))
+            {
+                var keys = _regions[region];
+                foreach (var key in keys)
+                {
+                    _memoryCache.Remove(key);
+                }
+            }
+        }
+
+        public void AddAndDelete(string key, T value, TimeSpan ttl, string region)
+        {
+            if (_memoryCache.TryGetValue(key, out object oldValue))
+                _memoryCache.Remove(key);
+
+            Add(key, value, ttl, region);
+        }
+
+        private void SetRegion(string region, string key)
+        {
+            if (_regions.ContainsKey(region))
+            {
+                var current = _regions[region];
+                if (!current.Contains(key))
+                {
+                    current.Add(key);
+                }
+            }
+            else
+            {
+                _regions.Add(region, new List<string> { key });
+            }
+        }
+    }
+}
diff --git a/src/Ocelot/DependencyInjection/OcelotBuilder.cs b/src/Ocelot/DependencyInjection/OcelotBuilder.cs
index 0336f691e..4109e6176 100644
--- a/src/Ocelot/DependencyInjection/OcelotBuilder.cs
+++ b/src/Ocelot/DependencyInjection/OcelotBuilder.cs
@@ -56,8 +56,12 @@ public OcelotBuilder(IServiceCollection services, IConfiguration configurationRo
             Services = services;
             Services.Configure<FileConfiguration>(configurationRoot);
 
-            Services.TryAddSingleton<IOcelotCache<FileConfiguration>, InMemoryCache<FileConfiguration>>();
-            Services.TryAddSingleton<IOcelotCache<CachedResponse>, InMemoryCache<CachedResponse>>();
+            //Services.TryAddSingleton<IOcelotCache<FileConfiguration>, InMemoryCache<FileConfiguration>>();
+            //Services.TryAddSingleton<IOcelotCache<CachedResponse>, InMemoryCache<CachedResponse>>();
+            Services.TryAddSingleton<IOcelotCache<FileConfiguration>, AspMemoryCache<FileConfiguration>>();
+            Services.TryAddSingleton<IOcelotCache<CachedResponse>, AspMemoryCache<CachedResponse>>();
+
+
             Services.TryAddSingleton<IHttpResponseHeaderReplacer, HttpResponseHeaderReplacer>();
             Services.TryAddSingleton<IHttpContextRequestHeaderReplacer, HttpContextRequestHeaderReplacer>();
             Services.TryAddSingleton<IHeaderFindAndReplaceCreator, HeaderFindAndReplaceCreator>();
diff --git a/test/Ocelot.UnitTests/Cache/AspMemoryCacheTests.cs b/test/Ocelot.UnitTests/Cache/AspMemoryCacheTests.cs
new file mode 100644
index 000000000..3019f8bf5
--- /dev/null
+++ b/test/Ocelot.UnitTests/Cache/AspMemoryCacheTests.cs
@@ -0,0 +1,83 @@
+using Microsoft.Extensions.Caching.Memory;
+
+namespace Ocelot.UnitTests.Cache
+{
+    using Ocelot.Cache;
+    using Shouldly;
+    using System;
+    using System.Threading;
+    using Xunit;
+
+    public class AspMemoryCacheTests
+    {
+        private readonly AspMemoryCache<Fake> _cache;
+
+        public AspMemoryCacheTests()
+        {
+            _cache = new AspMemoryCache<Fake>(new MemoryCache(new MemoryCacheOptions()));
+        }
+
+        [Fact]
+        public void should_cache()
+        {
+            var fake = new Fake(1);
+            _cache.Add("1", fake, TimeSpan.FromSeconds(100), "region");
+            var result = _cache.Get("1", "region");
+            result.ShouldBe(fake);
+            fake.Value.ShouldBe(1);
+        }
+
+        [Fact]
+        public void should_add_and_delete()
+        {
+            var fake = new Fake(1);
+            _cache.Add("1", fake, TimeSpan.FromSeconds(100), "region");
+            var newFake = new Fake(1);
+            _cache.AddAndDelete("1", newFake, TimeSpan.FromSeconds(100), "region");
+            var result = _cache.Get("1", "region");
+            result.ShouldBe(newFake);
+            newFake.Value.ShouldBe(1);
+        }
+
+        [Fact]
+        public void should_clear_region()
+        {
+            var fake = new Fake(1);
+            _cache.Add("1", fake, TimeSpan.FromSeconds(100), "region");
+            _cache.ClearRegion("region");
+            var result = _cache.Get("1", "region");
+            result.ShouldBeNull();
+        }
+
+        [Fact]
+        public void should_clear_key_if_ttl_expired()
+        {
+            var fake = new Fake(1);
+            _cache.Add("1", fake, TimeSpan.FromMilliseconds(50), "region");
+            Thread.Sleep(200);
+            var result = _cache.Get("1", "region");
+            result.ShouldBeNull();
+        }
+
+        [Theory]
+        [InlineData(0)]
+        [InlineData(-1)]
+        public void should_not_add_to_cache_if_timespan_empty(int ttl)
+        {
+            var fake = new Fake(1);
+            _cache.Add("1", fake, TimeSpan.FromSeconds(ttl), "region");
+            var result = _cache.Get("1", "region");
+            result.ShouldBeNull();
+        }
+
+        private class Fake
+        {
+            public Fake(int value)
+            {
+                Value = value;
+            }
+
+            public int Value { get; }
+        }
+    }
+}