@@ -5,7 +5,8 @@
55-----------------------------------------------------------------------------
66-- |
77-- Module : Database.PostgreSQL.Simple.Notification
8−-- Copyright : (c) 2011-2012 Leon P Smith
8+-- Copyright : (c) 2011-2015 Leon P Smith
9+-- (c) 2012 Joey Adams
910-- License : BSD3
1011--
1112-- Maintainer : leon@melding-monads.com
@@ -35,13 +36,23 @@ module Database.PostgreSQL.Simple.Notification
3536 , getBackendPID
3637 ) where
3738
38−import Control.Concurrent
39−import Control.Monad ( when )
40−import Control.Exception ( throwIO )
39+import Control.Monad ( join, void )
40+import Control.Exception ( throwIO, try )
4141import qualified Data.ByteString as B
42+import qualified Data.ByteString.Char8 as B8
4243import Database.PostgreSQL.Simple.Internal
4344import qualified Database.PostgreSQL.LibPQ as PQ
4445import System.Posix.Types ( CPid )
46+import GHC.IO.Exception ( ioe_location )
47+
48+#if defined(mingw32_HOST_OS)
49+import Control.Concurrent ( threadDelay )
50+#elif !MIN_VERSION_base(4,7,0)
51+import Control.Concurrent ( threadWaitRead )
52+#else
53+import GHC.Conc ( atomically )
54+import Control.Concurrent ( threadWaitReadSTM )
55+#endif
4556
4657data Notification = Notification
4758 { notificationPid :: !CPid
@@ -59,33 +70,69 @@ convertNotice PQ.Notify{..}
5970-- 'getNotification' blocks until one arrives.
6071
6172getNotification :: Connection -> IO Notification
62−getNotification conn = loop False
73+getNotification conn = join $ withConnection conn fetch
6374 where
6475 funcName = "Database.PostgreSQL.Simple.Notification.getNotification"
65− loop doConsume = do
66− res <- withConnection conn $ \c -> do
67− when doConsume (PQ.consumeInput c >> return ())
68− mmsg <- PQ.notifies c
69− case mmsg of
70− Nothing -> do
71− mfd <- PQ.socket c
72− case mfd of
73− Nothing -> throwIO $ fdError funcName
74− Just fd -> return (Left fd)
75− Just msg -> return (Right msg)
76− -- FIXME? what happens if the connection is closed/reset right here?
77− case res of
76+
77+ fetch c = do
78+ mmsg <- PQ.notifies c
79+ case mmsg of
80+ Just msg -> return (return $! convertNotice msg)
81+ Nothing -> do
82+ mfd <- PQ.socket c
83+ case mfd of
84+ Nothing -> return (throwIO $! fdError funcName)
7885#if defined(mingw32_HOST_OS)
79− -- threadWaitRead doesn't work for sockets on Windows, so just poll
80− -- for input every second (PQconsumeInput is non-blocking).
81− --
82− -- We could call select(), but FFI calls can't be interrupted with
83− -- async exceptions, whereas threadDelay can.
84− Left _fd -> threadDelay 1000000 >> loop True
86+ -- threadWaitRead doesn't work for sockets on Windows, so just
87+ -- poll for input every second (PQconsumeInput is non-blocking).
88+ --
89+ -- We could call select(), but FFI calls can't be interrupted
90+ -- with async exceptions, whereas threadDelay can.
91+ Just _fd -> do
92+ return (threadDelay 1000000 >> loop)
93+#elif !MIN_VERSION_base(4,7,0)
94+ -- Technically there's a race condition that is usually benign.
95+ -- If the connection is closed or reset after we drop the
96+ -- lock, and then the fd index is reallocated to a new
97+ -- descriptor before we call threadWaitRead, then
98+ -- we could end up waiting on the wrong descriptor.
99+ --
100+ -- Now, if the descriptor becomes readable promptly, then
101+ -- it's no big deal as we'll wake up and notice the change
102+ -- on the next iteration of the loop. But if are very
103+ -- unlucky, then we could end up waiting a long time.
104+ Just fd -> do
105+ return $ do
106+ merr <- try (threadWaitRead fd)
107+ case merr of
108+ Left err -> throwIO $! setLoc err
109+ Right _ -> loop
85110#else
86− Left fd -> threadWaitRead fd >> loop True
111+ -- This case fixes the race condition above. By registering
112+ -- our interest in the descriptor before we drop the lock,
113+ -- there is no opportunity for the descriptor index to be
114+ -- reallocated on us.
115+ --
116+ -- (That is, assuming there isn't concurrently executing
117+ -- code that manipulates the descriptor without holding
118+ -- the lock... but such a major bug is likely to exhibit
119+ -- itself in an at least somewhat more dramatic fashion.)
120+ Just fd -> do
121+ (waitRead, _) <- threadWaitReadSTM fd
122+ return $ do
123+ merr <- try (atomically waitRead)
124+ case merr of
125+ Left err -> throwIO $! setLoc err
126+ Right _ -> loop
87127#endif
88− Right msg -> return $! convertNotice msg
128+
129+ loop = join $ withConnection conn $ \c -> do
130+ void $ PQ.consumeInput c
131+ fetch c
132+
133+ setLoc :: IOError -> IOError
134+ setLoc err = err { ioe_location = B8.unpack funcName }
135+
89136
90137-- | Non-blocking variant of 'getNotification'. Returns a single notification,
91138-- if available. If no notifications are available, returns 'Nothing'.