fix: Reconnect the stream after non-HTTP read errors

src/LaunchDarkly/Server/Network/Streaming.hs

@@ -1,13 +1,13 @@
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE NumericUnderscores #-}
 
-module LaunchDarkly.Server.Network.Streaming (streamingThread) where
+module LaunchDarkly.Server.Network.Streaming (streamingThread, backoffCeiling) where
 
 import Control.Applicative (many, (<|>))
 import Control.Concurrent (threadDelay)
-import Control.Exception (throwIO)
+import Control.Exception (SomeAsyncException, SomeException, displayException, fromException, throwIO)
 import Control.Monad (mzero, void)
-import Control.Monad.Catch (Exception, MonadCatch, MonadMask, try)
+import Control.Monad.Catch (Exception, MonadCatch, MonadMask, MonadThrow, catch, throwM, try)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Monad.Logger (MonadLogger, logDebug, logError, logWarn)
 import Control.Monad.Loops (iterateUntilM)
@@ -16,9 +16,11 @@ import Data.Attoparsec.Text as P hiding (Result, try)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Lazy as L
+import Data.IORef (IORef, newIORef, readIORef, writeIORef)
+import Data.Maybe (isJust)
 import Data.Text (Text)
 import qualified Data.Text as T
-import Data.Text.Encoding (decodeUtf8, encodeUtf8)
+import Data.Text.Encoding (Decoding (Some), encodeUtf8, streamDecodeUtf8)
 import GHC.Generics (Generic)
 import GHC.Natural (Natural)
 import Network.HTTP.Client (HttpException (..), HttpExceptionContent (..), Manager, Request, Response (..), brRead)
@@ -188,17 +190,25 @@ tryReadE :: MonadCatch m => m a -> m (Either ReadE a)
 tryReadE = try
 
 -- heartbeat expected every 120 seconds
-readWithException :: IO ByteString -> IO Text
-readWithException body =
+-- A read can end partway through a multi-byte character, so the decoder keeps those bytes for the next read.
+-- `parseWith` treats empty text as end of input, so a read that decodes to nothing reads again.
+readWithException :: IORef (ByteString -> Decoding) -> IO ByteString -> IO Text
+readWithException decoder body =
     timeout (1_000_000 * 300) (brRead body) >>= \case
         Nothing -> throwIO ReadETimeout
-        Just bytes -> if bytes == B.empty then throwIO ReadEClosed else pure (decodeUtf8 bytes)
+        Just bytes ->
+            if bytes == B.empty
+                then throwIO ReadEClosed
+                else do
+                    Some text _ next <- ($ bytes) <$> readIORef decoder
+                    writeIORef decoder next
+                    if T.null text then readWithException decoder body else pure text
 
 readStream :: (MonadIO m, MonadLogger m, MonadMask m) => IO ByteString -> DataSourceUpdates -> m Bool
-readStream body dataSourceUpdates = loop "" False
+readStream body dataSourceUpdates = liftIO (newIORef streamDecodeUtf8) >>= \decoder -> loop decoder "" False
   where
-    loop initial processedEvent =
-        tryReadE (parseWith (liftIO $ readWithException body) parseEvent initial) >>= \case
+    loop decoder initial processedEvent =
+        tryReadE (parseWith (liftIO $ readWithException decoder body) parseEvent initial) >>= \case
             (Left ReadETimeout) -> do
                 $(logError) "streaming connection unexpectedly closed"
                 pure processedEvent
@@ -208,14 +218,24 @@ readStream body dataSourceUpdates = loop "" False
             (Right parsed) -> case parsed of
                 Done remaining event -> do
                     processed <- processEvent dataSourceUpdates (name event) (L.fromStrict $ encodeUtf8 $ buffer event)
-                    if processed then loop remaining True else pure processedEvent
+                    if processed then loop decoder remaining True else pure processedEvent
                 Fail _ context err -> do
                     $(logError) $ T.intercalate " " ["failed parsing SSE frame", T.pack $ show context, T.pack err]
                     pure processedEvent
                 Partial _ -> do
                     $(logError) "failed parsing SSE frame unexpected partial"
                     pure processedEvent
 
