LegionApp.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{- |
  This module contains the legion application definition.

  ToJSON instances are only to support a better 'Show' instance for
  debugging, they are not meant to be used as
-}
module Network.Legion.Discovery.LegionApp (
  -- * Types that make up the legion application definition.
  Input,
  Output,
  State,
  EntityName(..),
  ServiceAddr(..),
  InstanceInfo(..),
  Client(..),
  RequestInfo(..),
  Time(..),
  Service(..),

  -- * A more user-friendly interface to our legion application.
  getRange,
  ping,
  getService,
  getAllServices,

  -- * Utilities
  toKey,
) where

import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Aeson (ToJSON, toJSON, object, (.=))
import Data.Binary (Binary, get, put)
import Data.Conduit ((=$=), runConduit)
import Data.Default.Class (Default, def)
import Data.Digest.Pure.SHA (sha256)
import Data.Map (Map)
import Data.Monoid ((<>))
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8)
import Data.Time (UTCTime(UTCTime), Day(ModifiedJulianDay), addUTCTime,
  NominalDiffTime, getCurrentTime)
import Data.Version (showVersion)
import Distribution.Text (display)
import Distribution.Version (Version, VersionRange, withinRange)
import GHC.Generics (Generic)
import Network.Legion (Runtime, makeRequest, search, SearchTag(SearchTag),
  irKey, PartitionKey, Event, apply, Tag(Tag), Indexable, indexEntries)
import Web.HttpApiData (FromHttpApiData, parseUrlPiece)
import qualified Data.Aeson as A
import qualified Data.Binary as B
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Conduit.List as CL
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE


data Input
  = GetRange [Client] EntityName Time VersionRange
  | Ping Time EntityName Version ServiceAddr
  | GetService Time
  deriving (Show, Eq, Generic)
instance Binary Input


data Output
  = InstanceList (Map ServiceAddr InstanceInfo)
  | ServiceResponse (Maybe Service)
  | PingResponse
  deriving (Show, Generic)
instance Binary Output


newtype State = State {
    _unState :: Maybe Service
  }
  deriving (Binary)
instance Show State where
  show = T.unpack . decodeUtf8 . LBS.toStrict . A.encode
instance Default State where
  def = State Nothing
instance Indexable State where
  indexEntries (State Nothing) = Set.empty
  indexEntries (State (Just Service {name})) =
    (Set.singleton . Tag . TE.encodeUtf8 . unEntityName) name
instance ToJSON State where
  toJSON (State s) = object ["state" .= s]


{- | This type represents a discoverable service. -}
data Service = Service {
         name :: EntityName,
    instances :: Map ServiceAddr InstanceInfo,
     requests :: Map Client RequestInfo
  }
  deriving (Show, Generic)
instance Binary Service
instance ToJSON Service where
  toJSON (Service name_ instances_ requests_) = object [
      "name" .= name_,
      "instances" .= Map.mapKeys (T.unpack . unServiceAddr) instances_,
      "requests" .= Map.mapKeys show requests_
    ]


{- | Information about a request. -}
data RequestInfo = RequestInfo {
      riTime :: Time,
     riRange :: VersionRange
  }
  deriving (Show, Eq, Generic)
instance Binary RequestInfo
instance ToJSON RequestInfo where
  toJSON (RequestInfo time range) = object [
      "time" .= time,
      "range" .= display range
    ]


{- | Information about a query client. -}
data Client = Client {
       cName :: EntityName,
    cVersion :: Maybe Version
  }
  deriving (Eq, Generic, Ord)
instance Binary Client
instance Show Client where
  show (Client name version) =
    T.unpack (unEntityName name) <> maybe "" (("/" <>) . showVersion) version


{- | The address on which a service can be contacted.  -}
newtype ServiceAddr = ServiceAddr {unServiceAddr :: Text}
  deriving (Show, Eq, Ord, Binary)


{- | Additional information about the service instance. -}
data InstanceInfo = InstanceInfo {
     version :: Version,
    lastPing :: Time
  }
  deriving (Show, Generic)
instance Binary InstanceInfo
instance ToJSON InstanceInfo


{- | The main request handler. -}
handle :: Input -> State -> (Output, State)
handle (GetRange clients name now range) (State Nothing) =
  (
    InstanceList Map.empty,
    State (Just Service {
      name,
      instances = Map.empty,
      requests = Map.fromList [
          (client, RequestInfo {
              riTime = now,
              riRange = range
            })
          | client <- clients
        ]
    })
  )
