Skip to content

Commit

Permalink
64 bit key allocation (#283)
Browse files Browse the repository at this point in the history 
+semver: breaking
  • Loading branch information
@borland
borland authored May 21, 2024
1 parent e1efc75 commit 80ba0ad
Show file tree
Hide file tree
Showing 6 changed files with 122 additions and 89 deletions.
4 changes: 2 additions & 2 deletions source/Nevermore.IntegrationTests/Advanced/CustomKeyAllocatorFixture.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,12 +105,12 @@ public void Reset()
allocations.Clear();
}

public int NextId(string tableName)
public long NextId(string tableName)
{
return allocations.AddOrUpdate(tableName, (_) => 100, (_, prev) => prev + 100);
}

public ValueTask<int> NextIdAsync(string tableName, CancellationToken cancellationToken)
public ValueTask<long> NextIdAsync(string tableName, CancellationToken cancellationToken)
{
return ValueTask.FromResult(NextId(tableName));
}
Expand Down
57 changes: 54 additions & 3 deletions source/Nevermore.IntegrationTests/KeyAllocatorFixture.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ public void ShouldAllocateForDifferentCollections()
AssertNext(allocator, "Todos", 4);
AssertNext(allocator, "Todos", 5);
}

[Test]
public async Task NextIdAsync_ShouldAllocateForDifferentCollections()
{
Expand All @@ -128,6 +128,57 @@ public async Task NextIdAsync_ShouldAllocateForDifferentCollections()
await AssertNextAsync(allocator, "Todos", 4);
await AssertNextAsync(allocator, "Todos", 5);
}

[Test]
public void CanAllocateLongKeys()
{
var allocator = new KeyAllocator(Store, 10);

var collectionName = Guid.NewGuid().ToString("N"); // collection name doesn't matter, use a GUID to avoid colliding with other tests

using (var tx = Store.BeginWriteTransaction())
{
tx.ExecuteNonQuery("INSERT INTO [KeyAllocation] ([CollectionName], [Allocated]) VALUES (@collectionName, @allocated)",
new CommandParameterValues
{
{ "collectionName", collectionName },
{ "allocated", int.MaxValue - 1 }
});
tx.Commit();
}

// the allocator will try and allocate the next block of 10, from int.max-1 to int.max+9.
// the fact that it doesn't throw an exception is proof enough, but let's check the return values
// to be sure
AssertNext(allocator, collectionName, (long)int.MaxValue);
AssertNext(allocator, collectionName, (long)int.MaxValue + 1);
}

[Test]
public async Task NextIdAsync_CanAllocateLongKeys()
{
var allocator = new KeyAllocator(Store, 10);

var collectionName = Guid.NewGuid().ToString("N"); // collection name doesn't matter, use a GUID to avoid colliding with other tests

using (var tx = await Store.BeginWriteTransactionAsync())
{
await tx.ExecuteNonQueryAsync("INSERT INTO [KeyAllocation] ([CollectionName], [Allocated]) VALUES (@collectionName, @allocated)",
new CommandParameterValues
{
{ "collectionName", collectionName },
{ "allocated", int.MaxValue - 1 }
});
await tx.CommitAsync();
}

// the allocator will try and allocate the next block of 10, from int.max-1 to int.max+9.
// the fact that it doesn't throw an exception is proof enough, but let's check the return values
// to be sure
await AssertNextAsync(allocator, collectionName, (long)int.MaxValue);
await AssertNextAsync(allocator, collectionName, (long)int.MaxValue + 1);
}


[Test]
public void ShouldAllocateInParallel()
Expand Down Expand Up @@ -249,12 +300,12 @@ public async Task AllocateIdAsync_ShouldAllocateInParallel()
customerIdsAfter.Should().BeEquivalentTo(expectedProjectIds);
}

static void AssertNext(KeyAllocator allocator, string collection, int expected)
static void AssertNext(KeyAllocator allocator, string collection, long expected)
{
allocator.NextId(collection).Should().Be(expected);
}

