Utility function for modifying an IORef using a lens

.. and a supplied function that works inside the lens.
Signed-off-by: Petr Pudlak <pudlak@google.com>
Reviewed-by: Klaus Aehlig <aehlig@google.com>

Makefile.am

@@ -860,6 +860,7 @@ HS_LIB_SRCS = \
 	src/Ganeti/Utils.hs \
 	src/Ganeti/Utils/Atomic.hs \
 	src/Ganeti/Utils/AsyncWorker.hs \
+	src/Ganeti/Utils/IORef.hs \
 	src/Ganeti/Utils/Livelock.hs \
 	src/Ganeti/Utils/MonadPlus.hs \
 	src/Ganeti/VCluster.hs \

src/Ganeti/Utils/IORef.hs

+{-# LANGUAGE FlexibleContexts, RankNTypes #-}
+
+{-| Utility functions for working with IORefs. -}
+
+{-
+
+Copyright (C) 2014 Google Inc.
+
+This program is free software; you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation; either version 2 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful, but
+WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program; if not, write to the Free Software
+Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
+02110-1301, USA.
+
+-}
+
+module Ganeti.Utils.IORef
+  ( atomicModifyWithLens
+  ) where
+
+import Control.Monad.Base
+import Data.IORef.Lifted
+import Data.Tuple (swap)
+
+import Ganeti.BasicTypes
+import Ganeti.Lens
+
+-- | Atomically modifies an 'IORef' using a lens
+atomicModifyWithLens :: (MonadBase IO m)
+                     => IORef a -> Lens a a b c -> (b -> (r, c)) -> m r
+atomicModifyWithLens ref l f = atomicModifyIORef ref (swap . traverseOf l f)

src/Ganeti/WConfd/Monad.hs

-{-# LANGUAGE MultiParamTypeClasses, TypeFamilies, FlexibleContexts #-}
+{-# LANGUAGE MultiParamTypeClasses, TypeFamilies #-}
 {-# LANGUAGE TemplateHaskell #-}
 
 {-| All RPC calls are run within this monad.
@@ -58,7 +58,6 @@ import Control.Monad.Reader
 import Control.Monad.Trans.Control
 import Data.IORef.Lifted
 import qualified Data.Set as S
-import Data.Tuple (swap)
 import qualified Text.JSON as J
 
 import Ganeti.BasicTypes
@@ -69,6 +68,7 @@ import Ganeti.Locking.Locks
 import Ganeti.Locking.Waiting (getAllocation)
 import Ganeti.Logging
 import Ganeti.Utils.AsyncWorker
+import Ganeti.Utils.IORef
 import Ganeti.WConfd.ConfigState
 
 -- * Pure data types used in the monad
@@ -187,8 +187,7 @@ modifyConfigState f = do
   dh <- daemonHandle
   let modCS cs = let (cs', r) = f cs
                   in ((r, cs /= cs'), cs')
-  let mf = traverseOf dsConfigStateL modCS
-  (r, modified) <- atomicModifyIORef (dhDaemonState dh) (swap . mf)
+  (r, modified) <- atomicModifyWithLens (dhDaemonState dh) dsConfigStateL modCS
   when modified $ do
     -- trigger the config. saving worker and wait for it
     logDebug "Triggering config write"
@@ -209,10 +208,9 @@ modifyLockWaiting :: (GanetiLockWaiting -> ( GanetiLockWaiting
                      -> WConfdMonad a
 modifyLockWaiting f = do
   dh <- lift . WConfdMonadInt $ ask
-  let f' = swap . (fst &&& id) . f
-  (lockAlloc, (r, nfy)) <- atomicModifyIORef
-                             (dhDaemonState dh)
-                             (swap . traverseOf dsLockWaitingL f')
+  let f' = (id &&& fst) . f
+  (lockAlloc, (r, nfy)) <- atomicModifyWithLens
+                             (dhDaemonState dh) dsLockWaitingL f'
   logDebug $ "Current lock status: " ++ J.encode lockAlloc
   logDebug "Triggering lock state write"
   liftBase . triggerAndWait . dhSaveLocksWorker $ dh