WIP idea on read/write splitting wrapper over hedis

hedis_read_write.hs

{-# LANGUAGE ApplicativeDo     #-}
{-# LANGUAGE OverloadedStrings #-}
module Main where


-------------------------------------------------------------------------------
import           Control.Monad
import           Data.ByteString    (ByteString)
import           Data.List.NonEmpty
import           Data.Semigroup
import qualified Data.Vector        as V
import qualified Database.Redis     as R
import           Debug.Trace
import           System.Random
-------------------------------------------------------------------------------


main :: IO ()
main = do
  conn <- connect (R.defaultConnectInfo :| [R.defaultConnectInfo]) [R.defaultConnectInfo]
  res <- runRedis conn $ do
    a <- get "foo"
    b <- get "foo"
    c <- multiExec $ do
      bar <- get "bar"
      baz <- get "baz"
      quux <- set "quux" "xuuq"
      return $ (,,) <$> bar <*> baz <*> quux
    return (a, b, c)
  print res


-------------------------------------------------------------------------------
data Capability
  = ReadWrite
  | ReadOnly
  deriving (Show)


-- | Chooses the most powerful capability of any two
instance Semigroup Capability where
  ReadWrite <> _       = ReadWrite
  _ <> ReadWrite       = ReadWrite
  ReadOnly <> ReadOnly = ReadOnly


instance Monoid Capability where
  mempty = ReadOnly
  mappend = (<>)


-------------------------------------------------------------------------------
-- constructor not exported
data Connection = Connection
  { connRW :: !(V.Vector R.Connection)
  -- ^ Non empty via smart constructor
  , connRO :: !(V.Vector R.Connection)
  }


instance Semigroup Connection where
  Connection rw1 ro1 <> Connection rw2 ro2 =
    Connection (rw1 <> rw2) (ro1 <> ro2)


connect :: NonEmpty R.ConnectInfo -> [R.ConnectInfo] -> IO Connection
connect rws ros = Connection
  <$> (mapM R.connect (V.fromList (toList rws)))
  <*> (mapM R.connect (V.fromList ros))


-------------------------------------------------------------------------------
--TODO: figure out slot for applicative
data Redis m a = Redis Capability (m a)


instance (Functor m) => Functor (Redis m) where
  fmap f (Redis capability m) = Redis capability (fmap f m)


instance (Applicative m) => Applicative (Redis m) where
  pure a = Redis mempty (pure a)
  Redis c1 a <*> Redis c2 b = Redis (c1 <> c2) (a <*> b)


-------------------------------------------------------------------------------
runRedis :: Connection -> Redis R.Redis a -> IO a
runRedis conn (Redis capability m) = do
  rConn <- selectConnection conn capability
  R.runRedis rConn m


-------------------------------------------------------------------------------
selectConnection :: Connection -> Capability -> IO R.Connection
selectConnection (Connection rws _) ReadWrite  = trace "readwrite" $ pick rws
selectConnection (Connection rws ros) ReadOnly = trace "either readwrite or readonly" $ pick2 rws ros


-------------------------------------------------------------------------------
pick :: V.Vector a -> IO a
pick v = do
  idx <- randomRIO (0, V.length v - 1)
  return (v V.! idx)


-------------------------------------------------------------------------------
-- | Pick that spans 2 vectors
pick2 :: V.Vector a -> V.Vector a -> IO a
pick2 v1 v2 = do
  let v1l = V.length v1
  idx <- randomRIO (0, (v1l + V.length v2) - 1)
  return $ if idx >= v1l
    then v2 V.! (idx - v1l)
    else v1 V.! idx


-------------------------------------------------------------------------------
get :: (R.RedisCtx m f) => ByteString -> Redis m (f (Maybe ByteString))
get k = Redis ReadOnly (R.get k)


-------------------------------------------------------------------------------
set :: (R.RedisCtx m f) => ByteString -> ByteString -> Redis m (f R.Status)
set k v = Redis ReadWrite (R.set k v)


-------------------------------------------------------------------------------
multiExec :: Redis R.RedisTx (R.Queued a) -> Redis R.Redis (R.TxResult a)
multiExec (Redis capability q) = Redis capability (R.multiExec q)