Merge pull request #133 from Flow-Launcher/address_high_memory_issue

Limiting the number of ImageSources cached to address memory issue

Author: taooceros

Flow.Launcher.Infrastructure/Image/ImageCache.cs

@@ -2,53 +2,100 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading.Tasks;
 using System.Windows.Media;
 
 namespace Flow.Launcher.Infrastructure.Image
 {
     [Serializable]
+    public class ImageUsage
+    {
+
+        public int usage;
+        public ImageSource imageSource;
+
+        public ImageUsage(int usage, ImageSource image)
+        {
+            this.usage = usage;
+            imageSource = image;
+        }
+    }
+
     public class ImageCache
     {
-        private const int MaxCached = 5000;
-        public ConcurrentDictionary<string, int> Usage = new ConcurrentDictionary<string, int>();
-        private readonly ConcurrentDictionary<string, ImageSource> _data = new ConcurrentDictionary<string, ImageSource>();
+        private const int MaxCached = 50;
+        public ConcurrentDictionary<string, ImageUsage> Data { get; private set; } = new ConcurrentDictionary<string, ImageUsage>();
+        private const int permissibleFactor = 2;
 
+        public void Initialization(Dictionary<string, int> usage)
+        {
+            foreach (var key in usage.Keys)
+            {
+                Data[key] = new ImageUsage(usage[key], null);
+            }
+        }
 
         public ImageSource this[string path]
         {
             get
             {
-                Usage.AddOrUpdate(path, 1, (k, v) => v + 1);
-                var i = _data[path];
-                return i;
+                if (Data.TryGetValue(path, out var value))
+                {
+                    value.usage++;
+                    return value.imageSource;
+                }
+                
+                return null;
             }
-            set { _data[path] = value; }
-        }
+            set
+            {
+                Data.AddOrUpdate(
+                        path, 
+                        new ImageUsage(0, value), 
+                        (k, v) =>
+                            {
+                                v.imageSource = value;
+                                v.usage++;
+                                return v;
+                            }
+                );
 
-        public Dictionary<string, int> CleanupAndToDictionary() 
-            => Usage
-                .OrderByDescending(o => o.Value)
-                .Take(MaxCached)
-                .ToDictionary(i => i.Key, i => i.Value);
+                // To prevent the dictionary from drastically increasing in size by caching images, the dictionary size is not allowed to grow more than the permissibleFactor * maxCached size
+                // This is done so that we don't constantly perform this resizing operation and also maintain the image cache size at the same time
+                if (Data.Count > permissibleFactor * MaxCached)
+                {
+                    // To delete the images from the data dictionary based on the resizing of the Usage Dictionary.
+
+
+                    foreach (var key in Data.OrderBy(x => x.Value.usage).Take(Data.Count - MaxCached).Select(x => x.Key))
+                    {
+                        if (!(key.Equals(Constant.ErrorIcon) || key.Equals(Constant.DefaultIcon)))
+                        {
+                            Data.TryRemove(key, out _);
+                        }
+                    }
+                }
+            }
+        }
 
         public bool ContainsKey(string key)
         {
-            var contains = _data.ContainsKey(key);
+            var contains = Data.ContainsKey(key);
             return contains;
         }
 
         public int CacheSize()
         {
-            return _data.Count;
+            return Data.Count;
         }
 
         /// <summary>
         /// return the number of unique images in the cache (by reference not by checking images content)
         /// </summary>
         public int UniqueImagesInCache()
         {
-            return _data.Values.Distinct().Count();
+            return Data.Values.Select(x => x.imageSource).Distinct().Count();
         }
     }
 
-}
+}

Flow.Launcher.Infrastructure/Image/ImageLoader.cs

