RateLimitingPreExecutionInterceptor.java

package org.stagemonitor.tracing.sampling;

import org.stagemonitor.configuration.ConfigurationOption;
import org.stagemonitor.configuration.ConfigurationRegistry;
import org.stagemonitor.tracing.SpanContextInformation;
import org.stagemonitor.tracing.TracingPlugin;
import org.stagemonitor.tracing.utils.RateLimiter;

import java.util.HashMap;
import java.util.Map;

public class RateLimitingPreExecutionInterceptor extends PreExecutionSpanInterceptor {

	private RateLimiter serverSpanRateLimiter;
	private RateLimiter clientSpanRateLimiter;
	private Map<String, RateLimiter> clientSpanRateLimiterByType;

	@Override
	public void init(ConfigurationRegistry configuration) {
		TracingPlugin tracingPlugin = configuration.getConfig(TracingPlugin.class);
		serverSpanRateLimiter = getRateLimiter(tracingPlugin.getRateLimitServerSpansPerMinute());
		tracingPlugin.getRateLimitServerSpansPerMinuteOption().addChangeListener(new ConfigurationOption.ChangeListener<Double>() {
			@Override
			public void onChange(ConfigurationOption<?> configurationOption, Double oldValue, Double newValue) {
				serverSpanRateLimiter = getRateLimiter(newValue);
			}
		});

		clientSpanRateLimiter = getRateLimiter(tracingPlugin.getRateLimitClientSpansPerMinute());
		tracingPlugin.getRateLimitClientSpansPerMinuteOption().addChangeListener(new ConfigurationOption.ChangeListener<Double>() {
			@Override
			public void onChange(ConfigurationOption<?> configurationOption, Double oldValue, Double newValue) {
				clientSpanRateLimiter = getRateLimiter(newValue);
			}
		});

		setRateLimiterMap(tracingPlugin.getRateLimitClientSpansPerTypePerMinute());
		tracingPlugin.getRateLimitClientSpansPerTypePerMinuteOption().addChangeListener(new ConfigurationOption.ChangeListener<Map<String, Double>>() {
			@Override
			public void onChange(ConfigurationOption<?> configurationOption, Map<String, Double> oldValue, Map<String, Double> newValue) {
				setRateLimiterMap(newValue);
			}
		});
	}

	private void setRateLimiterMap(Map<String, Double> newValue) {
		Map<String, RateLimiter> rateLimiters = new HashMap<String, RateLimiter>();
		for (Map.Entry<String, Double> entry : newValue.entrySet()) {
			rateLimiters.put(entry.getKey(), getRateLimiter(entry.getValue()));
		}
		clientSpanRateLimiterByType = rateLimiters;
	}

	public static RateLimiter getRateLimiter(double creditsPerMinute) {
		if (creditsPerMinute >= 1000000) {
			return null;
		}
		final double maxTracesPerSecond = creditsPerMinute / 60;
		double maxBalance;
		if (maxTracesPerSecond <= 0) {
			maxBalance = 0.0;
		} else if (maxTracesPerSecond < 1.0) {
			maxBalance = 1.0;
		} else {
			maxBalance = maxTracesPerSecond;
		}
		return new RateLimiter(maxTracesPerSecond, maxBalance);
	}

	@Override
	public void interceptReport(PreExecutionInterceptorContext context) {
		final SpanContextInformation spanContext = context.getSpanContext();
		boolean rateExceeded = false;
		if (spanContext.isServerRequest()) {
			rateExceeded = isRateExceeded(serverSpanRateLimiter);
		} else if (spanContext.isExternalRequest()) {
			RateLimiter rateLimiter = clientSpanRateLimiter;
			if (clientSpanRateLimiterByType.containsKey(spanContext.getOperationType())) {
				rateLimiter = clientSpanRateLimiterByType.get(spanContext.getOperationType());
			}
			rateExceeded = isRateExceeded(rateLimiter);
		}
		if (rateExceeded) {
			context.shouldNotReport(getClass());
		}
	}

	public static boolean isRateExceeded(RateLimiter rateLimiter) {
		return rateLimiter != null && !rateLimiter.checkCredit(1.0);
	}

}