From 2989890b304a6b90253142b3f4407e923d68a682 Mon Sep 17 00:00:00 2001
From: Alexander Thiemann <mail@agrafix.net>
Date: Wed, 12 Aug 2015 14:55:13 +0200
Subject: [PATCH] add mapAllSessions

---
 Spock.cabal                              |  3 ++-
 src/Web/Spock/Internal/SessionManager.hs | 10 ++++++++++
 src/Web/Spock/Internal/SessionVault.hs   |  8 ++++++++
 src/Web/Spock/Internal/Types.hs          |  2 ++
 src/Web/Spock/Shared.hs                  | 10 +++++++++-
 5 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/Spock.cabal b/Spock.cabal
index e883d56..ac45b6a 100644
--- a/Spock.cabal
+++ b/Spock.cabal
@@ -1,5 +1,5 @@
 name:                Spock
-version:             0.8.0.0
+version:             0.8.1.0
 synopsis:            Another Haskell web framework for rapid development
 description:         This toolbox provides everything you need to get a quick start into web hacking with haskell:
                      .
@@ -57,6 +57,7 @@ library
                        case-insensitive >=1.1,
                        containers >=0.5,
                        directory >=1.2,
+                       focus >=0.1,
                        hashable >=1.2,
                        http-types >=0.8,
                        hvect >= 0.2,
diff --git a/src/Web/Spock/Internal/SessionManager.hs b/src/Web/Spock/Internal/SessionManager.hs
index 5a546fb..e6dd8af 100644
--- a/src/Web/Spock/Internal/SessionManager.hs
+++ b/src/Web/Spock/Internal/SessionManager.hs
@@ -55,6 +55,7 @@ createSessionManager cfg =
           , sm_writeSession = writeSessionImpl vaultKey cacheHM
           , sm_modifySession = modifySessionImpl vaultKey cacheHM
           , sm_clearAllSessions = clearAllSessionsImpl cacheHM
+          , sm_mapSessions = mapAllSessionsImpl cacheHM
           , sm_middleware = sessionMiddleware cfg vaultKey cacheHM
           , sm_addSafeAction = addSafeActionImpl vaultKey cacheHM
           , sm_lookupSafeAction = lookupSafeActionImpl vaultKey cacheHM
@@ -289,6 +290,15 @@ clearAllSessionsImpl :: SV.SessionVault (Session conn sess st)
 clearAllSessionsImpl sessionRef =
     liftIO $ atomically $ SV.filterSessions (const False) sessionRef
 
+mapAllSessionsImpl ::
+    SV.SessionVault (Session conn sess st)
+    -> (sess -> STM sess)
+    -> SpockAction conn sess st ()
+mapAllSessionsImpl sessionRef f =
+    liftIO $ atomically $ flip SV.mapSessions sessionRef $ \sess ->
+        do newData <- f (sess_data sess)
+           return $ sess { sess_data = newData }
+
 housekeepSessions :: SessionCfg sess
                   -> SV.SessionVault (Session conn sess st)
                   -> (HM.HashMap SessionId (Session conn sess st) -> IO ())
diff --git a/src/Web/Spock/Internal/SessionVault.hs b/src/Web/Spock/Internal/SessionVault.hs
index 7099d89..2e64c5d 100644
--- a/src/Web/Spock/Internal/SessionVault.hs
+++ b/src/Web/Spock/Internal/SessionVault.hs
@@ -12,6 +12,7 @@ import Control.Applicative
 import Control.Concurrent.STM (STM)
 import Control.Monad
 import Data.Hashable
+import Focus as F
 import qualified ListT as L
 import qualified STMContainers.Map as STMMap
 import qualified Data.Text as T
@@ -55,3 +56,10 @@ filterSessions cond sv =
                map getSessionKey $
                filter (not . cond) allVals
        forM_ deleteKeys $ flip deleteSession sv
+
+-- | Perform action on all sessions
+mapSessions :: IsSession s => (s -> STM s) -> SessionVault s -> STM ()
+mapSessions f sv@(SessionVault smap) =
+    do allVals <- toList sv
+       forM_ allVals $ \sess ->
+           STMMap.focus (F.adjustM f) (getSessionKey sess) smap
diff --git a/src/Web/Spock/Internal/Types.hs b/src/Web/Spock/Internal/Types.hs
index 7d77b95..865e94e 100644
--- a/src/Web/Spock/Internal/Types.hs
+++ b/src/Web/Spock/Internal/Types.hs
@@ -16,6 +16,7 @@ import Web.Spock.Internal.Wire
 #else
 import Control.Applicative
 #endif
+import Control.Concurrent.STM
 import Control.Monad.Base
 import Control.Monad.Reader
 import Control.Monad.Trans.Control
@@ -200,6 +201,7 @@ data SessionManager conn sess st
    , sm_readSession :: SpockAction conn sess st sess
    , sm_writeSession :: sess -> SpockAction conn sess st ()
    , sm_modifySession :: forall a. (sess -> (sess, a)) -> SpockAction conn sess st a
+   , sm_mapSessions :: (sess -> STM sess) -> SpockAction conn sess st ()
    , sm_clearAllSessions :: SpockAction conn sess st ()
    , sm_middleware :: Middleware
    , sm_addSafeAction :: PackedSafeAction conn sess st -> SpockAction conn sess st SafeActionHash
diff --git a/src/Web/Spock/Shared.hs b/src/Web/Spock/Shared.hs
index e75638e..b4c00ab 100644
--- a/src/Web/Spock/Shared.hs
+++ b/src/Web/Spock/Shared.hs
@@ -34,7 +34,7 @@ module Web.Spock.Shared
     , SessionPersistCfg(..), readShowSessionPersist
     , SessionId
     , getSessionId, readSession, writeSession
-    , modifySession, modifySession', modifyReadSession, clearAllSessions
+    , modifySession, modifySession', modifyReadSession, mapAllSessions, clearAllSessions
      -- * Internals for extending Spock
     , getSpockHeart, runSpockIO, WebStateM, WebState
     )
@@ -45,6 +45,7 @@ import Web.Spock.Internal.SessionManager
 import Web.Spock.Internal.Types
 import Web.Spock.Internal.CoreAction
 import Control.Monad
+import Control.Concurrent.STM (STM)
 import System.Directory
 import qualified Web.Spock.Internal.Wire as W
 import qualified Network.Wai as Wai
@@ -106,6 +107,13 @@ clearAllSessions =
     do mgr <- getSessMgr
        sm_clearAllSessions mgr
 
+-- | Apply a transformation to all sessions. Be careful with this, as this
+-- may cause many STM transaction retries.
+mapAllSessions :: (sess -> STM sess) -> SpockAction conn sess st ()
+mapAllSessions f =
+    do mgr <- getSessMgr
+       sm_mapSessions mgr f
+
 -- | Simple session persisting configuration. DO NOT USE IN PRODUCTION
 readShowSessionPersist :: (Read a, Show a) => FilePath -> SessionPersistCfg a
 readShowSessionPersist fp =