Permalink
Browse files

Merge pull request #16 from A1kmm/master

Add routedHost as an alternative to primary / secondaryOk that lets the user supply their own sort function
  • Loading branch information...
2 parents 5e015dc + d3f5492 commit a618ebbf957720ac3ecc59e2961bf64893f2e077 @TonyGen committed Feb 20, 2012
Showing with 38 additions and 2 deletions.
  1. +11 −2 Database/MongoDB/Connection.hs
  2. +27 −0 Database/MongoDB/Internal/Util.hs
@@ -12,7 +12,7 @@ module Database.MongoDB.Connection (
globalConnectTimeout, connect, connect',
-- * Replica Set
ReplicaSetName, openReplicaSet, openReplicaSet',
- ReplicaSet, primary, secondaryOk, closeReplicaSet, replSetName
+ ReplicaSet, primary, secondaryOk, routedHost, closeReplicaSet, replSetName
) where
import Prelude hiding (lookup)
@@ -29,7 +29,7 @@ import Control.Applicative ((<$>))
import Data.UString (UString, unpack)
import Data.Bson as D (Document, lookup, at, (=:))
import Database.MongoDB.Query (access, slaveOk, Failure(ConnectionFailure), Command, runCommand)
-import Database.MongoDB.Internal.Util (untilSuccess, liftIOE, runIOE, updateAssocs, shuffle)
+import Database.MongoDB.Internal.Util (untilSuccess, liftIOE, runIOE, updateAssocs, shuffle, mergesortM)
import Data.List as L (lookup, intersect, partition, (\\), delete)
import Data.IORef (IORef, newIORef, readIORef)
import System.Timeout (timeout)
@@ -146,6 +146,15 @@ secondaryOk rs = do
let hosts' = maybe hosts (\p -> delete p hosts ++ [p]) (statedPrimary info)
untilSuccess (connection rs Nothing) hosts'
+routedHost :: ((Host, Bool) -> (Host, Bool) -> IOE Ordering) -> ReplicaSet -> IOE 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
+ hosts <- lift $ 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
@@ -30,6 +30,33 @@ deriving instance Ord PortID
class (MonadIO m, Applicative m, Functor m) => MonadIO' m
instance (MonadIO m, Applicative m, Functor m) => MonadIO' m
+-- | A monadic sort implementation derived from the non-monadic one in ghc's Prelude
+mergesortM :: Monad m => (a -> a -> m Ordering) -> [a] -> m [a]
+mergesortM cmp = mergesortM' cmp . map wrap
+
+mergesortM' :: Monad m => (a -> a -> m Ordering) -> [[a]] -> m [a]
+mergesortM' _ [] = return []
+mergesortM' _ [xs] = return xs
+mergesortM' cmp xss = mergesortM' cmp =<< (merge_pairsM cmp xss)
+
+merge_pairsM :: Monad m => (a -> a -> m Ordering) -> [[a]] -> m [[a]]
+merge_pairsM _ [] = return []
+merge_pairsM _ [xs] = return [xs]
+merge_pairsM cmp (xs:ys:xss) = liftM2 (:) (mergeM cmp xs ys) (merge_pairsM cmp xss)
+
+mergeM :: Monad m => (a -> a -> m Ordering) -> [a] -> [a] -> m [a]
+mergeM _ [] ys = return ys
+mergeM _ xs [] = return xs
+mergeM cmp (x:xs) (y:ys)
+ = do
+ c <- x `cmp` y
+ case c of
+ GT -> liftM (y:) (mergeM cmp (x:xs) ys)
+ _ -> liftM (x:) (mergeM cmp xs (y:ys))
+
+wrap :: a -> [a]
+wrap x = [x]
+
shuffle :: [a] -> IO [a]
-- ^ Randomly shuffle items in list
shuffle list = shuffle' list (L.length list) <$> newStdGen

0 comments on commit a618ebb

Please sign in to comment.