Skip to content

Commit 02439f6

Browse files
tocsofttocsoft
committed
improve reconnect issue when multiple bridges running
1 parent ce45b60 commit 02439f6

File tree

6 files changed

+149
-18
lines changed

6 files changed

+149
-18
lines changed

‎KubeConnect/Args.cs

+4
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,9 @@ public Args(string[] args)
8686
case "working-directory":
8787
WorkingDirectory = Path.GetFullPath(argNext);
8888
break;
89+
case "trace-logs":
90+
EnableTraceLogs = true;
91+
break;
8992
default:
9093
// capture unknown args
9194
if (!string.IsNullOrWhiteSpace(arg))
@@ -180,6 +183,7 @@ private void ProcessArgs(string[] args, Func<string, string, string, bool> proce
180183
}
181184
}
182185

186+
public bool EnableTraceLogs { get; set; } = false;
183187
public string? Namespace { get; set; }
184188
public string? KubeconfigFile { get; private set; }
185189
public string? Context { get; private set; }

‎KubeConnect/Hubs/BridgeHub.cs

+9-3
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,19 @@
33
using System.Linq;
44
using System.Threading.Tasks;
55
using Microsoft.AspNetCore.SignalR;
6+
using Microsoft.Extensions.Hosting;
67

78
namespace KubeConnect.Hubs
89
{
910
public class BridgeHub : Hub
1011
{
1112
private readonly ServiceManager serviceManager;
13+
private readonly IHostApplicationLifetime hostApplicationLifetime;
1214

13-
public BridgeHub(ServiceManager serviceManager)
15+
public BridgeHub(ServiceManager serviceManager, IHostApplicationLifetime hostApplicationLifetime)
1416
{
1517
this.serviceManager = serviceManager;
18+
this.hostApplicationLifetime = hostApplicationLifetime;
1619
}
1720

1821
public override Task OnConnectedAsync()
@@ -22,7 +25,10 @@ public override Task OnConnectedAsync()
2225

2326
public override async Task OnDisconnectedAsync(Exception? exception)
2427
{
25-
await serviceManager.Release(Context.ConnectionId);
28+
if (!hostApplicationLifetime.ApplicationStopping.IsCancellationRequested)
29+
{
30+
await serviceManager.Release(Context.ConnectionId);
31+
}
2632

2733
// when disconnecting kill bridge
2834
await base.OnDisconnectedAsync(exception);
@@ -53,7 +59,7 @@ public async Task StartServiceBridge(string serviceName, Dictionary<int, int> po
5359
{
5460
throw new Exception($"Unable to find the service '{serviceName}'");
5561
}
56-
62+
5763
// bridge logging should be a write to the group `$"Bridge:{service.ServiceName}:{service.Namespace}"`
5864
await serviceManager.Intercept(service, ports.Select(x => (x.Key, x.Value)).ToList(), this.Context.ConnectionId, this.Clients.Client(this.Context.ConnectionId));
5965
}

‎KubeConnect/IConsoleLogProvider.cs

+13-3
Original file line numberDiff line numberDiff line change
@@ -6,28 +6,33 @@ namespace KubeConnect
66
public partial class IConsoleLogProvider : ILoggerProvider
77
{
88
private readonly IConsole console;
9+
private readonly Args args;
910

10-
public IConsoleLogProvider(IConsole console)
11+
public IConsoleLogProvider(IConsole console, Args args)
1112
{
1213
this.console = console;
14+
this.args = args;
1315
}
1416

1517
public ILogger CreateLogger(string categoryName)
16-
=> new IConsoleLoggger(categoryName, console);
18+
=> new IConsoleLoggger(categoryName, console, this.args.EnableTraceLogs);
1719

1820
public void Dispose()
1921
{
2022
}
2123

2224
public class IConsoleLoggger : ILogger, IDisposable
2325
{
24-
public IConsoleLoggger(string category, IConsole console)
26+
public IConsoleLoggger(string category, IConsole console, bool traceLogs)
2527
{
28+
this.traceLogs = traceLogs;
2629
Category = category;
2730
this.console = console;
2831
this.internalLogs = category.StartsWith("KubeConnect");
2932
}
3033

34+
private readonly bool traceLogs;
35+
3136
public string Category { get; }
3237

3338
private bool internalLogs;
@@ -45,6 +50,11 @@ public void Dispose()
4550

4651
public bool IsEnabled(LogLevel logLevel)
4752
{
53+
if (traceLogs)
54+
{
55+
return true;
56+
}
57+
4858
if (internalLogs)
4959
{
5060
return logLevel >= LogLevel.Debug;

‎KubeConnect/PortForwarding/PortForwardingConnectionHandler.cs

+4-3
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ public override async Task OnConnectedAsync(ConnectionContext connection)
2828
{
2929
var input = connection.Transport.Input;
3030
var output = connection.Transport.Output;
31-
31+
3232
var binding = connection.Features.Get<PortBinding>();
3333
if (binding == null)
3434
{
@@ -43,6 +43,7 @@ public override async Task OnConnectedAsync(ConnectionContext connection)
4343
if (pod == null)
4444
{
4545
connection.Abort();
46+
return;
4647
}
4748

4849
logger.LogInformation("[{ConnectionID}] Opening connection for {ServiceName}:{ServicePort} to {PodName}:{PodPort}", connection.ConnectionId, binding.Name, (connection.LocalEndPoint as IPEndPoint)?.Port, pod.Name(), binding.TargetPort);
@@ -67,7 +68,7 @@ async Task PushToPod()
6768
{
6869
while (true)
6970
{
70-
var result = await input.ReadAsync();
71+
var result = await input.ReadAsync(connection.ConnectionClosed);
7172

7273
foreach (var buffer in result.Buffer)
7374
{
@@ -82,7 +83,7 @@ async Task PushToClient()
8283
{
8384
while (true)
8485
{
85-
var result = await podOutput.ReadAsync();
86+
var result = await podOutput.ReadAsync(connection.ConnectionClosed);
8687

8788
foreach (var buffer in result.Buffer)
8889
{

‎KubeConnect/ServiceDetails.cs

+48
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
using System.Collections.Generic;
33
using System.Linq;
44
using System.Net;
5+
using System.Text;
6+
using System.Threading;
7+
using System.Threading.Tasks;
58
using Microsoft.AspNetCore.SignalR;
69

710
namespace KubeConnect
@@ -17,6 +20,51 @@ public class BridgeDetails
1720
public IReadOnlyList<(int remotePort, int localPort)> BridgedPorts { get; init; } = Array.Empty<(int, int)>();
1821

1922
public IClientProxy Client { get; internal set; }
23+
24+
object locker = new object();
25+
StringBuilder builder = new StringBuilder();
26+
StringBuilder builderAlt = new StringBuilder();
27+
internal void Log(string msg)
28+
{
29+
lock (locker)
30+
{
31+
if (builder.Length > 0)
32+
{
33+
builder.AppendLine();
34+
}
35+
builder.Append(msg);
36+
}
37+
38+
}
39+
SemaphoreSlim semaphore = new SemaphoreSlim(1);
40+
public async Task FlushLogs()
41+
{
42+
StringBuilder RotateLogBuffers()
43+
{
44+
lock (locker)
45+
{
46+
var current = builder;
47+
builder = builderAlt;
48+
builderAlt = current;
49+
return current;
50+
}
51+
}
52+
53+
await semaphore.WaitAsync();
54+
try
55+
{
56+
var currentBuffers = RotateLogBuffers();
57+
if (currentBuffers.Length > 0)
58+
{
59+
await Client.SendAsync("log", currentBuffers.ToString(), false);
60+
currentBuffers.Clear();
61+
}
62+
}
63+
finally
64+
{
65+
semaphore.Release();
66+
}
67+
}
2068
}
2169

2270
public class ServiceDetails

‎KubeConnect/ServiceManager.cs

+71-9
Original file line numberDiff line numberDiff line change
@@ -201,12 +201,52 @@ private async Task EnableDeployment(V1Deployment deployment)
201201
}
202202
}
203203

204+
Task logWriter;
205+
private void EnsureLogWriterRunning()
206+
{
207+
if (!BridgedServices.Any())
208+
{
209+
return;
210+
}
211+
212+
if (logWriter == null)
213+
{
214+
logWriter = Task.Run(async () =>
215+
{
216+
while (true)
217+
{
218+
try
219+
{
220+
foreach (var s in BridgedServices.ToList())
221+
{
222+
try
223+
{
224+
await s.FlushLogs();
225+
}
226+
catch
227+
{
228+
console.WriteErrorLine($"Failed to flush logs for {s.ServiceName} bridge");
229+
}
230+
}
231+
}
232+
catch
233+
{
234+
}
235+
await Task.Delay(250);
236+
}
237+
});
238+
}
239+
}
240+
204241
private async Task StartSshForward(BridgeDetails bridgeDetails, ServiceDetails service)
205242
{
243+
EnsureLogWriterRunning();
244+
206245
var logger = (string msg) =>
207246
{
208247
console.WriteLine(msg);
209-
_ = bridgeDetails.Client.SendAsync("log", msg, false);
248+
249+
bridgeDetails.Log(msg);
210250
};
211251

212252
var pod = await FindInterceptionPod(service);
@@ -221,6 +261,7 @@ private async Task StartSshForward(BridgeDetails bridgeDetails, ServiceDetails s
221261
ContainerPort = X.remotePort,
222262
Name = $"port-{X.remotePort}"
223263
}).ToList();
264+
224265
ports.Add(new V1ContainerPort
225266
{
226267
ContainerPort = 2222,
@@ -279,7 +320,14 @@ private async Task StartSshForward(BridgeDetails bridgeDetails, ServiceDetails s
279320
client.AddForwardedPort(port);
280321
port.RequestReceived += (object? sender, Renci.SshNet.Common.PortForwardEventArgs e) =>
281322
{
282-
logger($"Traffic redirected from {service.ServiceName}:{mappings.remotePort} to localhost:{mappings.localPort}");
323+
try
324+
{
325+
logger($"Traffic redirected from {service.ServiceName}:{mappings.remotePort} to localhost:{mappings.localPort}");
326+
}
327+
catch
328+
{
329+
330+
}
283331
};
284332
port.Start();
285333
}
@@ -368,12 +416,21 @@ public async Task Intercept(ServiceDetails service, IReadOnlyList<(int remotePor
368416
await StartSshForward(bridgeDetails, service);
369417
}
370418

419+
private SemaphoreSlim semaphore = new SemaphoreSlim(1);
371420
public async Task Release(string connectionId)
372421
{
373-
var services = this.BridgedServices.Where(x => x.ConnectionId == connectionId);
374-
foreach (var service in services)
422+
await semaphore.WaitAsync();
423+
try
424+
{
425+
var services = this.BridgedServices.Where(x => x.ConnectionId == connectionId).ToList();
426+
foreach (var service in services)
427+
{
428+
await Release(service);
429+
}
430+
}
431+
finally
375432
{
376-
await Release(service);
433+
semaphore.Release();
377434
}
378435
}
379436

@@ -390,11 +447,12 @@ public Task Release(BridgeDetails bridgeDetails)
390447

391448
public async Task Release(ServiceDetails service)
392449
{
393-
var details = BridgedServices.FirstOrDefault(x => x.ServiceName.Equals(service.ServiceName, StringComparison.OrdinalIgnoreCase) && x.Namespace.Equals(service.Namespace, StringComparison.OrdinalIgnoreCase));
394-
if (details != null)
450+
var details = BridgedServices.Where(x => x.ServiceName.Equals(service.ServiceName, StringComparison.OrdinalIgnoreCase) && x.Namespace.Equals(service.Namespace, StringComparison.OrdinalIgnoreCase)).ToList();
451+
452+
if (details.Any())
395453
{
396-
BridgedServices = BridgedServices.Where(x => x != details).ToArray();
397-
OnBridgedServicesChanged?.Invoke(this, BridgedServices);
454+
BridgedServices = BridgedServices.Where(x => !details.Contains(x)).ToArray();
455+
console.WriteLine($"Shutting down bridge for {service.ServiceName}");
398456
}
399457

400458
var pod = await FindInterceptionPod(service);
@@ -408,6 +466,10 @@ public async Task Release(ServiceDetails service)
408466
{
409467
await EnableDeployment(dep);
410468
}
469+
if (details != null)
470+
{
471+
OnBridgedServicesChanged?.Invoke(this, BridgedServices);
472+
}
411473
}
412474

413475
public async Task ReleaseAll()

0 commit comments

Comments
 (0)