Skip to content

Commit 24c035c

Browse files
LucHeartLucHeart
committed
rework lcg assigning logic, add integration tests for it
1 parent 0fb4499 commit 24c035c

File tree

8 files changed

+189
-12
lines changed

8 files changed

+189
-12
lines changed

‎API.IntegrationTests/API.IntegrationTests.csproj‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
<PackageReference Include="Testcontainers.Redis" />
1616
<PackageReference Include="TUnit" />
1717
</ItemGroup>
18+
<ItemGroup>
19+
<Folder Include="Seeders\" />
20+
</ItemGroup>
1821

1922
<!-- Git stuff -->
2023
<Target Name="SetHash" AfterTargets="InitializeSourceControlInformation">

‎API.IntegrationTests/Tests/LcgAssignmentTests.cs‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
using System.Net;
2+
using System.Net.Http.Json;
3+
using Microsoft.AspNetCore.Hosting;
4+
using Microsoft.EntityFrameworkCore;
5+
using Microsoft.Extensions.DependencyInjection;
6+
using Newtonsoft.Json;
7+
using OpenShock.API.Models.Response;
8+
using OpenShock.Common.Models;
9+
using OpenShock.Common.OpenShockDb;
10+
using OpenShock.Common.Redis;
11+
using OpenShock.Common.Utils;
12+
using Redis.OM.Contracts;
13+
using JsonSerializer = System.Text.Json.JsonSerializer;
14+
15+
namespace OpenShock.API.IntegrationTests.Tests;
16+
17+
public sealed class LcgAssignmentTests
18+
{
19+
[ClassDataSource<WebApplicationFactory>(Shared = SharedType.PerTestSession)]
20+
public required WebApplicationFactory WebApplicationFactory { get; init; }
21+
22+
private static readonly Guid UserId = Guid.Parse("11111111-1111-1111-1111-111111111111");
23+
private static readonly Guid HubId = Guid.Parse("11111111-1111-1111-1111-111111111111");
24+
private const string HubToken = "test";
25+
26+
[Before(Test)]
27+
public async Task Setup()
28+
{
29+
await using var context = WebApplicationFactory.Services.CreateAsyncScope();
30+
var db = context.ServiceProvider.GetRequiredService<OpenShockContext>();
31+
32+
var user = new User
33+
{
34+
Id = UserId,
35+
Name = "TestUser",
36+
Email = "test@test.org",
37+
PasswordHash = HashingUtils.HashPassword("password"),
38+
CreatedAt = DateTime.UtcNow,
39+
ActivatedAt = DateTime.UtcNow
40+
};
41+
42+
db.Users.Add(user);
43+
44+
var hub = new Device
45+
{
46+
Id = HubId,
47+
Name = "TestHub",
48+
OwnerId = UserId,
49+
Token = HubToken,
50+
CreatedAt = DateTime.UtcNow
51+
};
52+
53+
db.Devices.Add(hub);
54+
55+
await db.SaveChangesAsync();
56+
}
57+
58+
[After(Test)]
59+
public async Task Teardown()
60+
{
61+
await using var context = WebApplicationFactory.Services.CreateAsyncScope();
62+
var db = context.ServiceProvider.GetRequiredService<OpenShockContext>();
63+
await db.Devices.Where(x => x.Id == HubId).ExecuteDeleteAsync();
64+
await db.Users.Where(x => x.Id == UserId).ExecuteDeleteAsync();
65+
66+
var redisConnectionProvider = context.ServiceProvider.GetRequiredService<IRedisConnectionProvider>();
67+
var webHostEnvironment = context.ServiceProvider.GetRequiredService<IWebHostEnvironment>();
68+
var lcgNodesCollection = redisConnectionProvider.RedisCollection<LcgNode>(false);
69+
70+
var allLcg = await lcgNodesCollection.ToArrayAsync();
71+
await lcgNodesCollection.DeleteAsync(allLcg);
72+
}
73+
74+
[Test]
75+
[NotInParallel]
76+
[Arguments("US", "us1.example.com", new[] { "US|us1.example.com", "DE|de1.example.com", "AS|as1.example.com" })]
77+
[Arguments("DE", "de1.example.com", new[] { "US|us1.example.com", "DE|de1.example.com", "AS|as1.example.com" })]
78+
[Arguments("CA", "us1.example.com", new[] { "US|us1.example.com", "DE|de1.example.com", "AS|as1.example.com" })]
79+
[Arguments("CA", "us1.example.com", new[] { "US|us1.example.com", "DE|de1.example.com", "AS|as1.example.com" })]
80+
[Arguments("AT", "de1.example.com", new[] { "US|us1.example.com", "DE|de1.example.com", "AS|as1.example.com" })]
81+
[Arguments("FR", "de1.example.com", new[] { "US|us1.example.com", "DE|de1.example.com", "AS|as1.example.com" })]
82+
public async Task GetLcgAssignment(string requesterCountry, string expectedHost, string[] availableGateways)
83+
{
84+
using var client = WebApplicationFactory.CreateClient();
85+
86+
await using var context = WebApplicationFactory.Services.CreateAsyncScope();
87+
var redisConnectionProvider = context.ServiceProvider.GetRequiredService<IRedisConnectionProvider>();
88+
var webHostEnvironment = context.ServiceProvider.GetRequiredService<IWebHostEnvironment>();
89+
var lcgNodesCollection = redisConnectionProvider.RedisCollection<LcgNode>(false);
90+
91+
var testGateways = availableGateways.Select(x =>
92+
{
93+
var split = x.Split('|');
94+
if (split.Length != 2)
95+
throw new ArgumentException("Invalid gateway format");
96+
97+
return new LcgNode
98+
{
99+
Country = split[0],
100+
Fqdn = split[1],
101+
Load = 0,
102+
Environment = webHostEnvironment.EnvironmentName
103+
};
104+
});
105+
106+
await lcgNodesCollection.InsertAsync(testGateways);
107+
108+
var httpRequest = new HttpRequestMessage(HttpMethod.Get, "/2/device/assignLCG?version=1");
109+
httpRequest.Headers.Add("Device-Token", HubToken);
110+
httpRequest.Headers.Add("CF-IPCountry", requesterCountry);
111+
var response = await client.SendAsync(httpRequest);
112+
113+
await Assert.That(response.StatusCode).IsEqualTo(HttpStatusCode.OK);
114+
115+
var mediaType = response.Content.Headers.ContentType?.MediaType;
116+
await Assert.That(mediaType).IsEqualTo("application/json");
117+
118+
var data = await response.Content.ReadFromJsonAsync<LcgNodeResponseV2>();
119+
await Assert.That(data).IsNotNull();
120+
await Assert.That(data.Host).IsEqualTo(expectedHost);
121+
}
122+
123+
124+
}

