Fix underflow possibility on ColumnFamilyDbStore (#5975) (#6023)

* Fix underflow * Use compare exchange and make separate test * Use Max for count instead of trying to make atomic check
Azure · Feb 1, 2022 · ee79c13 · ee79c13
1 parent 20d7e11
commit ee79c13
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 1 deletion.
diff --git a/edge-util/src/Microsoft.Azure.Devices.Edge.Storage.RocksDb/ColumnFamilyDbStore.cs b/edge-util/src/Microsoft.Azure.Devices.Edge.Storage.RocksDb/ColumnFamilyDbStore.cs
@@ -141,7 +141,7 @@ public Task IterateBatch(int batchSize, Func<byte[], byte[], Task> callback, Can
             return this.IterateBatch(iterator => iterator.SeekToFirst(), batchSize, callback, cancellationToken);
         }
 
-        public Task<ulong> Count() => Task.FromResult((ulong)Interlocked.Read(ref this.count));
+        public Task<ulong> Count() => Task.FromResult((ulong)Math.Max(Interlocked.Read(ref this.count), 0));
 
         public Task<ulong> GetCountFromOffset(byte[] offset)
         {

diff --git a/edge-util/test/Microsoft.Azure.Devices.Edge.Storage.RocksDb.Test/ColumnFamilyStoreTest.cs b/edge-util/test/Microsoft.Azure.Devices.Edge.Storage.RocksDb.Test/ColumnFamilyStoreTest.cs
@@ -2,6 +2,7 @@
 namespace Microsoft.Azure.Devices.Edge.Storage.RocksDb.Test
 {
     using System.Linq;
+    using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Devices.Edge.Util;
     using Microsoft.Azure.Devices.Edge.Util.Test.Common;
@@ -103,5 +104,37 @@ public async Task MessageCountTest()
                 Assert.Equal(0ul, await columnFamilyDbStore.Count());
             }
         }
+
+        [Fact]
+        public async Task MessageCountUnderflowTest()
+        {
+            using (IDbStore columnFamilyDbStore = this.rocksDbStoreProvider.GetDbStore("test"))
+            {
+                Assert.Equal(0ul, await columnFamilyDbStore.Count());
+
+                for (int i = 0; i < 10; i++)
+                {
+                    string key = $"key{i}";
+                    string value = "$value{i}";
+                    await columnFamilyDbStore.Put(key.ToBytes(), value.ToBytes());
+                }
+
+                Assert.Equal(10ul, await columnFamilyDbStore.Count());
+            }
+
+            using (IDbStore columnFamilyDbStore = this.rocksDbStoreProvider.GetDbStore("test"))
+            {
+                Assert.Equal(10ul, await columnFamilyDbStore.Count());
+
+                // Using 11 (10 + 1) to make sure underflow is caught.
+                for (int i = 0; i < 11; i++)
+                {
+                    string key = $"key{i}";
+                    await columnFamilyDbStore.Remove(key.ToBytes());
+                }
+
+                Assert.Equal(0ul, await columnFamilyDbStore.Count());
+            }
+        }
     }
 }