Replicate counter identity rules

README.md

@@ -186,12 +186,17 @@ you shouldn’t have made this request because you are over the limit".
 
 ### Rate limiting models
 
-There are two models for rate limiting – the "bucket" one and the "rolling
-window" one. In the bucket model you have a single counter which is as
-granular as the rate-limiting unit – say, if the limit is "N requests per
-hour" then the counter will get incremented for all requests for
-X:00:00–X:59:59. If you have spent your limit at the beginning of the hour,
-you don't get any credit until the next hour starts and the counter resets.
+There are two models for rate limiting – the "bucket" one and the
+"rolling window" one. In the bucket model you have a single counter
+which is as granular as the rate-limiting unit – say, if the limit is
+"N requests per hour" then the counter will get incremented for all
+requests for X:00:00–X:59:59. If you have spent your limit at the
+beginning of the hour, you don't get any credit until the next hour
+starts and the counter resets. Every counter is indexed by a time unit
+so that if there is an update in the configuration and rules are
+reloaded as a result, a change in time limit yields a new reset
+counter; if there is an update in the request rate only, an existing
+counter is updated.
 
 The "rolling window" model is more interesting. In this model, you are
 limited based on how many requests you’ve made in the past hour, which

fencer.cabal

@@ -106,7 +106,7 @@ test-suite test-fencer
   type:
     exitcode-stdio-1.0
   main-is:
-    Tests.hs
+    Main.hs
   hs-source-dirs:
     test
   other-modules:
@@ -126,8 +126,8 @@ test-suite test-fencer
     , neat-interpolation
     , proto3-wire
     , proto3-suite
+    , stm-containers
     , tasty
-    , tasty-discover
     , tasty-hunit
     , temporary
     , text

lib/Fencer/AppState.hs

@@ -5,6 +5,8 @@
 -- | In-memory state of Fencer.
 module Fencer.AppState
     ( AppState
+    , appStateCounters
+    , appStateRules
     , initAppState
 
     -- * Methods for working with 'AppState'
@@ -27,10 +29,10 @@ import Named ((:!), arg)
 import qualified Focus as Focus
 import Control.Monad.Trans.Class (lift)
 
-import Fencer.Types
 import Fencer.Counter
-import Fencer.Time
 import Fencer.Rules
+import Fencer.Time
+import Fencer.Types
 
 -- | Fencer runtime context and in-memory state.
 --
@@ -154,6 +156,16 @@ getLimit appState domain descriptor =
         Just ruleTree -> pure (applyRules descriptor ruleTree)
 
 -- | Set 'appStateRules' and 'appStateRulesLoaded'.
+--
+-- The 'appStateCounters' field stays unchanged. This is in accordance
+-- with the behavior of @lyft/ratelimit@.
+--
+-- There might be a change in rules with the same descriptors that
+-- updates the value of 'requests_per_unit' (with the time unit left
+-- intact), which allows a different number of requests to be
+-- made. This is as expected. However, if there is a change in the
+-- rate limit time unit, a new counter will be created, regardless of
+-- how many requests the previous counter had used up.
 setRules :: AppState -> [(DomainId, RuleTree)] -> STM ()
 setRules appState rules = do
     writeTVar (appStateRulesLoaded appState) True

lib/Fencer/Counter.hs

@@ -24,6 +24,7 @@ import Fencer.Time
 data CounterKey = CounterKey
     { counterKeyDomain :: !DomainId
     , counterKeyDescriptor :: ![(RuleKey, RuleValue)]
+    , counterKeyUnit :: !TimeUnit
     }
     deriving stock (Eq, Generic)
     deriving anyclass (Hashable)
@@ -35,7 +36,7 @@ data Counter = Counter
       -- | Counter expiry date, inclusive (i.e. on 'counterExpiry' the
       -- counter is already expired).
     , counterExpiry :: !Timestamp
-    }
+    } deriving Eq
 
 data CounterStatus = CounterStatus
     { -- | How many hits can be taken before the limit is reached. Will be 0

lib/Fencer/Main.hs

@@ -89,7 +89,11 @@ reloadRules logger settings appState = do
                     show (map (unDomainId . domainDefinitionId) ruleDefinitions))
 
     -- Recreate 'appStateRules'
+    --
+    -- There is no need to remove old rate limiting rules
     atomically $
+        -- See the documentation of 'setRules' for details on what
+        -- happens with counters during rule reloading.
         setRules appState
             [ ( domainDefinitionId rule
               , definitionsToRuleTree (NE.toList . domainDefinitionDescriptors $ rule))

lib/Fencer/Server.hs

@@ -136,13 +136,13 @@ 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))
-  where
-    counterKey :: CounterKey
-    counterKey = CounterKey
-        { counterKeyDomain = domain
-        , counterKeyDescriptor = descriptor }
 
 ----------------------------------------------------------------------------
 -- Working with protobuf structures

test/Fencer/Rules/Test.hs

