// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.search.searchers;
import ai.vespa.metrics.ContainerMetrics;
import com.yahoo.component.annotation.Inject;
import com.yahoo.cloud.config.ClusterInfoConfig;
import com.yahoo.metrics.simple.MetricReceiver;
import com.yahoo.metrics.simple.Counter;
import com.yahoo.metrics.simple.Point;
import com.yahoo.processing.request.CompoundName;
import com.yahoo.search.Query;
import com.yahoo.search.Result;
import com.yahoo.search.Searcher;
import com.yahoo.search.config.RateLimitingConfig;
import com.yahoo.search.result.ErrorMessage;
import com.yahoo.search.searchchain.Execution;
import com.yahoo.yolean.chain.Provides;
import java.time.Clock;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
/**
* A simple rate limiter.
*
* This takes these query parameter arguments:
*
* - rate.id - (String) the id of the client from rate limiting perspective
*
- rate.cost - (Double) the cost Double of this query. This is read after executing the query and hence can be set
* by downstream searchers inspecting the result to allow differencing the cost of various queries. Default is 1.
*
- rate.quota - (Double) the cost per second a particular id is allowed to consume in this system.
*
- rate.idDimension - (String) the name of the rate-id dimension used when logging metrics.
* If this is not specified, the metric will be logged without dimensions.
*
- rate.dryRun - (Boolean) emit metrics on rejected requests but don't actually reject them
*
*
* Whenever quota is exceeded for an id this searcher will reject queries from that id by
* returning a result containing a status 429 error.
*
* If rate.id or rate.quota is not set in Query.properties this searcher will do nothing.
*
* Metrics: This will emit the count metric requestsOverQuota with the dimension [rate.idDimension=rate.id]
* counting rejected requests.
*
* Ordering: This searcher Provides rateLimiting
*
* @author bratseth
*/
@Provides(RateLimitingSearcher.RATE_LIMITING)
public class RateLimitingSearcher extends Searcher {
/** Constant containing the name this Provides - "rateLimiting", for ordering constraints */
public static final String RATE_LIMITING = "rateLimiting";
public static final CompoundName idKey = CompoundName.from("rate.id");
public static final CompoundName costKey = CompoundName.from("rate.cost");
public static final CompoundName quotaKey = CompoundName.from("rate.quota");
public static final CompoundName idDimensionKey = CompoundName.from("rate.idDimension");
public static final CompoundName dryRunKey = CompoundName.from("rate.dryRun");
private static final String requestsOverQuotaMetricName = ContainerMetrics.REQUESTS_OVER_QUOTA.baseName();
/** Used to divide quota by nodes. Assumption: All nodes get the same share of traffic. */
private final int nodeCount;
/** Shared capacity across all threads. Each thread will ask for more capacity from here when they run out. */
private final AvailableCapacity availableCapacity;
/** Capacity already allocated to this thread */
private final ThreadLocal