Skip to content
This repository
Browse code

connect timeout

  • Loading branch information...
commit 19da43a3484b4b6d20e9b16cb96d169fa462eb8b 1 parent 9d2f09a
authored July 21, 2011
59  Database/MongoDB/Connection.hs
@@ -4,13 +4,15 @@
4 4
 
5 5
 module Database.MongoDB.Connection (
6 6
 	-- * Util
7  
-	IOE, runIOE,
  7
+	IOE, runIOE, Secs,
8 8
 	-- * Connection
9 9
 	Pipe, close, isClosed,
10 10
 	-- * Server
11  
-	Host(..), PortID(..), defaultPort, host, showHostPort, readHostPort, readHostPortM, connect,
  11
+	Host(..), PortID(..), defaultPort, host, showHostPort, readHostPort, readHostPortM,
  12
+	globalConnectTimeout, connect, connect',
12 13
 	-- * Replica Set
13  
-	ReplicaSetName, openReplicaSet, ReplicaSet, primary, secondaryOk, closeReplicaSet
  14
+	ReplicaSetName, openReplicaSet, openReplicaSet',
  15
+	ReplicaSet, primary, secondaryOk, closeReplicaSet, replSetName
14 16
 ) where
15 17
 
16 18
 import Prelude hiding (lookup)
@@ -30,6 +32,9 @@ import Data.Bson as D (Document, lookup, at, (=:))
30 32
 import Database.MongoDB.Query (access, slaveOk, Failure(ConnectionFailure), Command, runCommand)
31 33
 import Database.MongoDB.Internal.Util (untilSuccess, liftIOE, runIOE, updateAssocs, shuffle)
32 34
 import Data.List as L (lookup, intersect, partition, (\\), delete)
  35
+import Data.IORef (IORef, newIORef, readIORef)
  36
+import System.Timeout (timeout)
  37
+import System.IO.Unsafe (unsafePerformIO)
33 38
 
34 39
 adminCommand :: Command -> Pipe -> IOE Document
35 40
 -- ^ Run command against admin database on server connected to pipe. Fail if connection fails.
@@ -80,10 +85,23 @@ readHostPort :: String -> Host
80 85
 -- ^ Read string \"hostname:port\" as @Host hostname (PortNumber port)@ or \"hostname\" as @host hostname@ (default port). Error if string does not match either syntax.
81 86
 readHostPort = runIdentity . readHostPortM
82 87
 
  88
+type Secs = Double
  89
+
  90
+globalConnectTimeout :: IORef Secs
  91
+-- ^ 'connect' (and 'openReplicaSet') fails if it can't connect within this many seconds (default is 6 seconds). Use 'connect\'' (and 'openReplicaSet\'') if you want to ignore this global and specify your own timeout. Note, this timeout only applies to initial connection establishment, not when reading/writing to the connection.
  92
+globalConnectTimeout = unsafePerformIO (newIORef 6)
  93
