fix stream reconnect logic

Author: eli-darkly

src/LaunchDarkly.EventSource/EventSource.cs

@@ -25,8 +25,9 @@ public class EventSource : IEventSource
         private string _eventName = Constants.MessageField;
         private string _lastEventId;
         private TimeSpan _retryDelay;
-        private CancellationTokenSource _pendingRequest;
         private readonly ExponentialBackoffWithDecorrelation _backOff;
+        private CancellationTokenSource _currentRequestToken;
+        private ReadyState _readyState;
 
         #endregion
 
@@ -65,8 +66,20 @@ public class EventSource : IEventSource
         /// </value>
         public ReadyState ReadyState
         {
-            get;
-            private set;
+            get
+            {
+                lock(this)
+                {
+                    return _readyState;
+                }
+            }
+            private set
+            {
+                lock(this)
+                {
+                    _readyState = value;
+                }
+            }
         }
 
         internal TimeSpan BackOffDelay
@@ -88,14 +101,12 @@ internal TimeSpan BackOffDelay
         /// configuration</exception>
         public EventSource(Configuration configuration)
         {
-            ReadyState = ReadyState.Raw;
+            _readyState = ReadyState.Raw;
 
             _configuration = configuration ?? throw new ArgumentNullException(nameof(configuration));
 
             _logger = _configuration.Logger ?? LogManager.GetLogger(typeof(EventSource));
 
-            _pendingRequest = new CancellationTokenSource();
-
             _retryDelay = _configuration.DelayRetryDuration;
 
             _backOff = new ExponentialBackoffWithDecorrelation(_retryDelay,
@@ -114,14 +125,22 @@ public EventSource(Configuration configuration)
         /// <exception cref="InvalidOperationException">The method was called after the connection <see cref="ReadyState"/> was Open or Connecting.</exception>
         public async Task StartAsync()
         {
-            var cancellationToken = _pendingRequest.Token;
-            while (!cancellationToken.IsCancellationRequested)
+            while (ReadyState != ReadyState.Shutdown)
             {
                 await MaybeWaitWithBackOff();
                 try
                 {
-                    await ConnectToEventSourceAsync(cancellationToken);
-                    _backOff.ResetReconnectAttemptCount();
+                    var newRequestTokenSource = new CancellationTokenSource();
+                    lock (this)
+                    {
+                        if (_readyState == ReadyState.Shutdown)
+                        {
+                            // in case Close() was called in between the previous ReadyState check and the creation of the new token
+                            return;
+                        }
+                        _currentRequestToken = newRequestTokenSource;
+                    }
+                    await ConnectToEventSourceAsync(newRequestTokenSource.Token);
                 }
                 catch (Exception e)
                 {
@@ -152,21 +171,24 @@ public void Close()
             if (ReadyState == ReadyState.Raw || ReadyState == ReadyState.Shutdown) return;
 
             Close(ReadyState.Shutdown);
-
-            // Cancel token to cancel requests.
-            CancelToken();
-
+            CancelCurrentRequest();
         }
 
         #endregion
 
         #region Private Methods
 
-        private void CancelToken()
+        private void CancelCurrentRequest()
         {
-            CancellationTokenSource cancellationTokenSource = Interlocked.Exchange(ref _pendingRequest, new CancellationTokenSource());
-            cancellationTokenSource.Cancel();
-            cancellationTokenSource.Dispose();
+            CancellationTokenSource requestTokenSource = null;
+            lock (this)
+            {
+                requestTokenSource = _currentRequestToken;
+            }
+            if (requestTokenSource != null)
+            {
+                requestTokenSource.Cancel();
+            }
         }
 
         internal virtual EventSourceService GetEventSourceService(Configuration configuration)
@@ -189,7 +211,10 @@ private async Task ConnectToEventSourceAsync(CancellationToken cancellationToken
 
                 var svc = GetEventSourceService(_configuration);
 
-                svc.ConnectionOpened += (o, e) => { SetReadyState(ReadyState.Open, OnOpened); };
+                svc.ConnectionOpened += (o, e) => {
+                    _backOff.ResetReconnectAttemptCount();
+                    SetReadyState(ReadyState.Open, OnOpened);
+                };
                 svc.ConnectionClosed += (o, e) => { SetReadyState(ReadyState.Closed, OnClosed); };
 
                 await svc.GetDataAsync(
@@ -199,7 +224,7 @@ await svc.GetDataAsync(
             }
             catch (EventSourceServiceCancelledException e)
             {
-                CancelToken();
+                CancelCurrentRequest();
 
                 CloseAndRaiseError(e);
             }
@@ -212,7 +237,6 @@ await svc.GetDataAsync(
 
                     throw;
                 }
-
             }
         }
 
@@ -252,12 +276,19 @@ private void ProcessResponseContent(string content)
 
         private void SetReadyState(ReadyState state, Action<StateChangedEventArgs> action = null)
         {
-            if (ReadyState == state) return;
-
-            ReadyState = state;
+            lock (this)
+            {
+                if (_readyState == state || _readyState == ReadyState.Shutdown)
+                {
+                    return;
+                }
+                _readyState = state;
+            }
 
-            if (action != null)
-                action(new StateChangedEventArgs(ReadyState));
+            if (action != null)
+            {
+                action(new StateChangedEventArgs(state));
+            }
         }
 
         private void ProcessField(string field, string value)

src/LaunchDarkly.EventSource/EventSourceService.cs

@@ -63,7 +63,6 @@ public async Task GetDataAsync(Action<string> processResponse, CancellationToken
             cancellationToken.ThrowIfCancellationRequested();
 
             await ConnectToEventSourceApi(processResponse, cancellationToken);
-
         }
 
         #endregion

test/LaunchDarkly.EventSource.Tests/EventSourceTests.cs

@@ -353,18 +353,15 @@ public async Task Given_content_type_not_equal_to_eventstream_when_the_http_resp
 
             var evt = new EventSource(config);
 
-            var wasErrorEventRaised = false;
-            evt.Error += (s, e) =>
-            {
-                wasErrorEventRaised = true;
-            };
+            var receiver = new ErrorReceiver(evt);
 
             //// Act
             await evt.StartAsync();
 
             //// Assert
-            Assert.True(wasErrorEventRaised);
-            Assert.True(evt.ReadyState == ReadyState.Closed);
+            Assert.NotNull(receiver.ErrorReceived);
+            Assert.Equal(ReadyState.Closed, receiver.SourceStateReceived);
+            Assert.Equal(ReadyState.Shutdown, evt.ReadyState);
         }
 
         [Fact]
@@ -377,21 +374,19 @@ public async Task Given_204_when_the_http_response_is_received_then_error_event_
 
             var evt = new EventSource(new Configuration(_uri, handler));
 
+            var receiver = new ErrorReceiver(evt);
+
             //Act
-            var raisedEvent = await Assert.RaisesAsync<ExceptionEventArgs>(
-                h => evt.Error += h,
-                h => evt.Error -= h,
-                () => evt.StartAsync());
+            await evt.StartAsync();
 
             //// Assert
-            Assert.NotNull(raisedEvent);
-            Assert.Equal(evt, raisedEvent.Sender);
-            Assert.IsType<EventSourceServiceUnsuccessfulResponseException>(raisedEvent.Arguments.Exception);
-            Assert.Equal(204, ((EventSourceServiceUnsuccessfulResponseException)raisedEvent.Arguments.Exception).StatusCode);
-            Assert.True(evt.ReadyState == ReadyState.Closed);
+            Assert.NotNull(receiver.ErrorReceived);
+            var ex = Assert.IsType<EventSourceServiceUnsuccessfulResponseException>(receiver.ErrorReceived);
+            Assert.Equal(204, ex.StatusCode);
+            Assert.Equal(ReadyState.Closed, receiver.SourceStateReceived);
+            Assert.Equal(ReadyState.Shutdown, evt.ReadyState);
         }
 
-
         [Theory]
         [InlineData(HttpStatusCode.InternalServerError)]
         [InlineData(HttpStatusCode.BadRequest)]
@@ -406,18 +401,16 @@ public async Task Given_status_code_when_the_http_response_is_received_then_erro
 
             var evt = new EventSource(new Configuration(_uri, handler));
 
-            //Act
-            var raisedEvent = await Assert.RaisesAsync<ExceptionEventArgs>(
-                h => evt.Error += h,
-                h => evt.Error -= h,
-                () => evt.StartAsync());
+            ErrorReceiver receiver = new ErrorReceiver(evt);
+
+            await evt.StartAsync();
 
             //// Assert
-            Assert.NotNull(raisedEvent);
-            Assert.Equal(evt, raisedEvent.Sender);
-            Assert.IsType<EventSourceServiceUnsuccessfulResponseException>(raisedEvent.Arguments.Exception);
-            Assert.Equal((int)statusCode, ((EventSourceServiceUnsuccessfulResponseException)raisedEvent.Arguments.Exception).StatusCode);
-            Assert.True(evt.ReadyState == ReadyState.Closed);
+            Assert.NotNull(receiver.ErrorReceived);
+            var ex = Assert.IsType<EventSourceServiceUnsuccessfulResponseException>(receiver.ErrorReceived);
+            Assert.Equal((int)statusCode, ex.StatusCode);
+            Assert.Equal(ReadyState.Closed, receiver.SourceStateReceived);
+            Assert.Equal(ReadyState.Shutdown, evt.ReadyState);
         }
 
         [Fact]
@@ -426,7 +419,8 @@ public async Task Given_bad_http_responses_then_retry_delay_durations_should_be_
             // Arrange
             var handler = new StubMessageHandler();
 
-            for (int i = 0; i < 2; i++)
+            var nAttempts = 2;
+            for (int i = 0; i < nAttempts; i++)
             {
                 var response = new HttpResponseMessageWithError();
 
@@ -436,7 +430,6 @@ public async Task Given_bad_http_responses_then_retry_delay_durations_should_be_
                     "text/event-stream");
 
                 handler.QueueResponse(response);
-
             }
 
             handler.QueueResponse(new HttpResponseMessage(HttpStatusCode.NoContent));