‎API.IntegrationTests/WebApplicationFactory.cs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
using Microsoft.Extensions.Http;
66
using OpenShock.API.IntegrationTests.Docker;
77
using OpenShock.API.IntegrationTests.HttpMessageHandlers;
8+
using Serilog;
9+
using Serilog.Events;
810
using TUnit.Core.Interfaces;
911

1012
namespace OpenShock.API.IntegrationTests;
@@ -67,12 +69,20 @@ protected override void ConfigureWebHost(IWebHostBuilder builder)
6769
{ "OPENSHOCK__LCG__FQDN", "de1-gateway.my-openshock-instance.net" },
6870
{ "OPENSHOCK__LCG__COUNTRYCODE", "DE" }
6971
};
70-
72+
7173
foreach (var envVar in environmentVariables)
7274
{
7375
Environment.SetEnvironmentVariable(envVar.Key, envVar.Value);
7476
}
75-
77+
78+
builder.ConfigureServices(services =>
79+
{
80+
services.AddSerilog(configuration =>
81+
{
82+
configuration.WriteTo.Console(LogEventLevel.Warning);
83+
});
84+
});
85+
7686
builder.ConfigureTestServices(services =>
7787
{
7888
services.AddTransient<HttpMessageHandlerBuilder, InterceptedHttpMessageHandlerBuilder>();

‎API/Controller/Device/AssignLCG.cs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
using OpenShock.API.Models.Response;
33
using System.Net.Mime;
44
using Asp.Versioning;
5+
using OpenShock.API.Services.LCGNodeProvisioner;
56
using OpenShock.Common.Errors;
67
using OpenShock.Common.Problems;
7-
using OpenShock.Common.Services.LCGNodeProvisioner;
88
using OpenShock.Common.Utils;
99
using OpenShock.Common.Models;
1010

@@ -28,6 +28,7 @@ public async Task<IActionResult> GetLiveControlGateway([FromServices] ILCGNodePr
2828
_logger.LogWarning("CF-IPCountry header could not be parsed into a alpha2 country code");
2929
}
3030

31+
try {
3132
var closestNode = await geoLocation.GetOptimalNodeAsync(countryCode);
3233
if (closestNode is null) return Problem(AssignLcgError.NoLcgNodesAvailable);
3334

@@ -36,5 +37,11 @@ public async Task<IActionResult> GetLiveControlGateway([FromServices] ILCGNodePr
3637
Fqdn = closestNode.Fqdn,
3738
Country = closestNode.Country
3839
});
40+
}
41+
catch (Exception ex)
42+
{
43+
_logger.LogError(ex, "Error while assigning LCG node");
44+
return Problem(AssignLcgError.NoLcgNodesAvailable);
45+
}
3946
}
4047
}

