Skip to content

Commit c47ddde

Browse files
rstamrstam
committed
CSHARP-1375: Refactored TcpStreamFactory Connect method to fix to problem with slow connect times, and refactored ConnectAsync to match.
1 parent fb10da0 commit c47ddde

File tree

3 files changed

+99
-53
lines changed

3 files changed

+99
-53
lines changed

‎src/MongoDB.Driver.Core.Tests/Core/Connections/TcpStreamFactoryTests.cs‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,50 @@ public void CreateStream_should_throw_a_SocketException_when_the_endpoint_could_
6060
act.ShouldThrow<SocketException>();
6161
}
6262

63+
[Test]
64+
public void CreateStream_should_throw_when_cancellation_is_requested(
65+
[Values(false, true)]
66+
bool async)
67+
{
68+
var subject = new TcpStreamFactory();
69+
var endPoint = new IPEndPoint(new IPAddress(0x01010101), 12345); // a non-existent host and port
70+
var cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromMilliseconds(20));
71+
72+
Action action;
73+
if (async)
74+
{
75+
action = () => subject.CreateStreamAsync(endPoint, cancellationTokenSource.Token).GetAwaiter().GetResult();
76+
}
77+
else
78+
{
79+
action = () => subject.CreateStream(endPoint, cancellationTokenSource.Token);
80+
}
81+
82+
action.ShouldThrow<OperationCanceledException>();
83+
}
84+
85+
[Test]
86+
public void CreateStream_should_throw_when_connect_timeout_has_expired(
87+
[Values(false, true)]
88+
bool async)
89+
{
90+
var settings = new TcpStreamSettings(connectTimeout: TimeSpan.FromMilliseconds(20));
91+
var subject = new TcpStreamFactory(settings);
92+
var endPoint = new IPEndPoint(new IPAddress(0x01010101), 12345); // a non-existent host and port
93+
94+
Action action;
95+
if (async)
96+
{
97+
action = () => subject.CreateStreamAsync(endPoint, CancellationToken.None).GetAwaiter().GetResult(); ;
98+
}
99+
else
100+
{
101+
action = () => subject.CreateStream(endPoint, CancellationToken.None);
102+
}
103+
104+
action.ShouldThrow<TimeoutException>();
105+
}
106+
63107
[Test]
64108
[RequiresServer]
65109
public void CreateStream_should_call_the_socketConfigurator(

‎src/MongoDB.Driver.Core/Core/Connections/TcpStreamFactory.cs‎

Lines changed: 50 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -74,74 +74,77 @@ private void ConfigureConnectedSocket(Socket socket)
7474

7575
private void Connect(Socket socket, EndPoint endPoint, CancellationToken cancellationToken)
7676
{
77-
IAsyncResult asyncResult;
78-
var dnsEndPoint = endPoint as DnsEndPoint;
79-
if (dnsEndPoint != null)
80-
{
81-
// mono doesn't support DnsEndPoint in its BeginConnect method.
82-
asyncResult = socket.BeginConnect(dnsEndPoint.Host, dnsEndPoint.Port, null, null);
83-
}
84-
else
85-
{
86-
asyncResult = socket.BeginConnect(endPoint, null, null);
87-
}
77+
var cancelled = false;
78+
var timedOut = false;
8879

89-
asyncResult.AsyncWaitHandle.WaitOne(_settings.ConnectTimeout);
90-
if (!asyncResult.IsCompleted)
80+
using (var registration = cancellationToken.Register(() => { cancelled = true; try { socket.Close(); } catch { } }))
81+
using (var timer = new Timer(_ => { timedOut = true; try { socket.Close(); } catch { } }, null, _settings.ConnectTimeout, Timeout.InfiniteTimeSpan))
9182
{
9283
try
9384
{
94-
socket.Dispose();
85+
var dnsEndPoint = endPoint as DnsEndPoint;
86+
if (dnsEndPoint != null)
87+
{
88+
// mono doesn't support DnsEndPoint in its BeginConnect method.
89+
socket.Connect(dnsEndPoint.Host, dnsEndPoint.Port);
90+
}
91+
else
92+
{
93+
socket.Connect(endPoint);
94+
}
9595
}
9696
catch
9797
{
98-
// ignore exceptions
98+
if (!cancelled && ! timedOut)
99+
{
100+
throw;
101+
}
99102
}
103+
}
100104

105+
cancellationToken.ThrowIfCancellationRequested();
106+
if (timedOut)
107+
{
101108
var message = string.Format("Timed out connecting to {0}. Timeout was {1}.", endPoint, _settings.ConnectTimeout);
102109
throw new TimeoutException(message);
103110
}
104-
105-
socket.EndConnect(asyncResult);
106111
}
107112

108113
private async Task ConnectAsync(Socket socket, EndPoint endPoint, CancellationToken cancellationToken)
109114
{
110-
Task connectTask;
111-
var dnsEndPoint = endPoint as DnsEndPoint;
112-
if (dnsEndPoint != null)
113-
{
114-
// mono doesn't support DnsEndPoint in its BeginConnect method.
115-
connectTask = Task.Factory.FromAsync(socket.BeginConnect(dnsEndPoint.Host, dnsEndPoint.Port, null, null), socket.EndConnect);
116-
}
117-
else
118-
{
119-
connectTask = Task.Factory.FromAsync(socket.BeginConnect(endPoint, null, null), socket.EndConnect);
120-
}
115+
var cancelled = false;
116+
var timedOut = false;
121117

122-
if (_settings.ConnectTimeout == Timeout.InfiniteTimeSpan)
118+
using (var registration = cancellationToken.Register(() => { cancelled = true; try { socket.Close(); } catch { } }))
119+
using (var timer = new Timer(_ => { timedOut = true; try { socket.Close(); } catch { } }, null, _settings.ConnectTimeout, Timeout.InfiniteTimeSpan))
123120
{
124-
await connectTask.ConfigureAwait(false);
125-
return;
126-
}
127-
128-
using (var delayCancellationTokenSource = new CancellationTokenSource())
129-
{
130-
var delayTask = Task.Delay(_settings.ConnectTimeout, delayCancellationTokenSource.Token);
131-
132-
var completedTask = await Task.WhenAny(connectTask, delayTask).ConfigureAwait(false);
133-
134-
// kill the delay timer as soon as possible
135-
delayCancellationTokenSource.Cancel();
136-
137-
if (completedTask == delayTask && !socket.Connected)
121+
try
122+
{
123+
var dnsEndPoint = endPoint as DnsEndPoint;
124+
if (dnsEndPoint != null)
125+
{
126+
// mono doesn't support DnsEndPoint in its BeginConnect method.
127+
await Task.Factory.FromAsync(socket.BeginConnect(dnsEndPoint.Host, dnsEndPoint.Port, null, null), socket.EndConnect).ConfigureAwait(false);
128+
}
129+
else
130+
{
131+
await Task.Factory.FromAsync(socket.BeginConnect(endPoint, null, null), socket.EndConnect).ConfigureAwait(false);
132+
}
133+
}
134+
catch
138135
{
139-
socket.Dispose();
140-
var message = string.Format("Timed out connecting to {0}. Timeout was {1}.", endPoint, _settings.ConnectTimeout);
141-
throw new TimeoutException(message);
136+
if (!cancelled && !timedOut)
137+
{
138+
throw;
139+
}
142140
}
141+
}
143142

144-
await connectTask.ConfigureAwait(false);
143+
cancellationToken.ThrowIfCancellationRequested();
144+
if (timedOut)
145+
{
146+
var message = string.Format("Timed out connecting to {0}. Timeout was {1}.", endPoint, _settings.ConnectTimeout);
147+
throw new TimeoutException(message);
145148
}
146149
}
147150

‎src/MongoDB.Driver.Core/Core/Misc/SemaphoreSlimRequest.cs‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ public sealed class SemaphoreSlimRequest : IDisposable
2929
{
3030
// private fields
3131
private readonly CancellationTokenSource _disposeCancellationTokenSource;
32+
private readonly CancellationTokenSource _linkedCancellationTokenSource;
3233
private readonly SemaphoreSlim _semaphore;
3334
private readonly Task _task;
3435

@@ -44,11 +45,8 @@ public SemaphoreSlimRequest(SemaphoreSlim semaphore, CancellationToken cancellat
4445
_semaphore = Ensure.IsNotNull(semaphore, nameof(semaphore));
4546

4647
_disposeCancellationTokenSource = new CancellationTokenSource();
47-
var disposeCancellationToken = _disposeCancellationTokenSource.Token;
48-
var linkedCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, disposeCancellationToken);
49-
var linkedCancellationToken = linkedCancellationTokenSource.Token;
50-
51-
_task = semaphore.WaitAsync(linkedCancellationToken);
48+
_linkedCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _disposeCancellationTokenSource.Token);
49+
_task = semaphore.WaitAsync(_linkedCancellationTokenSource.Token);
5250
}
5351

5452
// public properties
@@ -64,13 +62,14 @@ public SemaphoreSlimRequest(SemaphoreSlim semaphore, CancellationToken cancellat
6462
/// <inheritdoc/>
6563
public void Dispose()
6664
{
67-
_disposeCancellationTokenSource.Cancel(); // if we haven't gotten the lock by now we no longer want it
65+
_disposeCancellationTokenSource.Cancel(); // does nothing if we have the lock, otherwise cancels the request
6866
if (_task.Status == TaskStatus.RanToCompletion)
6967
{
7068
_semaphore.Release();
7169
}
7270

7371
_disposeCancellationTokenSource.Dispose();
72+
_linkedCancellationTokenSource.Dispose();
7473
}
7574
}
7675
}

0 commit comments

Comments
 (0)