App.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# 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.App (
  Input(..),
  Output(..),
  State,
  ServiceId(..),
  ServiceAddr(..),
  InstanceInfo(..),
  Client(..),
  RequestInfo(..),
  Time(..),
  Service(..),
  toKey,
) where

import Data.Aeson (ToJSON, toJSON, object, (.=))
import Data.Binary (Binary, get, put)
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)
import Data.Version (showVersion)
import Distribution.Text (display)
import Distribution.Version (Version, VersionRange, withinRange)
import GHC.Generics (Generic)
import Network.Legion (Event, apply, PartitionKey, Tag(Tag),
  Indexable, indexEntries)
import qualified Data.Aeson as A
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 as T
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 (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 . unServiceId) name
instance ToJSON State where
  toJSON (State s) = object ["state" .= s]


{- | 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
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 :: ServiceId,
    cVersion :: Maybe Version
  }
  deriving (Eq, Generic, Ord)
instance Binary Client
instance Show Client where
  show (Client name version) =
    T.unpack (unServiceId 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 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, ToJSON)


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 ServiceId into a partition key. -}
toKey :: ServiceId -> PartitionKey
toKey =
  B.decode
  . B.encode
  . sha256
  . LBS.fromStrict
  . TE.encodeUtf8
  . unServiceId


thirtySeconds :: NominalDiffTime
thirtySeconds = 30