Moonshine.hs

{-# LANGUAGE DeriveGeneric, OverloadedStrings, NamedFieldPuns #-}
module Web.Moonshine (
  Moonshine,
  runMoonshine,
  route,
  makeTimer,
  timerAdd,
  Timer,
  timed,
  liftSnap,
  liftIO,
  setBanner,
  timedIO,
  makeGauge,
  gaugeInc,
  gaugeDec,
  gaugeSet,
  gaugeAdd,
  gaugeSubtract,
  Gauge,
) where

import Canteven.Config (canteven)
import Control.Applicative (Applicative(pure, (<*>)), (<|>))
import Control.Monad (ap, liftM)
import Control.Monad.IO.Class (MonadIO(liftIO))
import Data.Aeson (Value(Object), (.:?), (.!=), (.!=))
import Data.ByteString (ByteString)
import Data.Int (Int64)
import Data.Maybe (fromMaybe)
import Data.Monoid (mempty, mconcat)
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8)
import Data.Time.Clock (getCurrentTime, diffUTCTime)
import Data.Yaml (FromJSON(parseJSON))
import GHC.Generics (Generic)
import Snap (Snap, httpServe, setPort, Config, setSSLCert, setSSLKey,
  setSSLPort, setResponseCode, getResponse, finishWith, setAccessLog,
  setErrorLog, ConfigLog(ConfigIoLog))
import System.Environment (setEnv, lookupEnv)
import System.Log.Logger (infoM, errorM)
import System.Metrics (registerGcMetrics)
import System.Metrics.Distribution (Distribution)
import System.Remote.Monitoring (Server, forkServerWith)
import qualified Data.Text as T (unpack)
import qualified Snap (route)
import qualified System.Log.Logger as L (debugM)
import qualified System.Metrics as EkgMetrics
import qualified System.Metrics.Distribution as D (add)
import qualified System.Metrics.Gauge as G (inc, Gauge, dec, set, add, subtract)

-- Public Types ---------------------------------------------------------------
-- Semi-Public Types ----------------------------------------------------------

{- |
  The moonshine monad.
-}
data Moonshine a =
  Moonshine (State -> IO (State , a))

instance Monad Moonshine where
  return a = Moonshine (\state -> return (state, a))
  (Moonshine m) >>= fun = Moonshine $ \state -> do
    (newState, val) <- m state
    let Moonshine m2 = fun val
    m2 newState

instance Applicative Moonshine where
  pure = return
  (<*>) = ap

instance Functor Moonshine where
  fmap = liftM

instance MonadIO Moonshine where
  liftIO io = Moonshine (\state -> do
      val <- io
      return (state, val)
    )


{- |
  The Timer type.
-}
newtype Timer = T Distribution

{- |
  The Gauge type.
-}
newtype Gauge = G {unGauge :: G.Gauge}


data ServerConfig =
  ServerConfig {
    applicationConnector :: [ConnectorConfig],
    adminConnector :: [ConnectorConfig]
      -- ^ Configuration for EKG.  Note that although this is a list,
      --   at present EKG makes it impossible to support more than one
      --   server. See https://github.com/tibbe/ekg/issues/32.
  } deriving (Generic)

instance FromJSON ServerConfig


data ConnectorConfig = ConnectorConfig {
    scheme :: Scheme
  , port :: Int
  , cert :: Maybe FilePath
  , key :: Maybe FilePath
  } deriving (Generic)

instance FromJSON ConnectorConfig


data Scheme = HTTP | HTTPS
  deriving (Generic)

instance FromJSON Scheme


-- Public Functions -----------------------------------------------------------

{- |
  Execute an insance of the Moonshine monad. This starts up a server
  and never returns.
-}
runMoonshine :: Moonshine a -> IO a
runMoonshine (Moonshine m) = do
    printBanner
    MoonshineConfig {mcServerConfig} <- canteven
    metricsStore <- EkgMetrics.newStore
    (State {snap}, val) <- m State {
        snap = return (),
        metricsStore
      }
    let snap2 = snap <|> notFound
    startServer mcServerConfig metricsStore snap2
    return val
  where
    notFound :: Snap a 
    notFound = do
      response <- fmap (setResponseCode 404) getResponse
      finishWith response

