Skip to content
This repository
Browse code

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...
commit a618ebbf957720ac3ecc59e2961bf64893f2e077 2 parents 5e015dc + d3f5492
Tony Hannan authored February 20, 2012
13  Database/MongoDB/Connection.hs
@@ -12,7 +12,7 @@ module Database.MongoDB.Connection (
12 12
 	globalConnectTimeout, connect, connect',
13 13
 	-- * Replica Set
14 14
 	ReplicaSetName, openReplicaSet, openReplicaSet',
15  
-	ReplicaSet, primary, secondaryOk, closeReplicaSet, replSetName
  15
+	ReplicaSet, primary, secondaryOk, routedHost, closeReplicaSet, replSetName
16 16
 ) where
17 17
 
18 18
 import Prelude hiding (lookup)
@@ -29,7 +29,7 @@ import Control.Applicative ((<$>))
29 29
 import Data.UString (UString, unpack)
30 30
 import Data.Bson as D (Document, lookup, at, (=:))
31 31
 import Database.MongoDB.Query (access, slaveOk, Failure(ConnectionFailure), Command, runCommand)
32  
-import Database.MongoDB.Internal.Util (untilSuccess, liftIOE, runIOE, updateAssocs, shuffle)
  32
+import Database.MongoDB.Internal.Util (untilSuccess, liftIOE, runIOE, updateAssocs, shuffle, mergesortM)
33 33
 import Data.List as L (lookup, intersect, partition, (\\), delete)
34 34
 import Data.IORef (IORef, newIORef, readIORef)
35 35
 import System.Timeout (timeout)
@@ -146,6 +146,15 @@ secondaryOk rs = do
146 146
 	let hosts' = maybe hosts (\p -> delete p hosts ++ [p]) (statedPrimary info)
147 147
 	untilSuccess (connection rs Nothing) hosts'
148 148
 
  149
+routedHost :: ((Host, Bool) -> (Host, Bool) -> IOE Ordering) -> ReplicaSet -> IOE Pipe
  150
+-- ^ 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.
  151
+routedHost f rs = do
  152
+  info <- updateMembers rs
  153
+  hosts <- lift $ shuffle (possibleHosts info)
  154
+  let addIsPrimary h = (h, if Just h == statedPrimary info then True else False)
  155
+  hosts' <- mergesortM (\a b -> f (addIsPrimary a) (addIsPrimary b)) hosts
  156
+  untilSuccess (connection rs Nothing) hosts'
  157
+
149 158
 type ReplicaInfo = (Host, Document)
150 159
 -- ^ Result of isMaster command on host in replica set. Returned fields are: setName, ismaster, secondary, hosts, [primary]. primary only present when ismaster = false
151 160
 
27  Database/MongoDB/Internal/Util.hs
@@ -30,6 +30,33 @@ deriving instance Ord PortID
30 30
 class (MonadIO m, Applicative m, Functor m) => MonadIO' m
31 31
 instance (MonadIO m, Applicative m, Functor m) => MonadIO' m
32 32
 
  33
+-- | A monadic sort implementation derived from the non-monadic one in ghc's Prelude
  34
+mergesortM :: Monad m => (a -> a -> m Ordering) -> [a] -> m [a]
  35
+mergesortM cmp = mergesortM' cmp . map wrap
  36
+
  37
+mergesortM' :: Monad m => (a -> a -> m Ordering) -> [[a]] -> m [a]
  38
+mergesortM' _  [] = return []
  39
+mergesortM' _  [xs] = return xs
  40
+mergesortM' cmp xss = mergesortM' cmp =<< (merge_pairsM cmp xss)
  41
+
  42
+merge_pairsM :: Monad m => (a -> a -> m Ordering) -> [[a]] -> m [[a]]
  43
+merge_pairsM _   [] = return []
  44
+merge_pairsM _   [xs] = return [xs]
  45
+merge_pairsM cmp (xs:ys:xss) = liftM2 (:) (mergeM cmp xs ys) (merge_pairsM cmp xss)
  46
+
  47
+mergeM :: Monad m => (a -> a -> m Ordering) -> [a] -> [a] -> m [a]
  48
+mergeM _   [] ys = return ys
  49
+mergeM _   xs [] = return xs
  50
+mergeM cmp (x:xs) (y:ys)
  51
+ = do
  52
+     c <- x `cmp` y
  53
+     case c of
  54
+        GT -> liftM (y:) (mergeM cmp (x:xs)   ys)
  55
+        _  -> liftM (x:) (mergeM cmp    xs (y:ys))
  56
+
  57
+wrap :: a -> [a]
  58
+wrap x = [x]
  59
+
33 60
 shuffle :: [a] -> IO [a]
34 61
 -- ^ Randomly shuffle items in list
35 62
 shuffle list = shuffle' list (L.length list) <$> newStdGen

0 notes on commit a618ebb

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