‎API/Controller/Device/AssignLCGV2.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
using OpenShock.API.Models.Response;
33
using System.Net.Mime;
44
using Asp.Versioning;
5+
using OpenShock.API.Services.LCGNodeProvisioner;
56
using OpenShock.Common.Errors;
67
using OpenShock.Common.Problems;
7-
using OpenShock.Common.Services.LCGNodeProvisioner;
88
using OpenShock.Common.Utils;
99

1010
namespace OpenShock.API.Controller.Device;

‎API/Program.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using OpenShock.API.Services.Account;
66
using OpenShock.API.Services.DeviceUpdate;
77
using OpenShock.API.Services.Email;
8+
using OpenShock.API.Services.LCGNodeProvisioner;
89
using OpenShock.API.Services.OAuthConnection;
910
using OpenShock.API.Services.Turnstile;
1011
using OpenShock.API.Services.UserService;
@@ -13,7 +14,6 @@
1314
using OpenShock.Common.Hubs;
1415
using OpenShock.Common.Services;
1516
using OpenShock.Common.Services.Device;
16-
using OpenShock.Common.Services.LCGNodeProvisioner;
1717
using OpenShock.Common.Services.Ota;
1818
using OpenShock.Common.Swagger;
1919
using Serilog;

Common/Services/LCGNodeProvisioner/ILCGNodeProvisioner.cs renamed to API/Services/LCGNodeProvisioner/ILCGNodeProvisioner.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
using OpenShock.Common.Geo;
22
using OpenShock.Common.Redis;
33

4-
namespace OpenShock.Common.Services.LCGNodeProvisioner;
4+
namespace OpenShock.API.Services.LCGNodeProvisioner;
55

66
public interface ILCGNodeProvisioner
77
{

Common/Services/LCGNodeProvisioner/LCGNodeProvisioner.cs renamed to API/Services/LCGNodeProvisioner/LCGNodeProvisioner.cs

Lines changed: 39 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
using Redis.OM.Contracts;
77
using Redis.OM.Searching;
88

9-
namespace OpenShock.Common.Services.LCGNodeProvisioner;
9+
namespace OpenShock.API.Services.LCGNodeProvisioner;
1010

1111
public sealed class LCGNodeProvisioner : ILCGNodeProvisioner
1212
{
@@ -41,16 +41,49 @@ public LCGNodeProvisioner(IRedisConnectionProvider redisConnectionProvider, IWeb
4141
return await GetOptimalNodeAsync();
4242
}
4343

44+
// Load all nodes for our environment
4445
var nodes = await _lcgNodes
4546
.Where(x => x.Environment == _environmentName)
4647
.ToArrayAsync();
48+
49+
if(nodes.Length < 1)
50+
{
51+
_logger.LogWarning("No LCG nodes available after filtering by environment [{Environment}]!", _environmentName);
52+
return null;
53+
}
54+
55+
// Precompute distances
56+
var withDistances = nodes
57+
.Select(x => new
58+
{
59+
Node = x,
60+
Distance = DistanceLookup.TryGetDistanceBetween(x.Country, countryCode, out var dist) ? dist : Distance.DistanceToAndromedaGalaxyInKm
61+
})
62+
.ToArray();
63+
64+
// 1) Find the closest region (min distance)
65+
var minDistance = withDistances.Min(x => x.Distance);
4766

48-
var node = nodes
49-
.OrderBy(x => DistanceLookup.TryGetDistanceBetween(x.Country, countryCode, out float distance) ? distance : Distance.DistanceToAndromedaGalaxyInKm) // Just a large number :3
50-
.ThenBy(x => x.Load)
51-
.FirstOrDefault();
67+
var closestRegionNodes = withDistances
68+
.Where(x => Math.Abs(x.Distance - minDistance) < 1)
69+
.Select(x => x.Node)
70+
.ToArray();
5271

53-
if (node is null) _logger.LogWarning("No LCG nodes available!");
72+
// 2) Among those, find minimal load
73+
var minLoad = closestRegionNodes.Min(x => x.Load);
74+
var loadCandidates = closestRegionNodes
75+
.Where(x => x.Load == minLoad)
76+
.ToArray();
77+
78+
if(loadCandidates.Length < 1)
79+
{
80+
_logger.LogWarning("No LCG nodes available after filtering by geo location and load!");
81+
return null;
82+
}
83+
84+
// 3) Randomly pick one of the tied nodes
85+
var node = loadCandidates[Random.Shared.Next(loadCandidates.Length)];
86+
5487
if (_logger.IsEnabled(LogLevel.Debug)) _logger.LogDebug("LCG node provisioned: {@LcgNode}", node);
5588

5689
return node;

0 commit comments

Comments
 (0)