data MoonshineConfig =
  MoonshineConfig {
    mcServerConfig :: ServerConfig
  }
instance FromJSON MoonshineConfig where
  parseJSON v =
    case v of
      Object topLevel ->  do
        mcServerConfig <- topLevel .:? "server" .!= defaultServerConfig
        return MoonshineConfig {mcServerConfig}
      _ -> 
        fail $ "Couldn't parse system config from value: " ++ show v


{- |
  Like `Snap.route`, but that automatically sets up metrics for the
  specified routes.
-}
route :: [(ByteString, Moonshine ())] -> Moonshine ()
route routes = Moonshine $ \state@State {snap, metricsStore} -> do
  snapRoutes <- mapM (getUnboundSnapRoute state) routes
  monitoredRoutes <- mapM (monitorRoute metricsStore) snapRoutes
  return (state {snap = snap >> Snap.route monitoredRoutes}, ())
  where
    getUnboundSnapRoute state (path, Moonshine m) = do
      (State {snap}, ()) <- m state {snap = return ()}
      return (path, snap)

    monitorRoute
      :: EkgMetrics.Store
      -> (ByteString, Snap ())
      -> IO (ByteString, Snap ())
    monitorRoute metricsStore (path, snap) = do -- IO monad
      timer <- getDistribution (decodeUtf8 path) metricsStore
      return (path, timedSnap timer snap)


{- |
  Lifts a snap action into a moonshine action.
-}
liftSnap :: Snap () -> Moonshine ()
liftSnap s = Moonshine (\state@State {snap} ->
    return (state {snap = snap >> s}, ())
  )


{- |
  Wrap a `Moonshine` in a timer, so that timing data for invocations
  will be logged to the metrics handler.
-}
timed
  :: Text
    -- ^ The name of the timer.
  -> Moonshine a
    -- ^ The action to be wrapped.
  -> Moonshine a
timed name (Moonshine m) = Moonshine $ \state@State {snap, metricsStore} -> do
  timer <- getDistribution name metricsStore
  m state {snap = timedSnap timer snap}


{- |
  Make a new timer with the given name.
-}
makeTimer :: Text -> Moonshine Timer
makeTimer name = Moonshine (\state@State {metricsStore} -> do
    dist <- getDistribution name metricsStore
    return (state, T dist)
  )


{- |
  Add a time to a timer.
-}
timerAdd :: (Real time, MonadIO io) => Timer -> time -> io ()
timerAdd (T timer) = liftIO . D.add timer . fromRational . toRational


{- |
  Time how long it takes to perform some kind of IO. This function is
  compatible with any `MonadIO`.
-}
timedIO :: (MonadIO io) => Timer -> io a -> io a
timedIO timer io = do
  start <- liftIO getCurrentTime
  result <- io
  end <- liftIO getCurrentTime
  timerAdd timer (diffUTCTime end start)
  return result


{- |
  Sets the banner that moonshine should display on startup.
-}
setBanner :: String -> IO ()
setBanner = setEnv bannerEnvName


{- |
  Add one to a gauge value.
-}
gaugeInc :: (MonadIO io) => Gauge -> io ()
gaugeInc = liftIO . G.inc . unGauge


{- |
  Subtract one from a gauge value.
-}
gaugeDec :: (MonadIO io) => Gauge -> io ()
gaugeDec = liftIO . G.dec . unGauge


{- |
  Set the gauge the a specific value one from a gauge value.
-}
gaugeSet :: (MonadIO io) => Gauge -> Int64 -> io ()
gaugeSet (G gauge) = liftIO . G.set gauge


{- |
  Subtrace a value from a gauge.
-}
gaugeSubtract :: (MonadIO io) => Gauge -> Int64 -> io ()
gaugeSubtract (G gauge) = liftIO . G.subtract gauge


