Skip to content

Commit

Permalink
Fixed data subscriptions not always closed in specific cases
Browse files Browse the repository at this point in the history 
We switched from manual resource handling to Exception.bracket with MVars to make sure that all data subscriptions are always closed when the underlying websocket connection is closed
  • Loading branch information
@mpscholten
mpscholten committed Feb 8, 2022
1 parent 2c5afc0 commit 501bbef
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 62 deletions.
134 changes: 75 additions & 59 deletions IHP/DataSync/Controller.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ instance (
initialState = DataSyncController

run = do
setState DataSyncReady { subscriptions = HashMap.empty, transactions = HashMap.empty }
setState DataSyncReady { subscriptions = HashMap.empty, transactions = HashMap.empty, asyncs = [] }

ensureRLSEnabled <- makeCachedEnsureRLSEnabled
installTableChangeTriggers <- ChangeNotifications.makeCachedInstallTableChangeTriggers
Expand All @@ -56,6 +56,8 @@ instance (
sendJSON DataSyncResult { result, requestId }

handleMessage CreateDataSubscription { query, requestId } = do
ensureBelowSubscriptionsLimit

tableNameRLS <- ensureRLSEnabled (get #table query)

subscriptionId <- UUID.nextRandom
Expand Down Expand Up @@ -111,22 +113,23 @@ instance (
when isWatchingRecord do
sendJSON DidDelete { subscriptionId, id }
let subscribe = PGListener.subscribeJSON (ChangeNotifications.channelName tableNameRLS) callback pgListener
let unsubscribe subscription = PGListener.unsubscribe subscription pgListener
channelSubscription <- pgListener
|> PGListener.subscribeJSON (ChangeNotifications.channelName tableNameRLS) callback
Exception.bracket subscribe unsubscribe \channelSubscription -> do
close <- MVar.newEmptyMVar
modifyIORef' ?state (\state -> state |> modify #subscriptions (HashMap.insert subscriptionId close))
modifyIORef' ?state (\state -> state |> modify #subscriptions (HashMap.insert subscriptionId Subscription { id = subscriptionId, channelSubscription }))
sendJSON DidCreateDataSubscription { subscriptionId, requestId, result }
sendJSON DidCreateDataSubscription { subscriptionId, requestId, result }
MVar.takeMVar close
handleMessage DeleteDataSubscription { requestId, subscriptionId } = do
DataSyncReady { subscriptions } <- getState
let maybeSubscription :: Maybe Subscription = HashMap.lookup subscriptionId subscriptions
let (Just closeSignalMVar) = HashMap.lookup subscriptionId subscriptions
-- Cancel table watcher
case maybeSubscription of
Just subscription -> pgListener |> PGListener.unsubscribe (get #channelSubscription subscription)
Nothing -> pure ()
MVar.putMVar closeSignalMVar ()
modifyIORef' ?state (\state -> state |> modify #subscriptions (HashMap.delete subscriptionId))
Expand Down Expand Up @@ -260,35 +263,49 @@ instance (
ensureBelowTransactionLimit
transactionId <- UUID.nextRandom
(connection, localPool) <- ?modelContext
|> get #connectionPool
|> Pool.takeResource
let transaction = DataSyncTransaction
{ id = transactionId
, connection
, releaseConnection = Pool.putResource localPool connection
}
let globalModelContext = ?modelContext
let ?modelContext = globalModelContext { transactionConnection = Just connection } in sqlExecWithRLS "BEGIN" ()
let takeConnection = ?modelContext
|> get #connectionPool
|> Pool.takeResource
let releaseConnection (connection, localPool) = do
PG.execute connection "ROLLBACK" () -- Make sure there's no pending transaction in case something went wrong
Pool.putResource localPool connection
Exception.bracket takeConnection releaseConnection \(connection, localPool) -> do
transactionSignal <- MVar.newEmptyMVar
let globalModelContext = ?modelContext
let ?modelContext = globalModelContext { transactionConnection = Just connection } in sqlExecWithRLS "BEGIN" ()
modifyIORef' ?state (\state -> state |> modify #transactions (HashMap.insert transactionId transaction))
let transaction = DataSyncTransaction
{ id = transactionId
, connection
, close = transactionSignal
}
sendJSON DidStartTransaction { requestId, transactionId }
modifyIORef' ?state (\state -> state |> modify #transactions (HashMap.insert transactionId transaction))
sendJSON DidStartTransaction { requestId, transactionId }
MVar.takeMVar transactionSignal
modifyIORef' ?state (\state -> state |> modify #transactions (HashMap.delete transactionId))
handleMessage RollbackTransaction { requestId, id } = do
sqlExecWithRLSAndTransactionId (Just id) "ROLLBACK" ()
DataSyncTransaction { id, close } <- findTransactionById id
closeTransaction id
sqlExecWithRLSAndTransactionId (Just id) "ROLLBACK" ()
MVar.putMVar close ()
sendJSON DidRollbackTransaction { requestId, transactionId = id }
handleMessage CommitTransaction { requestId, id } = do
sqlExecWithRLSAndTransactionId (Just id) "COMMIT" ()
DataSyncTransaction { id, close } <- findTransactionById id
closeTransaction id
sqlExecWithRLSAndTransactionId (Just id) "COMMIT" ()
MVar.putMVar close ()
sendJSON DidCommitTransaction { requestId, transactionId = id }
Expand All @@ -301,22 +318,24 @@ instance (
Right decodedMessage -> do
let requestId = get #requestId decodedMessage
-- Handle the messages in an async way
-- This increases throughput as multiple queries can be fetched
-- in parallel
async do
result <- Exception.try (handleMessage decodedMessage)
case result of
Left (e :: Exception.SomeException) -> do
let errorMessage = case fromException e of
Just (enhancedSqlError :: EnhancedSqlError) -> cs (get #sqlErrorMsg (get #sqlError enhancedSqlError))
Nothing -> cs (displayException e)
Log.error (tshow e)
sendJSON DataSyncError { requestId, errorMessage }
Right result -> pure ()
pure ()
Exception.mask \restore -> do
-- Handle the messages in an async way
-- This increases throughput as multiple queries can be fetched
-- in parallel
handlerProcess <- async $ restore do
result <- Exception.try (handleMessage decodedMessage)
case result of
Left (e :: Exception.SomeException) -> do
let errorMessage = case fromException e of
Just (enhancedSqlError :: EnhancedSqlError) -> cs (get #sqlErrorMsg (get #sqlError enhancedSqlError))
Nothing -> cs (displayException e)
Log.error (tshow e)
sendJSON DataSyncError { requestId, errorMessage }
Right result -> pure ()
modifyIORef' ?state (\state -> state |> modify #asyncs (handlerProcess:))
pure ()
Left errorMessage -> sendJSON FailedToDecodeMessageError { errorMessage = cs errorMessage }
onClose = cleanupAllSubscriptions
Expand All @@ -327,16 +346,7 @@ cleanupAllSubscriptions = do
let pgListener = ?applicationContext |> get #pgListener
case state of
DataSyncReady { subscriptions, transactions } -> do
let channelSubscriptions = subscriptions
|> HashMap.elems
|> map (get #channelSubscription)
forEach channelSubscriptions \channelSubscription -> do
pgListener |> PGListener.unsubscribe channelSubscription
forEach (HashMap.elems transactions) (get #releaseConnection)
pure ()
DataSyncReady { asyncs } -> forEach asyncs uninterruptibleCancel
_ -> pure ()
changesToValue :: [ChangeNotifications.Change] -> Value
Expand Down Expand Up @@ -369,11 +379,6 @@ findTransactionById transactionId = do
Just transaction -> pure transaction
Nothing -> error "No transaction with that id"
closeTransaction transactionId = do
DataSyncTransaction { releaseConnection } <- findTransactionById transactionId
modifyIORef' ?state (\state -> state |> modify #transactions (HashMap.delete transactionId))
releaseConnection
-- | Allow max 10 concurrent transactions per connection to avoid running out of database connections
--
-- Each transaction removes a database connection from the connection pool. If we don't limit the transactions,
Expand All @@ -389,6 +394,14 @@ ensureBelowTransactionLimit = do
when (transactionCount >= maxTransactionsPerConnection) do
error ("You've reached the transaction limit of " <> tshow maxTransactionsPerConnection <> " transactions")
ensureBelowSubscriptionsLimit :: (?state :: IORef DataSyncController) => IO ()
ensureBelowSubscriptionsLimit = do
subscriptions <- get #subscriptions <$> readIORef ?state
let subscriptionsCount = HashMap.size subscriptions
let maxSubscriptionsPerConnection = 128
when (subscriptionsCount >= maxSubscriptionsPerConnection) do
error ("You've reached the subscriptions limit of " <> tshow maxSubscriptionsPerConnection <> " subscriptions")
sqlQueryWithRLSAndTransactionId ::
( ?modelContext :: ModelContext
, PG.ToRow parameters
Expand Down Expand Up @@ -423,8 +436,11 @@ sqlExecWithRLSAndTransactionId transactionId theQuery theParams = runInModelCont
$(deriveFromJSON defaultOptions 'DataSyncQuery)
$(deriveToJSON defaultOptions 'DataSyncResult)
instance SetField "subscriptions" DataSyncController (HashMap UUID Subscription) where
instance SetField "subscriptions" DataSyncController (HashMap UUID (MVar.MVar ())) where
setField subscriptions record = record { subscriptions }
instance SetField "transactions" DataSyncController (HashMap UUID DataSyncTransaction) where
setField transactions record = record { transactions }
setField transactions record = record { transactions }
instance SetField "asyncs" DataSyncController [Async ()] where
setField asyncs record = record { asyncs }
7 changes: 4 additions & 3 deletions IHP/DataSync/Types.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import IHP.DataSync.DynamicQuery
import Data.HashMap.Strict (HashMap)
import qualified IHP.PGListener as PGListener
import qualified Database.PostgreSQL.Simple as PG
import Control.Concurrent.MVar as MVar

data DataSyncMessage
= DataSyncQuery { query :: !DynamicSQLQuery, requestId :: !Int, transactionId :: !(Maybe UUID) }
Expand Down Expand Up @@ -42,17 +43,17 @@ data DataSyncResponse
| DidRollbackTransaction { requestId :: !Int, transactionId :: !UUID }
| DidCommitTransaction { requestId :: !Int, transactionId :: !UUID }

data Subscription = Subscription { id :: !UUID, channelSubscription :: !PGListener.Subscription }
data DataSyncTransaction
= DataSyncTransaction
{ id :: !UUID
, connection :: !PG.Connection
, releaseConnection :: IO ()
, close :: MVar ()
}

data DataSyncController
= DataSyncController
| DataSyncReady
{ subscriptions :: !(HashMap UUID Subscription)
{ subscriptions :: !(HashMap UUID (MVar.MVar ()))
, transactions :: !(HashMap UUID DataSyncTransaction)
, asyncs :: ![Async ()]
}

0 comments on commit 501bbef

Please sign in to comment.