+{-# NOINLINE globalConnectTimeout #-}
  94
+
83 95
 connect :: Host -> IOE Pipe
84  
--- ^ Connect to Host returning pipelined TCP connection. Throw IOError if problem connecting.
85  
-connect (Host hostname port) = do
86  
-	handle <- ErrorT . E.try $ connectTo hostname port
  96
+-- ^ Connect to Host returning pipelined TCP connection. Throw IOError if connection refused or no response within 'globalConnectTimeout'.
  97
+connect h = lift (readIORef globalConnectTimeout) >>= flip connect' h
  98
+
  99
+connect' :: Secs -> Host -> IOE Pipe
  100
+-- ^ Connect to Host returning pipelined TCP connection. Throw IOError if connection refused or no response within given number of seconds.
  101
+connect' timeoutSecs (Host hostname port) = do
  102
+	handle <- ErrorT . E.try $ do
  103
+		mh <- timeout (round $ timeoutSecs * 1000000) (connectTo hostname port)
  104
+		maybe (ioError $ userError "connect timed out") return mh
87 105
 	lift $ newPipeline $ IOStream (writeMessage handle) (readMessage handle) (hClose handle)
88 106
 
89 107
 -- * Replica Set
@@ -91,22 +109,31 @@ connect (Host hostname port) = do
91 109
 type ReplicaSetName = UString
92 110
 
93 111
 -- | Maintains a connection (created on demand) to each server in the named replica set
94  
-data ReplicaSet = ReplicaSet ReplicaSetName (MVar [(Host, Maybe Pipe)])
  112
+data ReplicaSet = ReplicaSet ReplicaSetName (MVar [(Host, Maybe Pipe)]) Secs
  113
+
  114
+replSetName :: ReplicaSet -> UString
  115
+-- ^ name of connected replica set
  116
+replSetName (ReplicaSet rsName _ _) = rsName
95 117
 
96 118
 openReplicaSet :: (ReplicaSetName, [Host]) -> IOE ReplicaSet
97  
--- ^ Open connections (on demand) to servers in replica set. Supplied hosts is seed list. At least one of them must be a live member of the named replica set, otherwise fail.
98  
-openReplicaSet (rsName, seedList) = do
99  
-	rs <- ReplicaSet rsName <$> newMVar (map (, Nothing) seedList)
  119
+-- ^ Open connections (on demand) to servers in replica set. Supplied hosts is seed list. At least one of them must be a live member of the named replica set, otherwise fail. The value of 'globalConnectTimeout' at the time of this call is the timeout used for future member connect attempts. To use your own value call 'openReplicaSet\'' instead.
  120
+openReplicaSet rsSeed = lift (readIORef globalConnectTimeout) >>= flip openReplicaSet' rsSeed
  121
+
  122
+openReplicaSet' :: Secs -> (ReplicaSetName, [Host]) -> IOE ReplicaSet
  123
+-- ^ Open connections (on demand) to servers in replica set. Supplied hosts is seed list. At least one of them must be a live member of the named replica set, otherwise fail. Supplied seconds timeout is used for connect attempts to members.
  124
+openReplicaSet' timeoutSecs (rsName, seedList) = do
  125
+	vMembers <- newMVar (map (, Nothing) seedList)
  126
+	let rs = ReplicaSet rsName vMembers timeoutSecs
100 127
 	_ <- updateMembers rs
101 128
 	return rs
102 129
 
103 130
 closeReplicaSet :: ReplicaSet -> IO ()
104 131
 -- ^ Close all connections to replica set
105  
-closeReplicaSet (ReplicaSet _ vMembers) = withMVar vMembers $ mapM_ (maybe (return ()) close . snd)
  132
+closeReplicaSet (ReplicaSet _ vMembers _) = withMVar vMembers $ mapM_ (maybe (return ()) close . snd)
106 133
 
107 134
 primary :: ReplicaSet -> IOE Pipe
108 135
 -- ^ Return connection to current primary of replica set. Fail if no primary available.
109  
-primary rs@(ReplicaSet rsName _) = do
  136
+primary rs@(ReplicaSet rsName _ _) = do
110 137
 	mHost <- statedPrimary <$> updateMembers rs
111 138
 	case mHost of
112 139
 		Just host' -> connection rs Nothing host'
@@ -133,7 +160,7 @@ possibleHosts (_, info) = map readHostPort $ at "hosts" info
133 160
 
134 161
 updateMembers :: ReplicaSet -> IOE ReplicaInfo
135 162
 -- ^ Fetch replica info from any server and update members accordingly
136  
-updateMembers rs@(ReplicaSet _ vMembers) = do
  163
+updateMembers rs@(ReplicaSet _ vMembers _) = do
137 164
 	(host', info) <- untilSuccess (fetchReplicaInfo rs) =<< readMVar vMembers
138 165
 	modifyMVar vMembers $ \members -> do
139 166
 		let ((members', old), new) = intersection (map readHostPort $ at "hosts" info) members
@@ -147,7 +174,7 @@ updateMembers rs@(ReplicaSet _ vMembers) = do
147 174
 
148 175
 fetchReplicaInfo :: ReplicaSet -> (Host, Maybe Pipe) -> IOE ReplicaInfo
149 176
 -- Connect to host and fetch replica info from host creating new connection if missing or closed (previously failed). Fail if not member of named replica set.
150  
-fetchReplicaInfo rs@(ReplicaSet rsName _) (host', mPipe) = do
  177
+fetchReplicaInfo rs@(ReplicaSet rsName _ _) (host', mPipe) = do
151 178
 	pipe <- connection rs mPipe host'
152 179
 	info <- adminCommand ["isMaster" =: (1 :: Int)] pipe
153 180
 	case D.lookup "setName" info of
@@ -157,11 +184,11 @@ fetchReplicaInfo rs@(ReplicaSet rsName _) (host', mPipe) = do
157 184
 
158 185
 connection :: ReplicaSet -> Maybe Pipe -> Host -> IOE Pipe
159 186
 -- ^ Return new or existing connection to member of replica set. If pipe is already known for host it is given, but we still test if it is open.
160  
-connection (ReplicaSet _ vMembers) mPipe host' =
  187
+connection (ReplicaSet _ vMembers timeoutSecs) mPipe host' =
161 188
 	maybe conn (\p -> lift (isClosed p) >>= \bad -> if bad then conn else return p) mPipe
162 189
  where
163 190
  	conn = 	modifyMVar vMembers $ \members -> do
164  
-		let new = connect host' >>= \pipe -> return (updateAssocs host' (Just pipe) members, pipe)
  191
+		let new = connect' timeoutSecs host' >>= \pipe -> return (updateAssocs host' (Just pipe) members, pipe)
165 192
 		case L.lookup host' members of
166 193
 			Just (Just pipe) -> lift (isClosed pipe) >>= \bad -> if bad then new else return (members, pipe)
167 194
 			_ -> new
4  mongoDB.cabal
... ...
@@ -1,5 +1,5 @@
1 1
 name: mongoDB
2  
-version: 1.0.0
  2
+version: 1.0.1
3 3
 build-type: Simple
4 4
 license: OtherLicense
5 5
 license-file: LICENSE
@@ -58,7 +58,7 @@ install-includes:
58 58
 include-dirs:
59 59
 hs-source-dirs: .
60 60
 other-modules:
61  
-ghc-prof-options:
  61
+ghc-prof-options: -auto-all
62 62
 ghc-shared-options:
63 63
 ghc-options: -Wall -O2
64 64
 hugs-options:

0 notes on commit 19da43a

Please sign in to comment.
Something went wrong with that request. Please try again.