{- |
  Add a value to a gauge.
-}
gaugeAdd :: (MonadIO io) => Gauge -> Int64 -> io ()
gaugeAdd (G gauge) = liftIO . G.add gauge


{- |
  Create a new gauge with the given name.
-}
makeGauge :: Text -> Moonshine Gauge
makeGauge name = Moonshine (\state@State {metricsStore} -> do
    gauge <- EkgMetrics.createGauge name metricsStore
    return (state, G gauge)
  )


-- Private Types --------------------------------------------------------------

{- |
  Defines command line options.
-}
data Opt = Help
         | Config String
         deriving Show

{- |
  The internal running state used while executing the `Moonshine` monad.
-}
data State =
  State {
    snap :: Snap (),
    metricsStore :: EkgMetrics.Store
  }


-- Private Functions ----------------------------------------------------------

defaultServerConfig :: ServerConfig
defaultServerConfig = ServerConfig {
    adminConnector =
      [ConnectorConfig {scheme=HTTP, port=8001, cert=Nothing, key=Nothing}]
  , applicationConnector =
      [ConnectorConfig {scheme=HTTP, port=8000, cert=Nothing, key=Nothing}]
  }


{- |
  Find and output a banner.
-}
printBanner :: IO ()
printBanner = do
  banner <- lookupEnv bannerEnvName
  putStr $ fromMaybe
    "FIXME: failed to magic setBanner.  son, I am disappoint\n"
    banner


{- |
  This is the name o the magic banner environment variable.
-}
bannerEnvName :: String
bannerEnvName = "mclovin-banner-magic"


{- |
  A replacement for 'System.Remote.Monitoring.getDistribution' that uses
  an 'EkgMetrics.Store' instead of a 'Server'.
-}
getDistribution :: Text -> EkgMetrics.Store -> IO Distribution
getDistribution = EkgMetrics.createDistribution


{- |
  Start application listening on the ports given by the config.
-}
startServer :: ServerConfig -> EkgMetrics.Store -> Snap () -> IO ()
startServer
    ServerConfig { applicationConnector, adminConnector }
    metricsStore
    snap
  = do
    mapM_ startMetricsServer adminConnector
    httpServe snapConfig snap

  where
    startMetricsServer :: ConnectorConfig -> IO Server
    startMetricsServer ConnectorConfig { scheme=HTTP, port } = do
      registerGcMetrics metricsStore
      debugM ("Starting up ekg on port: " ++ show port)
      forkServerWith metricsStore "0.0.0.0" port
    startMetricsServer ConnectorConfig { scheme=HTTPS } =
      error "EKG does not support running on HTTPS"

    snapConfig :: Config m a
    snapConfig = mconcat $
        map toSnapConfig applicationConnector ++ [
          setAccessLog (ConfigIoLog (infoM "snap" . toString)) mempty,
          setErrorLog (ConfigIoLog (errorM "snap" . toString)) mempty
        ]
      where toString = T.unpack . decodeUtf8
        
    toSnapConfig :: ConnectorConfig -> Config m a
    toSnapConfig ConnectorConfig { scheme=HTTP, port } = setPort port mempty
    toSnapConfig ConnectorConfig { scheme=HTTPS, port, cert=mcert, key=mkey } =
      fromMaybe
        (error "You must provide cert and key in order to use HTTPS.")
        $ do -- Maybe monad
          cert <- mcert
          key <- mkey
          return $ setSSLCert cert $ setSSLKey key $ setSSLPort port mempty


{- |
  Add timing information to a snap action.
-}
timedSnap :: Distribution -> Snap () -> Snap ()
timedSnap timer snap = do -- snap monad
  start <- liftIO getCurrentTime
  result <- snap
  end <- liftIO getCurrentTime
  addTiming start end
  return result
  where
    addTiming start end = liftIO $
      D.add timer diff
      where
        diff = toDouble (diffUTCTime end start)
        toDouble = fromRational . toRational


{- |
  Shorthand logging
-}
debugM :: String -> IO ()
debugM = L.debugM "moonshine"