Use stm-containers Map in memory for concurrent read/writes

app/Main.hs

@@ -3,7 +3,7 @@
 
 module Main where
 
-import Soostone ( app, middleWare, AppState(AppState, sqliteFile) )
+import Soostone ( app, middleWare, AppState(AppState, sqliteFile, stmMap) )
 
 import Control.Monad ( void, when )
 import Data.Version (showVersion)
@@ -34,6 +34,7 @@ import Prometheus ( register )
 import Prometheus.Metric.GHC ( ghcMetrics )
 import Prometheus.Metric.Proc ( procMetrics )
 import Text.Read ( readEither )
+import qualified StmContainers.Map as SM
 
 
 versioner :: Parser (a -> a)
@@ -98,7 +99,8 @@ main = do
   case serverCommand of
     Run ServerConfig{..} -> do
       when (sqliteDB == ":memory:") $ fail "In-memory DB is not supported!"
-      let appState = AppState { sqliteFile = sqliteDB }
+      m <- SM.newIO
+      let appState = AppState { sqliteFile = sqliteDB, stmMap = m }
       void $ register ghcMetrics
       void $ register procMetrics
       run serverPort (middleWare $ app appState)

lib/Soostone/Handlers.hs

@@ -4,7 +4,7 @@
 module Soostone.Handlers where
 
 import Soostone.Types
-    ( AppState(AppState, sqliteFile), AppM, Count(..), Key )
+    ( AppState(AppState, sqliteFile, stmMap), AppM, Count(..), Key )
 
 import Data.Functor ( ($>) )
 import Database.SQLite.Simple
@@ -18,10 +18,20 @@ import Database.SQLite.Simple
 import Control.Monad.Catch ( finally )
 import Control.Monad.Reader   ( ask )
 import Control.Monad.IO.Class ( liftIO )
+import qualified StmContainers.Map as SM
+import GHC.Conc (atomically)
 
 
 submitKeyHandler :: Key -> AppM ()
-submitKeyHandler key = withSQLiteCon $ \con -> do
+submitKeyHandler key = do
+  AppState{..} <- ask
+  liftIO $ atomically $ do
+    SM.lookup key stmMap >>= \case
+      Just v  -> SM.insert (v + 1) key stmMap
+      Nothing -> SM.insert 1       key stmMap
+
+{--
+withSQLiteCon $ \con -> do
   r <- liftIO $ withImmediateTransaction con $ do
     queryNamed @Count con "SELECT count FROM keys WHERE key=:key" [":key" := key] >>= \case
       []  -> executeNamed con "INSERT INTO keys (key, count) VALUES (:key, :count)" [":key" := key, ":count" := Count 1] $> Right ()
@@ -30,14 +40,16 @@ submitKeyHandler key = withSQLiteCon $ \con -> do
   case r of
     Right _ -> pure ()
     Left e -> fail e -- internal error
+--}
 
 
 countKeyHandler :: Key -> AppM Count
-countKeyHandler key = withSQLiteCon $ \con -> do
+countKeyHandler key = do
-  liftIO (queryNamed @Count con "SELECT count FROM keys WHERE key=:key" [":key" := key]) >>= \case
+  AppState{..} <- ask
-    []  -> return (Count 0)
+  liftIO $ atomically $ do
-    [c] -> pure c
+    SM.lookup key stmMap >>= \case
-    _   -> fail "Oops" -- internal error
+      Just v  -> return v
+      Nothing -> return (Count 0)
 
 withSQLiteCon :: (Connection -> AppM a) -> AppM a
 withSQLiteCon action = do

lib/Soostone/Types.hs

@@ -17,10 +17,12 @@ import Database.SQLite.Simple.ToField ( ToField(..) )
 import Servant
     ( Handler, ServerError(errBody), FromHttpApiData, err500 )
 import Servant.API.Generic ( Generic )
+import qualified StmContainers.Map as SM
+import Data.Hashable
 
 
 newtype Key = Key Text
-  deriving (FromHttpApiData, Generic)
+  deriving (FromHttpApiData, Generic, Hashable, Eq)
 
 instance ToField Key where
   toField (Key k) = toField k
@@ -46,6 +48,7 @@ instance FromJSON Count
 
 data AppState = AppState {
     sqliteFile :: FilePath
+  , stmMap     :: SM.Map Key Count
   }
 
 newtype AppM a = AppM { runAppM :: ReaderT AppState Handler a }

soostone.cabal

@@ -21,6 +21,7 @@ common deps
     , base                       >=4.13   && <5.0
     , bytestring                 ^>=0.11
     , exceptions                 ^>=0.10
+    , hashable                   ^>=1.4
     , mtl                        ^>=2.2
     , optparse-applicative       ^>=0.16
     , prometheus-client          ^>=1.1
@@ -29,6 +30,7 @@ common deps
     , servant                    ^>=0.19
     , servant-server             ^>=0.19
     , sqlite-simple              ^>=0.4
+    , stm-containers             ^>=1.2
     , text                       ^>=2.0
     , transformers               ^>=0.5
     , wai                        ^>=3.2.3