Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve sampler #2761

Merged
merged 3 commits into from
Jul 17, 2024
Merged

Improve sampler #2761

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2b7101a
Improve sampler
estringana Jul 16, 2024
8a08961
Remove scope
estringana Jul 16, 2024
3a9942b
Address pr comments
estringana Jul 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 4 additions & 10 deletions appsec/src/helper/client.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -416,16 +416,10 @@ bool client::handle_command(network::request_shutdown::request &command)
auto free_ctx = defer([this]() { this->context_.reset(); });

auto sampler = service_->get_schema_sampler();
std::optional<sampler::scope> scope;
if (sampler) {
scope = sampler->get();
if (scope.has_value()) {
parameter context_processor = parameter::map();
context_processor.add(
"extract-schema", parameter::as_boolean(true));
command.data.add(
"waf.context.processor", std::move(context_processor));
}
if (sampler && sampler->picked()) {
parameter context_processor = parameter::map();
context_processor.add("extract-schema", parameter::as_boolean(true));
command.data.add("waf.context.processor", std::move(context_processor));
}

auto response = publish<network::request_shutdown>(command);
Expand Down
58 changes: 7 additions & 51 deletions appsec/src/helper/sampler.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,7 @@
#include <atomic>
#include <cmath>
#include <cstdint>
#include <iostream>
#include <mutex>
#include <optional>

namespace dds {
static const double min_rate = 0.0001;
Expand All @@ -29,63 +27,21 @@ class sampler {
}
// NOLINTEND(cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers)
}
class scope {
public:
explicit scope(std::atomic<bool> &concurrent) : concurrent_(&concurrent)
{
concurrent_->store(true, std::memory_order_relaxed);
}

scope(const scope &) = delete;
scope &operator=(const scope &) = delete;
scope(scope &&oth) noexcept
{
concurrent_ = oth.concurrent_;
oth.concurrent_ = nullptr;
}
scope &operator=(scope &&oth)
{
concurrent_ = oth.concurrent_;
oth.concurrent_ = nullptr;

return *this;
}

~scope()
{
if (concurrent_ != nullptr) {
concurrent_->store(false, std::memory_order_relaxed);
}
}

protected:
std::atomic<bool> *concurrent_;
};

std::optional<scope> get()
bool picked()
{
const std::lock_guard<std::mutex> lock_guard(mtx_);

std::optional<scope> result = std::nullopt;

if (!concurrent_ && floor(request_ * sample_rate_) !=
floor((request_ + 1) * sample_rate_)) {
result = {scope{concurrent_}};
}

if (request_ < std::numeric_limits<unsigned>::max()) {
request_++;
} else {
request_ = 1;
if (sample_rate_ == 1) {
return true;
}

return result;
auto old_request = request_.fetch_add(1, std::memory_order_relaxed);
return floor(old_request * sample_rate_) !=
floor((request_)*sample_rate_);
}

protected:
unsigned request_{1};
std::atomic<unsigned> request_{0};
double sample_rate_;
std::atomic<bool> concurrent_{false};
std::mutex mtx_;
};
} // namespace dds
31 changes: 4 additions & 27 deletions appsec/tests/helper/sampler_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ class sampler : public dds::sampler {
public:
sampler(double sample_rate) : dds::sampler(sample_rate) {}
void set_request(unsigned int i) { request_ = i; }
auto get_request() { return request_; }
unsigned int get_request() { return request_; }
};

} // namespace mock
Expand All @@ -25,8 +25,7 @@ std::atomic<int> picked = 0;
void count_picked(dds::sampler &sampler, int iterations)
{
for (int i = 0; i < iterations; i++) {
auto is_pick = sampler.get();
if (is_pick != std::nullopt) {
if (sampler.picked()) {
picked++;
}
}
Expand Down Expand Up @@ -198,29 +197,7 @@ TEST(SamplerTest, TestOverflow)
{
mock::sampler s(0);
s.set_request(UINT_MAX);
s.get();
EXPECT_EQ(1, s.get_request());
}

TEST(ScopeTest, TestConcurrent)
{
std::atomic<bool> concurrent = false;
{
auto s = sampler::scope(std::ref(concurrent));
EXPECT_TRUE(concurrent);
}
EXPECT_FALSE(concurrent);
}

TEST(ScopeTest, TestItDoesNotPickTokenUntilScopeReleased)
{
sampler sampler(1);
auto is_pick = sampler.get();
EXPECT_TRUE(is_pick != std::nullopt);
is_pick = sampler.get();
EXPECT_FALSE(is_pick != std::nullopt);
is_pick.reset();
is_pick = sampler.get();
EXPECT_TRUE(is_pick != std::nullopt);
s.picked();
EXPECT_EQ(0, s.get_request());
}
} // namespace dds
12 changes: 6 additions & 6 deletions appsec/tests/helper/service_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,15 +57,15 @@ TEST(ServiceTest, ServicePickSchemaExtractionSamples)
auto s = service(
engine, service_config, nullptr, {true, all_requests_are_picked});

EXPECT_TRUE(s.get_schema_sampler()->get().has_value());
EXPECT_TRUE(s.get_schema_sampler()->picked());
}

{ // Constructor. It does not pick based on rate
double no_request_is_picked = 0.0;
auto s = service(
engine, service_config, nullptr, {true, no_request_is_picked});

EXPECT_FALSE(s.get_schema_sampler()->get().has_value());
EXPECT_FALSE(s.get_schema_sampler()->picked());
}

{ // Constructor. It does not pick if disabled
Expand All @@ -74,7 +74,7 @@ TEST(ServiceTest, ServicePickSchemaExtractionSamples)
auto s = service(engine, service_config, nullptr,
{schema_extraction_disabled, all_requests_are_picked});

EXPECT_FALSE(s.get_schema_sampler()->get().has_value());
EXPECT_FALSE(s.get_schema_sampler()->picked());
}

{ // Static constructor. It picks based on rate
Expand All @@ -83,7 +83,7 @@ TEST(ServiceTest, ServicePickSchemaExtractionSamples)
auto service = service::from_settings(
service_identifier(sid), engine_settings, {}, meta, metrics, false);

EXPECT_TRUE(service->get_schema_sampler()->get().has_value());
EXPECT_TRUE(service->get_schema_sampler()->picked());
}

{ // Static constructor. It does not pick based on rate
Expand All @@ -92,7 +92,7 @@ TEST(ServiceTest, ServicePickSchemaExtractionSamples)
auto service = service::from_settings(
service_identifier(sid), engine_settings, {}, meta, metrics, false);

EXPECT_FALSE(service->get_schema_sampler()->get().has_value());
EXPECT_FALSE(service->get_schema_sampler()->picked());
}

{ // Static constructor. It does not pick if disabled
Expand All @@ -101,7 +101,7 @@ TEST(ServiceTest, ServicePickSchemaExtractionSamples)
auto service = service::from_settings(
service_identifier(sid), engine_settings, {}, meta, metrics, false);

EXPECT_FALSE(service->get_schema_sampler()->get().has_value());
EXPECT_FALSE(service->get_schema_sampler()->picked());
}
}

Expand Down
Loading