Skip to content

Commit adc96b1

Browse files
committed
Add HasKey method
1 parent 898640b commit adc96b1

File tree

8 files changed

+79
-7
lines changed

8 files changed

+79
-7
lines changed

benchmarks/RocksDb.Extensions.Benchmarks/RocksDbBenchmark.cs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,7 @@ public class RocksDbBenchmark : IDisposable
1313
private readonly CacheKey[] _keys;
1414
private readonly CacheValue[] _values;
1515
private readonly Random _random;
16-
17-
16+
1817
public RocksDbBenchmark()
1918
{
2019
var services = new ServiceCollection();

src/RocksDb.Extensions/IRocksDbAccessor.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ public interface IRocksDbAccessor<TKey, TValue>
1818
void PutRange(ReadOnlySpan<TKey> keys, ReadOnlySpan<TValue> values);
1919
void PutRange(ReadOnlySpan<TValue> values, Func<TValue, TKey> keySelector);
2020
IEnumerable<TValue> GetAll();
21+
bool HasKey(TKey key);
2122
}
2223

2324
#pragma warning restore CS1591

src/RocksDb.Extensions/IRocksDbStore.cs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,4 +57,11 @@ public abstract class RocksDbStore<TKey, TValue>
5757
/// </summary>
5858
/// <returns>An enumerable collection of all the values in the store.</returns>
5959
public IEnumerable<TValue> GetAll() => _rocksDbAccessor.GetAll();
60+
61+
/// <summary>
62+
/// Determines whether the store contains a value for a specific key.
63+
/// </summary>
64+
/// <param name="key">The key to check in the store for an associated value.</param>
65+
/// <returns><c>true</c> if the store contains an element with the specified key; otherwise, <c>false</c>.</returns>
66+
public bool HasKey(TKey key) => _rocksDbAccessor.HasKey(key);
6067
}

src/RocksDb.Extensions/RocksDbAccessor.cs

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,5 +297,45 @@ public IEnumerable<TValue> GetAll()
297297
_ = iterator.Next();
298298
}
299299
}
300+
301+
public bool HasKey(TKey key)
302+
{
303+
byte[]? rentedKeyBuffer = null;
304+
bool useSpan;
305+
306+
// ReSharper disable once AssignmentInConditionalExpression
307+
Span<byte> keyBuffer = (useSpan = _keySerializer.TryCalculateSize(ref key, out var keySize))
308+
? keySize < MaxStackSize
309+
? stackalloc byte[keySize]
310+
: (rentedKeyBuffer = ArrayPool<byte>.Shared.Rent(keySize)).AsSpan(0, keySize)
311+
: Span<byte>.Empty;
312+
313+
ReadOnlySpan<byte> keySpan = keyBuffer;
314+
ArrayPoolBufferWriter<byte>? keyBufferWriter = null;
315+
316+
try
317+
{
318+
if (useSpan)
319+
{
320+
_keySerializer.WriteTo(ref key, ref keyBuffer);
321+
}
322+
else
323+
{
324+
keyBufferWriter = new ArrayPoolBufferWriter<byte>();
325+
_keySerializer.WriteTo(ref key, keyBufferWriter);
326+
keySpan = keyBufferWriter.WrittenSpan;
327+
}
328+
329+
return _rocksDb.HasKey(keySpan, _columnFamilyHandle);
330+
}
331+
finally
332+
{
333+
keyBufferWriter?.Dispose();
334+
if (rentedKeyBuffer is not null)
335+
{
336+
ArrayPool<byte>.Shared.Return(rentedKeyBuffer);
337+
}
338+
}
339+
}
300340
}
301341

test/RocksDb.Extensions.Tests/RocksDbStoreWithJsonSerializerTests.cs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ public void should_put_and_retrieve_data_from_store()
2929
store.Put(cacheKey, cacheValue);
3030

3131
// Assert
32+
store.HasKey(cacheKey).ShouldBeTrue();
3233
store.TryGet(cacheKey, out var value).ShouldBeTrue();
3334
value.ShouldBeEquivalentTo(cacheValue);
3435
}
@@ -55,6 +56,7 @@ public void should_put_and_remove_data_from_store()
5556
store.Remove(cacheKey);
5657

5758
// Assert
59+
store.HasKey(cacheKey).ShouldBeFalse();
5860
store.TryGet(cacheKey, out _).ShouldBeFalse();
5961
}
6062

