This is an application I wrote that allows multiple TCP clients to share a single TCP connection to a remote server (hosted project, and a demo). Traffic generated by the server is forwarded to all clients. Traffic generated from any client is forwarded to the server. I don't concern the case where client traffics are interleaved, as in my domain traffics are spontaneous and it's affordable to have interleaved, invalid data.
The Goal
The goal for the code review is to examine if I am idiomatic on .NET socket programming and TPL programming, if I correctly handled task executions and cancellations, and if there is any performance concerns, e.g., unnecessarily blocking the socket communication, etc.
High Level Design
The project is hosted here. Feel free to read the README. Implementation wise, I designed it to be as responsive as possible - no socket connection is unnecessarily blocked.
- There is an outbound queue. Each outbound packet from each client is put into this queue.
- There is an inbound queue for each client. Each inbound packet is put into each of these queues.
For each client, there're two async tasks, one for dequeueing from the client's inbound queue (this can block when the queue is empty, which is desirable) and write the data to the client socket; one for reading data from the client socket and enqueueing to the outbound queue.
Similarly there're two async tasks for the remote connection, one for dequeueing the outbound queue and write to the remote socket; one for reading from the remote socket and enqueueing to each client's inbound queue.
There's another async task accepting commands from stdin for connecting/disconnecting to the server, dumping diagnostics info, etc. But that's not the main focus.
The entire code base is hosted here, but the most relevant classes are listed below.
Program.cs
namespace Multiplexer
{
using System.Threading.Tasks;
class Program
{
static void Main(string[] args)
{
var glob = new Global();
var ctrl = new ControlChannel(glob);
var clientServer = new ClientServer(3333, glob);
Task.WaitAny(
Task.Run(() => ctrl.Run(), glob.CancellationToken),
Task.Run(() => clientServer.Run(), glob.CancellationToken));
}
}
}
ClientServer.cs
namespace Multiplexer
{
using System;
using System.Net;
using System.Net.Sockets;
using System.Threading.Tasks;
/// <summary>
/// Listens to connection requests and manages client connections
/// </summary>
class ClientServer
{
/// <summary>
/// Port to listen on for clients connections
/// </summary>
readonly int port;
Global glob;
public ClientServer(
int port,
Global glob)
{
this.port = port;
this.glob = glob;
}
/// <summary>
/// Continuously listening for client connection requests
/// </summary>
public async Task Run()
{
var localserver = new TcpListener(IPAddress.Parse("127.0.0.1"), port);
localserver.Start();
while (true)
{
Console.WriteLine("Waiting for clients to connect...");
// AcceptTcpClientAsync() does not accept a cancellation token. But it's OK since
// in no case would I want the client listener loop to stop running during the entire
// multiplexer lifecycle. For the sake of completeness, if it is necessary to cancel
// this operation, one could use CancellationToken.Register(localserver.Stop).
// See: http://stackoverflow.com/a/30856169/695964
var client = await localserver.AcceptTcpClientAsync();
var clientWrapper = new Client(client, glob.CancellationToken, Upload);
Console.WriteLine($"Client connected: {clientWrapper}");
// Register client
glob.Clients[clientWrapper] = 0;
// Unregister client when it is terminating
clientWrapper.OnClose = () =>
{
Console.WriteLine($"Removing client from clients list: {clientWrapper}");
byte c;
glob.Clients.TryRemove(clientWrapper, out c);
};
// Start the client. This is fire-and-forget. We don't want to await on it. I
// t's OK because Start() has necessary logic to handle client termination and disposal.
var tsk = clientWrapper.Start();
}
}
/// <summary>
/// Implementation of upload delegate to be called when there's data to upload to remote server
/// </summary>
/// <param name="data">the outbound data</param>
void Upload(byte[] data)
{
// Do not enqueue data if remote is not connected (drop it)
if (glob.Remote.Connected)
{
glob.UploadQueue.TryAdd(data);
}
}
}
}
ControlChannel.cs
namespace Multiplexer
{
using Newtonsoft.Json;
using System;
using System.Linq;
using System.Net.Sockets;
using System.Threading.Tasks;
/// <summary>
/// Class to handle multiplexer management commands.
/// </summary>
class ControlChannel
{
Global glob;
public ControlChannel(Global glob)
{
this.glob = glob;
}
public void Run()
{
while (true)
{
var line = Console.ReadLine();
var toks = line.Split();
switch (toks[0])
{
// Connecting to remote server
// connect <remote_server> <port>
case "connect":
if (glob.Remote.Connected)
{
Console.WriteLine("Error: Remote is connected. Try disconnect first");
}
else
{
// Reset the upload queue so stale outbound data is not uploaded to the new
// connection
glob.ResetUploadQueue();
// This is a fire-and-forget task. It is the responsibility of
// the task to properly handle resource clean up.
Task.Run(() => StartServer(toks[1], int.Parse(toks[2])));
}
break;
// Disconnecting:
// disconnect
case "disconnect":
if (glob.Remote.Connected)
{
Console.WriteLine("Disconnecting");
// cancel the global cancellation token. This would disconnect the server and all
// clients. It is desirable to disconnect the clients to maintain the equivelency
// when a client is directly connected to the server.
glob.Cancel();
}
else
{
Console.WriteLine("Not connected");
}
break;
// Dump multiplexer status
// info|stats
case "info":
case "stats":
DumpStats();
break;
// Quit the multiplexer application
case "quit":
Console.WriteLine("Exiting...");
glob.Cancel();
return; // terminate the control channel loop
default:
Console.WriteLine("Unknown command: " + line);
break;
}
}
}
/// <summary>
/// Dump multiplxer info
/// </summary>
private void DumpStats()
{
var info = new
{
Remote = glob.Remote,
Clients = glob.Clients.Select(c => c.ToString()).ToArray(),
UploadQueue = glob.UploadQueue.Select(msg => System.Text.Encoding.UTF8.GetString(msg)).ToArray(),
};
Console.WriteLine(JsonConvert.SerializeObject(info, Formatting.Indented));
}
/// <summary>
/// Start the remote connection
/// </summary>
/// <param name="hostname">remote hostname</param>
/// <param name="port">remote port</param>
async Task StartServer(string hostname, int port)
{
Console.WriteLine($"Connecting to {hostname}:{port}");
var remote = new TcpClient(hostname, port);
var server = new Remote(remote, glob.UploadQueue, glob.CancellationToken, /* receive: */ data =>
{
// Implementation of receive() is to put inbound data to each of the client queues.
// Note that this is non-blocking. If any queue is full, the data is dropped from that
// queue.
foreach (var client in glob.Clients)
{
client.Key.DownlinkQueue.TryAdd(data);
}
});
// Register the remote connection globally so everyone is aware of the connection status. This is
// essential for several requirements, e.g., client data should not be added to outbound queue if
// there's no connection.
glob.RegisterRemote(server);
try
{
// Start and wait for the remote connection to terminate
await server.Start();
}
catch (Exception e)
{
Console.WriteLine(e);
}
finally
{
Console.WriteLine($"Disposing remote connection: {server}");
server.Dispose();
server = null;
// When remote connection is terminated. Also disconnects all the clients.
glob.Cancel();
}
}
}
}
Remote.cs
namespace Multiplexer
{
using System;
using System.Collections.Concurrent;
using System.Linq;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;
/// <summary>
/// An interface to expose read-only remote connection information
/// </summary>
public interface IRemoteInfo
{
/// <summary>
/// Whether connected to the remote service
/// </summary>
bool Connected { get; }
/// <summary>
/// Remote address
/// </summary>
string RemoteAddress { get; }
}
/// <summary>
/// Class to manage connection to the remote server
/// </summary>
class Remote : IDisposable, IRemoteInfo
{
readonly TcpClient client;
readonly NetworkStream stream;
/// <summary>
/// A delegate called on receiving a package. Implementation could be submitting the package to the queue.
/// </summary>
readonly Action<byte[]> receive;
/// <summary>
/// Queue for data to be uploaded to remote server
/// </summary>
readonly BlockingCollection<byte[]> uplinkQueue;
/// <summary>
/// a linked cancellation token that is cancelled when:
/// - external cancellation is requested, or
/// - the token of the linked CTS is cancelled
/// Note that the cancellation of the linked source won't propagate to the external token
/// </summary>
readonly CancellationTokenSource linkedCTS;
public bool Connected => client.Connected;
public string RemoteAddress => $"{client.Client.RemoteEndPoint}";
public Remote(
TcpClient client,
BlockingCollection<byte[]> uplinkQueue,
CancellationToken externalCancellationToken,
Action<byte[]> receive)
{
this.client = client;
this.uplinkQueue = uplinkQueue;
this.receive = receive;
linkedCTS = CancellationTokenSource.CreateLinkedTokenSource(externalCancellationToken);
stream = client.GetStream();
}
/// <summary>
/// Async task to handle downlink (remote -> multiplexer) traffic
///
/// This is to read from the socket and put data into the downlink queue (via receive())
/// </summary>
async Task HandleDownlink()
{
linkedCTS.Token.ThrowIfCancellationRequested();
int c;
byte[] buffer = new byte[256];
while ((c = await stream.ReadAsync(buffer, 0, buffer.Length, linkedCTS.Token)) > 0)
{
// Receive is non-blocking
receive(buffer.Take(c).ToArray());
}
}
/// <summary>
/// Async task to handle uplink (multiplexer -> remote) traffic
///
/// This is to take data from the uplink queue and write into the socket.
/// </summary>
async Task HandleUplink()
{
linkedCTS.Token.ThrowIfCancellationRequested();
byte[] data;
// Taking from the queue can be blocked if there's nothing in the queue for consumption
while (null != (data = uplinkQueue.Take(linkedCTS.Token)))
{
await stream.WriteAsync(data, 0, data.Length, linkedCTS.Token);
}
}
/// <summary>
/// Async task to start and wait for the uplink and downlink handlers
/// </summary>
public async Task Start()
{
try
{
var downlinkTask = Task.Run(HandleDownlink, linkedCTS.Token);
var uplinkTask = Task.Run(HandleUplink, linkedCTS.Token);
// If either task returns, the connection is considered to be terminated.
await await Task.WhenAny(downlinkTask, uplinkTask);
}
catch (Exception e)
{
Console.WriteLine(e);
}
finally
{
// Cancel the other task (uplink or downlink)
linkedCTS.Cancel();
Console.WriteLine("Remote connection exited.");
this.Dispose();
}
}
public void Dispose()
{
Console.WriteLine("Disposing of remote connection");
linkedCTS.Dispose();
stream.Dispose();
client.Close();
}
}
}
Client.cs
namespace Multiplexer
{
using System;
using System.Collections.Concurrent;
using System.Linq;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;
/// <summary>
/// A class to handle local client connections (client - multiplexer)
/// </summary>
class Client : IDisposable
{
readonly TcpClient client;
/// <summary>
/// A delegate to invoke when receiving data from the client socket that should be uploaded.
///
/// Implementation should put this to the outbound queue. This should be non-blocking.
/// </summary>
readonly Action<byte[]> upload;
readonly NetworkStream stream;
/// <summary>
/// A queue containing data from remote server that should be delivered to this client
/// </summary>
readonly BlockingCollection<byte[]> downlinkQueue = new BlockingCollection<byte[]>();
/// <summary>
/// A cancellation token source linked with an external token
/// </summary>
readonly CancellationTokenSource cts;
public BlockingCollection<byte[]> DownlinkQueue
{
get
{
return this.downlinkQueue;
}
}
/// <summary>
/// A delegate to be called when the client is closed. The <see cref="ClientServer"/> uses this to
/// properly remove the client from the clients list.
/// </summary>
public Action OnClose { get; set; }
public Client(TcpClient client, CancellationToken externalCancellationToken, Action<byte[]> upload)
{
this.client = client;
this.upload = upload;
this.cts = CancellationTokenSource.CreateLinkedTokenSource(externalCancellationToken);
this.stream = client.GetStream();
}
/// <summary>
/// Start the client traffic
/// </summary>
public async Task Start()
{
var uplinkTask = Task.Run(HandleUplink, cts.Token);
var downlinkTask = Task.Run(HandleDownlink, cts.Token);
try
{
// Await for either of the downlink or uplink task to finish
await await Task.WhenAny(uplinkTask, downlinkTask);
}
catch (Exception e)
{
Console.WriteLine(e);
}
finally
{
// Cancel the other task (uplink or downlink)
cts.Cancel();
Console.WriteLine("Client closing");
Dispose();
}
}
/// <summary>
/// Handle uplink traffic (client -> multiplexer -> remote)
/// </summary>
async Task HandleUplink()
{
cts.Token.ThrowIfCancellationRequested();
int c;
byte[] buffer = new byte[256];
while ((c = await stream.ReadAsync(buffer, 0, buffer.Length, cts.Token)) > 0)
{
upload(buffer.Take(c).ToArray());
}
}
/// <summary>
/// Handle downlink traffic (remote -> multiplexer -> client)
/// </summary>
async Task HandleDownlink()
{
cts.Token.ThrowIfCancellationRequested();
byte[] data;
// This would block if the downlink queue is empty
while (null != (data = downlinkQueue.Take(cts.Token)))
{
await stream.WriteAsync(data, 0, data.Length, cts.Token);
}
}
public override string ToString()
{
return client.Client.RemoteEndPoint.ToString();
}
public void Dispose()
{
Console.WriteLine($"Disposing of client: {this}");
OnClose();
cts.Dispose();
stream.Dispose();
client.Close();
}
}
}
Global.cs
namespace Multiplexer
{
using System;
using System.Collections.Concurrent;
using System.Threading;
/// <summary>
/// A class to hold common dependencies to other classes
/// </summary>
/// <remarks>
/// This used to be a singleton and referenced directly by other code, hence the name "Global".
/// I changed it to be dependencies passed as constructor parameters of other classes so it's
/// easier to write tests.
/// </remarks>
class Global
{
/// <summary>
/// Queue for data to be uploaded to remote server (client -> remote)
/// </summary>
public BlockingCollection<byte[]> UploadQueue => uploadQueue;
/// <summary>
/// Set of connected clients. This is used as a set (only keys are used), but there's no ConcurrentSet.
/// </summary>
public ConcurrentDictionary<Client, byte> Clients => clients;
/// <summary>
/// A cancellation token for disconnection. This is used to cancel the remote connection and client connections.
/// </summary>
public CancellationToken CancellationToken => cts.Token;
/// <summary>
/// A readonly object holding status for the remote connection
/// </summary>
public IRemoteInfo Remote => remote ?? new DummyRemote();
private BlockingCollection<byte[]> uploadQueue = new BlockingCollection<byte[]>();
private ConcurrentDictionary<Client, byte> clients = new ConcurrentDictionary<Client, byte>();
private CancellationTokenSource cts = new CancellationTokenSource();
private IRemoteInfo remote;
/// <summary>
/// Register remote connection
/// </summary>
public IRemoteInfo RegisterRemote(IRemoteInfo remote)
{
return Interlocked.Exchange(ref this.remote, remote);
}
/// <summary>
/// Cancel the global cancellation token. Used for disconnecting remote and clients
/// </summary>
public void Cancel()
{
if (!cts.Token.IsCancellationRequested)
{
lock (cts)
{
if (!cts.Token.IsCancellationRequested)
{
cts.Cancel();
cts.Dispose();
cts = new CancellationTokenSource();
}
}
}
}
/// <summary>
/// Clear the upload queue. Used at each new remote connection.
/// </summary>
public void ResetUploadQueue()
{
Console.WriteLine($"Resetting upload queue ({uploadQueue.Count})");
byte[] b;
while (uploadQueue.TryTake(out b)) { }
}
/// <summary>
/// A dummy remote info to avoid null referencing when no remote server is connected
/// </summary>
class DummyRemote : IRemoteInfo
{
public bool Connected => false;
public string RemoteAddress => "";
}
}
}
linkedCTS.Token.ThrowIfCancellationRequested()in methodsRemote.HandleUplink()andRemote.HandleDownlink()is called outside the while loop? \$\endgroup\$