diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,10 +5,30 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). -## Unreleased +## [1.2] - 2023-01-06 -# [1.1] - 2022-12-16 +### Added +- NARs are now streamed without invoking an external process, so if +you have a lot of small files, there should be some significant performance improvements + +- `cachix deploy activate` now by default waits for the agents to be deployed, displays the logs and exists if any deployments fail. + If you'd like to keep the old behaviour pass `--async` flag. + +## Changed + +- We no longer pin Nix to speed up version bumps of Nix + +## Fixed + +- A number of improvements to stability of the websocket connection used in cachix deploy. + +- Fixed a regression in `cachix deploy activate`, requiring `--agent` flag that should be optional. + +- Fixed a C++ crash that would sometimes happen on exceptions in rare conditions. + +## [1.1] - 2022-12-16 + ### Added - Use ZSTD compresion method by default and allow overriding it via `--compression-method` back to XZ. You can also change the default permanently on your binary cache settings page. @@ -21,13 +41,13 @@ - Generated NixOS module now uses the naming of Nix settings as introduced in NixOS 22.05. -# [1.0.1] - 2022-09-24 +## [1.0.1] - 2022-09-24 ### Added - `cachix config`: allow setting hostname -# [1.0.0] - 2022-09-06 +## [1.0.0] - 2022-09-06 - Cachix Deploy: auto rollback if the agent can't connect to the backend service anymore diff --git a/cachix.cabal b/cachix.cabal --- a/cachix.cabal +++ b/cachix.cabal @@ -1,6 +1,6 @@ cabal-version: 2.2 name: cachix -version: 1.1 +version: 1.2 license: Apache-2.0 license-file: LICENSE copyright: 2018 Domen Kozar @@ -102,9 +102,10 @@ , either , extra , filepath - , fsnotify + , fsnotify >=0.4.1 , hercules-ci-cnix-store , here + , hnix-store-core >=0.6.1.0 , http-client , http-client-tls , http-conduit @@ -147,8 +148,8 @@ , websockets , wuss - -- These versions should match those supported by hercules-ci-cnix-store - pkgconfig-depends: nix-main >=2.4 && <2.11, nix-store >=2.4 && <2.11 + -- These versions are pinned in hercules-ci-cnix-store + pkgconfig-depends: nix-main, nix-store executable cachix import: defaults diff --git a/src/Cachix/Client.hs b/src/Cachix/Client.hs --- a/src/Cachix/Client.hs +++ b/src/Cachix/Client.hs @@ -30,4 +30,4 @@ DeployCommand deployCommand -> case deployCommand of DeployOptions.Agent opts -> AgentCommand.run cachixOptions opts - DeployOptions.Activate opts -> ActivateCommand.run cachixOptions opts + DeployOptions.Activate opts -> ActivateCommand.run env opts diff --git a/src/Cachix/Client/Commands.hs b/src/Cachix/Client/Commands.hs --- a/src/Cachix/Client/Commands.hs +++ b/src/Cachix/Client/Commands.hs @@ -45,7 +45,7 @@ import qualified Data.Text as T import qualified Data.Text.IO as T.IO import GHC.IO.Handle (hDuplicate, hDuplicateTo) -import Hercules.CNix.Store (Store, StorePath, followLinksToStorePath, storePathToPath) +import Hercules.CNix.Store (Store, StorePath, followLinksToStorePath, storePathToPath, withStore) import Network.HTTP.Types (status401, status404) import Protolude hiding (toS) import Protolude.Conv @@ -153,22 +153,22 @@ -- that may or may not be written to by the caller. -- This is somewhat like the behavior of `cat` for example. (_, paths) -> return paths - pushParams <- getPushParams env opts name - normalized <- - liftIO $ - for inputStorePaths $ - \path -> do - storePath <- followLinksToStorePath (pushParamsStore pushParams) (encodeUtf8 path) - filterInvalidStorePath (pushParamsStore pushParams) storePath - pushedPaths <- - pushClosure - (mapConcurrentlyBounded (numJobs opts)) - pushParams - (catMaybes normalized) - case (length normalized, length pushedPaths) of - (0, _) -> putTextError "Nothing to push." - (_, 0) -> putTextError "Nothing to push - all store paths are already on Cachix." - _ -> putTextError "All done." + withPushParams env opts name $ \pushParams -> do + normalized <- + liftIO $ + for inputStorePaths $ + \path -> do + storePath <- followLinksToStorePath (pushParamsStore pushParams) (encodeUtf8 path) + filterInvalidStorePath (pushParamsStore pushParams) storePath + pushedPaths <- + pushClosure + (mapConcurrentlyBounded (numJobs opts)) + pushParams + (catMaybes normalized) + case (length normalized, length pushedPaths) of + (0, _) -> putTextError "Nothing to push." + (_, 0) -> putTextError "Nothing to push - all store paths are already on Cachix." + _ -> putTextError "All done." push _ _ = throwIO $ DeprecatedCommand "DEPRECATED: cachix watch-store has replaced cachix push --watch-store." @@ -178,12 +178,11 @@ watchStore :: Env -> PushOptions -> Text -> IO () watchStore env opts name = do - pushParams <- getPushParams env opts name - WatchStore.startWorkers (pushParamsStore pushParams) (numJobs opts) pushParams + withPushParams env opts name $ \pushParams -> + WatchStore.startWorkers (pushParamsStore pushParams) (numJobs opts) pushParams watchExec :: Env -> PushOptions -> Text -> Text -> [Text] -> IO () -watchExec env pushOpts name cmd args = do - pushParams <- getPushParams env pushOpts name +watchExec env pushOpts name cmd args = withPushParams env pushOpts name $ \pushParams -> do stdoutOriginal <- hDuplicate stdout let process = (System.Process.proc (toS cmd) (toS <$> args)) {System.Process.std_out = System.Process.UseHandle stdoutOriginal} watch = do @@ -221,10 +220,9 @@ Cachix.Client.Push.omitDeriver = Cachix.Client.OptionsParser.omitDeriver opts } -getPushParams :: Env -> PushOptions -> Text -> IO (PushParams IO ()) -getPushParams env pushOpts name = do +withPushParams :: Env -> PushOptions -> Text -> (PushParams IO () -> IO ()) -> IO () +withPushParams env pushOpts name m = do pushSecret <- findPushSecret (config env) name - store <- wait (storeAsync env) authToken <- Config.getAuthTokenMaybe (config env) compressionMethodBackend <- case pushSecret of PushSigningKey {} -> pure Nothing @@ -240,11 +238,12 @@ | otherwise -> throwM err Right binaryCache -> pure (Just $ BinaryCache.preferredCompressionMethod binaryCache) let compressionMethod = fromMaybe BinaryCache.ZSTD (head $ catMaybes [Cachix.Client.OptionsParser.compressionMethod pushOpts, compressionMethodBackend]) - return $ - PushParams - { pushParamsName = name, - pushParamsSecret = pushSecret, - pushParamsClientEnv = clientenv env, - pushParamsStrategy = pushStrategy store authToken pushOpts name compressionMethod, - pushParamsStore = store - } + withStore $ \store -> + m + PushParams + { pushParamsName = name, + pushParamsSecret = pushSecret, + pushParamsClientEnv = clientenv env, + pushParamsStrategy = pushStrategy store authToken pushOpts name compressionMethod, + pushParamsStore = store + } diff --git a/src/Cachix/Client/Env.hs b/src/Cachix/Client/Env.hs --- a/src/Cachix/Client/Env.hs +++ b/src/Cachix/Client/Env.hs @@ -11,6 +11,7 @@ import qualified Cachix.Client.OptionsParser as Options import Cachix.Client.URI (getBaseUrl) import Cachix.Client.Version (cachixVersion) +import qualified Hercules.CNix as CNix import Hercules.CNix.Store (Store, openStore) import Network.HTTP.Client ( ManagerSettings, @@ -28,13 +29,12 @@ data Env = Env { cachixoptions :: Config.CachixOptions, clientenv :: ClientEnv, - config :: Config, - storeAsync :: Async Store + config :: Config } mkEnv :: Options.Flags -> IO Env mkEnv flags = do - store <- async openStore + CNix.init -- make sure path to the config is passed as absolute to dhall logic canonicalConfigPath <- canonicalizePath (Options.configPath flags) cfg <- Config.getConfig canonicalConfigPath @@ -49,8 +49,7 @@ Env { cachixoptions = cachixOptions, clientenv = clientEnv, - config = cfg, - storeAsync = store + config = cfg } customManagerSettings :: ManagerSettings diff --git a/src/Cachix/Client/Push.hs b/src/Cachix/Client/Push.hs --- a/src/Cachix/Client/Push.hs +++ b/src/Cachix/Client/Push.hs @@ -47,7 +47,6 @@ import Data.Coerce (coerce) import Data.Conduit import qualified Data.Conduit.Lzma as Lzma (compress) -import Data.Conduit.Process hiding (env) import qualified Data.Conduit.Zstd as Zstd (compress) import Data.IORef import qualified Data.Set as Set @@ -67,6 +66,7 @@ import Servant.Conduit () import System.Environment (lookupEnv) import qualified System.Nix.Base32 +import System.Nix.Nar data PushSecret = PushToken Token @@ -164,17 +164,13 @@ -- This should be a noop because storePathText came from a StorePath normalized <- liftIO $ Store.followLinksToStorePath store $ toS storePathText pathinfo <- liftIO $ Store.queryPathInfo store normalized - -- stream store path as compressed nar file - let cmd = proc "nix-store" ["--dump", toS storePathText] - storePathSize :: Int64 + -- stream store path as xz compressed nar file + let storePathSize :: Int64 storePathSize = Store.validPathInfoNarSize pathinfo onAttempt strategy retrystatus storePathSize - -- create_group makes subprocess ignore signals such as ctrl-c that we handle in haskell main thread - -- see https://github.com/haskell/process/issues/198 - (ClosedStream, stdoutStream, Inherited, cph) <- liftIO $ streamingProcess (cmd {create_group = True}) withCompressor $ \compressor -> do let stream' = - stdoutStream + streamNarIO narEffectsIO (toS storePathText) Data.Conduit.yield .| passthroughSizeSink narSizeRef .| passthroughHashSink narHashRef .| compressor @@ -189,15 +185,10 @@ clientEnv { baseUrl = (baseUrl clientEnv) {baseUrlHost = subdomain <> "." <> baseUrlHost (baseUrl clientEnv)} } - (_ :: NoContent) <- - liftIO $ - (`withClientM` newClientEnv) - (API.createNar cachixClient (getCacheAuthToken (pushParamsSecret cache)) name (Just $ compressionMethod strategy) (mapOutput coerce stream')) - $ escalate - >=> \NoContent -> do - exitcode <- waitForStreamingProcess cph - when (exitcode /= ExitSuccess) $ throwM $ NarStreamingError exitcode $ show cmd - return NoContent + (_ :: NoContent) <- liftIO $ do + (`withClientM` newClientEnv) + (API.createNar cachixClient (getCacheAuthToken (pushParamsSecret cache)) name (Just $ compressionMethod strategy) (mapOutput coerce stream')) + escalate (_ :: NoContent) <- liftIO $ do narSize <- readIORef narSizeRef narHash <- ("sha256:" <>) . System.Nix.Base32.encode <$> readIORef narHashRef diff --git a/src/Cachix/Client/Servant.hs b/src/Cachix/Client/Servant.hs --- a/src/Cachix/Client/Servant.hs +++ b/src/Cachix/Client/Servant.hs @@ -1,6 +1,6 @@ {-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -O0 #-} -- TODO https://github.com/haskell-servant/servant/issues/986 @@ -8,12 +8,14 @@ module Cachix.Client.Servant ( isErr, cachixClient, - deployClient, + deployClientV1, + deployClientV2, ) where -import qualified Cachix.API -import qualified Cachix.API.Deploy as Cachix.Deploy.API +import qualified Cachix.API as API +import qualified Cachix.API.Deploy.V1 as API.Deploy.V1 +import qualified Cachix.API.Deploy.V2 as API.Deploy.V2 import Cachix.Types.ContentTypes () import Network.HTTP.Types (Status) import Protolude @@ -29,8 +31,11 @@ | Servant.Client.responseStatusCode resp == status = True isErr _ _ = False -cachixClient :: Cachix.API.BinaryCacheAPI (AsClientT ClientM) -cachixClient = fromServant $ client Cachix.API.api +cachixClient :: API.BinaryCacheAPI (AsClientT ClientM) +cachixClient = fromServant $ client (Proxy @API.API) -deployClient :: Cachix.Deploy.API.DeployAPI (AsClientT ClientM) -deployClient = fromServant $ client Cachix.Deploy.API.api +deployClientV1 :: API.Deploy.V1.DeployAPI (AsClientT ClientM) +deployClientV1 = fromServant $ client (Proxy @API.Deploy.V1.API) + +deployClientV2 :: API.Deploy.V2.DeployAPI (AsClientT ClientM) +deployClientV2 = fromServant $ client (Proxy @API.Deploy.V2.API) diff --git a/src/Cachix/Deploy/ActivateCommand.hs b/src/Cachix/Deploy/ActivateCommand.hs --- a/src/Cachix/Deploy/ActivateCommand.hs +++ b/src/Cachix/Deploy/ActivateCommand.hs @@ -1,34 +1,152 @@ module Cachix.Deploy.ActivateCommand where -import qualified Cachix.API.Deploy as API +import qualified Cachix.API.Deploy.V1 as API.V1 +import qualified Cachix.API.Deploy.V2 as API.V2 import Cachix.API.Error (escalate) +import qualified Cachix.API.WebSocketSubprotocol as WSS import qualified Cachix.Client.Config as Config import qualified Cachix.Client.Env as Env -import Cachix.Client.Servant (deployClient) +import qualified Cachix.Client.Retry as Retry +import Cachix.Client.Servant (deployClientV1, deployClientV2) +import qualified Cachix.Client.URI as URI +import Cachix.Client.Version (versionNumber) import qualified Cachix.Deploy.OptionsParser as DeployOptions +import qualified Cachix.Deploy.Websocket as WebSocket +import Cachix.Types.Deploy (Deploy) import qualified Cachix.Types.Deploy as Types import qualified Cachix.Types.DeployResponse as DeployResponse +import qualified Cachix.Types.Deployment as Deployment +import qualified Control.Concurrent.Async as Async import qualified Data.Aeson as Aeson import Data.HashMap.Strict (filterWithKey) import qualified Data.HashMap.Strict as HM +import qualified Data.Text as Text +import qualified Data.Text.IO as Text +import Data.UUID (UUID) +import qualified Data.UUID as UUID +import qualified Network.WebSockets as WS import Protolude hiding (toS) import Protolude.Conv import Servant.Auth.Client (Token (..)) -import Servant.Client.Streaming (runClientM) +import Servant.Client.Streaming (ClientEnv, runClientM) import Servant.Conduit () import System.Environment (getEnv) -run :: Config.CachixOptions -> DeployOptions.ActivateOptions -> IO () -run cachixOptions DeployOptions.ActivateOptions {DeployOptions.payloadPath, DeployOptions.agents} = do - agentToken <- getEnv "CACHIX_ACTIVATE_TOKEN" - clientEnv <- Env.createClientEnv cachixOptions - payloadEither <- Aeson.eitherDecodeFileStrict' payloadPath - case payloadEither of +run :: Env.Env -> DeployOptions.ActivateOptions -> IO () +run env DeployOptions.ActivateOptions {DeployOptions.payloadPath, DeployOptions.agents, DeployOptions.deployAsync} = do + -- TODO: improve the error message here + agentToken <- toS <$> getEnv "CACHIX_ACTIVATE_TOKEN" + Aeson.eitherDecodeFileStrict' payloadPath >>= \case Left err -> do - hPutStrLn stderr $ "Error while parsing JSON: " <> err + hPutStrLn stderr $ "Error parsing the deployment spec: " <> err exitFailure - Right payload -> do - let deploy = payload {Types.agents = filterWithKey (\k _ -> k `elem` agents) (Types.agents payload)} - response <- escalate <=< (`runClientM` clientEnv) $ API.activate deployClient (Token $ toS agentToken) deploy - for_ (HM.toList $ DeployResponse.agents response) $ - \(_, url) -> putStrLn url + Right deploySpec -> do + activate env deployAsync agentToken (filterAgents agents deploySpec) + where + filterAgents [] deploySpec = deploySpec + filterAgents chosenAgents deploySpec = + deploySpec + { Types.agents = filterWithKey (\k _ -> k `elem` chosenAgents) (Types.agents deploySpec) + } + +-- TODO: use prettyprinter +activate :: Env.Env -> Bool -> ByteString -> Deploy -> IO () +activate Env.Env {cachixoptions, clientenv} deployAsync agentToken payload = do + deployResponse <- + escalate <=< (`runClientM` clientenv) $ + API.V2.activate deployClientV2 (Token agentToken) payload + + let agents = HM.toList (DeployResponse.agents deployResponse) + + Text.putStr (renderOverview agents) + + when deployAsync exitSuccess + + Text.putStr "\n\n" + deployments <- Async.mapConcurrently watchDeployments agents + + Text.putStr "\n" + Text.putStr (renderSummary deployments) + + if all isSuccessfulDeployment deployments + then exitSuccess + else exitFailure + where + isSuccessfulDeployment = (==) Deployment.Succeeded . Deployment.status . snd + + watchDeployments (agentName, details) = do + let deploymentID = DeployResponse.id details + host = Config.host cachixoptions + hostname = URI.getHostname host + port = fromMaybe (URI.Port 80) (URI.getPortFor (URI.getScheme host)) + path = "/api/v1/deploy/log/" <> UUID.toText deploymentID <> "?view=true" + useSSL = URI.requiresSSL (URI.getScheme host) + headers = [("Authorization", "Bearer " <> agentToken)] + identifier = unwords ["cachix", versionNumber] + options = + WebSocket.Options + { WebSocket.host = hostname, + WebSocket.port = port, + WebSocket.path = path, + WebSocket.useSSL = useSSL, + WebSocket.headers = headers, + WebSocket.identifier = identifier + } + + deployment <- + Async.withAsync (printLogsToTerminal options agentName) $ \_ -> + pollDeploymentStatus clientenv (Token agentToken) deploymentID + + pure (agentName, deployment) + +pollDeploymentStatus :: ClientEnv -> Token -> UUID -> IO Deployment.Deployment +pollDeploymentStatus clientEnv token deploymentID = loop + where + loop = do + deployment <- + Retry.retryAll . const $ + escalate <=< (`runClientM` clientEnv) $ + API.V1.getDeployment deployClientV1 token deploymentID + + case Deployment.status deployment of + Deployment.Cancelled -> pure deployment + Deployment.Failed -> pure deployment + Deployment.Succeeded -> pure deployment + _ -> do + threadDelay (2 * 1000 * 1000) + loop + +printLogsToTerminal :: WebSocket.Options -> Text -> IO a +printLogsToTerminal options agentName = + WebSocket.runClientWith options WS.defaultConnectionOptions $ \connection -> + forever $ do + message <- WS.receiveData connection + case Aeson.eitherDecodeStrict' message of + Left error -> Text.putStrLn $ "Error parsing the log message: " <> show error + Right msg -> Text.putStrLn $ unwords [inBrackets agentName, WSS.line msg] + +renderOverview :: [(Text, DeployResponse.Details)] -> Text +renderOverview agents = + Text.intercalate "\n" $ + "Deploying agents:" : + [ inBrackets agentName <> " " <> DeployResponse.url details + | (agentName, details) <- agents + ] + +renderSummary :: [(Text, Deployment.Deployment)] -> Text +renderSummary results = + Text.intercalate "\n" $ + "Deployment summary:" : + [ inBrackets agentName <> " " <> renderStatus (Deployment.status deployment) + | (agentName, deployment) <- results + ] + where + renderStatus = \case + Deployment.Succeeded -> "Deployed successfully" + Deployment.Failed -> "Failed to deploy" + Deployment.Cancelled -> "Deployment cancelled" + Deployment.InProgress -> "Still deploying" + Deployment.Pending -> "Deployment not started" + +inBrackets :: (Semigroup a, IsString a) => a -> a +inBrackets s = "[" <> s <> "]" diff --git a/src/Cachix/Deploy/Agent.hs b/src/Cachix/Deploy/Agent.hs --- a/src/Cachix/Deploy/Agent.hs +++ b/src/Cachix/Deploy/Agent.hs @@ -219,16 +219,14 @@ -- sending deployment progress updates. connectToService :: ServiceWebSocket -> IO (IO ()) connectToService websocket = do - initialConnection <- MVar.newEmptyMVar close <- MVar.newEmptyMVar thread <- Async.async $ WebSocket.runConnection websocket $ do - MVar.putMVar initialConnection () WebSocket.handleJSONMessages websocket (MVar.readMVar close) - -- Block until the connection has been established - void $ MVar.takeMVar initialConnection + -- Block until the initial connection is established + void $ MVar.readMVar (WebSocket.connection websocket) once $ MVar.putMVar close () >> Async.wait thread diff --git a/src/Cachix/Deploy/Log.hs b/src/Cachix/Deploy/Log.hs --- a/src/Cachix/Deploy/Log.hs +++ b/src/Cachix/Deploy/Log.hs @@ -70,7 +70,7 @@ .| sendLog connection streamLine :: LogStream -> ByteString -> IO () -streamLine logStream msg = Conduit.connect (Conduit.yield $ "\n" <> msg <> "\n") logStream +streamLine logStream msg = Conduit.connect (Conduit.yield $ msg <> "\n") logStream logLocal :: WithLog -> Conduit.ConduitT ByteString ByteString IO () logLocal logger = diff --git a/src/Cachix/Deploy/OptionsParser.hs b/src/Cachix/Deploy/OptionsParser.hs --- a/src/Cachix/Deploy/OptionsParser.hs +++ b/src/Cachix/Deploy/OptionsParser.hs @@ -34,7 +34,8 @@ data ActivateOptions = ActivateOptions { payloadPath :: FilePath, - agents :: [Text] + agents :: [Text], + deployAsync :: Bool } deriving (Show) @@ -68,3 +69,4 @@ <> help "Deploy only specific agent(s)." ) ) + <*> switch (long "async" <> help "Skip waiting for the agents to deploy") diff --git a/src/Cachix/Deploy/Websocket.hs b/src/Cachix/Deploy/Websocket.hs --- a/src/Cachix/Deploy/Websocket.hs +++ b/src/Cachix/Deploy/Websocket.hs @@ -18,7 +18,6 @@ import qualified Control.Retry as Retry import qualified Data.Aeson as Aeson import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy as BL import Data.Conduit ((.|)) import qualified Data.Conduit as Conduit import qualified Data.Conduit.Combinators as Conduit @@ -92,31 +91,32 @@ read = atomically . TMChan.readTMChan -- | Read incoming data messages, ignoring everything else. --- --- Stops reading when either a close request is received or the channel is --- closed. readDataMessages :: Receive rx -> (rx -> IO ()) -> IO () readDataMessages channel action = loop where loop = read channel >>= \case Just (DataMessage message) -> action message *> loop - Just (ControlMessage (WS.Close _ _)) -> pure () Just (ControlMessage _) -> loop Nothing -> pure () -- | Close the outgoing queue. -drainQueue :: WebSocket tx rx -> IO () -drainQueue WebSocket {tx} = atomically $ TBMQueue.closeTBMQueue tx +drainQueue :: WebSocket tx rx -> Async () -> IO () +drainQueue WebSocket {tx} outgoingThread = do + atomically $ TBMQueue.closeTBMQueue tx + Async.wait outgoingThread -shutdownNow :: WebSocket tx rx -> Word16 -> BL.ByteString -> IO () -shutdownNow websocket@WebSocket {rx} code msg = do - -- Close the outgoing queue - drainQueue websocket - -- Signal to all receivers that the socket is closed - atomically (TMChan.closeTMChan rx) - throwIO $ WS.CloseRequest code msg +closeGracefully :: WebSocket tx rx -> Async () -> IO () +closeGracefully websocket incomingThread = do + repsonseToCloseRequest <- startGracePeriod $ do + MVar.tryReadMVar (connection websocket) >>= \case + Just activeConnection -> do + WS.sendClose activeConnection ("Peer initiated a close request" :: ByteString) + Async.wait incomingThread + Nothing -> pure () + when (isNothing repsonseToCloseRequest) throwNoResponseToCloseRequest + -- | Run an app inside a new WebSocket connection. withConnection :: Log.WithLog -> Options -> (WebSocket tx rx -> IO ()) -> IO () withConnection withLog options app = do @@ -188,26 +188,25 @@ -- dropping messages. handleJSONMessages :: (Aeson.ToJSON tx, Aeson.FromJSON rx) => WebSocket tx rx -> IO () -> IO () handleJSONMessages websocket app = - Async.withAsync (handleIncomingJSON websocket) $ \incomingThread -> - Async.withAsync (handleOutgoingJSON websocket) $ \outgoingThread -> - Async.withAsync - ( do + handleJust unwrapThreadExceptions throwIO $ + mask $ \restore -> do + incomingThread <- Async.async (handleIncomingJSON websocket) + outgoingThread <- Async.async (handleOutgoingJSON websocket) + + let threads = [incomingThread, outgoingThread] + cancelThreads = mapM_ Async.uninterruptibleCancel threads + mapM_ Async.link threads + + let runApp = do app - drainQueue websocket - Async.wait outgoingThread - closeGracefully incomingThread - ) - (\appThread -> mapM_ wait [appThread, outgoingThread, incomingThread]) + drainQueue websocket outgoingThread + closeGracefully websocket incomingThread + restore (runApp `Safe.finally` cancelThreads) where - closeGracefully incomingThread = do - repsonseToCloseRequest <- startGracePeriod $ do - MVar.tryReadMVar (connection websocket) >>= \case - Just activeConnection -> do - WS.sendClose activeConnection ("Peer initiated a close request" :: ByteString) - Async.wait incomingThread - Nothing -> pure () - - when (isNothing repsonseToCloseRequest) throwNoResponseToCloseRequest + unwrapThreadExceptions :: SomeException -> Maybe SomeException + unwrapThreadExceptions e + | Just (Async.ExceptionInLinkedThread _ e') <- fromException e = Just e' + | otherwise = Nothing handleIncomingJSON :: (Aeson.FromJSON rx) => WebSocket tx rx -> IO () handleIncomingJSON websocket@WebSocket {connection, rx, withLog} = do @@ -223,12 +222,14 @@ Right pMsg -> broadcast (DataMessage pMsg) WS.ControlMessage controlMsg -> do case controlMsg of - WS.Ping pl -> send websocket (ControlMessage (WS.Pong pl)) - WS.Pong _ -> WS.connectionOnPong (WS.Connection.connectionOptions activeConnection) + WS.Ping pl -> + send websocket (ControlMessage (WS.Pong pl)) + WS.Pong _ -> + WS.connectionOnPong (WS.Connection.connectionOptions activeConnection) WS.Close code closeMsg -> do hasSentClose <- IORef.readIORef $ WS.Connection.connectionSentClose activeConnection unless hasSentClose $ WS.send activeConnection msg - shutdownNow websocket code closeMsg + throwIO $ WS.CloseRequest code closeMsg broadcast (ControlMessage controlMsg)