static async Task AssertNextAsync(KeyAllocator allocator, string collection, int expected)
static async Task AssertNextAsync(KeyAllocator allocator, string collection, long expected)
{
(await allocator.NextIdAsync(collection, CancellationToken.None)).Should().Be(expected);
}
Expand Down
4 changes: 2 additions & 2 deletions source/Nevermore/Mapping/IKeyAllocator.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ namespace Nevermore.Mapping
public interface IKeyAllocator
{
void Reset();
int NextId(string tableName);
ValueTask<int> NextIdAsync(string tableName, CancellationToken cancellationToken);
long NextId(string tableName);
ValueTask<long> NextIdAsync(string tableName, CancellationToken cancellationToken);
}
}
136 changes: 59 additions & 77 deletions source/Nevermore/Mapping/KeyAllocator.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,31 +21,26 @@ public KeyAllocator(IRelationalStore store, int blockSize)
}

public void Reset()
{
allocations.Clear();
}
=> allocations.Clear();

public int NextId(string tableName)
{
var allocation = allocations.GetOrAdd(tableName, _ => new Allocation(store, tableName, blockSize));
return allocation.Next();
}
Allocation GetAllocation(string tableName)
=> allocations.GetOrAdd(tableName, t => new Allocation(store, t, blockSize));

public async ValueTask<int> NextIdAsync(string tableName, CancellationToken cancellationToken)
{
var allocation = allocations.GetOrAdd(tableName, _ => new Allocation(store, tableName, blockSize));
return await allocation.NextAsync(cancellationToken).ConfigureAwait(false);
}
public long NextId(string tableName)
=> GetAllocation(tableName).Next();

public ValueTask<long> NextIdAsync(string tableName, CancellationToken cancellationToken)
=> GetAllocation(tableName).NextAsync(cancellationToken);

class Allocation
{
readonly IRelationalStore store;
readonly string collectionName;
readonly int blockSize;
readonly SemaphoreSlim sync = new(1, 1);
int blockStart;
int blockNext;
int blockFinish;
long blockStart;
long blockNext;
long blockFinish;

public Allocation(IRelationalStore store, string collectionName, int blockSize)
{
Expand All @@ -54,78 +49,68 @@ public Allocation(IRelationalStore store, string collectionName, int blockSize)
this.blockSize = blockSize;
}

public async ValueTask<int> NextAsync(CancellationToken cancellationToken)
public async ValueTask<long> NextAsync(CancellationToken cancellationToken)
{
using (await sync.LockAsync(cancellationToken))
using var releaseLock = await sync.LockAsync(cancellationToken);

async Task<long> GetNextMaxValue(CancellationToken ct)
{
async Task<int> GetNextMaxValue(CancellationToken ct)
using var transaction = await store.BeginWriteTransactionAsync(IsolationLevel.Serializable, name: $"{nameof(KeyAllocator)}.{nameof(Allocation)}.{nameof(GetNextMaxValue)}", cancellationToken: ct).ConfigureAwait(false);
var parameters = new CommandParameterValues
{
using var transaction = await store.BeginWriteTransactionAsync(IsolationLevel.Serializable, name: $"{nameof(KeyAllocator)}.{nameof(Allocation)}.{nameof(GetNextMaxValue)}", cancellationToken: ct).ConfigureAwait(false);
var parameters = new CommandParameterValues
{
{ "collectionName", collectionName },
{ "blockSize", blockSize }
};
parameters.CommandType = CommandType.StoredProcedure;

var result = await transaction.ExecuteScalarAsync<int>("GetNextKeyBlock", parameters, cancellationToken: ct).ConfigureAwait(false);
await transaction.CommitAsync(ct).ConfigureAwait(false);
return result;
}

async Task ExtendAllocation(CancellationToken ct)
{ "collectionName", collectionName },
{ "blockSize", blockSize }
};
parameters.CommandType = CommandType.StoredProcedure;

var result = await transaction.ExecuteScalarAsync<object>("GetNextKeyBlock", parameters, cancellationToken: ct).ConfigureAwait(false);
await transaction.CommitAsync(ct).ConfigureAwait(false);
// Older versions of the GetNextKeyBlock stored proc and KeyAllocation table might be using 32-bit ID's
// The type-check here lets us remain compatible with that while supporting 64-bit ID's as well
return result is int i ? i : (long)result;
}

if (blockNext == blockFinish)
{
await GetRetryPolicy().ExecuteActionAsync(async ct =>
{
var max = await GetNextMaxValue(ct).ConfigureAwait(false);
SetRange(max);
}

if (blockNext == blockFinish)
{
await GetRetryPolicy().ExecuteActionAsync(ExtendAllocation, cancellationToken).ConfigureAwait(false);
}

var result = blockNext;
blockNext++;

return result;
}, cancellationToken).ConfigureAwait(false);
}

