#include "common.h"

#include <balancer/kernel/balancing/pessimization.h>
#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/log/log.h>
#include <library/cpp/pop_count/popcount.h>
#include <library/cpp/testing/unittest/registar.h>
#include <util/random/random.h>


using namespace NSrvKernel::NDynamicBalancing;
using NSrvKernel::TError;
using NSrvKernel::TNotReadyBackendSet;

// If you want to see debug output, use following trick:
//   TLog log{"/dev/stderr"};
//   TDynamicLog nullLog{"test"};
//   nullLog.SetLog(&log);


Y_UNIT_TEST_SUITE(PessimizationTracker) {
    Y_UNIT_TEST(TestUnhealthyBackendsArePessimized) {
        TConfig config;
        config.set_max_pessimized_share(0.5);

        TDynamicLog nullLog{"test"};
        TPessimizationTracker tracker{config, nullLog};
        auto backends = FillBackends(tracker, 2);
        backends[0]->SetEnabled(false);
        UNIT_ASSERT(!tracker.IsOverPessimized());
        auto state = tracker.Step();
        UNIT_ASSERT(!tracker.IsOverPessimized());
        UNIT_ASSERT(state.size() == 1);
        UNIT_ASSERT(state[0].Get() == backends[1]);
    }

    Y_UNIT_TEST(TestFailRateCheck) {
        TConfig config;
        config.set_max_pessimized_share(1);

        {
            TFailRateCheckConfig failRateCheck;
            failRateCheck.set_min_requests(10);
            failRateCheck.set_threshold(0.5);
            config.set_fail_rate_check(std::move(failRateCheck));
            TDynamicLog nullLog{"test"};
            TPessimizationTracker tracker{config, nullLog};

            auto backends = FillBackends(tracker, 1);
            for (size_t i = 0; i < 9; i++) {
                TError error = Y_MAKE_ERROR(yexception{} << "test");
                backends[0]->OnFailRequest(error, TDuration::Seconds(1));
            }
            bool hasSkippedCheck = false;
            bool hasFiredCheck = false;
            UNIT_ASSERT_VALUES_EQUAL(tracker.RunBackendChecks(*backends[0], TStateAggregate{}, hasFiredCheck, hasSkippedCheck), TDynamicBackend::TStatus::Healthy);
            UNIT_ASSERT_VALUES_EQUAL(hasSkippedCheck, true);
            UNIT_ASSERT_VALUES_EQUAL(hasFiredCheck, false);

            backends[0]->OnFailRequest(Y_MAKE_ERROR(yexception{} << "test"), TDuration::Seconds(1));
            UNIT_ASSERT_VALUES_EQUAL(tracker.RunBackendChecks(*backends[0], TStateAggregate{}, hasFiredCheck, hasSkippedCheck), TDynamicBackend::TStatus::Critical);
            UNIT_ASSERT_VALUES_EQUAL(hasSkippedCheck, false);
            UNIT_ASSERT_VALUES_EQUAL(hasFiredCheck, true);
        }
    }

    Y_UNIT_TEST(TestConsecutiveFailsCheck) {
        TConfig config;
        config.set_max_pessimized_share(1);

        {
            TConsecutiveFailsCheckConfig consecutiveFailsCheckConfig;
            consecutiveFailsCheckConfig.set_threshold(5);
            config.set_consecutive_fails_check(std::move(consecutiveFailsCheckConfig));
            TDynamicLog nullLog{"test"};
            TPessimizationTracker tracker{config, nullLog};

            auto backends = FillBackends(tracker, 1);
            for (size_t i = 0; i < 10; i++) {
                TError error = Y_MAKE_ERROR(yexception{} << "test");
                backends[0]->OnFailRequest(error, TDuration::Seconds(1));
                backends[0]->OnCompleteRequest(TDuration::Seconds(1));
            }
            bool hasSkippedCheck = false;
            bool hasFiredCheck = false;
            UNIT_ASSERT_VALUES_EQUAL(tracker.RunBackendChecks(*backends[0], TStateAggregate{}, hasFiredCheck, hasSkippedCheck), TDynamicBackend::TStatus::Healthy);
            UNIT_ASSERT_VALUES_EQUAL(hasSkippedCheck, false);
            UNIT_ASSERT_VALUES_EQUAL(hasFiredCheck, false);

            for (size_t i = 0; i < 10; i++) {
                TError error = Y_MAKE_ERROR(yexception{} << "test");
                backends[0]->OnFailRequest(error, TDuration::Seconds(1));
                backends[0]->OnFailRequest(error, TDuration::Seconds(1));
                backends[0]->OnFailRequest(error, TDuration::Seconds(1));
                backends[0]->OnFailRequest(error, TDuration::Seconds(1));
                backends[0]->OnFailRequest(error, TDuration::Seconds(1));
                backends[0]->OnCompleteRequest(TDuration::Seconds(1));
            }
            UNIT_ASSERT_VALUES_EQUAL(tracker.RunBackendChecks(*backends[0], TStateAggregate{}, hasFiredCheck, hasSkippedCheck), TDynamicBackend::TStatus::Critical);
            UNIT_ASSERT_VALUES_EQUAL(hasSkippedCheck, false);
            UNIT_ASSERT_VALUES_EQUAL(hasFiredCheck, true);
        }
    }

    Y_UNIT_TEST(TestUnhealthyBackendsSuperceedCritical) {
        TConfig config;
        config.set_max_pessimized_share(0.5);
        TFailRateCheckConfig failRateCheck;
        failRateCheck.set_min_requests(10);
        failRateCheck.set_threshold(0.5);
        config.set_fail_rate_check(std::move(failRateCheck));

        const int backendCount = 10;
        for (size_t mask = 0; mask < (1 << backendCount); mask++) {
            TDynamicLog nullLog{TStringBuilder{} << "test_" << mask};
            TPessimizationTracker tracker{config, nullLog};
            auto backends = FillBackends(tracker, backendCount);

            for (size_t i = 0; i < 10; i++) {
                if (mask & (1 << i)) {
                    backends[i]->SetEnabled(true);
                    MakeBackendCritical(*backends[i], config);
                } else {
                    backends[i]->SetEnabled(false);
                }
            }

            UNIT_ASSERT(!tracker.IsOverPessimized());
            auto state = tracker.Step();
            UNIT_ASSERT(tracker.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 5);
            size_t criticalCount = PopCount(mask);
            size_t unhealthyCount = backendCount - criticalCount;
            if (unhealthyCount > 5) {
                UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Unhealthy), unhealthyCount - 5);
                UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Critical), criticalCount);
            } else {
                UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Critical), 5);
            }
        }
    }

    Y_UNIT_TEST(TestBackendRecovering) {
        TConfig config;
        config.set_max_pessimized_share(0.5);
        config.set_weight_increase_step(0.1);
        TFailRateCheckConfig failRateCheck;
        failRateCheck.set_min_requests(10);
        failRateCheck.set_threshold(0.5);
        config.set_fail_rate_check(std::move(failRateCheck));

        const int backendCount = 10;
        TDynamicLog nullLog{"test"};
        TPessimizationTracker tracker{config, nullLog};
        auto backends = FillBackends(tracker, backendCount);

        Sleep(config.history_interval() + TDuration::MilliSeconds(500));
        MakeBackendCritical(*backends[0], config);

        UNIT_ASSERT(!tracker.IsOverPessimized());
        auto state = tracker.Step();
        UNIT_ASSERT(!tracker.IsOverPessimized());
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 9);
        UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 9);
        UNIT_ASSERT(backends[0]->Pessimized);

        // Test that we can start recovery without stats update
        state = tracker.Step();
        UNIT_ASSERT(!tracker.IsOverPessimized());
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
        UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 9);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 0.1);

        // Test that we go back to Critical if stats with fails are still actual
        // and pessimization steps increase in this case
        for (size_t i = 0; i < 2; i++) {
            auto state = tracker.Step();
            UNIT_ASSERT(!tracker.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 9);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 9);
            UNIT_ASSERT(backends[0]->Pessimized);
        }

        // Test that we can start recovery without stats update
        // and do not go to Critical if stats with fails are outdated
        Sleep(config.history_interval() + TDuration::MilliSeconds(500));
        state = tracker.Step();
        UNIT_ASSERT(!tracker.IsOverPessimized());
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
        UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 9);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 0.1);

        // Test that recovery has no progress without new requests
        for (size_t i = 1; i < 10; i++) {
            state = tracker.Step();
            UNIT_ASSERT(!tracker.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 9);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
            UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 0.1);
        }

        // Test that recovery has no progress with too few new requests
        AddSuccessfulRequests(*backends[0], config.fail_rate_check().min_requests() / 2);
        for (size_t i = 1; i < 10; i++) {
            state = tracker.Step();
            UNIT_ASSERT(!tracker.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 9);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
            UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 0.1);
        }

        //Test that recovery has progress if we passed any checks (even if it took some time to save enough requests)
        Sleep(config.history_interval() + TDuration::MilliSeconds(500));
        AddSuccessfulRequests(*backends[0], config.fail_rate_check().min_requests() / 2 + 1);
        for (size_t i = 2; i < 10; i++) {
            state = tracker.Step();
            UNIT_ASSERT(!tracker.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 9);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
            UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), i * 0.1);
        }

        // Test that on final step in recovery we directly transit to healthy state
        state = tracker.Step();
        UNIT_ASSERT(!tracker.IsOverPessimized());
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
        UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 10);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Healthy);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 1);
    }

    Y_UNIT_TEST(TestMaxPessimizedShareExtremeValues) {
        auto test = [](float mps) {
            TConfig config;
            config.set_max_pessimized_share(mps);
            TFailRateCheckConfig failRateCheck;
            failRateCheck.set_min_requests(10);
            failRateCheck.set_threshold(0.5);
            config.set_fail_rate_check(std::move(failRateCheck));

            const int backendCount = 10;
            TDynamicLog nullLog{"test"};
            TPessimizationTracker tracker{config, nullLog};
            auto backends = FillBackends(tracker, backendCount);

            for (size_t i = 0; i < 10; i++) {
                backends[i]->SetEnabled(false);
            }

            UNIT_ASSERT(!tracker.IsOverPessimized());
            auto state = tracker.Step();
            UNIT_ASSERT_VALUES_EQUAL(tracker.IsOverPessimized(), state.size() > 0);
            UNIT_ASSERT_VALUES_EQUAL(state.size(), static_cast<int>(backendCount * (1 - mps)));
        };

        test(0);
        test(1);
    }

    Y_UNIT_TEST(TestWeightsFromBackend) {
        TConfig config;
        config.set_weight_increase_step(0.1);
        config.set_max_pessimized_share(0.5);
        THealthCheckConfig healthCheckConfig;
        healthCheckConfig.set_use_backend_weight(true);
        healthCheckConfig.set_weight_normalization_coeff(10);
        config.set_active(std::move(healthCheckConfig));
        TFailRateCheckConfig failRateCheck;
        failRateCheck.set_min_requests(10);
        failRateCheck.set_threshold(0.5);
        config.set_fail_rate_check(std::move(failRateCheck));

        TVector<double> weightsFromPing = { 1, 5, 10, 20, 50, 75, 100, 1000 };
        TDynamicLog nullLog{"test"};
        TPessimizationTracker tracker{config, nullLog};
        auto backends = FillBackends(tracker, weightsFromPing.size());
        for (size_t i = 0; i < backends.size(); i++) {
            backends[i]->SetEnabled(true);
            backends[i]->SetWeightNormalizationCoeff(config.active().weight_normalization_coeff());
        }

        for (size_t i = 0; i < backends.size(); i++) {
            const double weightFromPing = weightsFromPing[i];
            backends[i]->SetWeightFromPing(weightFromPing, true);

            double finalWeight = Min<double>(weightFromPing / config.active().weight_normalization_coeff(), 10);
            if (finalWeight < 1) {
                auto state = tracker.Step();
                UNIT_ASSERT_VALUES_EQUAL(state.size(), backends.size());
                UNIT_ASSERT_VALUES_EQUAL(backends[i]->FinalWeight(), finalWeight);
            } else {
                double currentWeight = 1;
                if (currentWeight < finalWeight) {
                    // Test that weight does not increase if we didn't pass any checks
                    auto state = tracker.Step();
                    UNIT_ASSERT_VALUES_EQUAL(state.size(), backends.size());
                    UNIT_ASSERT_VALUES_EQUAL(backends[i]->FinalWeight(), currentWeight);
                }
                AddSuccessfulRequests(*backends[i], config.fail_rate_check().min_requests());
                while (currentWeight < finalWeight) {
                    auto state = tracker.Step();
                    currentWeight = Min<double>(currentWeight + config.weight_increase_step(), finalWeight);
                    UNIT_ASSERT_VALUES_EQUAL(state.size(), backends.size());
                    UNIT_ASSERT_VALUES_EQUAL(backends[i]->FinalWeight(), currentWeight);
                }
                UNIT_ASSERT_VALUES_EQUAL(backends[i]->FinalWeight(), finalWeight);
            }
        }
    }

    Y_UNIT_TEST(TestMaxPessimizedShareAppliesToBackendsWithZeroWeight) {
        TConfig config;
        config.set_max_pessimized_share(0.5);
        THealthCheckConfig healthCheckConfig;
        healthCheckConfig.set_use_backend_weight(true);
        healthCheckConfig.set_weight_normalization_coeff(100);
        config.set_active(std::move(healthCheckConfig));

        const int backendCount = 10;
        TDynamicLog nullLog{"test"};
        TPessimizationTracker tracker{config, nullLog};
        auto backends = FillBackends(tracker, backendCount);

        for (size_t i = 0; i < backendCount; i++) {
            backends[i]->SetEnabled(true);
            backends[i]->SetWeightFromPing(0, true);
        }

        UNIT_ASSERT(!tracker.IsOverPessimized());
        auto state = tracker.Step();
        UNIT_ASSERT(tracker.IsOverPessimized());
        UNIT_ASSERT_VALUES_EQUAL(state.size(), backendCount / 2);
        for (size_t i = 0; i < state.size(); i++) {
            UNIT_ASSERT(state[i]->FinalWeight() > 0);
        }
    }

    Y_UNIT_TEST(TestBlacklist) {
        TConfig config;
        config.set_max_pessimized_share(0.5);

        const int backendCount = 10;
        TDynamicLog nullLog{"test"};
        TPessimizationTracker tracker{config, nullLog};
        auto backends = FillBackends(tracker, backendCount);

        for (size_t mask = 0; mask < (1 << backendCount); mask++) {
            TStringBuilder blacklistRaw;
            for (size_t i = 0; i < backendCount; i++) {
                if (mask & (1 << i)) {
                    blacklistRaw << backends[i]->Name() << "\n";
                }
            }
            auto blacklist = TBlacklist::Parse(blacklistRaw);
            tracker.UpdateBlacklist(std::move(blacklist));
            UNIT_ASSERT(!tracker.IsOverPessimized());
            auto state = tracker.Step();
            UNIT_ASSERT(!tracker.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), backendCount - PopCount(mask));
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), backendCount - PopCount(mask));
        }
    }

    Y_UNIT_TEST(TestNotReadySet) {
        TConfig config;
        config.set_max_pessimized_share(0.5);

        const size_t backendCount = 10;
        TDynamicLog nullLog{"test"};
        TPessimizationTracker tracker{config, nullLog};
        auto backends = FillBackends(tracker, backendCount);

        for (size_t mask = 0; mask < (1 << backendCount); mask++) {
            auto notReadySet = MakeIntrusive<TNotReadyBackendSet>();
            for (size_t i = 0; i < backendCount; i++) {
                if (mask & (1 << i)) {
                    notReadySet->insert(backends[i]->Name());
                }
            }
            tracker.UpdateNotReadySet(std::move(notReadySet));
            auto state = tracker.Step();
            auto readyCount = backendCount - PopCount(mask);
            UNIT_ASSERT_VALUES_EQUAL(state.size(), Max<size_t>(readyCount, 5));
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::NotReady), Max<int>(0, 5 - readyCount));
        }
    }

    Y_UNIT_TEST(TestUnhealthyIsUnpessimizedAsFastAsPossible) {
        TConfig config;
        config.set_weight_increase_step(0.1);
        config.set_max_pessimized_share(0.5);
        THealthCheckConfig healthCheckConfig;
        healthCheckConfig.set_weight_normalization_coeff(1);
        config.set_active(std::move(healthCheckConfig));

        TDynamicLog nullLog{"test"};
        TPessimizationTracker tracker{config, nullLog};
        auto backends = FillBackends(tracker, 10);
        for (size_t i = 0; i < backends.size(); i++) {
            backends[i]->SetEnabled(true);
        }

        backends[0]->SetEnabled(false);
        for (size_t i = 0; i < 30; i++) {
            auto state = tracker.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 9);
            UNIT_ASSERT(backends[0]->Pessimized);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Unhealthy);
        }

        backends[0]->SetEnabled(true);
        auto state = tracker.Step();
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 0.1);

        //Test that recovery has progress when we have no checks configured
        for (size_t i = 2; i < 10; i++) {
            state = tracker.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
            UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), i * 0.1);
        }

        // Test that on final step in recovery we directly transit to healthy state
        state = tracker.Step();
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
        UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 10);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Healthy);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 1);
    }

    Y_UNIT_TEST(TestDegradedIsUnpessimizedAsFastAsPossible) {
        TConfig config;
        config.set_weight_increase_step(0.1);
        config.set_max_pessimized_share(0.5);
        THealthCheckConfig healthCheckConfig;
        healthCheckConfig.set_use_backend_degraded_notification(true);
        config.set_active(std::move(healthCheckConfig));

        TDynamicLog nullLog{"test"};
        TPessimizationTracker tracker{config, nullLog};
        auto backends = FillBackends(tracker, 10);
        for (size_t i = 0; i < backends.size(); i++) {
            backends[i]->SetEnabled(true);
            backends[i]->SetDegraded(false);
        }

        backends[0]->SetDegraded(true);
        for (size_t i = 0; i < 30; i++) {
            auto state = tracker.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 9);
            UNIT_ASSERT(backends[0]->Pessimized);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Degraded);
        }

        backends[0]->SetDegraded(false);
        auto state = tracker.Step();
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 0.1);

        //Test that recovery has progress when we have no checks configured
        for (size_t i = 2; i < 10; i++) {
            state = tracker.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
            UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), i * 0.1);
        }

        // Test that on final step in recovery we directly transit to healthy state
        state = tracker.Step();
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
        UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 10);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Healthy);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 1);
    }

    Y_UNIT_TEST(TestCriticalBackendsSuperceedDegraded) {
        TConfig config;
        config.set_max_pessimized_share(0.5);
        TFailRateCheckConfig failRateCheck;
        failRateCheck.set_min_requests(10);
        failRateCheck.set_threshold(0.5);
        config.set_fail_rate_check(std::move(failRateCheck));
        THealthCheckConfig healthCheckConfig;
        healthCheckConfig.set_use_backend_degraded_notification(true);
        config.set_active(std::move(healthCheckConfig));

        const int backendCount = 10;
        for (size_t mask = 0; mask < (1 << backendCount); mask++) {
            TDynamicLog nullLog{TStringBuilder{} << "test_" << mask};
            TPessimizationTracker tracker{config, nullLog};
            auto backends = FillBackends(tracker, backendCount);

            for (size_t i = 0; i < 10; i++) {
                if (mask & (1 << i)) {
                    backends[i]->SetDegraded(false);
                    MakeBackendCritical(*backends[i], config);
                } else {
                    backends[i]->SetDegraded(true);
                }
            }

            UNIT_ASSERT(!tracker.IsOverPessimized());
            auto state = tracker.Step();
            UNIT_ASSERT(tracker.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 5);
            size_t criticalCount = PopCount(mask);
            size_t degradedCount = backendCount - criticalCount;
            if (criticalCount > 5) {
                UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Critical), criticalCount - 5);
                UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Degraded), degradedCount);
            } else {
                UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Degraded), 5);
            }
        }
    }
}