handle (GetRange clients _name now range) (State (Just service)) =
    let
      output =
        InstanceList . Map.fromAscList $ [
          (addr, info)
          | (addr, info) <- Map.toAscList (instances service)
          , addUTCTime thirtySeconds (unTime (lastPing info)) >= unTime now
          , withinRange (version info) range
        ]
    in (output, State (Just (addRequest service)))
  where
    addRequest :: Service -> Service
    addRequest s@Service {requests} = s {
        requests =
          Map.union
            (Map.fromList [
                (client, RequestInfo {
                    riTime = now,
                    riRange = range
                  })
                | client <- clients
              ])
          requests
      }
handle (Ping now name version addr) (State Nothing) =
  let
    state = State (Just Service {
        name = name,
        requests = Map.empty,
        instances = Map.singleton addr InstanceInfo {
          version,
          lastPing = now
        }
      })
  in (PingResponse, state)
handle (Ping now name version addr) (State (Just service)) =
  let
    state = State (Just service {
        name = name,
        instances = Map.alter doPing addr (instances service)
      })
    doPing Nothing = Just InstanceInfo {
        version,
        lastPing = now
      }
    doPing (Just info) = Just info {
        version,
        lastPing = now
      }
  in (PingResponse, state)
handle (GetService _ ) (State Nothing) =
  (ServiceResponse Nothing, State Nothing)
handle (GetService now) state@(State (Just service)) =
    let
      i = Map.filter notExpired (instances service)
      output = ServiceResponse $ if Map.null i
        then Nothing
        else Just service {
            instances = i
          }
    in (output, state)
  where
    notExpired = (unTime now <=) . addUTCTime thirtySeconds . unTime . lastPing


{- | The name of an entity (client, or service, or both). -}
newtype EntityName = EntityName {
    unEntityName :: Text
  }
  deriving (Show, Eq, Binary, Ord, ToJSON)
instance FromHttpApiData EntityName where
  parseUrlPiece = Right . EntityName


instance Event Input Output State where
  apply = handle


{- |
  Wrapper for UTCTime. Useful beacuse UTCTime doesn't have a `Binary`
  instane.
-}
newtype Time = Time {unTime :: UTCTime}
  deriving (Show, Eq, Ord, ToJSON)
instance Binary Time where
  put (Time (UTCTime (ModifiedJulianDay day) time)) = put (day, toRational time)
  get = do
    (day, time) <- get
    return (Time (UTCTime (ModifiedJulianDay day) (fromRational time)))


{- | Convert a EntityName into a partition key. -}
toKey :: EntityName -> PartitionKey
toKey =
  B.decode
  . B.encode
  . sha256
  . LBS.fromStrict
  . TE.encodeUtf8
  . unEntityName


thirtySeconds :: NominalDiffTime
thirtySeconds = 30


{- | Get the instances of a service that match a version range. -}
getRange :: (MonadIO io)
  => Runtime Input Output
  -> [Client] {- ^ The clients making the requests. -}
  -> EntityName {- ^ The name of the target service. -}
  -> VersionRange {- ^ The allowable instance versions. -}
  -> io (Map ServiceAddr InstanceInfo)
getRange runtime clients name range = do
  now <- Time <$> liftIO getCurrentTime
  let req = GetRange clients name now range
  makeRequest runtime (toKey name) req >>= \case
    InstanceList instances -> return instances
    r -> fail $ "Bad runtime response to req [" ++ show req ++ "]: " ++ show r


{- | Register a service instance. -}
ping :: (MonadIO io)
  => Runtime Input Output
  -> EntityName {- ^ The name of the registering service. -}
  -> Version {- ^ The version of the registering service. -}
  -> ServiceAddr
    {- ^ The service address on which the service can be contacted. -}
  -> io ()
ping runtime name version addy = do
  now <- Time <$> liftIO getCurrentTime
  let req = Ping now name version addy
  makeRequest runtime (toKey name) req >>= \case
    PingResponse -> return ()
    r -> fail $ "Bad runtime response to req [" ++ show req ++ "]: " ++ show r


{- | Returns all current information about a service. -}
getService :: (MonadIO io)
  => Runtime Input Output
  -> EntityName {- ^ The name of the service to retrieve. -}
  -> io (Maybe Service)
getService runtime name = getServiceByKey runtime (toKey name)


{- | Like `getService`, but using the precomputed partition key. -}
getServiceByKey :: (MonadIO io)
  => Runtime Input Output
  -> PartitionKey
  -> io (Maybe Service)
getServiceByKey runtime key = do
  now <- Time <$> liftIO getCurrentTime
  let req = GetService now
  makeRequest runtime key req >>= \case
    ServiceResponse s -> return s
    r -> fail $ "Bad runtime response to req [" ++ show req ++ "]: " ++ show r


{- | Returns all services. -}
getAllServices :: (MonadIO io)
  => Runtime Input Output
  -> io [Service]
getAllServices runtime =
  runConduit (
      search runtime (SearchTag "" Nothing)
      =$= CL.mapMaybeM (getServiceByKey runtime . irKey)
      =$= CL.consume
    )