From ed15e2b05564bfe62089ce468bb9dd94dd93bb8a Mon Sep 17 00:00:00 2001
From: Denis Tereshkin <denis@kasan.ws>
Date: Wed, 28 Sep 2016 14:49:42 +0700
Subject: [PATCH] Refactoring

---
 libatrade.cabal             |  1 +
 src/ATrade/Broker/Server.hs | 13 +++++++------
 src/ATrade/Util.hs          |  9 +++++++++
 test/TestBrokerServer.hs    |  5 +++--
 4 files changed, 20 insertions(+), 8 deletions(-)
 create mode 100644 src/ATrade/Util.hs

diff --git a/libatrade.cabal b/libatrade.cabal
index f37f4a0..0c9ea96 100644
--- a/libatrade.cabal
+++ b/libatrade.cabal
@@ -20,6 +20,7 @@ library
                      , ATrade.QuoteSource.Server
                      , ATrade.Broker.Protocol
                      , ATrade.Broker.Server
+                     , ATrade.Util
   build-depends:       base >= 4.7 && < 5
                      , Decimal
                      , time
diff --git a/src/ATrade/Broker/Server.hs b/src/ATrade/Broker/Server.hs
index 6bffc6c..452d117 100644
--- a/src/ATrade/Broker/Server.hs
+++ b/src/ATrade/Broker/Server.hs
@@ -22,6 +22,7 @@ import Control.Concurrent
 import Control.Exception
 import Control.Monad
 import System.Log.Logger
+import ATrade.Util
 
 newtype OrderIdGenerator = IO OrderId
 type PeerId = B.ByteString
@@ -75,10 +76,10 @@ notificationCallback state n = do
     Nothing -> warningM "Broker.Server" "Notification: unknown order"
 
   where
-    addNotification peerId n = atomicModifyIORef' state (\s ->
+    addNotification peerId n = atomicMapIORef state (\s ->
       case M.lookup peerId . pendingNotifications $ s of
-        Just ns -> (s { pendingNotifications = M.insert peerId (n : ns) (pendingNotifications s)}, ())
-        Nothing -> (s { pendingNotifications = M.insert peerId [n] (pendingNotifications s)}, ()))
+        Just ns -> s { pendingNotifications = M.insert peerId (n : ns) (pendingNotifications s)}
+        Nothing -> s { pendingNotifications = M.insert peerId [n] (pendingNotifications s)})
 
 brokerServerThread state = finally brokerServerThread' cleanup
   where
@@ -103,9 +104,9 @@ brokerServerThread state = finally brokerServerThread' cleanup
           case findBrokerForAccount (orderAccountId order) bros of
             Just bro -> do
               oid <- nextOrderId
-              atomicModifyIORef' state (\s -> (s {
+              atomicMapIORef state (\s -> s {
                 orderToBroker = M.insert oid bro (orderToBroker s),
-                orderMap = M.insert oid peerId (orderMap s) }, ()))
+                orderMap = M.insert oid peerId (orderMap s) })
               submitOrder bro order { orderId = oid }
               return $ ResponseOrderSubmitted oid
 
@@ -121,7 +122,7 @@ brokerServerThread state = finally brokerServerThread' cleanup
           maybeNs <- M.lookup peerId . pendingNotifications <$> readIORef state
           case maybeNs of
             Just ns -> do
-              atomicModifyIORef' state (\s -> (s { pendingNotifications = M.insert peerId [] (pendingNotifications s)}, ()))
+              atomicMapIORef state (\s -> s { pendingNotifications = M.insert peerId [] (pendingNotifications s)})
               return $ ResponseNotifications ns
             Nothing -> return $ ResponseNotifications []
         Nothing -> return $ ResponseError "Unable to parse request"
diff --git a/src/ATrade/Util.hs b/src/ATrade/Util.hs
new file mode 100644
index 0000000..597364b
--- /dev/null
+++ b/src/ATrade/Util.hs
@@ -0,0 +1,9 @@
+
+module ATrade.Util (
+  atomicMapIORef
+) where
+
+import Data.IORef
+
+atomicMapIORef :: IORef a -> (a -> a) -> IO ()
+atomicMapIORef ioref f = atomicModifyIORef' ioref (\s -> (f s, ()))
diff --git a/test/TestBrokerServer.hs b/test/TestBrokerServer.hs
index 3406876..da48b4e 100644
--- a/test/TestBrokerServer.hs
+++ b/test/TestBrokerServer.hs
@@ -14,6 +14,7 @@ import qualified Data.ByteString as B
 import qualified Data.ByteString.Lazy as BL
 import ATrade.Broker.Server
 import ATrade.Broker.Protocol
+import ATrade.Util
 import qualified Data.Text as T
 import Control.Monad
 import Control.Monad.Loops
@@ -39,7 +40,7 @@ data MockBrokerState = MockBrokerState {
 
 mockSubmitOrder :: IORef MockBrokerState -> Order -> IO ()
 mockSubmitOrder state order = do
-  atomicModifyIORef' state (\s -> (s { orders = submittedOrder : orders s }, ()))
+  atomicMapIORef state (\s -> s { orders = submittedOrder : orders s })
   maybeCb <- notificationCallback <$> readIORef state
   case maybeCb of
     Just cb -> cb $ OrderNotification (orderId order) Submitted
@@ -67,7 +68,7 @@ mkMockBroker accs = do
 
   return (BrokerInterface {
     accounts = accs,
-    setNotificationCallback = \cb -> atomicModifyIORef' state (\s -> (s { notificationCallback = cb }, ())),
+    setNotificationCallback = \cb -> atomicMapIORef state (\s -> s { notificationCallback = cb }),
     submitOrder = mockSubmitOrder state,
     cancelOrder = mockCancelOrder state,
     stopBroker = mockStopBroker state