/*
 * TagThrottle.actor.cpp
 *
 * This source file is part of the FoundationDB open source project
 *
 * Copyright 2013-2020 Apple Inc. and the FoundationDB project authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "fdbclient/NativeAPI.actor.h"
#include "fdbclient/TagThrottle.h"
#include "fdbserver/TesterInterface.actor.h"
#include "fdbserver/workloads/workloads.actor.h"
#include "fdbrpc/simulator.h"
#include "flow/actorcompiler.h" // This must be the last #include.

struct TagThrottleApiWorkload : TestWorkload {
	bool autoThrottleEnabled;
	double testDuration;

	constexpr static const char* NAME = "TagThrottleApi";

	TagThrottleApiWorkload(WorkloadContext const& wcx) : TestWorkload(wcx) {
		testDuration = getOption(options, LiteralStringRef("testDuration"), 10.0);
		autoThrottleEnabled = SERVER_KNOBS->AUTO_TAG_THROTTLING_ENABLED;
	}

	std::string description() const override { return TagThrottleApiWorkload::NAME; }

	Future<Void> setup(Database const& cx) override {
		DatabaseContext::debugUseTags = true;
		return Void();
	}

	Future<Void> start(Database const& cx) override {
		if (this->clientId != 0)
			return Void();
		return timeout(runThrottleApi(this, cx), testDuration, Void());
	}

	Future<bool> check(Database const& cx) override { return true; }

	void getMetrics(vector<PerfMetric>& m) override {}

	static Optional<TagThrottleType> randomTagThrottleType() {
		Optional<TagThrottleType> throttleType;
		switch (deterministicRandom()->randomInt(0, 3)) {
		case 0:
			throttleType = TagThrottleType::AUTO;
			break;
		case 1:
			throttleType = TagThrottleType::MANUAL;
			break;
		default:
			break;
		}

		return throttleType;
	}

	ACTOR Future<Void> throttleTag(
	    Database cx,
	    std::map<std::pair<TransactionTag, TransactionPriority>, TagThrottleInfo>* manuallyThrottledTags) {
		state TransactionTag tag =
		    TransactionTagRef(deterministicRandom()->randomChoice(DatabaseContext::debugTransactionTagChoices));
		state TransactionPriority priority = deterministicRandom()->randomChoice(allTransactionPriorities);
		state double rate = deterministicRandom()->random01() * 20;
		state double duration = 1 + deterministicRandom()->random01() * 19;

		TagSet tagSet;
		tagSet.addTag(tag);

		try {
			wait(ThrottleApi::throttleTags(cx,
			                               tagSet,
			                               rate,
			                               duration,
			                               TagThrottleType::MANUAL,
			                               priority,
			                               Optional<double>(),
			                               TagThrottledReason::MANUAL));
		} catch (Error& e) {
			state Error err = e;
			if (e.code() == error_code_too_many_tag_throttles) {
				ASSERT(manuallyThrottledTags->size() >= SERVER_KNOBS->MAX_MANUAL_THROTTLED_TRANSACTION_TAGS);
				return Void();
			}

			throw err;
		}

		manuallyThrottledTags->insert_or_assign(
		    std::make_pair(tag, priority),
		    TagThrottleInfo(
		        tag, TagThrottleType::MANUAL, priority, rate, now() + duration, duration, TagThrottledReason::MANUAL));

		return Void();
	}

	ACTOR Future<Void> unthrottleTag(
	    Database cx,
	    std::map<std::pair<TransactionTag, TransactionPriority>, TagThrottleInfo>* manuallyThrottledTags) {
		state TransactionTag tag =
		    TransactionTagRef(deterministicRandom()->randomChoice(DatabaseContext::debugTransactionTagChoices));
		TagSet tagSet;
		tagSet.addTag(tag);

		state Optional<TagThrottleType> throttleType = TagThrottleApiWorkload::randomTagThrottleType();
		Optional<TransactionPriority> priority = deterministicRandom()->coinflip()
		                                             ? Optional<TransactionPriority>()
		                                             : deterministicRandom()->randomChoice(allTransactionPriorities);

		state bool erased = false;
		state double maxExpiration = 0;
		if (!throttleType.present() || throttleType.get() == TagThrottleType::MANUAL) {
			for (auto p : allTransactionPriorities) {
				if (!priority.present() || priority.get() == p) {
					auto itr = manuallyThrottledTags->find(std::make_pair(tag, p));
					if (itr != manuallyThrottledTags->end()) {
						maxExpiration = std::max(maxExpiration, itr->second.expirationTime);
						erased = true;
						manuallyThrottledTags->erase(itr);
					}
				}
			}
		}

		bool removed = wait(ThrottleApi::unthrottleTags(cx, tagSet, throttleType, priority));
		if (removed) {
			ASSERT(erased || !throttleType.present() || throttleType.get() == TagThrottleType::AUTO);
		} else {
			ASSERT(maxExpiration < now());
		}

		return Void();
	}

	ACTOR Future<Void> getTags(
	    TagThrottleApiWorkload* self,
	    Database cx,
	    std::map<std::pair<TransactionTag, TransactionPriority>, TagThrottleInfo> const* manuallyThrottledTags) {
		std::vector<TagThrottleInfo> tags = wait(ThrottleApi::getThrottledTags(cx, CLIENT_KNOBS->TOO_MANY));

		int manualThrottledTags = 0;
		int activeAutoThrottledTags = 0;
		for (auto& tag : tags) {
			if (tag.throttleType == TagThrottleType::AUTO) {
				ASSERT(self->autoThrottleEnabled);
			} else if (tag.throttleType == TagThrottleType::MANUAL) {
				ASSERT(manuallyThrottledTags->find(std::make_pair(tag.tag, tag.priority)) !=
				       manuallyThrottledTags->end());
				++manualThrottledTags;
			} else if (tag.expirationTime > now()) {
				++activeAutoThrottledTags;
			}
		}

		ASSERT(manualThrottledTags <= SERVER_KNOBS->MAX_MANUAL_THROTTLED_TRANSACTION_TAGS);
		ASSERT(activeAutoThrottledTags <= SERVER_KNOBS->MAX_AUTO_THROTTLED_TRANSACTION_TAGS);

		int minManualThrottledTags = 0;
		int maxManualThrottledTags = 0;
		for (auto& tag : *manuallyThrottledTags) {
			if (tag.second.expirationTime > now()) {
				++minManualThrottledTags;
			}
			++maxManualThrottledTags;
		}

		ASSERT(manualThrottledTags >= minManualThrottledTags && manualThrottledTags <= maxManualThrottledTags);
		return Void();
	}

	ACTOR Future<Void> getRecommendedTags(TagThrottleApiWorkload* self, Database cx) {
		std::vector<TagThrottleInfo> tags = wait(ThrottleApi::getRecommendedTags(cx, CLIENT_KNOBS->TOO_MANY));

		for (auto& tag : tags) {
			ASSERT(tag.throttleType == TagThrottleType::AUTO);
		}
		return Void();
	}

	ACTOR Future<Void> unthrottleTagGroup(
	    Database cx,
	    std::map<std::pair<TransactionTag, TransactionPriority>, TagThrottleInfo>* manuallyThrottledTags) {
		state Optional<TagThrottleType> throttleType = TagThrottleApiWorkload::randomTagThrottleType();
		state Optional<TransactionPriority> priority =
		    deterministicRandom()->coinflip() ? Optional<TransactionPriority>()
		                                      : deterministicRandom()->randomChoice(allTransactionPriorities);

		bool unthrottled = wait(ThrottleApi::unthrottleAll(cx, throttleType, priority));
		if (!throttleType.present() || throttleType.get() == TagThrottleType::MANUAL) {
			bool unthrottleExpected = false;
			bool empty = manuallyThrottledTags->empty();
			for (auto itr = manuallyThrottledTags->begin(); itr != manuallyThrottledTags->end();) {
				if (!priority.present() || priority.get() == itr->first.second) {
					if (itr->second.expirationTime > now()) {
						unthrottleExpected = true;
					}

					itr = manuallyThrottledTags->erase(itr);
				} else {
					++itr;
				}
			}

			if (throttleType.present()) {
				ASSERT((unthrottled && !empty) || (!unthrottled && !unthrottleExpected));
			} else {
				ASSERT(unthrottled || !unthrottleExpected);
			}
		}

		return Void();
	}

	ACTOR Future<Void> enableAutoThrottling(TagThrottleApiWorkload* self, Database cx) {
		if (deterministicRandom()->coinflip()) {
			wait(ThrottleApi::enableAuto(cx, true));
			self->autoThrottleEnabled = true;
			if (deterministicRandom()->coinflip()) {
				bool unthrottled =
				    wait(ThrottleApi::unthrottleAll(cx, TagThrottleType::AUTO, Optional<TransactionPriority>()));
			}
		} else {
			wait(ThrottleApi::enableAuto(cx, false));
			self->autoThrottleEnabled = false;
		}

		return Void();
	}

	ACTOR Future<Void> runThrottleApi(TagThrottleApiWorkload* self, Database cx) {
		state std::map<std::pair<TransactionTag, TransactionPriority>, TagThrottleInfo> manuallyThrottledTags;
		loop {
			double delayTime = deterministicRandom()->random01() * 5;
			wait(delay(delayTime));

			state int action = deterministicRandom()->randomInt(0, 6);

			if (action == 0) {
				wait(self->throttleTag(cx, &manuallyThrottledTags));
			} else if (action == 1) {
				wait(self->unthrottleTag(cx, &manuallyThrottledTags));
			} else if (action == 2) {
				wait(self->getTags(self, cx, &manuallyThrottledTags));
			} else if (action == 3) {
				wait(self->unthrottleTagGroup(cx, &manuallyThrottledTags));
			} else if (action == 4) {
				wait(self->enableAutoThrottling(self, cx));
			} else if (action == 5) {
				wait(self->getRecommendedTags(self, cx));
			}
		}
	}
};

WorkloadFactory<TagThrottleApiWorkload> TagThrottleApiWorkloadFactory(TagThrottleApiWorkload::NAME);