return blockNext++;
}

public int Next()
public long Next()
{
using (sync.Lock())
using var releaseLock = sync.Lock();

long GetNextMaxValue()
{
int GetNextMaxValue()
using var transaction = store.BeginWriteTransaction(IsolationLevel.Serializable, name: $"{nameof(KeyAllocator)}.{nameof(Allocation)}.{nameof(GetNextMaxValue)}");
var parameters = new CommandParameterValues
{
using var transaction = store.BeginWriteTransaction(IsolationLevel.Serializable, name: $"{nameof(KeyAllocator)}.{nameof(Allocation)}.{nameof(GetNextMaxValue)}");
var parameters = new CommandParameterValues
{
{"collectionName", collectionName},
{"blockSize", blockSize}
};
parameters.CommandType = CommandType.StoredProcedure;

var result = transaction.ExecuteScalar<int>("GetNextKeyBlock", parameters);
transaction.Commit();
return result;
}

void ExtendAllocation()
{ "collectionName", collectionName },
{ "blockSize", blockSize }
};
parameters.CommandType = CommandType.StoredProcedure;

var result = transaction.ExecuteScalar<object>("GetNextKeyBlock", parameters);
transaction.Commit();
return result is int i ? i : (long)result; // 32/64-bit compatibility, see NextAsync() for explanation
}

if (blockNext == blockFinish)
{
GetRetryPolicy().ExecuteAction(() =>
{
var max = GetNextMaxValue();
SetRange(max);
}

if (blockNext == blockFinish)
{
GetRetryPolicy().ExecuteAction(ExtendAllocation);
}

var result = blockNext;
blockNext++;

return result;
});
}

return blockNext++;
}

RetryPolicy GetRetryPolicy()
Expand All @@ -138,18 +123,15 @@ RetryPolicy GetRetryPolicy()
/// <remarks>
/// Must only ever be executed while protected by the <see cref="sync"/> mutex!
/// </remarks>
void SetRange(int max)
void SetRange(long max)
{
var first = (max - blockSize) + 1;
blockStart = first;
blockNext = first;
blockFinish = max + 1;
}

public override string ToString()
{
return $"{blockStart} to {blockNext} (next: {blockFinish})";
}
public override string ToString() => $"{blockStart} to {blockNext} (next: {blockFinish})";
}
}
}
4 changes: 2 additions & 2 deletions source/Nevermore/Mapping/StringPrimaryKeyHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@ namespace Nevermore.Mapping
{
public sealed class StringPrimaryKeyHandler : AsyncPrimaryKeyHandler<string>
{
readonly Func<(string idPrefix, int key), string> format;
readonly Func<(string idPrefix, long key), string> format;

public StringPrimaryKeyHandler(string? idPrefix = null, Func<(string idPrefix, int key), string>? format = null)
public StringPrimaryKeyHandler(string? idPrefix = null, Func<(string idPrefix, long key), string>? format = null)
{
IdPrefix = idPrefix;
this.format = format ?? (x => $"{x.idPrefix}-{x.key}");
Expand Down
6 changes: 3 additions & 3 deletions source/Nevermore/UpgradeScripts/Script0001-KeyAllocation.sql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
CREATE TABLE dbo.KeyAllocation
(
CollectionName nvarchar(50) constraint PK_KeyAllocation_CollectionName primary key,
Allocated int not null
Allocated bigint not null
)
GO

Expand All @@ -18,7 +18,7 @@ CREATE PROCEDURE dbo.GetNextKeyBlock
AS
BEGIN
SET NOCOUNT ON
DECLARE @result int
DECLARE @result bigint

UPDATE KeyAllocation
SET @result = Allocated = (Allocated + @blockSize)
Expand All @@ -27,7 +27,7 @@ BEGIN
if (@@ROWCOUNT = 0)
begin
INSERT INTO KeyAllocation (CollectionName, Allocated) values (@collectionName, @blockSize)
SELECT @blockSize
SELECT CAST(@blockSize as bigint) -- type must be the same as @result
end

SELECT @result
Expand Down

0 comments on commit 80ba0ad

Please sign in to comment.