mongodb/Database/MongoDB/Connection.hs

227 lines
11 KiB
Haskell
Raw Normal View History

-- | Connect to a single server or a replica set of servers
{-# LANGUAGE CPP, OverloadedStrings, ScopedTypeVariables, TupleSections #-}
2016-05-20 04:44:42 +00:00
#if (__GLASGOW_HASKELL__ >= 706)
{-# LANGUAGE RecursiveDo #-}
#else
{-# LANGUAGE DoRec #-}
#endif
module Database.MongoDB.Connection (
2013-12-26 14:57:33 +00:00
-- * Util
2013-12-26 15:23:02 +00:00
Secs,
2013-12-26 14:57:33 +00:00
-- * Connection
Pipe, close, isClosed,
-- * Server
Host(..), PortID(..), defaultPort, host, showHostPort, readHostPort,
2012-06-10 19:47:14 +00:00
readHostPortM, globalConnectTimeout, connect, connect',
2013-12-26 14:57:33 +00:00
-- * Replica Set
ReplicaSetName, openReplicaSet, openReplicaSet',
ReplicaSet, primary, secondaryOk, routedHost, closeReplicaSet, replSetName
) where
2011-03-14 20:24:28 +00:00
import Prelude hiding (lookup)
2012-06-10 19:47:14 +00:00
import Data.IORef (IORef, newIORef, readIORef)
import Data.List (intersect, partition, (\\), delete)
2019-10-04 16:10:24 +00:00
import Data.Maybe (fromJust)
#if !MIN_VERSION_base(4,8,0)
2012-06-10 19:47:14 +00:00
import Control.Applicative ((<$>))
#endif
2012-06-10 19:47:14 +00:00
import Control.Monad (forM_)
2019-10-04 16:10:24 +00:00
import Control.Monad.Fail(MonadFail)
2012-06-10 19:47:14 +00:00
import System.IO.Unsafe (unsafePerformIO)
import System.Timeout (timeout)
import Text.ParserCombinators.Parsec (parse, many1, letter, digit, char, eof,
spaces, try, (<|>))
import qualified Data.List as List
import Control.Monad.Identity (runIdentity)
import Control.Monad.Except (throwError)
2012-06-10 19:47:14 +00:00
import Control.Concurrent.MVar.Lifted (MVar, newMVar, withMVar, modifyMVar,
readMVar)
import Data.Bson (Document, at, (=:))
2012-05-08 15:13:25 +00:00
import Data.Text (Text)
2012-06-10 19:47:14 +00:00
import qualified Data.Bson as B
2012-05-08 15:13:25 +00:00
import qualified Data.Text as T
2012-06-10 19:47:14 +00:00
import Database.MongoDB.Internal.Network (HostName, PortID(..), connectTo)
import Database.MongoDB.Internal.Protocol (Pipe, newPipe, close, isClosed)
2013-12-26 15:23:02 +00:00
import Database.MongoDB.Internal.Util (untilSuccess, liftIOE,
2012-06-10 19:47:14 +00:00
updateAssocs, shuffle, mergesortM)
import Database.MongoDB.Query (Command, Failure(ConnectionFailure), access,
2016-05-20 04:44:42 +00:00
slaveOk, runCommand, retrieveServerData)
2013-12-26 15:23:02 +00:00
adminCommand :: Command -> Pipe -> IO Document
-- ^ Run command against admin database on server connected to pipe. Fail if connection fails.
adminCommand cmd pipe =
2013-12-26 15:23:02 +00:00
liftIOE failureToIOError $ access pipe slaveOk "admin" $ runCommand cmd
where
2013-12-26 14:57:33 +00:00
failureToIOError (ConnectionFailure e) = e
failureToIOError e = userError $ show e
-- * Host
data Host = Host HostName PortID deriving (Show, Eq, Ord)
defaultPort :: PortID
-- ^ Default MongoDB port = 27017
defaultPort = PortNumber 27017
host :: HostName -> Host
-- ^ Host on 'defaultPort'
host hostname = Host hostname defaultPort
showHostPort :: Host -> String
-- ^ Display host as \"host:port\"
2011-02-22 15:48:14 +00:00
-- TODO: Distinguish Service and UnixSocket port
showHostPort (Host hostname port) = hostname ++ ":" ++ portname where
2013-12-26 14:57:33 +00:00
portname = case port of
PortNumber p -> show p
2019-10-04 16:10:24 +00:00
readHostPortM :: (MonadFail m) => String -> m Host
-- ^ Read string \"hostname:port\" as @Host hosthame (PortNumber port)@ or \"hostname\" as @host hostname@ (default port). Fail if string does not match either syntax.
2011-02-22 15:48:14 +00:00
-- TODO: handle Service and UnixSocket port
readHostPortM = either (fail . show) return . parse parser "readHostPort" where
2013-12-26 14:57:33 +00:00
hostname = many1 (letter <|> digit <|> char '-' <|> char '.')
parser = do
spaces
h <- hostname
try (spaces >> eof >> return (host h)) <|> do
_ <- char ':'
port :: Int <- read <$> many1 digit
spaces >> eof
return $ Host h (PortNumber $ fromIntegral port)
readHostPort :: String -> Host
-- ^ Read string \"hostname:port\" as @Host hostname (PortNumber port)@ or \"hostname\" as @host hostname@ (default port). Error if string does not match either syntax.
2019-10-04 16:10:24 +00:00
readHostPort = fromJust . readHostPortM
2011-07-21 15:27:19 +00:00
type Secs = Double
globalConnectTimeout :: IORef Secs
-- ^ '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.
globalConnectTimeout = unsafePerformIO (newIORef 6)
{-# NOINLINE globalConnectTimeout #-}
2013-12-26 15:23:02 +00:00
connect :: Host -> IO Pipe
2011-07-21 15:27:19 +00:00
-- ^ Connect to Host returning pipelined TCP connection. Throw IOError if connection refused or no response within 'globalConnectTimeout'.
2013-12-26 15:23:02 +00:00
connect h = readIORef globalConnectTimeout >>= flip connect' h
2011-07-21 15:27:19 +00:00
2013-12-26 15:23:02 +00:00
connect' :: Secs -> Host -> IO Pipe
2011-07-21 15:27:19 +00:00
-- ^ Connect to Host returning pipelined TCP connection. Throw IOError if connection refused or no response within given number of seconds.
connect' timeoutSecs (Host hostname port) = do
2013-12-26 15:23:02 +00:00
mh <- timeout (round $ timeoutSecs * 1000000) (connectTo hostname port)
handle <- maybe (ioError $ userError "connect timed out") return mh
2016-05-20 04:44:42 +00:00
rec
p <- newPipe sd handle
sd <- access p slaveOk "admin" retrieveServerData
return p
-- * Replica Set
2012-05-08 15:13:25 +00:00
type ReplicaSetName = Text
-- | Maintains a connection (created on demand) to each server in the named replica set
2011-07-21 15:27:19 +00:00
data ReplicaSet = ReplicaSet ReplicaSetName (MVar [(Host, Maybe Pipe)]) Secs
2012-05-08 15:13:25 +00:00
replSetName :: ReplicaSet -> Text
2011-07-21 15:27:19 +00:00
-- ^ name of connected replica set
replSetName (ReplicaSet rsName _ _) = rsName
2013-12-26 15:23:02 +00:00
openReplicaSet :: (ReplicaSetName, [Host]) -> IO ReplicaSet
2011-07-21 15:27:19 +00:00
-- ^ 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.
2013-12-26 15:23:02 +00:00
openReplicaSet rsSeed = readIORef globalConnectTimeout >>= flip openReplicaSet' rsSeed
2011-07-21 15:27:19 +00:00
2013-12-26 15:23:02 +00:00
openReplicaSet' :: Secs -> (ReplicaSetName, [Host]) -> IO ReplicaSet
2011-07-21 15:27:19 +00:00
-- ^ 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.
openReplicaSet' timeoutSecs (rsName, seedList) = do
2013-12-26 14:57:33 +00:00
vMembers <- newMVar (map (, Nothing) seedList)
let rs = ReplicaSet rsName vMembers timeoutSecs
_ <- updateMembers rs
return rs
closeReplicaSet :: ReplicaSet -> IO ()
-- ^ Close all connections to replica set
2011-07-21 15:27:19 +00:00
closeReplicaSet (ReplicaSet _ vMembers _) = withMVar vMembers $ mapM_ (maybe (return ()) close . snd)
2013-12-26 15:23:02 +00:00
primary :: ReplicaSet -> IO Pipe
-- ^ Return connection to current primary of replica set. Fail if no primary available.
2011-07-21 15:27:19 +00:00
primary rs@(ReplicaSet rsName _ _) = do
2013-12-26 14:57:33 +00:00
mHost <- statedPrimary <$> updateMembers rs
case mHost of
Just host' -> connection rs Nothing host'
Nothing -> throwError $ userError $ "replica set " ++ T.unpack rsName ++ " has no primary"
2013-12-26 15:23:02 +00:00
secondaryOk :: ReplicaSet -> IO Pipe
-- ^ Return connection to a random secondary, or primary if no secondaries available.
secondaryOk rs = do
2013-12-26 14:57:33 +00:00
info <- updateMembers rs
2013-12-26 15:23:02 +00:00
hosts <- shuffle (possibleHosts info)
2013-12-26 14:57:33 +00:00
let hosts' = maybe hosts (\p -> delete p hosts ++ [p]) (statedPrimary info)
untilSuccess (connection rs Nothing) hosts'
2013-12-26 15:23:02 +00:00
routedHost :: ((Host, Bool) -> (Host, Bool) -> IO Ordering) -> ReplicaSet -> IO Pipe
-- ^ Return a connection to a host using a user-supplied sorting function, which sorts based on a tuple containing the host and a boolean indicating whether the host is primary.
routedHost f rs = do
info <- updateMembers rs
2013-12-26 15:23:02 +00:00
hosts <- shuffle (possibleHosts info)
let addIsPrimary h = (h, if Just h == statedPrimary info then True else False)
hosts' <- mergesortM (\a b -> f (addIsPrimary a) (addIsPrimary b)) hosts
untilSuccess (connection rs Nothing) hosts'
type ReplicaInfo = (Host, Document)
-- ^ Result of isMaster command on host in replica set. Returned fields are: setName, ismaster, secondary, hosts, [primary]. primary only present when ismaster = false
statedPrimary :: ReplicaInfo -> Maybe Host
-- ^ Primary of replica set or Nothing if there isn't one
2012-06-10 19:47:14 +00:00
statedPrimary (host', info) = if (at "ismaster" info) then Just host' else readHostPort <$> B.lookup "primary" info
possibleHosts :: ReplicaInfo -> [Host]
-- ^ Non-arbiter, non-hidden members of replica set
possibleHosts (_, info) = map readHostPort $ at "hosts" info
2013-12-26 15:23:02 +00:00
updateMembers :: ReplicaSet -> IO ReplicaInfo
-- ^ Fetch replica info from any server and update members accordingly
2011-07-21 15:27:19 +00:00
updateMembers rs@(ReplicaSet _ vMembers _) = do
2013-12-26 14:57:33 +00:00
(host', info) <- untilSuccess (fetchReplicaInfo rs) =<< readMVar vMembers
modifyMVar vMembers $ \members -> do
let ((members', old), new) = intersection (map readHostPort $ at "hosts" info) members
2013-12-26 15:23:02 +00:00
forM_ old $ \(_, mPipe) -> maybe (return ()) close mPipe
2013-12-26 14:57:33 +00:00
return (members' ++ map (, Nothing) new, (host', info))
where
2013-12-26 14:57:33 +00:00
intersection :: (Eq k) => [k] -> [(k, v)] -> (([(k, v)], [(k, v)]), [k])
intersection keys assocs = (partition (flip elem inKeys . fst) assocs, keys \\ inKeys) where
assocKeys = map fst assocs
inKeys = intersect keys assocKeys
2013-12-26 15:23:02 +00:00
fetchReplicaInfo :: ReplicaSet -> (Host, Maybe Pipe) -> IO ReplicaInfo
-- 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.
2011-07-21 15:27:19 +00:00
fetchReplicaInfo rs@(ReplicaSet rsName _ _) (host', mPipe) = do
2013-12-26 14:57:33 +00:00
pipe <- connection rs mPipe host'
info <- adminCommand ["isMaster" =: (1 :: Int)] pipe
case B.lookup "setName" info of
Nothing -> throwError $ userError $ show host' ++ " not a member of any replica set, including " ++ T.unpack rsName ++ ": " ++ show info
Just setName | setName /= rsName -> throwError $ userError $ show host' ++ " not a member of replica set " ++ T.unpack rsName ++ ": " ++ show info
Just _ -> return (host', info)
2013-12-26 15:23:02 +00:00
connection :: ReplicaSet -> Maybe Pipe -> Host -> IO Pipe
-- ^ 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.
2011-07-21 15:27:19 +00:00
connection (ReplicaSet _ vMembers timeoutSecs) mPipe host' =
2013-12-26 15:23:02 +00:00
maybe conn (\p -> isClosed p >>= \bad -> if bad then conn else return p) mPipe
where
2013-12-26 14:57:33 +00:00
conn = modifyMVar vMembers $ \members -> do
let new = connect' timeoutSecs host' >>= \pipe -> return (updateAssocs host' (Just pipe) members, pipe)
case List.lookup host' members of
2013-12-26 15:23:02 +00:00
Just (Just pipe) -> isClosed pipe >>= \bad -> if bad then new else return (members, pipe)
2013-12-26 14:57:33 +00:00
_ -> new
{- Authors: Tony Hannan <tony@10gen.com>
Copyright 2011 10gen Inc.
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at: http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. -}