Merge pull request stratisproject#35 from quantumagi/kbdispose

Refactor KeyValueStore disposal

src/Stratis.Bitcoin/FullNode.cs

@@ -297,6 +297,9 @@ public void Dispose()
             this.logger.LogInformation("Disposing settings.");
             this.Settings.Dispose();
 
+            // Dispose the node storage.
+            this.DisposeNodeStorage();
+
             // Fire INodeLifetime.Stopped.
             this.logger.LogInformation("Notify application has stopped.");
             this.nodeLifetime.NotifyStopped();
@@ -305,5 +308,31 @@ public void Dispose()
 
             this.State = FullNodeState.Disposed;
         }
+
+        /// <summary>
+        /// Due to the potential shared nature of node storage it can't be up to any given
+        /// feature to deal with the disposal thereof. Instead we do it here for all
+        /// singleton storage objects during shutdown.
+        /// </summary>
+        private void DisposeNodeStorage()
+        {
+            // Identify classes that support the IKeyValueStore interface.
+            var storageClasses = AppDomain.CurrentDomain.GetAssemblies()
+                .SelectMany(x => x.GetTypes())
+                .Where(x => typeof(IKeyValueStore).IsAssignableFrom(x) && !x.IsInterface && !x.IsAbstract);
+
+            // Find the class interfaces that are derived from IKeyValueStore...
+            foreach (Type type in storageClasses.SelectMany(x => x.GetInterfaces().Where(i => i.GetInterfaces().Any(i2 => i2.UnderlyingSystemType == typeof(IKeyValueStore)))))
+            {
+                // ...that are disposable...
+                var obj = this.Services.ServiceProvider.GetService(type) as IDisposable;
+                if (obj == null)
+                    continue;
+
+                // ...and dispose them.
+                this.logger.LogInformation("Disposing node storage '{0}'.", obj.GetType().Name);
+                obj.Dispose();
+            }
+        }
     }
 }

src/Stratis.Bitcoin/Interfaces/IKeyValueStoreRepository.cs

@@ -1,12 +1,13 @@
-using System.Collections.Generic;
+using System;
+using System.Collections.Generic;
 using Stratis.Bitcoin.KeyValueStore;
 
 namespace Stratis.Bitcoin.Interfaces
 {
     /// <summary>
     /// Represents a glue-layer containing the basic methods that all key-value databases should support.
     /// </summary>
-    public interface IKeyValueStoreRepository
+    public interface IKeyValueStoreRepository : IDisposable
     {
         /// <summary>
         /// Initialize the underlying database / glue-layer.

src/Stratis.Bitcoin/NodeStorage/KeyValueStore/KeyValueStore.cs

@@ -96,7 +96,7 @@ protected override void Dispose(bool disposing)
 
             if (disposing)
             {
-                this.Repository.Close();
+                this.Repository.Dispose();
             }
 
             this.disposed = true;

src/Stratis.Bitcoin/NodeStorage/KeyValueStore/KeyValueStoreRepository.cs

@@ -119,5 +119,20 @@ public virtual T Deserialize<T>(byte[] objBytes)
 
         /// <inheritdoc />
         public abstract void Close();
+
+        // Public implementation of Dispose pattern callable by consumers.
+        public void Dispose()
+        {
+            this.Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        /// <summary>Protected implementation of Dispose pattern.</summary>
+        /// <param name="disposing">Indicates whether disposing.</param>
+        protected virtual void Dispose(bool disposing)
+        {
+            if (disposing)
+                this.Close();
+        }
     }
 }

src/Stratis.Bitcoin/NodeStorage/KeyValueStoreDBreeze/KeyValueStoreDBreeze.cs

@@ -53,6 +53,7 @@ public override byte[] Serialize<T>(T obj)
 
         public override void Init(string rootPath)
         {
+            this.Close();
             this.storage = new DBreezeEngine(rootPath);
         }
 
@@ -219,6 +220,8 @@ public override void OnRollback(KeyValueStoreTransaction keyValueStoreTransactio
 
         public override void Close()
         {
+            this.storage?.Dispose();
+            this.storage = null;
         }
     }
 }

src/Stratis.Bitcoin/NodeStorage/KeyValueStoreLevelDB/KeyValueStoreLevelDB.cs

@@ -64,8 +64,11 @@ public override void Init(string rootPath)
                 CreateIfMissing = true,
             };
 
+            this.Close();
             this.Storage = new DB(options, rootPath);
 
+            Guard.NotNull(this.Storage, nameof(this.Storage));
+
             for (this.nextTablePrefix = 1; ; this.nextTablePrefix++)
             {
                 byte[] tableNameBytes = this.Storage.Get(new byte[] { 0, (byte)this.nextTablePrefix });
@@ -246,7 +249,8 @@ public override void OnRollback(KeyValueStoreTransaction keyValueStoreTransactio
 
         public override void Close()
         {
-            this.Storage.Close();
+            this.Storage?.Dispose();
+            this.Storage = null;
         }
     }
 }