@@ -447,14 +440,18 @@ public async Task Given_bad_http_responses_then_retry_delay_durations_should_be_
             evt.Error += (_, e) =>
             {
                 backoffs.Add(evt.BackOffDelay);
+                if (backoffs.Count >= nAttempts)
+                {
+                    evt.Close();
+                }
             };
 
             //Act
             await evt.StartAsync();
 
             //// Assert
             Assert.NotEmpty(backoffs);
-            Assert.True(backoffs.Distinct().Count() == backoffs.Count());
+            Assert.Equal(backoffs.Distinct().Count(), backoffs.Count());
         }
 
 
@@ -471,25 +468,14 @@ public async Task When_response_exceeds_read_timeout_then_read_timeout_exception
 
             var evt = new StubEventSource(new Configuration(_uri, handler, readTimeout: readTimeout), (int)timeout.TotalMilliseconds);
 
-            var exceptionMessage = string.Empty;
-
-            try
-            {
-
-                evt.Error += (_, e) =>
-                {
-                    exceptionMessage = e.Exception.Message;
-                    evt.Close();
-                };
-
-                await evt.StartAsync();
-
-            }
-            catch (TaskCanceledException) {}
+            var receiver = new ErrorReceiver(evt);
 
-            Assert.Contains(exceptionMessage, Resources.EventSourceService_Read_Timeout);
-            Assert.True(evt.ReadyState == ReadyState.Shutdown);
+            await evt.StartAsync();
 