Y_UNIT_TEST_SUITE(PessimizationTrackerDonor) {
    Y_UNIT_TEST(TestUnhealthyTransition) {
        TConfig config;
        config.set_max_pessimized_share(0.5);

        TDynamicLog nullLog{"test"};
        TPessimizationTracker donor{config, nullLog};
        TVector<TAtomicSharedPtr<NSrvKernel::TBackendDescriptor>> backendRefs;
        for (size_t i = 0; i < 2; i++) {
            backendRefs.emplace_back(MakeAtomicShared<NSrvKernel::TBackendDescriptor>(ToString(i), nullptr, 1.0));
        }
        donor.DiscoverNewBackends(backendRefs, nullptr);
        donor.Backends().at("0")->Get()->SetEnabled(false);
        backendRefs.emplace_back(MakeAtomicShared<NSrvKernel::TBackendDescriptor>("2", nullptr, 1.0));
        TPessimizationTracker tracker{config, nullLog};
        tracker.DiscoverNewBackends(backendRefs, &donor);
        auto state = tracker.Step();
        UNIT_ASSERT(!tracker.IsOverPessimized());
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 2);
        for( auto const& backend : state) {
            UNIT_ASSERT(backend->Name() != "0");
        }
    }

    Y_UNIT_TEST(TestStatusTransition) {
        TConfig config;
        config.set_max_pessimized_share(1);

        TFailRateCheckConfig failRateCheck;
        failRateCheck.set_min_requests(10);
        failRateCheck.set_threshold(0.5);
        config.set_fail_rate_check(std::move(failRateCheck));
        TDynamicLog nullLog{"test"};
        TPessimizationTracker donor{config, nullLog};

        TVector<TAtomicSharedPtr<NSrvKernel::TBackendDescriptor>> backendRefs;
        backendRefs.emplace_back(MakeAtomicShared<NSrvKernel::TBackendDescriptor>("0", nullptr, 1.0));
        donor.DiscoverNewBackends(backendRefs, nullptr);
        for (size_t i = 0; i < 10; i++) {
            TError error = Y_MAKE_ERROR(yexception{} << "test");
            donor.Backends().at("0")->Get()->OnFailRequest(error, TDuration::Seconds(1));
        }
        bool hasSkippedCheck = false;
        bool hasFiredCheck = false;
        UNIT_ASSERT_VALUES_EQUAL(donor.RunBackendChecks(*donor.Backends().at("0")->Get(), TStateAggregate{}, hasFiredCheck, hasSkippedCheck), TDynamicBackend::TStatus::Critical);
        UNIT_ASSERT_VALUES_EQUAL(hasSkippedCheck, false);
        UNIT_ASSERT_VALUES_EQUAL(hasFiredCheck, true);

        TPessimizationTracker tracker{config, nullLog};
        backendRefs.emplace_back(MakeAtomicShared<NSrvKernel::TBackendDescriptor>("1", nullptr, 1.0));
        tracker.DiscoverNewBackends(backendRefs, &donor);
        UNIT_ASSERT_VALUES_EQUAL(tracker.RunBackendChecks(*tracker.Backends().at("0")->Get(), TStateAggregate{}, hasFiredCheck, hasSkippedCheck), TDynamicBackend::TStatus::Critical);
        UNIT_ASSERT_VALUES_EQUAL(hasSkippedCheck, false);
        UNIT_ASSERT_VALUES_EQUAL(hasFiredCheck, true);
        UNIT_ASSERT_VALUES_EQUAL(tracker.RunBackendChecks(*tracker.Backends().at("1")->Get(), TStateAggregate{}, hasFiredCheck, hasSkippedCheck), TDynamicBackend::TStatus::Healthy);
        UNIT_ASSERT_VALUES_EQUAL(hasSkippedCheck, true);
        UNIT_ASSERT_VALUES_EQUAL(hasFiredCheck, false);
        auto state = tracker.Step();
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 1);
        UNIT_ASSERT_VALUES_EQUAL(state[0]->Name(), "1");
    }

    Y_UNIT_TEST(TestDonorRecovery) {
        TConfig config;
        config.set_weight_increase_step(0.1);
        config.set_max_pessimized_share(0.5);
        THealthCheckConfig healthCheckConfig;
        healthCheckConfig.set_use_backend_degraded_notification(true);
        config.set_active(std::move(healthCheckConfig));

        TDynamicLog nullLog{"test"};
        TPessimizationTracker donor{config, nullLog};
        TVector<TAtomicSharedPtr<NSrvKernel::TBackendDescriptor>> backendRefs;
        auto backends = FillBackends(donor, 10, &backendRefs);
        for (size_t i = 0; i < backends.size(); i++) {
            backends[i]->SetEnabled(true);
            backends[i]->SetDegraded(false);
        }

        backends[0]->SetDegraded(true);
        for (size_t i = 0; i < 30; i++) {
            auto state = donor.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 9);
            UNIT_ASSERT(backends[0]->Pessimized);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Degraded);
        }

        backends[0]->SetDegraded(false);
        auto state = donor.Step();
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 10);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 0.1);

        TPessimizationTracker tracker{config, nullLog};
        backendRefs.pop_back();
        tracker.DiscoverNewBackends(backendRefs, &donor);

        //Test that recovery has progress when we have no checks configured
        for (size_t i = 2; i < 10; i++) {
            state = tracker.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 9);
            UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Recovering);
            UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), i * 0.1);
        }

        // Test that on final step in recovery we directly transit to healthy state
        state = tracker.Step();
        UNIT_ASSERT_VALUES_EQUAL(state.size(), 9);
        UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 9);
        UNIT_ASSERT_VALUES_EQUAL(backends[0]->Status.load(), TDynamicBackend::TStatus::Healthy);
        UNIT_ASSERT_EQUAL_EPS(backends[0]->FinalWeight(), 1);
    }

    Y_UNIT_TEST(TestMaxPessimizedShareOverflow) {
        TConfig config;
        config.set_max_pessimized_share(0.5);

        TFailRateCheckConfig failRateCheck;
        failRateCheck.set_min_requests(10);
        failRateCheck.set_threshold(0.5);
        config.set_fail_rate_check(std::move(failRateCheck));
        TDynamicLog nullLog{"test"};
        TPessimizationTracker donor{config, nullLog};

        TVector<TAtomicSharedPtr<NSrvKernel::TBackendDescriptor>> backendRefs;
        for (size_t i = 0; i < 4; i++) {
            backendRefs.emplace_back(MakeAtomicShared<NSrvKernel::TBackendDescriptor>(ToString(i), nullptr, 1.0));
        }
        donor.DiscoverNewBackends(backendRefs, nullptr);
        for (size_t i = 0; i < 10; i++) {
            TError error = Y_MAKE_ERROR(yexception{} << "test");
            donor.Backends().at("0")->Get()->OnFailRequest(error, TDuration::Seconds(1));
            donor.Backends().at("1")->Get()->OnFailRequest(error, TDuration::Seconds(1));
        }
        {
            UNIT_ASSERT(!donor.IsOverPessimized());
            auto state = donor.Step();
            UNIT_ASSERT(!donor.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 2);
            for( auto const& backend : state) {
                UNIT_ASSERT(backend->Name() != "0");
                UNIT_ASSERT(backend->Name() != "1");
            }
        }

        TPessimizationTracker tracker{config, nullLog};
        backendRefs.pop_back();
        tracker.DiscoverNewBackends(backendRefs, &donor);
        {
            UNIT_ASSERT(tracker.IsOverPessimized());
            auto state = tracker.Step();
            UNIT_ASSERT(tracker.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 2);
        }
    }

    Y_UNIT_TEST(TestNotCheckedBackends) {
        TConfig config;
        config.set_max_pessimized_share(0.5);

        TFailRateCheckConfig failRateCheck;
        failRateCheck.set_min_requests(10);
        failRateCheck.set_threshold(0.5);
        config.set_fail_rate_check(std::move(failRateCheck));
        THealthCheckConfig healthCheckConfig;
        healthCheckConfig.set_tcp_check(true);
        config.set_active(healthCheckConfig);
        TDynamicLog nullLog{"test"};
        TPessimizationTracker donor{config, nullLog};

        TVector<TAtomicSharedPtr<NSrvKernel::TBackendDescriptor>> backendRefs;
        for (size_t i = 0; i < 4; i++) {
            backendRefs.emplace_back(MakeAtomicShared<NSrvKernel::TBackendDescriptor>(ToString(i), nullptr, 1.0));
        }
        donor.DiscoverNewBackends(backendRefs, nullptr);

        // At start all backends considered healthy
        {
            auto state = donor.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 4);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 4);
        }

        // On update new backends considered not checked
        TPessimizationTracker tracker1{config, nullLog};
        for (size_t i = 0; i < 2; i++) {
            backendRefs.emplace_back(MakeAtomicShared<NSrvKernel::TBackendDescriptor>(ToString(backendRefs.size()), nullptr, 1.0));
        }
        tracker1.DiscoverNewBackends(backendRefs, &donor);
        {
            auto state = tracker1.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 4);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 4);
        }

        // On first check they become healthy
        tracker1.Backends().at("4")->Get()->ActiveChecked = true;
        tracker1.Backends().at("5")->Get()->ActiveChecked = true;
        {
            auto state = tracker1.Step();
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 6);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 6);
        }

        // Not checked backends respect max_pessimized_share
        TPessimizationTracker tracker2{config, nullLog};
        for (size_t i = 0; i < 10; i++) {
            backendRefs.emplace_back(MakeAtomicShared<NSrvKernel::TBackendDescriptor>(ToString(backendRefs.size()), nullptr, 1.0));
        }
        tracker2.DiscoverNewBackends(backendRefs, &tracker1);
        {
            auto state = tracker2.Step();
            UNIT_ASSERT(tracker2.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 8);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 6);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::NotChecked), 2);
        }

        // Not checked backends have minimal pessimization priority
        for (size_t i = 0; i < 10; i++) {
            TError error = Y_MAKE_ERROR(yexception{} << "test");
            tracker2.Backends().at("0")->Get()->OnFailRequest(error, TDuration::Seconds(1));
            tracker2.Backends().at("1")->Get()->OnFailRequest(error, TDuration::Seconds(1));
        }
        {
            auto state = tracker2.Step();
            UNIT_ASSERT(tracker2.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 8);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 4);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::NotChecked), 4);
        }

        // Not checked status transits from donor, updated max_pessimized_share count is respected
        TPessimizationTracker tracker3{config, nullLog};
        backendRefs.pop_back();
        backendRefs.pop_back();
        tracker3.DiscoverNewBackends(backendRefs, &tracker2);
        {
            // backends state is preserved in DiscoverNewBackends, so "0" and "1" move to recovering state on next step
            auto state = tracker3.Step();
            UNIT_ASSERT(tracker3.IsOverPessimized());
            UNIT_ASSERT_VALUES_EQUAL(state.size(), 9);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Healthy), 4);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::NotChecked), 3);
            UNIT_ASSERT_VALUES_EQUAL(CountBackendsByStatus(state, TDynamicBackend::TStatus::Recovering), 2);
        }
    }
}
