App.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{- | This module contains the legion application definition.  -}
module Network.Legion.Discovery.App (
  newApp,
  Input(..),
  Output(..),
  State,
  ServiceId(..),
  ServiceAddr(..),
  InstanceInfo(..),
  Time(..),
  Service(..),
  toKey,
) where

import Data.Binary (Binary, get, put)
import Data.Default.Class (Default, def)
import Data.Digest.Pure.SHA (sha256)
import Data.Map (Map)
import Data.Text (Text)
import Data.Time (UTCTime(UTCTime), Day(ModifiedJulianDay), addUTCTime,
  NominalDiffTime)
import Distribution.Version (Version, VersionRange, withinRange)
import GHC.Generics (Generic)
import Network.Legion (Legionary(Legionary), ApplyDelta, apply, PartitionKey,
  newMemoryPersistence, handleRequest, persistence, index, Tag(Tag))
import qualified Data.Binary as B
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text.Encoding as TE


data Input
  = GetRange Time VersionRange
  | GetAll Time
  | Ping Time ServiceId 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 (Show, Binary)
instance Default State where
  def = State Nothing


{- | This type represents a discoverable service. -}
data Service = Service {
         name :: ServiceId,
    instances :: Map ServiceAddr InstanceInfo
  }
  deriving (Show, Generic)
instance Binary Service


{- | 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


newApp :: IO (Legionary Input Output State)
newApp = do
    persist <- newMemoryPersistence
    return Legionary {
        handleRequest = \i s -> let (o, _) = handle i s in o,
        persistence = persist,
        index
      }
  where
    index (State Nothing) = Set.empty
    index (State (Just Service {name})) =
      (Set.singleton . Tag . TE.encodeUtf8 . unServiceId) name


{- | The main request handler. -}
handle :: Input -> State -> (Output, State)
handle (GetRange _ _) (State Nothing) =
  (InstanceList Map.empty, State Nothing)
handle (GetRange 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 service))
handle (GetAll _) (State Nothing) =
  (InstanceList Map.empty, State Nothing)
handle (GetAll now) (State (Just service)) =
  let
    output =
      InstanceList . Map.fromAscList $ [
        (addr, info)
        | (addr, info) <- Map.toAscList (instances service)
        , addUTCTime thirtySeconds (unTime (lastPing info)) >= unTime now
      ]
  in (output, State (Just service))
handle (Ping now name version addr) (State Nothing) =
  let
    state = State (Just Service {
        name = name,
        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 ping addr (instances service)
      })
    ping Nothing = Just InstanceInfo {
        version,
        lastPing = now
      }
    ping (Just info) = Just info {
        version,
        lastPing = now
      }
  in (PingResponse, state)
handle (GetService _ ) (State Nothing) =
  (ServiceResponse Nothing, State Nothing)
handle (GetService now) (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 (Just service))
  where
    notExpired = (unTime now <=) . addUTCTime thirtySeconds . unTime . lastPing


{- | A service id. -}
newtype ServiceId = ServiceId {
    unServiceId :: Text
  }
  deriving (Show, Eq, Binary)


instance ApplyDelta Input State where
  apply i s = let (_, s2) = handle i s in s2


{- |
  Wrapper for UTCTime. Useful beacuse UTCTime doesn't have a `Binary`
  instane.
-}
newtype Time = Time {unTime :: UTCTime}
  deriving (Show, Eq)
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 ServiceId into a partition key. -}
toKey :: ServiceId -> PartitionKey
toKey =
  B.decode
  . B.encode
  . sha256
  . LBS.fromStrict
  . TE.encodeUtf8
  . unServiceId


thirtySeconds :: NominalDiffTime
thirtySeconds = 30