@@ -13,12 +13,11 @@ namespace Flow.Launcher.Infrastructure.Image
 {
     public static class ImageLoader
     {
-        private static readonly ImageCache _imageCache = new ImageCache();
-        private static readonly ConcurrentDictionary<string, string> _guidToKey = new ConcurrentDictionary<string, string>();
-        private static readonly bool _enableHashImage = true;
-
+        private static readonly ImageCache ImageCache = new ImageCache();
         private static BinaryStorage<Dictionary<string, int>> _storage;
+        private static readonly ConcurrentDictionary<string, string> GuidToKey = new ConcurrentDictionary<string, string>();
         private static IImageHashGenerator _hashGenerator;
+        private static bool EnableImageHash = true;
 
         private static readonly string[] ImageExtensions =
         {
@@ -36,39 +35,39 @@ public static void Initialize()
             _storage = new BinaryStorage<Dictionary<string, int>>("Image");
             _hashGenerator = new ImageHashGenerator();
 
-            _imageCache.Usage = LoadStorageToConcurrentDictionary();
+            var usage = LoadStorageToConcurrentDictionary();
 
             foreach (var icon in new[] { Constant.DefaultIcon, Constant.MissingImgIcon })
             {
                 ImageSource img = new BitmapImage(new Uri(icon));
                 img.Freeze();
-                _imageCache[icon] = img;
+                ImageCache[icon] = img;
             }
 
             Task.Run(() =>
             {
                 Stopwatch.Normal("|ImageLoader.Initialize|Preload images cost", () =>
                 {
-                    _imageCache.Usage.AsParallel().ForAll(x =>
+                    ImageCache.Data.AsParallel().ForAll(x =>
                     {
                         Load(x.Key);
                     });
                 });
-                Log.Info($"|ImageLoader.Initialize|Number of preload images is <{_imageCache.Usage.Count}>, Images Number: {_imageCache.CacheSize()}, Unique Items {_imageCache.UniqueImagesInCache()}");
+                Log.Info($"|ImageLoader.Initialize|Number of preload images is <{ImageCache.CacheSize()}>, Images Number: {ImageCache.CacheSize()}, Unique Items {ImageCache.UniqueImagesInCache()}");
             });
         }
 
         public static void Save()
         {
             lock (_storage)
             {
-                _storage.Save(_imageCache.CleanupAndToDictionary());
+                _storage.Save(ImageCache.Data.Select(x => (x.Key, x.Value.usage)).ToDictionary(x => x.Key, y => y.usage));
             }
         }
 
         private static ConcurrentDictionary<string, int> LoadStorageToConcurrentDictionary()
         {
-            lock(_storage)
+            lock (_storage)
             {
                 var loaded = _storage.TryLoad(new Dictionary<string, int>());
 
@@ -106,11 +105,11 @@ private static ImageResult LoadInternal(string path, bool loadFullImage = false)
             {
                 if (string.IsNullOrEmpty(path))
                 {
-                    return new ImageResult(_imageCache[Constant.MissingImgIcon], ImageType.Error);
+                    return new ImageResult(ImageCache[Constant.MissingImgIcon], ImageType.Error);
                 }
-                if (_imageCache.ContainsKey(path))
+                if (ImageCache.ContainsKey(path))
                 {
-                    return new ImageResult(_imageCache[path], ImageType.Cache);
+                    return new ImageResult(ImageCache[path], ImageType.Cache);
                 }
 
                 if (path.StartsWith("data:", StringComparison.OrdinalIgnoreCase))
@@ -139,8 +138,8 @@ private static ImageResult LoadInternal(string path, bool loadFullImage = false)
                     Log.Exception($"|ImageLoader.Load|Failed to get thumbnail for {path} on first try", e);
                     Log.Exception($"|ImageLoader.Load|Failed to get thumbnail for {path} on second try", e2);
 
-                    ImageSource image = _imageCache[Constant.MissingImgIcon];
-                    _imageCache[path] = image;
+                    ImageSource image = ImageCache[Constant.MissingImgIcon];
+                    ImageCache[path] = image;
                     imageResult = new ImageResult(image, ImageType.Error);
                 }
             }
@@ -191,7 +190,7 @@ private static ImageResult GetThumbnailResult(ref string path, bool loadFullImag
             }
             else
             {
-                image = _imageCache[Constant.MissingImgIcon];
+                image = ImageCache[Constant.MissingImgIcon];
                 path = Constant.MissingImgIcon;
             }
 
@@ -218,27 +217,26 @@ public static ImageSource Load(string path, bool loadFullImage = false)
 
             var img = imageResult.ImageSource;
             if (imageResult.ImageType != ImageType.Error && imageResult.ImageType != ImageType.Cache)
-            { 
-                // we need to get image hash
-                string hash = _enableHashImage ? _hashGenerator.GetHashFromImage(img) : null;
+            { // we need to get image hash
+                string hash = EnableImageHash ? _hashGenerator.GetHashFromImage(img) : null;
                 if (hash != null)
                 {
-                    if (_guidToKey.TryGetValue(hash, out string key))
-                    { 
-                        // image already exists
-                        img = _imageCache[key];
+                    
+                    if (GuidToKey.TryGetValue(hash, out string key))
+                    { // image already exists
+                        img = ImageCache[key] ?? img;
                     }
                     else
-                    { 
-                        // new guid
-                        _guidToKey[hash] = path;
+                    { // new guid
+                        GuidToKey[hash] = path;
                     }
                 }
 
                 // update cache
-                _imageCache[path] = img;
+                ImageCache[path] = img;
             }
 
+
             return img;
         }