@@ -78,6 +80,7 @@ public void should_put_range_of_data_to_store()
7880
for (var index = 0; index < cacheKeys.Length; index++)
7981
{
8082
var cacheKey = cacheKeys[index];
83+
store.HasKey(cacheKey).ShouldBeTrue();
8184
store.TryGet(cacheKey, out var cacheValue).ShouldBeTrue();
8285
cacheValue.ShouldBeEquivalentTo(cacheValues[index]);
8386
}
@@ -99,7 +102,9 @@ public void should_put_range_of_data_to_store_when_key_is_derived_from_value()
99102
// Assert
100103
foreach (var expectedCacheValue in cacheValues)
101104
{
102-
store.TryGet(new ProtoNetCacheKey { Id = expectedCacheValue.Id }, out var cacheValue).ShouldBeTrue();
105+
var key = new ProtoNetCacheKey { Id = expectedCacheValue.Id };
106+
store.HasKey(key).ShouldBeTrue();
107+
store.TryGet(key, out var cacheValue).ShouldBeTrue();
103108
cacheValue.ShouldBeEquivalentTo(expectedCacheValue);
104109
}
105110
}

test/RocksDb.Extensions.Tests/RocksDbStoreWithPrimitiveSerializerTests.cs

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,11 @@ public void should_put_and_retrieve_data_from_store_using_int_32_types()
1717
store.Put(1, 2);
1818

1919
// Assert
20+
store.HasKey(1).ShouldBeTrue();
2021
store.TryGet(1, out var value).ShouldBeTrue();
2122
value.ShouldBe(2);
2223
}
2324

24-
2525
[Test]
2626
public void should_put_and_remove_data_from_store_using_int32_types()
2727
{
@@ -34,6 +34,7 @@ public void should_put_and_remove_data_from_store_using_int32_types()
3434
store.Remove(1);
3535

3636
// Assert
37+
store.HasKey(1).ShouldBeFalse();
3738
store.TryGet(1, out _).ShouldBeFalse();
3839
}
3940

@@ -55,6 +56,7 @@ public void should_put_range_of_data_to_store_using_int_32_types()
5556
for (var index = 0; index < cacheKeys.Length; index++)
5657
{
5758
var cacheKey = cacheKeys[index];
59+
store.HasKey(cacheKey).ShouldBeTrue();
5860
store.TryGet(cacheKey, out var cacheValue).ShouldBeTrue();
5961
cacheValue.ShouldBe(cacheValues[index]);
6062
}
@@ -71,6 +73,7 @@ public void should_put_and_retrieve_data_from_store_using_int_64_types()
7173
store.Put(1, 2);
7274

7375
// Assert
76+
store.HasKey(1).ShouldBeTrue();
7477
store.TryGet(1, out var value).ShouldBeTrue();
7578
value.ShouldBe(2);
7679
}
@@ -87,6 +90,7 @@ public void should_put_and_remove_data_from_store_using_int64_types()
8790
store.Remove(1);
8891

8992
// Assert
93+
store.HasKey(1).ShouldBeFalse();
9094
store.TryGet(1, out _).ShouldBeFalse();
9195
}
9296

@@ -110,6 +114,7 @@ public void should_put_range_of_data_to_store_using_int_64_types()
110114
for (var index = 0; index < cacheKeys.Length; index++)
111115
{
112116
var cacheKey = cacheKeys[index];
117+
store.HasKey(cacheKey).ShouldBeTrue();
113118
store.TryGet(cacheKey, out var cacheValue).ShouldBeTrue();
114119
cacheValue.ShouldBe(cacheValues[index]);
115120
}
@@ -129,6 +134,7 @@ public void should_put_and_retrieve_data_from_store_using_string_types(int paylo
129134
store.Put(key, value);
130135

131136
// Assert
137+
store.HasKey(key);
132138
store.TryGet(key, out var retrievedValue).ShouldBeTrue();
133139
retrievedValue.ShouldBe(value);
134140
}
@@ -149,6 +155,7 @@ public void should_put_and_remove_data_from_store_using_string_types(int payload
149155

150156

151157
// Assert
158+
store.HasKey(key).ShouldBeFalse();
152159
store.TryGet(key, out _).ShouldBeFalse();
153160
}
154161

