diff options
Diffstat (limited to 'src/mongo/db')
-rw-r--r-- | src/mongo/db/query/SConscript | 13 | ||||
-rw-r--r-- | src/mongo/db/query/rate_limiting.cpp | 92 | ||||
-rw-r--r-- | src/mongo/db/query/rate_limiting.h | 109 | ||||
-rw-r--r-- | src/mongo/db/query/rate_limiting_test.cpp | 77 |
4 files changed, 291 insertions, 0 deletions
diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript index f3cce404031..1c320503176 100644 --- a/src/mongo/db/query/SConscript +++ b/src/mongo/db/query/SConscript @@ -357,6 +357,17 @@ env.Library( ) env.Library( + target='rate_limiting', + source=[ + 'rate_limiting.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/util/clock_sources', + ], +) + +env.Library( target='telemetry', source=[ 'telemetry.cpp', @@ -423,6 +434,7 @@ env.CppUnitTest( "query_request_test.cpp", "query_settings_test.cpp", "query_solution_test.cpp", + "rate_limiting_test.cpp", "sbe_and_hash_test.cpp", "sbe_and_sorted_test.cpp", "sbe_stage_builder_accumulator_test.cpp", @@ -459,6 +471,7 @@ env.CppUnitTest( "query_planner_test_fixture", "query_request", "query_test_service_context", + "rate_limiting", ], ) diff --git a/src/mongo/db/query/rate_limiting.cpp b/src/mongo/db/query/rate_limiting.cpp new file mode 100644 index 00000000000..44c94ead313 --- /dev/null +++ b/src/mongo/db/query/rate_limiting.cpp @@ -0,0 +1,92 @@ +/** + * Copyright (C) 2022-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "rate_limiting.h" + +namespace mongo { +RateLimiting::RateLimiting(RequestCount samplingRate, Milliseconds timePeriod) + : _clockSource(SystemClockSource::get()), + _samplingRate(samplingRate), + _timePeriod(timePeriod), + _windowStart(_clockSource->now()), + _prevCount(0), + _currentCount(0) {} + +Date_t RateLimiting::tickWindow() { + Date_t currentTime = _clockSource->now(); + + // Elapsed time since window start exceeds the time period. Start a new window. + if (currentTime - _windowStart > _timePeriod) { + _windowStart = currentTime; + _prevCount = _currentCount; + _currentCount = 0; + } + return currentTime; +} + +bool RateLimiting::handleRequestFixedWindow() { + stdx::unique_lock windowLock{_windowMutex}; + tickWindow(); + + if (_currentCount < _samplingRate) { + _currentCount += 1; + return true; + } + return false; +} + +bool RateLimiting::handleRequestSlidingWindow() { + stdx::unique_lock windowLock{_windowMutex}; + + Date_t currentTime = tickWindow(); + auto windowStart = _windowStart; + auto prevCount = _prevCount; + + // Sliding window is implemented over fixed size time periods/blocks as follows. Instead of + // making the decision to limit the rate using only the current time period, we look to the rate + // of the previous period to predicate the rate of the current. This smooths the "sampling" of + // the events by predicting a constant rate and limiting accordingly. + + // Percentage of time remaining in current window. + double percentRemainingOfCurrentWindow = + ((double)(_timePeriod.count() - (currentTime - windowStart).count())) / _timePeriod.count(); + // Estimate the number of requests remaining in the current period. We assume the requests in + // the previous time block occurred at a constant rate. We multiply the total number of requests + // in the previous period by the percentage of time remaining in the current period. + double estimatedRemaining = prevCount * percentRemainingOfCurrentWindow; + // Add this estimate to the requests we know have taken place within the current time block. + double estimatedCount = _currentCount + estimatedRemaining; + + if (estimatedCount < _samplingRate) { + _currentCount += 1; + return true; + } + return false; +} +} // namespace mongo diff --git a/src/mongo/db/query/rate_limiting.h b/src/mongo/db/query/rate_limiting.h new file mode 100644 index 00000000000..959ac52692e --- /dev/null +++ b/src/mongo/db/query/rate_limiting.h @@ -0,0 +1,109 @@ +/** + * Copyright (C) 2022-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/util/clock_source.h" +#include "mongo/util/system_clock_source.h" + +namespace mongo { + +/** + * Rate limiting is used to put a bound on the number of requests to a certain resource over a fixed + * time window. This implementation is approximate in the sense that it may permit the bound to + * exceeded. The bound is approximate as a trade off to reduce contention on internal resources. + */ +class RateLimiting { + using RequestCount = uint32_t; + +public: + /* + * Constructor for a rate limiter. Specify the number of requests you want to take place, as + * well as the time period in milliseconds. + */ + RateLimiting(RequestCount samplingRate, Milliseconds timePeriod = Seconds{1}); + + /* + * A simple method for rate limiting. Returns false if we have reached the request limit for the + * current time window; otherwise, returns true and adds the request to the count for the + * current window. If we have passed the end of the previous window, the slate is wiped clean. + */ + bool handleRequestFixedWindow(); + + /* + * A method that ensures a more steady rate of requests. Rather than only looking at the current + * time block, this method simulates a sliding window to estimate how many requests occurred in + * the last full time period. Like the above, returns whether the request should be handled, and + * resets the window if enough time has passed. + */ + bool handleRequestSlidingWindow(); + +private: + /* + * Resets the current window if it has ended. Returns the current time. This must be called in + * the beginning of each handleRequest...() method. + */ + Date_t tickWindow(); + + /* + * Clock source used to track time. + */ + ClockSource* const _clockSource; + + /* + * Sampling rate is the bound on the number of requests we want to admit per window. + */ + const RequestCount _samplingRate; + + /* + * Time period is the window size in ms. + */ + const Milliseconds _timePeriod; + + /* + * Window start. + */ + Date_t _windowStart; + + /* + * Count of requests handled in the previous window. + */ + RequestCount _prevCount; + + /* + * Count of requests handled in the current window. + */ + RequestCount _currentCount; + + /* + * Mutex used when reading/writing the window. + */ + stdx::recursive_mutex _windowMutex; +}; +} // namespace mongo diff --git a/src/mongo/db/query/rate_limiting_test.cpp b/src/mongo/db/query/rate_limiting_test.cpp new file mode 100644 index 00000000000..2d9ef35647a --- /dev/null +++ b/src/mongo/db/query/rate_limiting_test.cpp @@ -0,0 +1,77 @@ +/** + * Copyright (C) 2022-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/db/query/rate_limiting.h" +#include "mongo/unittest/unittest.h" +#include "mongo/util/time_support.h" + +namespace mongo { +TEST(RateLimitingTest, FixedWindowSucceeds) { + auto rl = RateLimiting(1); + ASSERT_TRUE(rl.handleRequestFixedWindow()); +} + +TEST(RateLimitingTest, SlidingWindowSucceeds) { + auto rl = RateLimiting(1); + ASSERT_TRUE(rl.handleRequestSlidingWindow()); +} + +TEST(RateLimitingTest, FixedWindowFails) { + auto rl = RateLimiting(0); + ASSERT_FALSE(rl.handleRequestFixedWindow()); +} + +TEST(RateLimitingTest, SlidingWindowFails) { + auto rl = RateLimiting(0); + ASSERT_FALSE(rl.handleRequestSlidingWindow()); +} + +TEST(RateLimitingTest, FixedWindowSucceedsThenFails) { + auto rl = RateLimiting(1, Hours{1}); + ASSERT_TRUE(rl.handleRequestFixedWindow()); + ASSERT_FALSE(rl.handleRequestFixedWindow()); + ASSERT_FALSE(rl.handleRequestFixedWindow()); +} + +TEST(RateLimitingTest, SlidingWindowSucceedsThenFails) { + auto rl = RateLimiting(1, Hours{1}); + ASSERT_TRUE(rl.handleRequestSlidingWindow()); + ASSERT_FALSE(rl.handleRequestSlidingWindow()); + ASSERT_FALSE(rl.handleRequestSlidingWindow()); +} + +TEST(RateLimitingTest, FixedWindowPermitsRequestAfterWindowExpires) { + auto rl = RateLimiting(1, Milliseconds{10}); + ASSERT_TRUE(rl.handleRequestFixedWindow()); + ASSERT_FALSE(rl.handleRequestFixedWindow()); + sleepmillis(11); + ASSERT_TRUE(rl.handleRequestFixedWindow()); +} + +} // namespace mongo |