server: configurable websocket keep alive interval (#6092)

Accept new server flag --websocket-keepalive to control
websockets keep-alive interval

Co-authored-by: Auke Booij <auke@hasura.io>
Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -120,6 +120,8 @@ This release contains the [PDV refactor (#4111)](https://github.com/hasura/graph
 - server: accept only non-negative integers for batch size and refetch interval (close #5653) (#5759)
 - server: fix bug which arised when renaming a table which had a manual relationship defined (close #4158)
 - server: limit the length of event trigger names (close #5786)
+- server: Configurable websocket keep-alive interval. Add `--websocket-keepalive` command-line flag
+          and handle `HASURA_GRAPHQL_WEBSOCKET_KEEPALIVE` env variable (fix #3539) 
 **NOTE:** If you have event triggers with names greater than 42 chars, then you should update their names to avoid running into Postgres identifier limit bug (#5786)
 - server: validate remote schema queries (fixes #4143)
 - server: fix issue with tracking custom functions that return `SETOF` materialized view (close #5294) (#5945)

server/src-lib/Hasura/App.hs

@@ -339,7 +339,6 @@ runHGEServer env ServeOptions{..} InitCtx{..} pgExecCtx initTime shutdownApp pos
 
   _idleGCThread <- C.forkImmortal "ourIdleGC" logger $ liftIO $
     ourIdleGC logger (seconds 0.3) (seconds 10) (seconds 60)
-
   HasuraApp app cacheRef cacheInitTime stopWsServer <- flip onException (flushLogger loggerCtx) $
     mkWaiApp env
              soTxIso
@@ -364,6 +363,7 @@ runHGEServer env ServeOptions{..} InitCtx{..} pgExecCtx initTime shutdownApp pos
              _icSchemaCache
              ekgStore
              soConnectionOptions
+             soWebsocketKeepAlive
 
   -- log inconsistent schema objects
   inconsObjs <- scInconsistentObjs <$> liftIO (getSCFromRef cacheRef)
@@ -429,7 +429,7 @@ runHGEServer env ServeOptions{..} InitCtx{..} pgExecCtx initTime shutdownApp pos
                         , eventQueueThread
                         , scheduledEventsThread
                         , cronEventsThread
-                        ] <> maybe [] pure telemetryThread
+                        ] <> onNothing telemetryThread []
 
   finishTime <- liftIO Clock.getCurrentTime
   let apiInitTime = realToFrac $ Clock.diffUTCTime finishTime initTime

server/src-lib/Hasura/GraphQL/Transport/WebSocket.hs

@@ -73,6 +73,7 @@ import qualified Hasura.GraphQL.Transport.WebSocket.Server   as WS
 import qualified Hasura.Logging                              as L
 import qualified Hasura.Server.Telemetry.Counters            as Telem
 import qualified Hasura.Tracing                              as Tracing
+import           Hasura.Server.Init.Config (KeepAliveDelay (..))
 
 -- | 'LQ.LiveQueryId' comes from 'Hasura.GraphQL.Execute.LiveQuery.State.addLiveQuery'. We use
 -- this to track a connection's operations so we can remove them from 'LiveQueryState', and
@@ -228,11 +229,12 @@ data WSServerEnv
   -- , _wseQueryCache      :: !E.PlanCache -- See Note [Temporarily disabling query plan caching]
   , _wseServer          :: !WSServer
   , _wseEnableAllowlist :: !Bool
+  , _wseKeepAliveDelay  :: !KeepAliveDelay
   }
 
-onConn :: (MonadIO m)
-       => L.Logger L.Hasura -> CorsPolicy -> WS.OnConnH m WSConnData
-onConn (L.Logger logger) corsPolicy wsId requestHead ipAddress = do
+onConn :: (MonadIO m, MonadReader WSServerEnv m)
+       => WS.OnConnH m WSConnData
+onConn wsId requestHead ipAddress = do
   res <- runExceptT $ do
     (errType, queryType) <- checkPath
     let reqHdrs = WS.requestHeaders requestHead
@@ -241,9 +243,10 @@ onConn (L.Logger logger) corsPolicy wsId requestHead ipAddress = do
   either reject accept res
 
   where
-    keepAliveAction wsConn = liftIO $ forever $ do
-      sendMsg wsConn SMConnKeepAlive
-      sleep $ seconds 5
+    keepAliveAction keepAliveDelay wsConn = do
+      liftIO $ forever $ do
+        sendMsg wsConn SMConnKeepAlive
+        sleep $ seconds (unKeepAliveDelay keepAliveDelay)
 
     tokenExpiryHandler wsConn = do
       expTime <- liftIO $ STM.atomically $ do
@@ -256,6 +259,8 @@ onConn (L.Logger logger) corsPolicy wsId requestHead ipAddress = do
       sleep $ convertDuration $ TC.diffUTCTime expTime currTime
 
     accept (hdrs, errType, queryType) = do
+      (L.Logger logger) <- asks _wseLogger
+      keepAliveDelay <- asks _wseKeepAliveDelay
       logger $ mkWsInfoLog Nothing (WsConnInfo wsId Nothing Nothing) EAccepted
       connData <- liftIO $ WSConnData
                   <$> STM.newTVarIO (CSNotInitialised hdrs ipAddress)
@@ -264,9 +269,9 @@ onConn (L.Logger logger) corsPolicy wsId requestHead ipAddress = do
                   <*> pure queryType
       let acceptRequest = WS.defaultAcceptRequest
                           { WS.acceptSubprotocol = Just "graphql-ws"}
-      return $ Right $ WS.AcceptWith connData acceptRequest keepAliveAction tokenExpiryHandler
-
+      return $ Right $ WS.AcceptWith connData acceptRequest (keepAliveAction keepAliveDelay) tokenExpiryHandler
     reject qErr = do
+      (L.Logger logger) <- asks _wseLogger
       logger $ mkWsErrorLog Nothing (WsConnInfo wsId Nothing Nothing) (ERejected qErr)
       return $ Left $ WS.RejectRequest
         (H.statusCode $ qeStatus qErr)
@@ -283,21 +288,24 @@ onConn (L.Logger logger) corsPolicy wsId requestHead ipAddress = do
     getOrigin =
       find ((==) "Origin" . fst) (WS.requestHeaders requestHead)
 
-    enforceCors origin reqHdrs = case cpConfig corsPolicy of
-      CCAllowAll -> return reqHdrs
-      CCDisabled readCookie ->
-        if readCookie
-        then return reqHdrs
-        else do
-          lift $ logger $ mkWsInfoLog Nothing (WsConnInfo wsId Nothing (Just corsNote)) EAccepted
-          return $ filter (\h -> fst h /= "Cookie") reqHdrs
-      CCAllowedOrigins ds
-        -- if the origin is in our cors domains, no error
-        | bsToTxt origin `elem` dmFqdns ds   -> return reqHdrs
-        -- if current origin is part of wildcard domain list, no error
-        | inWildcardList ds (bsToTxt origin) -> return reqHdrs
-        -- otherwise error
-        | otherwise                          -> corsErr
+    enforceCors origin reqHdrs = do
+      (L.Logger logger) <- asks _wseLogger
+      corsPolicy <- asks _wseCorsPolicy
+      case cpConfig corsPolicy of
+        CCAllowAll -> return reqHdrs
+        CCDisabled readCookie ->
+          if readCookie
+          then return reqHdrs
+          else do
+            lift $ logger $ mkWsInfoLog Nothing (WsConnInfo wsId Nothing (Just corsNote)) EAccepted
+            return $ filter (\h -> fst h /= "Cookie") reqHdrs
+        CCAllowedOrigins ds
+          -- if the origin is in our cors domains, no error
+          | bsToTxt origin `elem` dmFqdns ds   -> return reqHdrs
+          -- if current origin is part of wildcard domain list, no error
+          | inWildcardList ds (bsToTxt origin) -> return reqHdrs
+          -- otherwise error
+          | otherwise                          -> corsErr
 
     filterWsHeaders hdrs = flip filter hdrs $ \(n, _) ->
       n `notElem` [ "sec-websocket-key"
@@ -444,7 +452,7 @@ onStart env serverEnv wsConn (StartMsg opId q) = catchAndIgnore $ do
       return $ ResultsFragment telemTimeIO_DT Telem.Remote (JO.toEncJSON value) []
 
     WSServerEnv logger pgExecCtx lqMap getSchemaCache httpMgr _ sqlGenCtx {- planCache -}
-      _ enableAL = serverEnv
+      _ enableAL _keepAliveDelay = serverEnv
 
     WSConnData userInfoR opMap errRespTy queryType = WS.getData wsConn
 
@@ -690,14 +698,15 @@ createWSServerEnv
   -> CorsPolicy
   -> SQLGenCtx
   -> Bool
+  -> KeepAliveDelay
   -- -> E.PlanCache
   -> m WSServerEnv
 createWSServerEnv logger isPgCtx lqState getSchemaCache httpManager
-  corsPolicy sqlGenCtx enableAL {- planCache -} = do
+  corsPolicy sqlGenCtx enableAL keepAliveDelay {- planCache -} = do
   wsServer <- liftIO $ STM.atomically $ WS.createWSServer logger
   return $
     WSServerEnv logger isPgCtx lqState getSchemaCache httpManager corsPolicy
-    sqlGenCtx {- planCache -} wsServer enableAL
+    sqlGenCtx {- planCache -} wsServer enableAL keepAliveDelay
 
 createWSServerApp
   :: ( HasVersion
@@ -723,7 +732,7 @@ createWSServerApp env authMode serverEnv = \ !ipAddress !pendingConn ->
     handlers =
       WS.WSHandlers
       -- Mask async exceptions during event processing to help maintain integrity of mutable vars:
-      (\rid rh ip -> mask_ $ onConn (_wseLogger serverEnv) (_wseCorsPolicy serverEnv) rid rh ip)
+      (\rid rh ip -> mask_ $ flip runReaderT serverEnv $ onConn rid rh ip)
       (\conn bs -> mask_ $ onMessage env authMode serverEnv conn bs)
       (mask_ . onClose (_wseLogger serverEnv) (_wseLiveQMap serverEnv))
 

server/src-lib/Hasura/Server/App.hs

@@ -333,7 +333,7 @@ mkSpockAction serverCtx qErrEncoder qErrModifier apiHandler = do
       possiblyCompressedLazyBytes userInfo reqId waiReq req qTime respBytes respHeaders reqHeaders = do
         let (compressedResp, mEncodingHeader, mCompressionType) =
               compressResponse (Wai.requestHeaders waiReq) respBytes
-            encodingHeader = maybe [] pure mEncodingHeader
+            encodingHeader = onNothing mEncodingHeader []
             reqIdHeader = (requestIdHeader, txtToBs $ unRequestId reqId)
             allRespHeaders = pure reqIdHeader <> encodingHeader <> respHeaders
         lift $ logHttpSuccess logger userInfo reqId waiReq req respBytes compressedResp qTime mCompressionType reqHeaders
@@ -602,9 +602,10 @@ mkWaiApp
   -> (RebuildableSchemaCache Run, Maybe UTCTime)
   -> EKG.Store
   -> WS.ConnectionOptions
+  -> KeepAliveDelay
   -> m HasuraApp
 mkWaiApp env isoLevel logger sqlGenCtx enableAL pool pgExecCtxCustom ci httpManager mode corsCfg enableConsole consoleAssetsDir
-         enableTelemetry instanceId apis lqOpts _ {- planCacheOptions -} responseErrorsConfig liveQueryHook (schemaCache, cacheBuiltTime) ekgStore connectionOptions  = do
+         enableTelemetry instanceId apis lqOpts _ {- planCacheOptions -} responseErrorsConfig liveQueryHook (schemaCache, cacheBuiltTime) ekgStore connectionOptions keepAliveDelay = do
 
     -- See Note [Temporarily disabling query plan caching]
     -- (planCache, schemaCacheRef) <- initialiseCache
@@ -617,7 +618,7 @@ mkWaiApp env isoLevel logger sqlGenCtx enableAL pool pgExecCtxCustom ci httpMana
 
     lqState <- liftIO $ EL.initLiveQueriesState lqOpts pgExecCtx postPollHook
     wsServerEnv <- WS.createWSServerEnv logger pgExecCtx lqState getSchemaCache httpManager
-                                        corsPolicy sqlGenCtx enableAL {- planCache -}
+                                        corsPolicy sqlGenCtx enableAL keepAliveDelay {- planCache -}
 
     let serverCtx = ServerCtx
                     { scPGExecCtx       =  pgExecCtx

server/src-lib/Hasura/Server/Init.hs

@@ -180,13 +180,15 @@ mkServeOptions rso = do
                                 then WS.PermessageDeflateCompression WS.defaultPermessageDeflate
                                 else WS.NoCompression
                           }
+  webSocketKeepAlive <- KeepAliveDelay . fromIntegral . fromMaybe 5
+      <$> withEnv (rsoWebSocketKeepAlive rso) (fst webSocketKeepAliveEnv)
 
   return $ ServeOptions port host connParams txIso adminScrt authHook jwtSecret
                         unAuthRole corsCfg enableConsole consoleAssetsDir
                         enableTelemetry strfyNum enabledAPIs lqOpts enableAL
                         enabledLogs serverLogLevel planCacheOptions
                         internalErrorsConfig eventsHttpPoolSize eventsFetchInterval
-                        logHeadersFromEnv connectionOptions
+                        logHeadersFromEnv connectionOptions webSocketKeepAlive
   where
 #ifdef DeveloperAPIs
     defaultAPIs = [METADATA,GRAPHQL,PGDUMP,CONFIG,DEVELOPER]
@@ -325,7 +327,7 @@ serveCmdFooter =
       , jwtSecretEnv, unAuthRoleEnv, corsDomainEnv, corsDisableEnv, enableConsoleEnv
       , enableTelemetryEnv, wsReadCookieEnv, stringifyNumEnv, enabledAPIsEnv
       , enableAllowlistEnv, enabledLogsEnv, logLevelEnv, devModeEnv
-      , adminInternalErrorsEnv
+      , adminInternalErrorsEnv, webSocketKeepAliveEnv
       ]
 
     eventEnvs = [ eventsHttpPoolSizeEnv, eventsFetchIntervalEnv ]
@@ -943,6 +945,7 @@ serveOptsToLog so =
       , "log_level" J..= soLogLevel so
       , "plan_cache_options" J..= soPlanCacheOptions so
       , "websocket_compression_options" J..= show (WS.connectionCompressionOptions . soConnectionOptions $ so)
+      , "websocket_keep_alive" J..= show (soWebsocketKeepAlive so)
       ]
 
 mkGenericStrLog :: L.LogLevel -> Text -> String -> StartupLog
@@ -989,6 +992,7 @@ serveOptionsParser =
   <*> parseGraphqlEventsFetchInterval
   <*> parseLogHeadersFromEnv
   <*> parseWebSocketCompression
+  <*> parseWebSocketKeepAlive
 
 -- | This implements the mapping between application versions
 -- and catalog schema versions.
@@ -1035,3 +1039,17 @@ parseWebSocketCompression =
   switch ( long "websocket-compression" <>
            help (snd webSocketCompressionEnv)
          )
+
+webSocketKeepAliveEnv :: (String, String)
+webSocketKeepAliveEnv =
+  ( "HASURA_GRAPHQL_WEBSOCKET_KEEPALIVE"
+  , "Control websocket keep-alive timeout (default 5 seconds)"
+  )
+
+parseWebSocketKeepAlive :: Parser (Maybe Int)
+parseWebSocketKeepAlive =
+  optional $
+  option (eitherReader readEither)
+         ( long "websocket-keepalive" <>
+           help (snd webSocketKeepAliveEnv)
+         )

server/src-lib/Hasura/Server/Init/Config.hs

@@ -65,7 +65,8 @@ data RawServeOptions impl
   , rsoEventsHttpPoolSize    :: !(Maybe Int)
   , rsoEventsFetchInterval   :: !(Maybe Milliseconds)
   , rsoLogHeadersFromEnv     :: !Bool
-  , rsoWebSocketCompression :: !Bool
+  , rsoWebSocketCompression  :: !Bool
+  , rsoWebSocketKeepAlive    :: !(Maybe Int)
   }
 
 -- | @'ResponseInternalErrorsConfig' represents the encoding of the internal
@@ -83,6 +84,11 @@ shouldIncludeInternal role = \case
   InternalErrorsAdminOnly   -> isAdmin role
   InternalErrorsDisabled    -> False
 
+newtype KeepAliveDelay
+  = KeepAliveDelay
+      { unKeepAliveDelay :: Seconds
+      } deriving (Eq, Show)
+
 data ServeOptions impl
   = ServeOptions
   { soPort                         :: !Int
@@ -109,6 +115,7 @@ data ServeOptions impl
   , soEventsFetchInterval          :: !(Maybe Milliseconds)
   , soLogHeadersFromEnv            :: !Bool
   , soConnectionOptions            :: !WS.ConnectionOptions
+  , soWebsocketKeepAlive           :: !KeepAliveDelay
   }
 
 data DowngradeOptions