1. import java.util.HashMap;
  2. import java.util.Map;
  3. import java.util.concurrent.ConcurrentHashMap;
  4. import java.util.concurrent.atomic.AtomicInteger;
  6. public class RequestHeaderTracker {
  8. private final Map<String, AtomicInteger> headerCounts = new ConcurrentHashMap<>(); // Track header counts
  9. private final long rateLimit; // Define the rate limit
  10. private final long windowMs; // Define the time window for rate limiting
  12. public RequestHeaderTracker(long rateLimit, long windowMs) {
  13. this.rateLimit = rateLimit;
  14. this.windowMs = windowMs;
  15. }
  17. public boolean allowRequest(String requestId, Map<String, String> headers) {
  18. // Iterate over the request headers
  19. for (String header : headers.keySet()) {
  20. // Get the current count for the header
  21. AtomicInteger count = headerCounts.get(header);
  22. if (count == null) {
  23. count = new AtomicInteger(0); // Initialize if header not seen before
  24. }
  26. // Check if the header exceeds the rate limit
  27. if (count.getAndIncrement() > rateLimit) {
  28. return false; // Rate limit exceeded
  29. }
  30. }
  31. return true; // Request allowed
  32. }
  34. public void resetHeaderCounts() {
  35. headerCounts.clear(); // Reset counts after a period
  36. }
  38. public static void main(String[] args) {
  39. RequestHeaderTracker tracker = new RequestHeaderTracker(5, 60000); // Rate limit 5 per minute
  41. // Example usage
  42. for (int i = 0; i < 10; i++) {
  43. Map<String, String> headers = new HashMap<>();
  44. headers.put("X-Request-ID", "123");
  45. headers.put("X-Custom-Header", "value");
  47. if (tracker.allowRequest("req-" + i, headers)) {
  48. System.out.println("Request " + i + " allowed");
  49. } else {
  50. System.out.println("Request " + i + " rate limited");
  51. }
  52. }
  53. tracker.resetHeaderCounts();
  54. }
  55. }

Add your comment