|
1 | 1 | using Microsoft.AspNetCore.Http;
|
2 | 2 | using Ocelot.Infrastructure;
|
| 3 | +using Ocelot.Middleware; |
3 | 4 | using Ocelot.Responses;
|
4 | 5 | using Ocelot.Values;
|
5 | 6 |
|
6 |
| -namespace Ocelot.LoadBalancer.LoadBalancers |
| 7 | +namespace Ocelot.LoadBalancer.LoadBalancers; |
| 8 | + |
| 9 | +public class CookieStickySessions : ILoadBalancer |
7 | 10 | {
|
8 |
| - public class CookieStickySessions : ILoadBalancer |
| 11 | + private readonly int _keyExpiryInMs; |
| 12 | + private readonly string _cookieName; |
| 13 | + private readonly ILoadBalancer _loadBalancer; |
| 14 | + private readonly IBus<StickySession> _bus; |
| 15 | + |
| 16 | + private static readonly object Locker = new(); |
| 17 | + private static readonly Dictionary<string, StickySession> Stored = new(); // TODO Inject instead of static sharing |
| 18 | + |
| 19 | + public CookieStickySessions(ILoadBalancer loadBalancer, string cookieName, int keyExpiryInMs, IBus<StickySession> bus) |
9 | 20 | {
|
10 |
| - private readonly int _keyExpiryInMs; |
11 |
| - private readonly string _key; |
12 |
| - private readonly ILoadBalancer _loadBalancer; |
13 |
| - private readonly ConcurrentDictionary<string, StickySession> _stored; |
14 |
| - private readonly IBus<StickySession> _bus; |
15 |
| - private readonly object _lock = new(); |
| 21 | + _bus = bus; |
| 22 | + _cookieName = cookieName; |
| 23 | + _keyExpiryInMs = keyExpiryInMs; |
| 24 | + _loadBalancer = loadBalancer; |
| 25 | + _bus.Subscribe(CheckExpiry); |
| 26 | + } |
16 | 27 |
|
17 |
| - public CookieStickySessions(ILoadBalancer loadBalancer, string key, int keyExpiryInMs, IBus<StickySession> bus) |
| 28 | + private void CheckExpiry(StickySession sticky) |
| 29 | + { |
| 30 | + // TODO Get test coverage for this |
| 31 | + lock (Locker) |
18 | 32 | {
|
19 |
| - _bus = bus; |
20 |
| - _key = key; |
21 |
| - _keyExpiryInMs = keyExpiryInMs; |
22 |
| - _loadBalancer = loadBalancer; |
23 |
| - _stored = new ConcurrentDictionary<string, StickySession>(); |
24 |
| - _bus.Subscribe(ss => |
| 33 | + if (!Stored.TryGetValue(sticky.Key, out var session) || session.Expiry >= DateTime.UtcNow) |
25 | 34 | {
|
26 |
| - //todo - get test coverage for this. |
27 |
| - if (_stored.TryGetValue(ss.Key, out var stickySession)) |
28 |
| - { |
29 |
| - lock (_lock) |
30 |
| - { |
31 |
| - if (stickySession.Expiry < DateTime.UtcNow) |
32 |
| - { |
33 |
| - _stored.TryRemove(stickySession.Key, out _); |
34 |
| - _loadBalancer.Release(stickySession.HostAndPort); |
35 |
| - } |
36 |
| - } |
37 |
| - } |
38 |
| - }); |
| 35 | + return; |
| 36 | + } |
| 37 | + |
| 38 | + Stored.Remove(session.Key); |
| 39 | + _loadBalancer.Release(session.HostAndPort); |
39 | 40 | }
|
| 41 | + } |
40 | 42 |
|
41 |
| - public async Task<Response<ServiceHostAndPort>> Lease(HttpContext httpContext) |
| 43 | + public Task<Response<ServiceHostAndPort>> Lease(HttpContext httpContext) |
| 44 | + { |
| 45 | + var route = httpContext.Items.DownstreamRoute(); |
| 46 | + var serviceName = route.LoadBalancerKey; |
| 47 | + var cookie = httpContext.Request.Cookies[_cookieName]; |
| 48 | + var key = $"{serviceName}:{cookie}"; // strong key name because of static store |
| 49 | + lock (Locker) |
42 | 50 | {
|
43 |
| - var key = httpContext.Request.Cookies[_key]; |
44 |
| - |
45 |
| - lock (_lock) |
| 51 | + if (!string.IsNullOrEmpty(key) && Stored.TryGetValue(key, out StickySession cached)) |
46 | 52 | {
|
47 |
| - if (!string.IsNullOrEmpty(key) && _stored.ContainsKey(key)) |
48 |
| - { |
49 |
| - var cached = _stored[key]; |
50 |
| - |
51 |
| - var updated = new StickySession(cached.HostAndPort, DateTime.UtcNow.AddMilliseconds(_keyExpiryInMs), key); |
52 |
| - |
53 |
| - _stored[key] = updated; |
54 |
| - |
55 |
| - _bus.Publish(updated, _keyExpiryInMs); |
56 |
| - |
57 |
| - return new OkResponse<ServiceHostAndPort>(updated.HostAndPort); |
58 |
| - } |
| 53 | + var updated = new StickySession(cached.HostAndPort, DateTime.UtcNow.AddMilliseconds(_keyExpiryInMs), key); |
| 54 | + Update(key, updated); |
| 55 | + return Task.FromResult<Response<ServiceHostAndPort>>(new OkResponse<ServiceHostAndPort>(updated.HostAndPort)); |
59 | 56 | }
|
60 | 57 |
|
61 |
| - var next = await _loadBalancer.Lease(httpContext); |
62 |
| - |
| 58 | + // There is no value in the store, so lease it now! |
| 59 | + var next = _loadBalancer.Lease(httpContext).GetAwaiter().GetResult(); // unfortunately the operation must be synchronous |
63 | 60 | if (next.IsError)
|
64 | 61 | {
|
65 |
| - return new ErrorResponse<ServiceHostAndPort>(next.Errors); |
66 |
| - } |
67 |
| - |
68 |
| - lock (_lock) |
69 |
| - { |
70 |
| - if (!string.IsNullOrEmpty(key) && !_stored.ContainsKey(key)) |
71 |
| - { |
72 |
| - var ss = new StickySession(next.Data, DateTime.UtcNow.AddMilliseconds(_keyExpiryInMs), key); |
73 |
| - _stored[key] = ss; |
74 |
| - _bus.Publish(ss, _keyExpiryInMs); |
75 |
| - } |
| 62 | + return Task.FromResult<Response<ServiceHostAndPort>>(new ErrorResponse<ServiceHostAndPort>(next.Errors)); |
76 | 63 | }
|
77 | 64 |
|
78 |
| - return new OkResponse<ServiceHostAndPort>(next.Data); |
| 65 | + var ss = new StickySession(next.Data, DateTime.UtcNow.AddMilliseconds(_keyExpiryInMs), key); |
| 66 | + Update(key, ss); |
| 67 | + return Task.FromResult<Response<ServiceHostAndPort>>(new OkResponse<ServiceHostAndPort>(next.Data)); |
79 | 68 | }
|
| 69 | + } |
80 | 70 |
|
81 |
| - public void Release(ServiceHostAndPort hostAndPort) |
| 71 | + protected void Update(string key, StickySession value) |
| 72 | + { |
| 73 | + lock (Locker) |
82 | 74 | {
|
| 75 | + Stored[key] = value; |
| 76 | + _bus.Publish(value, _keyExpiryInMs); |
83 | 77 | }
|
84 | 78 | }
|
| 79 | + |
| 80 | + public void Release(ServiceHostAndPort hostAndPort) |
| 81 | + { |
| 82 | + } |
85 | 83 | }
|
0 commit comments