+-- `tryHTTP` only catches `HttpException`. Any other exception from the body read (a socket reset, a TLS error,
+-- invalid UTF-8) would end the streaming thread with no log and no reconnect, so it is logged and rethrown as an
+-- `HttpException` to get the usual backoff and reconnect. Async exceptions pass through so `close` can stop the thread.
+logReadError :: (MonadLogger m, MonadThrow m) => Request -> SomeException -> m a
+logReadError request e
+    | isJust (fromException @SomeAsyncException e) || isJust (fromException @HttpException e) = throwM e
+    | otherwise = do
+        $(logError) $ T.append "streaming connection failed, will retry: " (T.pack $ displayException e)
+        throwM $ HttpExceptionRequest request (InternalException e)
+
 -- This function is responsible for consuming a streaming connection.
 --
 -- It is responsible for opening the connection, parsing the body for as long as it is able, and then handling shut down
@@ -227,7 +247,7 @@ startNewConnection manager request dataSourceUpdates state@(StreamingState {acti
     status <- tryHTTP $ withResponseGeneric request manager $ \response -> do
         checkAuthorization response
             >> throwIfNot200 request response
-            >> readStream (responseBody response) dataSourceUpdates
+            >> (readStream (responseBody response) dataSourceUpdates `catch` logReadError request)
     now <- liftIO $ getTime Monotonic
     handleResponse state now status
   where
@@ -270,10 +290,19 @@ startNewConnection manager request dataSourceUpdates state@(StreamingState {acti
     calculateDelay StreamingState {initialRetryDelay, attempt = att} = do
         liftIO $
             newStdGen >>= \gen ->
-                let timespan = min (30 * 1_000_000) ((initialRetryDelay * 1_000) * (2 ^ (att - 1)))
+                let timespan = backoffCeiling initialRetryDelay att
                     jitter = fst $ randomR (0, timespan `div` 2) gen
                  in pure $ (timespan - jitter)
 
+-- The retry delay in microseconds before jitter, given the initial delay in milliseconds and the attempt number.
+--
+-- The product is an `Integer` because it overflows `Int` on 32-bit runtimes at the default 1 s initial delay, and
+-- on any runtime for a large enough initial delay. Any initial delay of 1 ms or more reaches the 30 s cap by 2^15,
+-- so the exponent cap keeps the `Integer` small without changing the result.
+backoffCeiling :: Int -> Int -> Int
+backoffCeiling initialRetryDelay att =
+    fromInteger $ min (30 * 1_000_000) (toInteger initialRetryDelay * 1_000 * 2 ^ min 15 (att - 1))
+
 data StreamingState = StreamingState
     { initialConnection :: Bool -- Marker used to determine the first time the streamer connects.
     , initialRetryDelay :: Int -- The base duration used for the retry delay calculation

test/Spec/Streaming.hs

@@ -1,6 +1,199 @@
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE NumericUnderscores #-}
+
 module Spec.Streaming (allTests) where
 
+import Control.Concurrent (forkFinally, killThread, threadDelay)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, readMVar, tryPutMVar, tryReadMVar)
+import Control.Exception (AsyncException (ThreadKilled), SomeException, fromException, throwIO, toException)
+import Control.Monad (forever, void)
+import Control.Monad.Logger (LogLevel (..), fromLogStr, runLoggingT)
+import Data.Aeson (encode, object, (.=))
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+import qualified Data.ByteString.Char8 as B8
+import qualified Data.ByteString.Lazy as L
+import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)
+import Data.Maybe (isJust)
+import Data.Text (Text)
+import qualified Data.Text as T
+import Data.Text.Encoding (decodeUtf8With)
+import Data.Text.Encoding.Error (lenientDecode)
+import Network.HTTP.Client (Manager, defaultManagerSettings, managerRawConnection, newManager, responseTimeoutNone)
+import Network.HTTP.Client.Internal (makeConnection)
+import Numeric (showHex)
+import System.Timeout (timeout)
 import Test.HUnit
 
+import LaunchDarkly.AesonCompat (KeyMap, lookupKey)
+import LaunchDarkly.Server.Config.ClientContext (ClientContext (..))
+import LaunchDarkly.Server.Config.HttpConfiguration (HttpConfiguration (..))
+import LaunchDarkly.Server.DataSource.Internal (DataSourceUpdates (..))
+import LaunchDarkly.Server.Features (Flag)
+import LaunchDarkly.Server.Network.Streaming (backoffCeiling, streamingThread)
+import Util.Features (makeTestFlag)
+
 allTests :: Test
-allTests = TestList []
+allTests =
+    TestList
+        [ TestLabel "healthy stream initializes" testHealthyStream
+        , TestLabel "reconnects after an HttpException mid-body" testReconnectAfterHttpException
+        , TestLabel "reconnects after a socket reset mid-body" testReconnectAfterSocketReset
+        , TestLabel "decodes UTF-8 split across reads" testUtf8SplitAcrossReads
+        , TestLabel "killThread stops the streaming thread" testKillThreadStopsStream
+        , TestLabel "backoff stays at the cap after many attempts" testBackoffCeiling
+        ]
+
+-- | One socket read. 'Fail' throws from inside the read, the way a reset socket does.
+data ReadStep = Send ByteString | Fail SomeException
+
+-- | The reads one TCP connection returns, in order. Once they run out, reads block.
+type ConnectionScript = [ReadStep]
+
+data StreamRun = StreamRun
+    { connections :: Int
+    -- ^ How many scripted connections the SDK opened
+    , initialized :: Maybe (KeyMap Flag)
+    -- ^ The flags passed to 'dataSourceUpdatesInit', if it was called
+    , logs :: [(LogLevel, Text)]
+    , exited :: Maybe (Either SomeException ())
+    -- ^ How the streaming thread ended. 'Nothing' when it was still running 1 s after 'killThread'.
+    }
+
+-- | Serves each new connection from the next script. Once the scripts run out, opening a connection blocks.
+scriptedManager :: IORef [ConnectionScript] -> IO Manager
+scriptedManager pending =
+    newManager defaultManagerSettings {managerRawConnection = pure $ \_ _ _ -> openNext}
+  where
+    openNext =
+        pop pending >>= \case
+            Nothing -> blockForever
+            Just script -> do
+                steps <- newIORef script
+                makeConnection (readStep steps) (const $ pure ()) (pure ())
+    readStep steps =
+        pop steps >>= \case
+            Nothing -> blockForever
+            Just (Send bytes) -> pure bytes
+            Just (Fail e) -> throwIO e
+
+pop :: IORef [a] -> IO (Maybe a)
+pop ref = atomicModifyIORef' ref $ \case
+    [] -> ([], Nothing)
+    x : rest -> (rest, Just x)
+
+blockForever :: IO a
+blockForever = forever $ threadDelay 1_000_000_000
+
+-- | Runs 'streamingThread' against the scripts until it initializes or exits (at most 5 s), then kills it.
+runStream :: [ConnectionScript] -> IO StreamRun
+runStream scripts = do
+    pending <- newIORef scripts
+    logRef <- newIORef []
+    initVar <- newEmptyMVar
+    settled <- newEmptyMVar
+    done <- newEmptyMVar
+    manager <- scriptedManager pending
+    let logTo _ _ level msg = atomicModifyIORef' logRef (\xs -> ((level, decodeUtf8With lenientDecode $ fromLogStr msg) : xs, ()))
+        context =
+            ClientContext
+                { runLogger = (`runLoggingT` logTo)
+                , httpConfiguration =
+                    HttpConfiguration
+                        { defaultRequestHeaders = []
+                        , defaultRequestTimeout = responseTimeoutNone
+                        , tlsManager = manager
+                        }
+                }
+        updates =
+            DataSourceUpdates
+                { dataSourceUpdatesInit = \flags _ -> tryPutMVar initVar flags >> tryPutMVar settled () >> pure (Right ())
+                , dataSourceUpdatesInsertFlag = const $ pure (Right ())
+                , dataSourceUpdatesInsertSegment = const $ pure (Right ())
+                , dataSourceUpdatesDeleteFlag = \_ _ -> pure (Right ())
+                , dataSourceUpdatesDeleteSegment = \_ _ -> pure (Right ())
+                , dataSourceUpdatesSetStatus = const $ pure ()
+                }
+    thread <-
+        forkFinally
+            (runLogger context $ streamingThread "http://localhost" 1 context updates)
+            (\result -> putMVar done result >> void (tryPutMVar settled ()))
+    void $ timeout 5_000_000 $ readMVar settled
+    killThread thread
+    exited <- timeout 1_000_000 $ readMVar done
+    initialized <- tryReadMVar initVar
+    remaining <- readIORef pending
+    logs <- reverse <$> readIORef logRef
+    pure StreamRun {connections = length scripts - length remaining, initialized, logs, exited}
+
+responseHead :: ReadStep
+responseHead = Send "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n"
+
+-- | Frames the bytes as one HTTP chunk, which the SDK receives from a single 'brRead'.
+chunk :: ByteString -> ReadStep
+chunk bytes = Send $ B8.pack (showHex (B.length bytes) "") <> "\r\n" <> bytes <> "\r\n"
+
+-- | A put event whose only flag key, @café@, is the first non-ASCII text in the payload.
+putEvent :: ByteString
+putEvent = "event: put\ndata: " <> L.toStrict (encode payload) <> "\n\n"
+  where
+    payload =
+        object
+            [ "path" .= ("/" :: Text)
+            , "data" .= object ["flags" .= object ["café" .= makeTestFlag "café" 1], "segments" .= object []]
+            ]
+
+-- | 'putEvent' split at the first byte of the two-byte UTF-8 encoding of @é@ (0xC3 0xA9).
+beforeAccent, fromAccent :: ByteString
+(beforeAccent, fromAccent) = B.break (== 0xC3) putEvent
+
+healthyConnection :: ConnectionScript
+healthyConnection = [responseHead, chunk putEvent]
+
+assertInitialized :: StreamRun -> Assertion
+assertInitialized run =
+    assertBool
+        ("stream never initialized; thread exit: " ++ show (exited run) ++ "; logs: " ++ show (logs run))
+        (isJust $ initialized run >>= lookupKey "café")
+
+testHealthyStream :: Test
+testHealthyStream = TestCase $ do
+    run <- runStream [healthyConnection]
+    assertInitialized run
+    assertEqual "connections" 1 (connections run)
+
+testReconnectAfterHttpException :: Test
+testReconnectAfterHttpException = TestCase $ do
+    run <- runStream [[responseHead, Send "zz\r\n"], healthyConnection]
+    assertInitialized run
+    assertEqual "connections" 2 (connections run)
+
+testReconnectAfterSocketReset :: Test
+testReconnectAfterSocketReset = TestCase $ do
+    run <- runStream [[responseHead, chunk beforeAccent, Fail $ toException $ userError "Connection reset by peer"], healthyConnection]
+    assertInitialized run
+    assertEqual "connections" 2 (connections run)
+    assertBool "logs the read error" $
+        any (\(level, msg) -> level == LevelError && T.isPrefixOf "streaming connection failed, will retry" msg) (logs run)
+
+testUtf8SplitAcrossReads :: Test
+testUtf8SplitAcrossReads = TestCase $ do
+    assertBool "payload contains 0xC3" (not $ B.null fromAccent)
+    run <- runStream [[responseHead, chunk beforeAccent, chunk (B.take 1 fromAccent), chunk (B.drop 1 fromAccent)]]
+    assertInitialized run
+    assertEqual "connections" 1 (connections run)
+
+testKillThreadStopsStream :: Test
+testKillThreadStopsStream = TestCase $ do
+    run <- runStream [healthyConnection]
+    assertInitialized run
+    assertBool ("thread stopped by ThreadKilled, got: " ++ show (exited run)) $
+        case exited run of
+            Just (Left e) -> fromException e == Just ThreadKilled
+            _ -> False
+
+testBackoffCeiling :: Test
+testBackoffCeiling = TestCase $ do
+    assertEqual "first attempt waits the initial delay" 1_000_000 (backoffCeiling 1000 1)
+    assertEqual "attempt 46 stays at the 30 s cap" 30_000_000 (backoffCeiling 1000 46)
+    assertEqual "an initial delay too large for Int arithmetic stays at the 30 s cap" 30_000_000 (backoffCeiling maxBound 1)