diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2015-2016, Soostone Inc + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Ozgun Ataman nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,24 @@ +# Katip Elasticsearch [![Build Status](https://travis-ci.org/Soostone/katip.svg?branch=master)](https://travis-ci.org/Soostone/katip) + +katip-elasticsearch is a scribe for the Katip logging framework that +sends structured logs to ElasticSearch. + +## Features + +* Built in bounded buffering. + +* Configurable pool of logging workers to help with high write + volume. + +* Optional field type annotation to avoid mistyping values. + +* Optional automatic date sharding, so logs can be filed into monthly, + weekly, daily, hourly, every minute indices. You can even specify + your own index routing logic. This pattern can be seen in the ELK + stack as a way of keeping indexes reasonably sized and easy to + optimize, rotate, and manage. + +* Customizable retry policy for temporary outages and errors. + +* Automatic index and mapping setup. + diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/bench/Main.hs b/bench/Main.hs new file mode 100644 --- /dev/null +++ b/bench/Main.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Main + ( main + ) where + +------------------------------------------------------------------------------- +import Control.DeepSeq +import Control.Monad +import Criterion.Main +import Data.Aeson +import qualified Data.HashMap.Strict as HM +import Data.RNG +import qualified Data.Text as T +import Database.Bloodhound.Types +import Numeric +------------------------------------------------------------------------------- +import Katip.Scribes.ElasticSearch +import Katip.Scribes.ElasticSearch.Annotations +------------------------------------------------------------------------------- + +main :: IO () +main = do + rng <- mkRNG + defaultMain + [ + mkDocIdBenchmark rng + , deannotateValueBenchmark + ] + +------------------------------------------------------------------------------- +mkDocIdBenchmark :: RNG -> Benchmark +mkDocIdBenchmark rng = bgroup "mkDocId" + [ + bench "mkDocId (randomIO)" $ nfIO mkDocId + , bench "mkDocId' (shared )" $ nfIO $ mkDocId' rng + ] + + +------------------------------------------------------------------------------- +deannotateValueBenchmark :: Benchmark +deannotateValueBenchmark = bgroup "deannotateValue" + [ + bench "deannotateValue" $ nf deannotateValue annotatedValue + ] + + +------------------------------------------------------------------------------- +annotatedValue :: Value +annotatedValue = Object $ HM.fromList [ ("a::string", String "whatever") + , ("b::double", Number 4.5) + , ("c::long", Number 4) + , ("d::boolean", Bool True) + , ("e::null", Null) + ] + +------------------------------------------------------------------------------- +mkDocId' :: RNG -> IO DocId +mkDocId' rng = do + is <- withRNG rng $ \gen -> replicateM len $ mk gen + return . DocId . T.pack . concatMap (`showHex` "") $ is + where + len = 32 + mk :: GenIO -> IO Int + mk = uniformR (0,15) + +deriving instance NFData DocId diff --git a/changelog.md b/changelog.md new file mode 100644 --- /dev/null +++ b/changelog.md @@ -0,0 +1,4 @@ +0.1.0.0 +============== + +Initial release diff --git a/katip-elasticsearch.cabal b/katip-elasticsearch.cabal new file mode 100644 --- /dev/null +++ b/katip-elasticsearch.cabal @@ -0,0 +1,103 @@ +name: katip-elasticsearch +synopsis: ElasticSearch scribe for the Katip logging framework. +description: See README.md for more details. +version: 0.1.0.0 +license: BSD3 +license-file: LICENSE +author: Ozgun Ataman +maintainer: ozgun.ataman@soostone.com +copyright: Soostone Inc, 2015-2016 +category: Data +build-type: Simple +cabal-version: >=1.10 +extra-source-files: + README.md + changelog.md + bench/Main.hs + test/Main.hs + +flag lib-Werror + default: False + manual: True + +library + exposed-modules: + Katip.Scribes.ElasticSearch + Katip.Scribes.ElasticSearch.Annotations + build-depends: + base >=4.6 && <5 + , katip >= 0.1.0.0 && < 0.2 + , bloodhound >= 0.11.0.0 && < 0.12 + , random >= 1.0 && < 1.2 + , uuid >= 1.3 && < 1.4 + , aeson >=0.6 && <0.12 + , stm >= 2.4.3 && < 2.5 + , stm-chans >= 3.0.0.2 && < 3.1 + , async >= 2.0.1.0 && < 2.2 + , enclosed-exceptions >= 1.0.0 && < 1.1 + , http-client + , text + , retry >= 0.7 && < 0.8 + , exceptions + , scientific + , unordered-containers + , transformers + , http-types + , time + hs-source-dirs: src + default-language: Haskell2010 + hs-source-dirs: src + default-language: Haskell2010 + if flag(lib-Werror) + ghc-options: -Wall -Werror + +benchmark bench + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: bench + default-language: Haskell2010 + ghc-options: -O2 + build-depends: + base + , katip-elasticsearch + , bloodhound + , deepseq + , criterion >= 1.1.0.0 + , rng-utils + , uuid + , text + , unordered-containers + , aeson + +test-suite test + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: test + default-language: Haskell2010 + build-depends: + katip-elasticsearch + , katip + , http-client + , bloodhound + , base + , tasty >= 0.11 + , tasty-hunit + , tasty-quickcheck + , quickcheck-instances + , containers + , transformers + , lens + , lens-aeson + , text + , http-types + , aeson + , vector + , unordered-containers + , scientific + , time + , stm + + if flag(lib-Werror) + ghc-options: -Wall -Werror + + ghc-options: -threaded -rtsopts diff --git a/src/Katip/Scribes/ElasticSearch.hs b/src/Katip/Scribes/ElasticSearch.hs new file mode 100644 --- /dev/null +++ b/src/Katip/Scribes/ElasticSearch.hs @@ -0,0 +1,390 @@ +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +module Katip.Scribes.ElasticSearch + (-- * Building a scribe + mkEsScribe + -- * Scribe configuration + , EsScribeSetupError(..) + , EsQueueSize + , mkEsQueueSize + , EsPoolSize + , mkEsPoolSize + , EsScribeCfg + , essRetryPolicy + , essQueueSize + , essPoolSize + , essAnnotateTypes + , essIndexSettings + , essIndexSharding + , IndexShardingPolicy(..) + , IndexNameSegment(..) + , defaultEsScribeCfg + -- * Utilities + , mkDocId + , module Katip.Scribes.ElasticSearch.Annotations + , roundToSunday + ) where + + +------------------------------------------------------------------------------- +import Control.Applicative as A +import Control.Concurrent +import Control.Concurrent.Async +import Control.Concurrent.STM.TBMQueue +import Control.Exception.Base +import Control.Exception.Enclosed +import Control.Monad +import Control.Monad.Catch +import Control.Monad.IO.Class +import Control.Monad.STM +import Control.Retry (RetryPolicy, + exponentialBackoff, + limitRetries, + recovering) +import Data.Aeson +import Data.Monoid ((<>)) +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import Data.Time +import Data.Time.Calendar.WeekDate +import Data.Typeable +import Data.UUID +import Database.Bloodhound +import Network.HTTP.Client +import Network.HTTP.Types.Status +import System.Random +------------------------------------------------------------------------------- +import Katip.Core +import Katip.Scribes.ElasticSearch.Annotations +------------------------------------------------------------------------------- + + +data EsScribeCfg = EsScribeCfg { + essRetryPolicy :: RetryPolicy + -- ^ Retry policy when there are errors sending logs to the server + , essQueueSize :: EsQueueSize + -- ^ Maximum size of the bounded log queue + , essPoolSize :: EsPoolSize + -- ^ Worker pool size limit for sending data to the + , essAnnotateTypes :: Bool + -- ^ Different payload items coexist in the "data" attribute in + -- ES. It is possible for different payloads to have different + -- types for the same key, e.g. an "id" key that is sometimes a + -- number and sometimes a string. If you're having ES do dynamic + -- mapping, the first log item will set the type and any that + -- don't conform will be *discarded*. If you set this to True, + -- keys will recursively be appended with their ES core + -- type. e.g. "id" would become "id::l" and "id::s" + -- automatically, so they won't conflict. When this library + -- exposes a querying API, we will try to make deserialization and + -- querying transparently remove the type annotations if this is + -- enabled. + , essIndexSettings :: IndexSettings + , essIndexSharding :: IndexShardingPolicy + } deriving (Typeable) + + +-- | Reasonable defaults for a config: +-- +-- * defaultManagerSettings +-- +-- * exponential backoff with 25ms base delay up to 5 retries +-- +-- * Queue size of 1000 +-- +-- * Pool size of 2 +-- +-- * Annotate types set to False +-- +-- * DailyIndexSharding +defaultEsScribeCfg :: EsScribeCfg +defaultEsScribeCfg = EsScribeCfg { + essRetryPolicy = exponentialBackoff 25 <> limitRetries 5 + , essQueueSize = EsQueueSize 1000 + , essPoolSize = EsPoolSize 2 + , essAnnotateTypes = False + , essIndexSettings = defaultIndexSettings + , essIndexSharding = DailyIndexSharding + } + + +------------------------------------------------------------------------------- +-- | How should katip store your log data? +-- +-- * NoIndexSharding will store all logs in one index name. This is +-- the simplest option but is not advised in production. In practice, +-- the index will grow very large and will get slower to +-- search. Deleting records based on some sort of retention period is +-- also extremely slow. +-- +-- * MonthlyIndexSharding, DailyIndexSharding, HourlyIndexSharding, +-- EveryMinuteIndexSharding will generate indexes based on the time of +-- the log. Index name is treated as a prefix. So if your index name +-- is @foo@ and DailySharding is used, logs will be stored in +-- @foo-2016-2-25@, @foo-2016-2-26@ and so on. Index templating will +-- be used to set up mappings automatically. Deletes based on date are +-- very fast and queries can be restricted to date ranges for better +-- performance. Queries against all dates should use @foo-*@ as an +-- index name. Note that index aliasing's glob feature is not suitable +-- for these date ranges as it matches index names as they are +-- declared, so new dates will be excluded. DailyIndexSharding is a +-- reasonable choice. Changing index sharding strategies is not +-- advisable. +-- +-- * CustomSharding: supply your own function that decomposes an item +-- into its index name heirarchy which will be appended to the index +-- name. So for instance if your function return ["arbitrary", +-- "prefix"], the index will be @foo-arbitrary-prefix@ and the index +-- template will be set to match @foo-*@. In general, you want to use +-- segments of increasing granularity (like year, month, day for +-- dates). This makes it easier to address groups of indexes +-- (e.g. @foo-2016-*@). +data IndexShardingPolicy = NoIndexSharding + | MonthlyIndexSharding + | WeeklyIndexSharding + -- ^ A special case of daily which shards to sunday + | DailyIndexSharding + | HourlyIndexSharding + | EveryMinuteIndexSharding + | CustomIndexSharding (forall a. Item a -> [IndexNameSegment]) + + +instance Show IndexShardingPolicy where + show NoIndexSharding = "NoIndexSharding" + show MonthlyIndexSharding = "MonthlyIndexSharding" + show WeeklyIndexSharding = "WeeklyIndexSharding" + show DailyIndexSharding = "DailyIndexSharding" + show HourlyIndexSharding = "HourlyIndexSharding" + show EveryMinuteIndexSharding = "EveryMinuteIndexSharding" + show (CustomIndexSharding _) = "CustomIndexSharding λ" + + +------------------------------------------------------------------------------- +newtype IndexNameSegment = IndexNameSegment { + indexNameSegment :: Text + } deriving (Show, Eq, Ord) + + +------------------------------------------------------------------------------- +shardPolicySegs :: IndexShardingPolicy -> Item a -> [IndexNameSegment] +shardPolicySegs NoIndexSharding _ = [] +shardPolicySegs MonthlyIndexSharding Item {..} = [sis y, sis m] + where + (y, m, _) = toGregorian (utctDay _itemTime) +shardPolicySegs WeeklyIndexSharding Item {..} = [sis y, sis m, sis d] + where + (y, m, d) = toGregorian (roundToSunday (utctDay _itemTime)) +shardPolicySegs DailyIndexSharding Item {..} = [sis y, sis m, sis d] + where + (y, m, d) = toGregorian (utctDay _itemTime) +shardPolicySegs HourlyIndexSharding Item {..} = [sis y, sis m, sis d, sis h] + where + (y, m, d) = toGregorian (utctDay _itemTime) + (h, _) = splitTime (utctDayTime _itemTime) +shardPolicySegs EveryMinuteIndexSharding Item {..} = [sis y, sis m, sis d, sis h, sis mn] + where + (y, m, d) = toGregorian (utctDay _itemTime) + (h, mn) = splitTime (utctDayTime _itemTime) +shardPolicySegs (CustomIndexSharding f) i = f i + + +------------------------------------------------------------------------------- +-- | If the given day is sunday, returns the input, otherwise returns +-- the previous sunday +roundToSunday :: Day -> Day +roundToSunday d + | dow == 7 = d + | w > 1 = fromWeekDate y (w - 1) 7 + | otherwise = fromWeekDate (y - 1) 53 7 + where + (y, w, dow) = toWeekDate d + + +------------------------------------------------------------------------------- +chooseIxn :: IndexName -> IndexShardingPolicy -> Item a -> IndexName +chooseIxn (IndexName ixn) p i = + IndexName (T.intercalate "-" (ixn:segs)) + where + segs = indexNameSegment A.<$> shardPolicySegs p i + + +------------------------------------------------------------------------------- +sis :: Show a => a -> IndexNameSegment +sis = IndexNameSegment . T.pack . show + + +------------------------------------------------------------------------------- +splitTime :: DiffTime -> (Int, Int) +splitTime t = asMins `divMod` 60 + where + asMins = floor t `div` 60 + + +------------------------------------------------------------------------------- +data EsScribeSetupError = CouldNotCreateIndex !Reply + | CouldNotCreateMapping !Reply deriving (Typeable, Show) + + +instance Exception EsScribeSetupError + +------------------------------------------------------------------------------- +mkEsScribe + :: EsScribeCfg + -> BHEnv + -> IndexName + -- ^ Treated as a prefix if index sharding is enabled + -> MappingName + -> Severity + -> Verbosity + -> IO (Scribe, IO ()) + -- ^ Returns a finalizer that will gracefully flush all remaining logs before shutting down workers +mkEsScribe cfg@EsScribeCfg {..} env ix mapping sev verb = do + q <- newTBMQueueIO $ unEsQueueSize essQueueSize + endSig <- newEmptyMVar + + runBH env $ do + chk <- indexExists ix + -- note that this doesn't update settings. That's not available + -- through the Bloodhound API yet + unless chk $ void $ do + r1 <- createIndex essIndexSettings ix + unless (statusIsSuccessful (responseStatus r1)) $ + liftIO $ throwIO (CouldNotCreateIndex r1) + r2 <- if shardingEnabled + then putTemplate tpl tplName + else putMapping ix mapping (baseMapping mapping) + unless (statusIsSuccessful (responseStatus r2)) $ + liftIO $ throwIO (CouldNotCreateMapping r2) + + workers <- replicateM (unEsPoolSize essPoolSize) $ async $ + startWorker cfg env mapping q + + _ <- async $ do + takeMVar endSig + atomically $ closeTBMQueue q + mapM_ waitCatch workers + putMVar endSig () + + let scribe = Scribe $ \ i -> + when (_itemSeverity i >= sev) $ + void $ atomically $ tryWriteTBMQueue q (chooseIxn ix essIndexSharding i, itemJson' i) + let finalizer = putMVar endSig () >> takeMVar endSig + return (scribe, finalizer) + where + tplName = TemplateName ixn + shardingEnabled = case essIndexSharding of + NoIndexSharding -> False + _ -> True + tpl = IndexTemplate (TemplatePattern (ixn <> "-*")) (Just essIndexSettings) [toJSON (baseMapping mapping)] + IndexName ixn = ix + itemJson' i + | essAnnotateTypes = itemJson verb (TypeAnnotated <$> i) + | otherwise = itemJson verb i + + +------------------------------------------------------------------------------- +baseMapping :: MappingName -> Value +baseMapping (MappingName mn) = + object [ mn .= object ["properties" .= object prs] ] + where prs = [ str "thread" + , str "sev" + , str "pid" + , str "ns" + , str "msg" + , "loc" .= locType + , str "host" + , str "env" + , "at" .= dateType + , str "app" + ] + str k = k .= object ["type" .= String "string"] + locType = object ["properties" .= object locPairs] + locPairs = [ str "loc_pkg" + , str "loc_mod" + , str "loc_ln" + , str "loc_fn" + , str "loc_col" + ] + dateType = object [ "format" .= esDateFormat + , "type" .= String "date" + ] + + +------------------------------------------------------------------------------- +-- | Handle both old-style aeson and picosecond-level precision +esDateFormat :: Text +esDateFormat = "yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ss.SSSZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSSSSSSSSZ" + + +------------------------------------------------------------------------------- +mkDocId :: IO DocId +mkDocId = (DocId . T.decodeUtf8 . toASCIIBytes) `fmap` randomIO + + +------------------------------------------------------------------------------- +newtype EsQueueSize = EsQueueSize { + unEsQueueSize :: Int + } deriving (Show, Eq, Ord) + + +instance Bounded EsQueueSize where + minBound = EsQueueSize 1 + maxBound = EsQueueSize maxBound + + +mkEsQueueSize :: Int -> Maybe EsQueueSize +mkEsQueueSize = mkNonZero EsQueueSize + + +------------------------------------------------------------------------------- +newtype EsPoolSize = EsPoolSize { + unEsPoolSize :: Int + } deriving (Show, Eq, Ord) + + +instance Bounded EsPoolSize where + minBound = EsPoolSize 1 + maxBound = EsPoolSize maxBound + + +mkEsPoolSize :: Int -> Maybe EsPoolSize +mkEsPoolSize = mkNonZero EsPoolSize + + +------------------------------------------------------------------------------- +mkNonZero :: (Int -> a) -> Int -> Maybe a +mkNonZero ctor n + | n > 0 = Just $ ctor n + | otherwise = Nothing + + +------------------------------------------------------------------------------- +startWorker + :: EsScribeCfg + -> BHEnv + -> MappingName + -> TBMQueue (IndexName, Value) + -> IO () +startWorker EsScribeCfg {..} env mapping q = go + where + go = do + popped <- atomically $ readTBMQueue q + case popped of + Just (ixn, v) -> do + sendLog ixn v `catchAny` eat + go + Nothing -> return () + sendLog :: IndexName -> Value -> IO () + sendLog ixn v = void $ recovering essRetryPolicy [handler] $ const $ do + did <- mkDocId + res <- runBH env $ indexDocument ixn mapping defaultIndexDocumentSettings v did + return res + eat _ = return () + handler _ = Handler $ \e -> + case fromException e of + Just (_ :: AsyncException) -> return False + _ -> return True diff --git a/src/Katip/Scribes/ElasticSearch/Annotations.hs b/src/Katip/Scribes/ElasticSearch/Annotations.hs new file mode 100644 --- /dev/null +++ b/src/Katip/Scribes/ElasticSearch/Annotations.hs @@ -0,0 +1,111 @@ +{-# LANGUAGE OverloadedStrings #-} +module Katip.Scribes.ElasticSearch.Annotations + ( TypeAnnotated(..) + -- * Exported for benchmarking + , deannotateValue + ) where + + +------------------------------------------------------------------------------- +import Control.Applicative as A +import Data.Aeson +import qualified Data.Foldable as FT +import qualified Data.HashMap.Strict as HM +import Data.Maybe +import Data.Monoid +import Data.Scientific (isFloating) +import Data.Text (Text) +import qualified Data.Text as T +------------------------------------------------------------------------------- +import Katip +------------------------------------------------------------------------------- + + +-- | Represents a value that can be converted to and from JSON that +-- will type annotate object keys when serializing and strip them out when deserializating +newtype TypeAnnotated a = TypeAnnotated { + typeAnnotatedValue :: a + } + +instance ToJSON a => ToJSON (TypeAnnotated a) where + toJSON = annotateValue . toJSON . typeAnnotatedValue + + +instance ToObject a => ToObject (TypeAnnotated a) + + +instance FromJSON a => FromJSON (TypeAnnotated a) where + parseJSON v = TypeAnnotated A.<$> parseJSON (deannotateValue v) + + +instance LogItem a => LogItem (TypeAnnotated a) where + payloadKeys v (TypeAnnotated x) = case payloadKeys v x of + AllKeys -> AllKeys + -- Take the key selection, overlap it with the actual keys + -- produced and annotate them + SomeKeys ks -> let o = toObject x + oInFocus = HM.fromList $ zip ks (repeat Null) + final = annotateKeys $ HM.intersection o oInFocus + in SomeKeys $ HM.keys final + +------------------------------------------------------------------------------- +-- Conversion Functions +------------------------------------------------------------------------------- + + +annotateValue :: Value -> Value +annotateValue (Object o) = Object $ annotateKeys o +annotateValue (Array a) = Array (annotateValue <$> a) +annotateValue x = x + + +annotateKeys :: Object -> Object +annotateKeys = HM.fromList . map go . HM.toList + where + go (k, Object o) = (k, Object $ annotateKeys o) + go (k, Array a) = (k, Array (annotateValue <$> a)) + go (k, s@(String _)) = (k <> stringAnn, s) + go (k, n@(Number sci)) = if isFloating sci + then (k <> doubleAnn, n) + else (k <> longAnn, n) + go (k, b@(Bool _)) = (k <> booleanAnn, b) + go (k, Null) = (k <> nullAnn, Null) + + +deannotateValue :: Value -> Value +deannotateValue (Object o) = Object $ deannotateKeys o +deannotateValue (Array a) = Array (deannotateValue <$> a) +deannotateValue x = x + + +deannotateKeys :: Object -> Object +deannotateKeys = HM.fromList . map go . HM.toList + where + go (k, Object o) = (k, Object $ deannotateKeys o) + go (k, Array a) = (k, Array (deannotateValue <$> a)) + go (k, v) = (fromMaybe k k', v) + where + k' = FT.asum (stripSuffix <$> suffixes) + suffixes = [stringAnn, doubleAnn, longAnn, booleanAnn, nullAnn] + stripSuffix suffix = T.stripSuffix suffix k + + +------------------------------------------------------------------------------- +-- Annotation Constants +------------------------------------------------------------------------------- + + +stringAnn :: Text +stringAnn = "::s" + +doubleAnn :: Text +doubleAnn = "::d" + +longAnn :: Text +longAnn = "::l" + +booleanAnn :: Text +booleanAnn = "::b" + +nullAnn :: Text +nullAnn = "::n" diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,308 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Main + ( main + ) where + + +------------------------------------------------------------------------------- +import Control.Applicative as A +import Control.Concurrent.STM +import Control.Lens hiding (mapping, (.=)) +import Control.Monad +import Control.Monad.IO.Class +import Data.Aeson +import Data.Aeson.Lens +import Data.Aeson.Types +import qualified Data.HashMap.Strict as HM +import qualified Data.Map as M +import Data.Monoid +import Data.Scientific +import Data.Time +import Data.Time.Calendar.WeekDate +import qualified Data.Vector as V +import Database.Bloodhound hiding (key) +import Network.HTTP.Client +import Network.HTTP.Types.Status +import Test.QuickCheck.Instances () +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck +------------------------------------------------------------------------------- +import Katip +import Katip.Scribes.ElasticSearch +------------------------------------------------------------------------------- + + +main :: IO () +main = defaultMain $ testGroup "katip-elasticsearch" + [ + esTests + , typeAnnotatedTests + , roundToSundayTests + ] + + +------------------------------------------------------------------------------- +setupSearch :: (EsScribeCfg -> EsScribeCfg) -> IO (Scribe, IO ()) +setupSearch modScribeCfg = do + bh dropESSchema + mgr <- newManager defaultManagerSettings + mkEsScribe cfg (mkBHEnv svr mgr) ixn mn DebugS V3 + where + cfg = modScribeCfg (defaultEsScribeCfg { essAnnotateTypes = True + , essIndexSettings = ixs + }) + + +------------------------------------------------------------------------------- +teardownSearch :: (Scribe, IO ()) -> IO () +teardownSearch (_, finalizer) = do + finalizer + bh $ do + when False $ dropESSchema + when False $ dropESSTemplate --TODO: drop + + +------------------------------------------------------------------------------- +withSearch :: (IO (Scribe, IO ()) -> TestTree) -> TestTree +withSearch = withSearch' id + + +------------------------------------------------------------------------------- +withSearch' :: (EsScribeCfg -> EsScribeCfg) -> (IO (Scribe, IO ()) -> TestTree) -> TestTree +withSearch' modScribeCfg = withResource (setupSearch modScribeCfg) teardownSearch + + +------------------------------------------------------------------------------- +esTests :: TestTree +esTests = testGroup "elasticsearch scribe" + [ + withSearch' (\c -> c { essIndexSharding = NoIndexSharding}) $ \setup -> testCase "it flushes to elasticsearch" $ withTestLogging setup $ \done -> do + $(logT) (ExampleCtx True) mempty InfoS "A test message" + liftIO $ do + void done + logs <- getLogs + length logs @?= 1 + let l = head logs + l ^? key "_source" . key "msg" . _String @?= Just "A test message" + l ^? key "_source" . key "data" . key "whatever::b" . _Bool @?= Just True + , withSearch $ \setup -> testCase "date-based index sharding" $ do + let t1 = mkTime 2016 1 2 3 4 5 + fakeClock <- newTVarIO t1 + withTestLogging' (set logEnvTimer (readTVarIO fakeClock)) setup $ \done -> do + $(logT) (ExampleCtx True) mempty InfoS "today" + let t2 = mkTime 2016 1 3 3 4 5 + liftIO (atomically (writeTVar fakeClock t2)) + $(logT) (ExampleCtx True) mempty InfoS "tomorrow" + liftIO $ do + void done + todayLogs <- getLogsByIndex (IndexName "katip-elasticsearch-tests-2016-1-2") + tomorrowLogs <- getLogsByIndex (IndexName "katip-elasticsearch-tests-2016-1-3") + assertBool ("todayLogs has " <> show (length todayLogs) <> " items") (length todayLogs == 1) + assertBool ("tomorrowLogs has " <> show (length tomorrowLogs) <> " items") (length tomorrowLogs == 1) + let logToday = head todayLogs + let logTomorrow = head tomorrowLogs + logToday ^? key "_source" . key "msg" . _String @?= Just "today" + logTomorrow ^? key "_source" . key "msg" . _String @?= Just "tomorrow" + , withSearch' (\c -> c { essIndexSharding = WeeklyIndexSharding}) $ \setup -> testCase "weekly index sharding rounds to previous sunday" $ do + let t1 = mkTime 2016 3 5 0 0 0 -- saturday, march 5th + fakeClock <- newTVarIO t1 + withTestLogging' (set logEnvTimer (readTVarIO fakeClock)) setup $ \done -> do + $(logT) (ExampleCtx True) mempty InfoS "today" + let t2 = mkTime 2016 3 6 0 0 0 -- sunday march 6th + liftIO (atomically (writeTVar fakeClock t2)) + $(logT) (ExampleCtx True) mempty InfoS "tomorrow" + liftIO $ do + void done + todayLogs <- getLogsByIndex (IndexName "katip-elasticsearch-tests-2016-2-28") -- rounds back to previous sunday + tomorrowLogs <- getLogsByIndex (IndexName "katip-elasticsearch-tests-2016-3-6") -- is on sunday, so uses current date + assertBool ("todayLogs has " <> show (length todayLogs) <> " items") (length todayLogs == 1) + assertBool ("tomorrowLogs has " <> show (length tomorrowLogs) <> " items") (length tomorrowLogs == 1) + let logToday = head todayLogs + let logTomorrow = head tomorrowLogs + logToday ^? key "_source" . key "msg" . _String @?= Just "today" + logTomorrow ^? key "_source" . key "msg" . _String @?= Just "tomorrow" + ] + + +------------------------------------------------------------------------------- +mkTime :: Integer -> Int -> Int -> DiffTime -> DiffTime -> DiffTime -> UTCTime +mkTime y m d hr minute s = UTCTime day dt + where + day = mkDay y m d + dt = s + 60 * minute + 60 * 60 * hr + + +------------------------------------------------------------------------------- +mkDay :: Integer -> Int -> Int -> Day +mkDay y m d = day + where + Just day = fromGregorianValid y m d + + +------------------------------------------------------------------------------- +data ExampleCtx = ExampleCtx { + ecBool :: Bool + } + +instance ToJSON ExampleCtx where + toJSON c = object ["whatever" .= ecBool c] + + +instance ToObject ExampleCtx + + +instance LogItem ExampleCtx where + payloadKeys _ _ = AllKeys + +------------------------------------------------------------------------------- +typeAnnotatedTests :: TestTree +typeAnnotatedTests = testGroup "TypeAnnotated" + [ + testCase "annotates values on toJSON" $ do + toJSON (TypeAnnotated exampleValue) @?= annotatedExampleValue + , testCase "deannotates on parseJSON" $ do + parseEither parseJSON (toJSON exampleValue) @?= Right exampleValue + , testProperty "roundtrips the same as raw" $ \(v :: Value) -> + let res = typeAnnotatedValue <$> (parseEither parseJSON (toJSON (TypeAnnotated v))) + in res === Right v + ] + + +------------------------------------------------------------------------------- +roundToSundayTests :: TestTree +roundToSundayTests = testGroup "roundToSunday" + [ + testProperty "always returns a sunday" $ \d -> + getDOW (roundToSunday d) === 7 + , testProperty "returns input on sunday" $ \d -> getDOW d == 7 ==> + roundToSunday d === d + , testProperty "goes back a week when not sunday" $ \d -> getDOW d /= 7 ==> + roundToSunday d < d + ] + where + getDOW = view _3 . toWeekDate + +------------------------------------------------------------------------------- +exampleValue :: Value +exampleValue = Array $ V.fromList [Null, Object ob] + where + ob = HM.fromList [ ("a bool", Bool False) + , ("a long", Number 24) + , ("a double", Number 52.3) + , ("a string", String "s") + , ("a null", Null) + , ("a map", Object (HM.singleton "baz" (Bool True))) + ] + + +------------------------------------------------------------------------------- +annotatedExampleValue :: Value +annotatedExampleValue = Array $ V.fromList + [ + Null + , Object $ HM.fromList + [ + ("a map",Object $ HM.fromList [("baz::b", Bool True)]) + , ("a bool::b", Bool False) + , ("a null::n", Null) + , ("a string::s", String "s") + , ("a double::d", Number 52.3) + , ("a long::l", Number 24.0) + ] + ] + + +------------------------------------------------------------------------------- +getLogs :: IO [Value] +getLogs = getLogsByIndex ixn + + +------------------------------------------------------------------------------- +getLogsByIndex :: IndexName -> IO [Value] +getLogsByIndex i = do + r <- bh $ do + void (refreshIndex i) + searchByIndex i (mkSearch Nothing Nothing) + let actualCode = statusCode (responseStatus r) + assertBool ("search by " <> show i <> " " <> show actualCode <> " /= 200") (actualCode == 200) + return $ responseBody r ^.. key "hits" . key "hits" . values + + +------------------------------------------------------------------------------- +bh :: BH IO a -> IO a +bh = withBH defaultManagerSettings svr + + +------------------------------------------------------------------------------- +withTestLogging + :: IO (Scribe, IO a) -> (IO Reply -> KatipT IO b) -> IO b +withTestLogging = withTestLogging' id + + +------------------------------------------------------------------------------- +withTestLogging' + :: (LogEnv -> LogEnv) + -> IO (Scribe, IO a) + -> (IO Reply -> KatipT IO b) + -> IO b +withTestLogging' modEnv setup f = do + (scr, done) <- setup + le <- modEnv <$> initLogEnv ns env + let done' = done >> bh (refreshIndex ixn) + runKatipT le { _logEnvScribes = M.singleton "es" scr} (f done') + where + ns = Namespace ["katip-test"] + env = Environment "test" + + +------------------------------------------------------------------------------- +svr :: Server +svr = Server "http://localhost:9200" + + +------------------------------------------------------------------------------- +ixn :: IndexName +ixn = IndexName "katip-elasticsearch-tests" + + +------------------------------------------------------------------------------- +ixs :: IndexSettings +ixs = defaultIndexSettings { indexShards = ShardCount 1 + , indexReplicas = ReplicaCount 1} + +------------------------------------------------------------------------------- +mn :: MappingName +mn = MappingName "logs" + + +------------------------------------------------------------------------------- +dropESSchema :: BH IO () +dropESSchema = void $ deleteIndex (IndexName "katip-elasticsearch-tests*") + + +------------------------------------------------------------------------------- +dropESSTemplate :: BH IO () +dropESSTemplate = void $ deleteTemplate (TemplateName "katip-elasticsearch-tests") + + +------------------------------------------------------------------------------- +instance Arbitrary Value where + arbitrary = oneof + [ Object <$> reduceSize arbitrary + , Array . V.fromList <$> reduceSize arbitrary + , String <$> arbitrary + , Number <$> (scientific <$> arbitrary <*> arbitrary) + , Bool <$> arbitrary + , A.pure Null + ] + + +------------------------------------------------------------------------------- +-- | Reduce the size of Arbitrary input for the given generator +reduceSize :: Gen a -> Gen a +reduceSize f = sized $ \ n -> resize (n `div` 2) f