@@ -172,6 +179,7 @@ public void should_put_range_of_data_to_store_using_string_types()
172179
for (var index = 0; index < cacheKeys.Length; index++)
173180
{
174181
var cacheKey = cacheKeys[index];
182+
store.HasKey(cacheKey).ShouldBeTrue();
175183
store.TryGet(cacheKey, out var cacheValue).ShouldBeTrue();
176184
cacheValue.ShouldBe(cacheValues[index]);
177185
}
@@ -193,7 +201,9 @@ public void should_put_range_of_data_when_key_is_derived_from_value()
193201
// Assert
194202
foreach (var expectedCacheValue in cacheValues)
195203
{
196-
store.TryGet($"{expectedCacheValue}+", out var cacheValue).ShouldBeTrue();
204+
var key = $"{expectedCacheValue}+";
205+
store.HasKey(key).ShouldBeTrue();
206+
store.TryGet(key, out var cacheValue).ShouldBeTrue();
197207
cacheValue.ShouldBeEquivalentTo(expectedCacheValue);
198208
}
199209
}

test/RocksDb.Extensions.Tests/RocksDbStoreWithProtoBufNetSerializerTests.cs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ public void should_put_and_retrieve_data_from_store()
3030
store.Put(cacheKey, cacheValue);
3131

3232
// Assert
33+
store.HasKey(cacheKey).ShouldBeTrue();
3334
store.TryGet(cacheKey, out var value).ShouldBeTrue();
3435
value.ShouldBeEquivalentTo(cacheValue);
3536
}
@@ -56,6 +57,7 @@ public void should_put_and_remove_data_from_store()
5657
store.Remove(cacheKey);
5758

5859
// Assert
60+
store.HasKey(cacheKey).ShouldBeFalse();
5961
store.TryGet(cacheKey, out _).ShouldBeFalse();
6062
}
6163

@@ -79,6 +81,7 @@ public void should_put_range_of_data_to_store()
7981
for (var index = 0; index < cacheKeys.Length; index++)
8082
{
8183
var cacheKey = cacheKeys[index];
84+
store.HasKey(cacheKey).ShouldBeTrue();
8285
store.TryGet(cacheKey, out var cacheValue).ShouldBeTrue();
8386
cacheValue.ShouldBeEquivalentTo(cacheValues[index]);
8487
}
@@ -100,7 +103,9 @@ public void should_put_range_of_data_when_key_is_derived_from_value()
100103
// Assert
101104
foreach (var expectedCacheValue in cacheValues)
102105
{
103-
store.TryGet(new ProtoNetCacheKey { Id = expectedCacheValue.Id }, out var cacheValue).ShouldBeTrue();
106+
var key = new ProtoNetCacheKey { Id = expectedCacheValue.Id };
107+
store.HasKey(key).ShouldBeTrue();
108+
store.TryGet(key, out var cacheValue).ShouldBeTrue();
104109
cacheValue.ShouldBeEquivalentTo(expectedCacheValue);
105110
}
106111
}

test/RocksDb.Extensions.Tests/RocksDbStoreWithProtobufSerializerTests.cs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ public void should_put_and_retrieve_data_from_store(int payloadSize)
3131
// Assert
3232
for (int i = 0; i < cacheKeys.Length; i++)
3333
{
34+
store.HasKey(cacheKeys[i]).ShouldBeTrue();
3435
store.TryGet(cacheKeys[i], out var value).ShouldBeTrue();
3536
value.ShouldBe(cacheValues[i]);
3637
}
@@ -61,6 +62,7 @@ public void should_put_and_remove_data_from_store(int payloadSize)
6162
// Assert
6263
foreach (var cacheKey in cacheKeys)
6364
{
65+
store.HasKey(cacheKey).ShouldBeFalse();
6466
store.TryGet(cacheKey, out _).ShouldBeFalse();
6567
}
6668
}
@@ -85,6 +87,7 @@ public void should_put_range_of_data_to_store()
8587
for (var index = 0; index < cacheKeys.Length; index++)
8688
{
8789
var cacheKey = cacheKeys[index];
90+
store.HasKey(cacheKey).ShouldBeTrue();
8891
store.TryGet(cacheKey, out var cacheValue).ShouldBeTrue();
8992
cacheValue.ShouldBe(cacheValues[index]);
9093
}
@@ -130,7 +133,9 @@ public void should_put_range_of_data_when_key_is_derived_from_value()
130133
// Assert
131134
foreach (var expectedCacheValue in cacheValues)
132135
{
133-
store.TryGet(new CacheKey { Id = expectedCacheValue.Id }, out var cacheValue).ShouldBeTrue();
136+
var key = new CacheKey { Id = expectedCacheValue.Id };
137+
store.HasKey(key).ShouldBeTrue();
138+
store.TryGet(key, out var cacheValue).ShouldBeTrue();
134139
cacheValue.ShouldBeEquivalentTo(expectedCacheValue);
135140
}
136141
}

0 commit comments

Comments
 (0)