Rework counter logic test for rule reloading (#69)

* Use the server's getLimit instead of doing a custom lookup
* Move and rename shouldRateLimitDescriptor to Fencer.AppState
* Rename the Fencer.AppState module to Fencer.Logic
* Use updateLimitCounter instead of recordHits
* Move a counter logic test into a separate module

Co-Authored-By: Artyom Kazak <[email protected]>

.hlint.yaml

@@ -38,3 +38,11 @@
 - ignore: { name: 'Use const' }
 - ignore: { name: 'Move brackets to avoid $' }
 - ignore: { name: 'Use ?~' }
+
+# Custom hints
+- warn:
+    lhs: "assertBool x (a == b)"
+    rhs: "assertEqual x a b"
+    note: >-
+      'assertEqual' takes arguments in order /expected, actual/;
+      make sure that the order is right.

fencer.cabal

@@ -45,7 +45,7 @@ library
     common
   exposed-modules:
     Fencer.Main
-    Fencer.AppState
+    Fencer.Logic
     Fencer.Server
     Fencer.Counter
     Fencer.Rules
@@ -110,6 +110,7 @@ test-suite test-fencer
   hs-source-dirs:
     test
   other-modules:
+    Fencer.Logic.Test
     Fencer.Types.Test
     Fencer.Rules.Test
     Fencer.Server.Test
@@ -126,7 +127,6 @@ test-suite test-fencer
     , neat-interpolation
     , proto3-wire
     , proto3-suite
-    , stm-containers
     , tasty
     , tasty-hunit
     , temporary

lib/Fencer/AppState.hs → lib/Fencer/Logic.hs

@@ -2,20 +2,18 @@
 {-# LANGUAGE OverloadedLabels #-}
 {-# LANGUAGE RecordWildCards #-}
 
--- | In-memory state of Fencer.
-module Fencer.AppState
+-- | In-memory state of Fencer and the logic for managing it.
+module Fencer.Logic
     ( AppState
-    , appStateCounters
-    , appStateRules
     , initAppState
 
     -- * Methods for working with 'AppState'
-    , recordHits
     , getLimit
     , setRules
     , getAppStateRulesLoaded
     , updateCurrentTime
     , deleteCountersWithExpiry
+    , updateLimitCounter
     )
 where
 
@@ -155,6 +153,31 @@ getLimit appState domain descriptor =
         Nothing -> pure Nothing
         Just ruleTree -> pure (applyRules descriptor ruleTree)
 
+-- | Handle a single descriptor in a 'shouldRateLimit' request.
+--
+-- Returns the current limit and response.
+--
+-- 'updateLimitCounter' will create a new counter if the counter does
+-- not exist, or update an existing counter otherwise. The counter will be
+-- reset if it has expired, and 'appStateCounterExpiry' will be updated.
+updateLimitCounter
+    :: AppState
+    -> "hits" :! Word
+    -> DomainId
+    -> [(RuleKey, RuleValue)]
+    -> STM (Maybe (RateLimit, CounterStatus))
+updateLimitCounter appState (arg #hits -> hits) domain descriptor =
+    getLimit appState domain descriptor >>= \case
+        Nothing -> pure Nothing
+        Just limit -> do
+            let counterKey :: CounterKey
+                counterKey = CounterKey
+                  { counterKeyDomain = domain
+                  , counterKeyDescriptor = descriptor
+                  , counterKeyUnit = rateLimitUnit limit }
+            status <- recordHits appState (#hits hits) (#limit limit) counterKey
+            pure (Just (limit, status))
+
 -- | Set 'appStateRules' and 'appStateRulesLoaded'.
 --
 -- The 'appStateCounters' field stays unchanged. This is in accordance

lib/Fencer/Main.hs

@@ -19,7 +19,7 @@ import qualified System.Logger as Logger
 import System.Logger (Logger)
 
 import Fencer.Types
-import Fencer.AppState
+import Fencer.Logic
 import Fencer.Server
 import Fencer.Rules
 import Fencer.Watch

lib/Fencer/Server.hs

@@ -13,7 +13,6 @@ where
 import BasePrelude hiding ((+++))
 
 import           Control.Concurrent.STM (atomically)
-import           Named ((:!), arg)
 import qualified Data.ByteString.Char8 as B
 import           Data.Text (Text)
 import qualified Data.Text.Lazy as TL
@@ -24,7 +23,7 @@ import qualified System.Logger as Logger
 import           System.Logger (Logger)
 import           System.Logger.Message ((+++))
 
-import           Fencer.AppState
+import           Fencer.Logic
 import           Fencer.Counter
 import qualified Fencer.Proto as Proto
 import           Fencer.Settings (defaultGRPCPort)
@@ -107,7 +106,7 @@ shouldRateLimit logger appState (Grpc.ServerNormalRequest serverCall request) =
      statuses :: [Proto.RateLimitResponse_DescriptorStatus]) <-
         fmap (unzip . map (maybe ruleNotFoundResponse counterStatusToProto)) $
         atomically $ forM descriptors $ \descriptor ->
-            shouldRateLimitDescriptor appState (#hits hits) domain descriptor
+            updateLimitCounter appState (#hits hits) domain descriptor
 
     -- Return server response.
     let overallCode =
@@ -129,31 +128,6 @@ shouldRateLimit logger appState (Grpc.ServerNormalRequest serverCall request) =
     let statusDetails = ""
     pure $ Grpc.ServerNormalResponse answer metadata Grpc.StatusOk statusDetails
 
--- | Handle a single descriptor in a 'shouldRateLimit' request.
---
--- Returns the current limit and response.
---
--- 'shouldRateLimitDescriptor' will create a new counter if the counter does
--- not exist, or update an existing counter otherwise. The counter will be
--- reset if it has expired, and 'appStateCounterExpiry' will be updated.
-shouldRateLimitDescriptor
-    :: AppState
-    -> "hits" :! Word
-    -> DomainId
-    -> [(RuleKey, RuleValue)]
-    -> STM (Maybe (RateLimit, CounterStatus))
-shouldRateLimitDescriptor appState (arg #hits -> hits) domain descriptor =
-    getLimit appState domain descriptor >>= \case
-        Nothing -> pure Nothing
-        Just limit -> do
-            let counterKey :: CounterKey
-                counterKey = CounterKey
-                  { counterKeyDomain = domain
-                  , counterKeyDescriptor = descriptor
-                  , counterKeyUnit = rateLimitUnit limit }
-            status <- recordHits appState (#hits hits) (#limit limit) counterKey
-            pure (Just (limit, status))
-
 ----------------------------------------------------------------------------
 -- Working with protobuf structures
 ----------------------------------------------------------------------------

test/Fencer/Logic/Test.hs

@@ -0,0 +1,105 @@
+{-# LANGUAGE OverloadedLabels  #-}
+{-# LANGUAGE OverloadedStrings #-}
+
+-- | Tests for "Fencer.Logic".
+module Fencer.Logic.Test
+  ( test_logicLimitUnitChange
+  ) where
+
+import           BasePrelude
+
+import           Data.List.NonEmpty (NonEmpty((:|)))
+import qualified Data.List.NonEmpty as NE
+import           Test.Tasty (TestTree)
+import           Test.Tasty.HUnit (assertEqual, testCase)
+
+import           Fencer.Counter (CounterStatus, counterRemainingLimit)
+import           Fencer.Logic (AppState, setRules, updateLimitCounter)
+import           Fencer.Rules (definitionsToRuleTree)
+import           Fencer.Server.Test (withServer, serverAppState)
+import           Fencer.Types
+
+
+-- | Test that a rule limit unit change adds a new counter and leaves
+-- the old one intact.
+test_logicLimitUnitChange :: TestTree
+test_logicLimitUnitChange =
+  withServer $ \serverIO ->
+    testCase "A rule limit unit change on rule reloading" $ do
+      state <- serverAppState <$> serverIO
+      void $ atomically $ setRules state (mapRuleDefs definitions1)
+
+      -- Record a hit and get the remaining limit
+      st1 <- getRemainingLimit <$> makeAHit state
+      assertEqual
+        "The remaining rate limit was not updated!"
+        (limit - hits)
+        st1
+
+      -- Set the new rules and the rules reloaded flag
+      atomically $ setRules state (mapRuleDefs definitions2)
+      -- Record a hit and get a remaining limit
+      st2 <- getRemainingLimit <$> makeAHit state
+      assertEqual
+        "The remaining rate was affected by a different counter!"
+        (limit - hits)
+        st2
+
+      -- Set the old rules again
+      void $ atomically $ setRules state (mapRuleDefs definitions1)
+      -- Record a hit and get a remaining limit
+      st1' <- getRemainingLimit <$> makeAHit state
+      assertEqual
+        "The old counter did not persist!"
+        (st1 - hits)
+        st1'
+ where
+  getRemainingLimit :: Maybe (RateLimit, CounterStatus) -> Word
+  getRemainingLimit = counterRemainingLimit . snd . fromMaybe (error "")
+
+  makeAHit :: AppState -> IO (Maybe (RateLimit, CounterStatus))
+  makeAHit st = atomically $
+    updateLimitCounter st (#hits hits) domainId ruleList
+
+  mapRuleDefs :: [DomainDefinition] -> [(DomainId, RuleTree)]
+  mapRuleDefs defs =
+    [ ( domainDefinitionId rule
+      , definitionsToRuleTree (NE.toList . domainDefinitionDescriptors $ rule))
+    | rule <- defs
+    ]
+
+  -- rate limit and hits in the test
+  limit = 4 :: Word
+  hits  = 1 :: Word
+
+  ruleKey   = RuleKey   "generic_key"
+  ruleValue = RuleValue "dream11_order_create"
+  ruleList  = [(ruleKey, ruleValue)]
+  domainId  = DomainId  "merchant_rate_limits"
+
+  descriptor :: DescriptorDefinition
+  descriptor = DescriptorDefinition
+        { descriptorDefinitionKey         = ruleKey
+        , descriptorDefinitionValue       = Just ruleValue
+        , descriptorDefinitionRateLimit   = Just $ RateLimit Minute limit
+        , descriptorDefinitionDescriptors = Nothing
+        }
+
+  definition1 :: DomainDefinition
+  definition1 = DomainDefinition
+    { domainDefinitionId          = domainId
+    , domainDefinitionDescriptors = descriptor :| []
+    }
+
+  definitions1 :: [DomainDefinition]
+  definitions1 = [definition1]
+
+  definition2 = definition1 {
+    domainDefinitionDescriptors =
+      (descriptor
+        { descriptorDefinitionRateLimit = Just $ RateLimit Hour limit }
+      ) :| []
+    }
+
+  definitions2 :: [DomainDefinition]
+  definitions2 = [definition2]