Client.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{- |
  This module contains a client library for communicating with the
  'legion-discovery' service discovery program.
-}
module Network.Legion.Discovery.Client (
  -- * Establishing a Connection.
  connect,
  -- * Performing Queries.
  query,
  -- * Registering a Service.
  withService,
  -- * HTTP Utilities.
  newLB,
  withResponse,
  -- * Other Types.
  LBHttp,
  Discovery,
  Name(..),
  ServiceAddr(..),
) where

import Control.Concurrent (newEmptyMVar, MVar, putMVar, tryTakeMVar,
  forkIO, threadDelay)
import Control.Concurrent.LoadDistribution (withResource,
  evenlyDistributed, LoadBalanced)
import Control.Exception (bracket)
import Control.Monad (void)
import Data.Aeson (eitherDecode, Value, encode, object, (.=))
import Data.ByteString.Lazy (fromChunks)
import Data.Default.Class (def)
import Data.Map (Map)
import Data.Monoid ((<>))
import Data.Set (Set)
import Data.String (IsString)
import Data.Text (Text, unpack, pack)
import Data.Text.Encoding (encodeUtf8)
import Distribution.Text (display)
import Distribution.Version (VersionRange, Version)
import Network.HTTP.Client (Request, Response, BodyReader,
  parseRequest, host, secure, port, Manager, requestHeaders,
  checkStatus, responseStatus, brConsume, responseBody, path, method,
  RequestBody(RequestBodyLBS), requestBody)
import Network.HTTP.Types (urlEncode, statusIsSuccessful)
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Network.HTTP.Client as C

{- | This type represents a handle on the discovery service. -}
data Discovery = D {
       dName :: Name,
    dVersion :: Version,
         dLb :: LBHttp
  }


{- | Create a connection to the discovery service. -}
connect
  :: Name
    {- ^
      The name of the local program. This is used to track and display
      a network graph of service dependencies.
    -}
  -> Version
    {- ^ The version of the local program. -}
  -> Set ServiceAddr
    {- ^ The well-known set of urls where the discovery service lives. -}
  -> Manager
    {- ^
      The http manager used to manage communication to the discovery
      service, and also for communication with any services accessed
      via the 'withResponse' function.
    -}
  -> IO Discovery
connect name version urls manager =
  D name version . LB manager <$> evenlyDistributed (return urls)


{- |
  Run a registered service, making sure to unregister upon ternination.
-}
withService
  :: ServiceAddr
    {- ^ The service address on which your service instance can be contacted. -}
  -> Discovery
    {- ^ A handle on the discovery service, obtained via `connect`. -}
  -> IO a
    {- ^ The IO action to perform while registered as a service. -}
  -> IO a
withService addy d =
    bracket launchPing stopPing . const
  where
    launchPing :: IO (MVar ())
    launchPing = do
      stop <- newEmptyMVar
      void . forkIO $
        let
          loop =
            tryTakeMVar stop >>= \case
              Nothing -> ping >> threadDelay tenSeconds >> loop
              Just () -> return ()
        in loop
      return stop

    stopPing :: MVar () -> IO ()
    stopPing stop = putMVar stop ()

    ping :: IO ()
    ping = do
      let
        userAgent = encodeUtf8
          $ unName (dName d) <> "/" <> pack (display (dVersion d))
        req = def {
            path = 
              "/v1/ping/"
              <> urlEncode False (encodeUtf8 (unName (dName d)))
              <> "/"
              <> urlEncode False (encodeUtf8 (pack (display (dVersion d)))),
            method = "POST",
            requestHeaders = [
                ("user-agent", userAgent),
                ("content-type", pingRequestCT)
              ],
            checkStatus = const . const . const $ Nothing,
            requestBody = RequestBodyLBS . encode . object $ [
                "serviceAddress" .= unServiceAddr addy
              ]
          }
      {- TODO figure out what we want to do when the ping failes. -}
      withResponse req (dLb d) (const (return ()))

    tenSeconds :: Int
    tenSeconds = 10000000 {- in microseconds. -}


{- | Query the discovery service. -}
query
  :: Name
    {- ^ The name of the service you are looking for. -}
  -> VersionRange
    {- ^ The range of service versions with which you are compatible. -}
  -> Discovery
    {- ^ A handle on the discovery service. -}
  -> IO (Set ServiceAddr)
query name range d = do
  let
    userAgent = encodeUtf8
      $ unName (dName d) <> "-" <> pack (display (dVersion d))
    req = def {
        path = 
          "/v1/services/"
          <> urlEncode False (encodeUtf8 (unName name))
          <> "/"
          <> urlEncode False (encodeUtf8 (pack (display range))),
        method = "GET",
        requestHeaders = [("user-agent", userAgent)],
        checkStatus = const . const . const $ Nothing
      }
  withResponse req (dLb d) (\case
      Nothing -> fail "No Discovery instances available."
      Just resp -> do
        body <- brConsume (responseBody resp)
        case responseStatus resp of
          status | statusIsSuccessful status ->
            case eitherDecode (fromChunks body) of
              Left err -> fail
                $ "Couldn't decode Discovery response: " ++ show err
              Right instances -> return (
                  Set.fromList
                    (ServiceAddr <$> Map.keys (instances :: Map Text Value))
                )
          status -> fail $ "Bad response from Discovery:" ++ show (status, body)
    )


{- | The name of a service. -}
newtype Name = Name {unName :: Text} deriving (IsString)


{- |
  The type of a service address.
  
  A 'ServiceAddr' is a wrapper around an unstructured text value. The
  meaning of the text value is dependent on the registered service,
  as it is the one responsible for publishing the service address. We
  __highly__ recommend that registered services choose to publish a
  fully qualified URL as their service address.
-}
newtype ServiceAddr = ServiceAddr {unServiceAddr :: Text}
  deriving (IsString, Eq, Ord, Show)


{- |
  Analog of 'C.withResponse', but automatically replaces the host, port,
  and scheme portions of the 'Request' with values obtained from the
  load balancer.

  If a 'Nothing' value is passed to the response handler, that means
  that there are no available instances that match the query params that were
  passed to 'newLB'.
-}
withResponse
  :: Request
  -> LBHttp
  -> (Maybe (Response BodyReader) -> IO a)
  -> IO a
withResponse r lb f = withResource (lbLb lb) (\case
    Nothing -> f Nothing
    Just url -> do
      lbReq <- parseRequest (unpack (unServiceAddr url))
      let
        req = r {
            host = host lbReq,
            secure = secure lbReq,
            port = port lbReq
          }
      C.withResponse req (lbManager lb) (f . Just)
  )


{- | A handle on the load balancing context. -}
data LBHttp = LB {
    lbManager :: Manager,
         lbLb :: LoadBalanced ServiceAddr
  }


{- | Create a new load balanced http client, for use with 'withResponse'. -}
newLB
  :: Discovery
    {- ^ A handle on the discovery service, obtained via 'connect'. -}
  -> Name
    {- ^ The name of the target service. -}
  -> VersionRange
    {- ^ The range of service versions with which you are compatible. -}
  -> IO LBHttp
    {- ^ Returns a load balanced http client, for use with 'withResponse'. -}
newLB d n r = LB (lbManager (dLb d)) <$> evenlyDistributed (query n r d)


{- | The content type of a ping request. -}
pingRequestCT :: (IsString a) => a
pingRequestCT = "application/vnd.legion-discovery.ping-request+json"