App.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{- | This module contains the legion application definition.  -}
module Network.Legion.Discovery.App (
  Input(..),
  Output(..),
  State,
  ServiceId(..),
  ServiceAddr(..),
  InstanceInfo(..),
  Client(..),
  RequestInfo(..),
  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 (Event, apply, PartitionKey, Tag(Tag),
  Indexable, indexEntries)
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 [Client] ServiceId Time VersionRange
  | Ping Time ServiceId Version ServiceAddr
  | GetService Time
  | GetRequests
  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

instance Indexable State where
  indexEntries (State Nothing) = Set.empty
  indexEntries (State (Just Service {name})) =
    (Set.singleton . Tag . TE.encodeUtf8 . unServiceId) name

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


{- | Information about a request. -}
data RequestInfo = RequestInfo {
      riTime :: Time,
     riRange :: VersionRange
  }
  deriving (Show, Eq, Generic)
instance Binary RequestInfo


{- | Information about a query client. -}
data Client = Client {
       cName :: ServiceId,
    cVersion :: Maybe Version
  }
  deriving (Show, Eq, Generic, Ord)
instance Binary Client


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


{- | 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 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@(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
handle GetRequests (State Nothing) =
  (ServiceResponse Nothing, State Nothing)
handle GetRequests state@(State (Just service)) =
  (ServiceResponse (Just service), state)


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


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