folly/experimental/logging/RateLimiter.h

   1 /*
   2  * Copyright 2004-present Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 #pragma once
  17
  18 #include <atomic>
  19 #include <chrono>
  20 #include <cstdint>
  21
  22 namespace folly {
  23 namespace logging {
  24
  25 /**
  26  * An interface for rate limiting checkers.
  27  */
  28 class RateLimiter {
  29  public:
  30   virtual ~RateLimiter() {}
  31   virtual bool check() = 0;
  32 };
  33
  34 /**
  35  * A rate limiter that can rate limit events to N events per M milliseconds.
  36  *
  37  * It is intended to be fast to check when messages are not being rate limited.
  38  * When messages are being rate limited it is slightly slower, as it has to
  39  * check the clock each time check() is called in this case.
  40  */
  41 class IntervalRateLimiter : public RateLimiter {
  42  public:
  43   IntervalRateLimiter(
  44       uint64_t maxPerInterval,
  45       std::chrono::steady_clock::duration interval);
  46
  47   bool check() override final {
  48     auto origCount = count_.fetch_add(1, std::memory_order_acq_rel);
  49     if (origCount < maxPerInterval_) {
  50       return true;
  51     }
  52     return checkSlow();
  53   }
  54
  55  private:
  56   bool checkSlow();
  57
  58   const uint64_t maxPerInterval_;
  59   const std::chrono::steady_clock::time_point::duration interval_;
  60
  61   std::atomic<uint64_t> count_{0};
  62   // Ideally timestamp_ would be a
  63   // std::atomic<std::chrono::steady_clock::time_point>, but this does not work
  64   // since time_point's constructor is not noexcept
  65   std::atomic<std::chrono::steady_clock::rep> timestamp_;
  66 };
  67 }
  68 }