+            Assert.NotNull(receiver.ErrorReceived);
+            Assert.Contains(receiver.ErrorReceived.Message, Resources.EventSourceService_Read_Timeout);
+            Assert.Equal(ReadyState.Closed, receiver.SourceStateReceived);
+            Assert.Equal(ReadyState.Shutdown, evt.ReadyState);
         }
 
         [Fact]
@@ -520,7 +506,36 @@ public async Task When_response_does_not_exceed_read_timeout_then_expected_messa
 
             Assert.Equal("put", eventName);
             Assert.True(wasMessageReceivedEventRaised);
+        }
+
+        [Fact]
+        public async Task When_server_returns_HTTP_error_a_reconnect_attempt_is_made()
+        {
+            var messageData = "hello";
 
+            var handler = new StubMessageHandler();
+            handler.QueueResponse(new HttpResponseMessage(HttpStatusCode.Unauthorized));
+            handler.QueueStringResponse("event: put\ndata: " + messageData + "\n\n");
+
+            var evt = new EventSource(new Configuration(_uri, handler));
+
+            ErrorReceiver receiver = new ErrorReceiver(evt);
+            receiver.CloseEventSourceOnError = false;
+
+            string messageReceived = null;
+            evt.MessageReceived += (_, e) =>
+            {
+                messageReceived = e.Message.Data;
+                evt.Close();
+            };
+
+            await evt.StartAsync();
+
+            Assert.Equal(2, handler.GetRequests().Count());
+            Assert.NotNull(receiver.ErrorReceived);
+            var ex = Assert.IsType<EventSourceServiceUnsuccessfulResponseException>(receiver.ErrorReceived);
+            Assert.Equal((int)HttpStatusCode.Unauthorized, ex.StatusCode);
+            Assert.Equal(messageData, messageReceived);
         }
 
         [Fact]
@@ -531,14 +546,35 @@ public async Task When_error_handler_closes_event_source_no_reconnect_attempt_is
 
             var evt = new EventSource(new Configuration(_uri, handler));
 
-            evt.Error += (_, e) =>
-            {
-                evt.Close();
-            };
+            var receiver = new ErrorReceiver(evt);
 
             await evt.StartAsync();
 
             Assert.Equal(1, handler.GetRequests().Count());
         }
     }
+
+    class ErrorReceiver
+    {
+        public bool CloseEventSourceOnError = true;
+        public Exception ErrorReceived = null;
+        public ReadyState SourceStateReceived;
+        private readonly EventSource _source;
+        
+        public ErrorReceiver(EventSource source)
+        {
+            _source = source;
+            source.Error += HandleError;
+        }
+
+        public void HandleError(object sender, ExceptionEventArgs e)
+        {
+            ErrorReceived = e.Exception;
+            SourceStateReceived = _source.ReadyState;
+            if (CloseEventSourceOnError)
+            {
+                _source.Close();
+            }
+        }
+    }
 }