@@ -4,30 +4,40 @@
 
 -- | Tests for "Fencer.Rules".
 module Fencer.Rules.Test
-  ( test_loadRulesYaml
-  , test_loadRulesNonYaml
-  , test_loadRulesRecursively
+  ( test_rulesLoadRulesYaml
+  , test_rulesLoadRulesNonYaml
+  , test_rulesLoadRulesRecursively
+  , test_rulesLimitUnitChange
   )
 where
 
 import           BasePrelude
 
+import qualified Data.HashMap.Strict as HM
+import           Data.List (sortOn)
+import qualified Data.List.NonEmpty as NE
+import           Data.Maybe (fromMaybe)
 import           Data.Text (Text)
 import qualified Data.Text.IO as TIO
-import           Test.Tasty (TestTree)
-import           Test.Tasty.HUnit (assertEqual, testCase)
-import qualified System.IO.Temp as Temp
 import           NeatInterpolation (text)
+import qualified StmContainers.Map as StmMap
+import qualified System.IO.Temp as Temp
 import           System.FilePath ((</>))
 import           System.Directory (createDirectoryIfMissing)
-import           Data.List (sortOn)
+import           Test.Tasty (TestTree, withResource)
+import           Test.Tasty.HUnit (assertBool, assertEqual, testCase)
 
-import           Fencer.Types
+import           Fencer.AppState (appStateCounters, appStateRules, recordHits, setRules)
+import           Fencer.Counter (CounterKey(..), counterHits)
 import           Fencer.Rules
+import           Fencer.Types
+
+import           Fencer.Server.Test (createServerAppState, destroyServerAppState)
+
 
 -- | Test that 'loadRulesFromDirectory' loads rules from YAML files.
-test_loadRulesYaml :: TestTree
-test_loadRulesYaml =
+test_rulesLoadRulesYaml :: TestTree
+test_rulesLoadRulesYaml =
   testCase "Rules are loaded from YAML files" $ do
     Temp.withSystemTempDirectory "fencer-config" $ \tempDir -> do
       TIO.writeFile (tempDir </> "config1.yml") domain1Text
@@ -42,8 +52,8 @@ test_loadRulesYaml =
 -- YAML files.
 --
 -- This counterintuitive behavior matches the behavior of @lyft/ratelimit@.
-test_loadRulesNonYaml :: TestTree
-test_loadRulesNonYaml =
+test_rulesLoadRulesNonYaml :: TestTree
+test_rulesLoadRulesNonYaml =
   testCase "Rules are loaded from non-YAML files" $ do
     Temp.withSystemTempDirectory "fencer-config" $ \tempDir -> do
       TIO.writeFile (tempDir </> "config1.bin") domain1Text
@@ -57,8 +67,8 @@ test_loadRulesNonYaml =
 -- | Test that 'loadRulesFromDirectory' loads rules recursively.
 --
 -- This matches the behavior of @lyft/ratelimit@.
-test_loadRulesRecursively :: TestTree
-test_loadRulesRecursively =
+test_rulesLoadRulesRecursively :: TestTree
+test_rulesLoadRulesRecursively =
   testCase "Rules are loaded recursively" $ do
     Temp.withSystemTempDirectory "fencer-config" $ \tempDir -> do
       createDirectoryIfMissing True (tempDir </> "domain1")
@@ -71,6 +81,78 @@ test_loadRulesRecursively =
         (sortOn domainDefinitionId [domain1, domain2])
         (sortOn domainDefinitionId definitions)
 
+-- | Test that a rule limit unit change adds a new counter and leaves
+-- the old one intact.
+test_rulesLimitUnitChange :: TestTree
+test_rulesLimitUnitChange =
+  withResource createServerAppState destroyServerAppState $ \ioLogIdState ->
+    testCase "A rule limit unit change on rule reloading" $ do
+      Temp.withSystemTempDirectory "fencer-config-unit" $ \tempDir -> do
+        createDirectoryIfMissing True (tempDir </> dir)
+
+        definitions1 <- writeLoad tempDir merchantLimitsText1
+        (_, _, state) <- ioLogIdState
+
+        atomically $ setRules state (mapRuleDefs definitions1)
+
+        ruleTree :: RuleTree <- atomically $
+          fromMaybe' <$> StmMap.lookup domainId (appStateRules state)
+        let ruleBranch = fromMaybe' $ HM.lookup (ruleKey, Just ruleValue) ruleTree
+        let rateLimit =  fromMaybe' $ ruleBranchRateLimit ruleBranch
+
+        -- Record a hit
+        void $ atomically $ recordHits state (#hits 1) (#limit rateLimit) counterKey1
+
+        mV1 <- atomically $ StmMap.lookup counterKey1 $ appStateCounters state
+
+        -- Change rules in the configuration
+        definitions2 <- writeLoad tempDir merchantLimitsText2
+
+        -- Set the new rules and the rules reloaded flag
+        atomically $ setRules state (mapRuleDefs definitions2)
+
+        mV1' <- atomically $ StmMap.lookup counterKey1 $ appStateCounters state
+        mV2  <- atomically $ StmMap.lookup counterKey2 $ appStateCounters state
+
+        assertBool
+          "The original counter was not updated after recording a hit!"
+          ((counterHits <$> mV1) == Just 1)
+        assertBool
+          "The original counter was mistakenly updated in the meantime!"
+          (mV1 == mV1')
+        assertBool "The secondary counter was set!" (mV2 == Nothing)
+ where
+  mapRuleDefs :: [DomainDefinition] -> [(DomainId, RuleTree)]
+  mapRuleDefs defs =
+    [ ( domainDefinitionId rule
+      , definitionsToRuleTree (NE.toList . domainDefinitionDescriptors $ rule))
+    | rule <- defs
+    ]
+
+  dir     = "d11-ratelimits"
+  cfgFile = "d11-ratelimits1.yaml"
+
+  writeLoad :: FilePath -> Text -> IO [DomainDefinition]
+  writeLoad tempDir txt = do
+    TIO.writeFile (tempDir </> dir </> cfgFile) txt
+    loadRulesFromDirectory (#directory tempDir) (#ignoreDotFiles True)
+
+  ruleKey   = RuleKey   "generic_key"
+  ruleValue = RuleValue "dream11_order_create"
+  domainId  = DomainId  "merchant_rate_limits"
+
+  counterKey1 = CounterKey
+    { counterKeyDomain     = domainId
+    , counterKeyDescriptor = [ (ruleKey, ruleValue) ]
+    , counterKeyUnit       = Minute }
+
+  counterKey2 :: CounterKey
+  counterKey2 = counterKey1 { counterKeyUnit = Hour }
+
+  fromMaybe' :: Maybe a -> a
+  fromMaybe' = fromMaybe (error "")
+
+
 ----------------------------------------------------------------------------
 -- Sample definitions
 ----------------------------------------------------------------------------
@@ -117,3 +199,25 @@ domain2Text = [text|
   descriptors:
     - key: some key 2
   |]
+
+merchantLimitsText1 :: Text
+merchantLimitsText1 = [text|
+  domain: merchant_rate_limits
+  descriptors:
+  - key: generic_key
+    value: dream11_order_create
+    rate_limit:
+      unit: minute
+      requests_per_unit: 400000
+  |]
+
+merchantLimitsText2 :: Text
+merchantLimitsText2 = [text|
+  domain: merchant_rate_limits
+  descriptors:
+  - key: generic_key
+    value: dream11_order_create
+    rate_limit:
+      unit: hour
+      requests_per_unit: 400000
+  |]

test/Fencer/Server/Test.hs

@@ -5,7 +5,9 @@
 
 -- | Tests for "Fencer.Server".
 module Fencer.Server.Test
-  ( test_responseNoRules
+  ( test_serverResponseNoRules
+  , createServerAppState
+  , destroyServerAppState
   )
 where
 
@@ -28,8 +30,8 @@ import qualified Fencer.Proto as Proto
 -- 'reloadRules' has never been ran), requests to Fencer will error out.
 --
 -- This behavior matches @lyft/ratelimit@.
-test_responseNoRules :: TestTree
-test_responseNoRules =
+test_serverResponseNoRules :: TestTree
+test_serverResponseNoRules =
   withResource createServer destroyServer $ \_ ->
     testCase "When no rules have been loaded, all requests error out" $ do
       Grpc.withGRPCClient clientConfig $ \grpcClient -> do
@@ -72,20 +74,30 @@ test_responseNoRules =
 -- | Start Fencer on port 50051.
 createServer :: IO (Logger.Logger, ThreadId)
 createServer = do
+  (logger, threadId, _) <- createServerAppState
+  pure (logger, threadId)
+
+-- | Start Fencer on port 50051.
+createServerAppState :: IO (Logger.Logger, ThreadId, AppState)
+createServerAppState = do
   -- TODO: not the best approach. Ideally we should use e.g.
   -- https://hackage.haskell.org/package/tasty-hunit/docs/Test-Tasty-HUnit.html#v:testCaseSteps
   -- but we can't convince @tinylog@ to use the provided step function.
   logger <- Logger.create (Logger.Path "/dev/null")
   appState <- initAppState
   threadId <- forkIO $ runServer logger appState
-  pure (logger, threadId)
+  pure (logger, threadId, appState)
 
 -- | Kill Fencer.
 destroyServer :: (Logger.Logger, ThreadId) -> IO ()
 destroyServer (logger, threadId) = do
   Logger.close logger
   killThread threadId
 
+-- | Kill Fencer.
+destroyServerAppState :: (Logger.Logger, ThreadId, AppState) -> IO ()
+destroyServerAppState (logger, threadId, _) = destroyServer (logger, threadId)
+
 ----------------------------------------------------------------------------
 -- gRPC client
 ----------------------------------------------------------------------------