Source code

001package org.nasdanika.models.gitlab.util;
002
003import java.time.ZonedDateTime;
004import java.time.format.DateTimeFormatter;
005import java.util.concurrent.atomic.AtomicInteger;
006import java.util.concurrent.atomic.AtomicLong;
007import java.util.logging.Handler;
008import java.util.logging.LogRecord;
009import java.util.regex.Matcher;
010import java.util.regex.Pattern;
011
012/**
013 * A {@link Handler} for throttling API calls based on rate limit HTTP response headers
014 */
015public class ThrottlingHandler extends Handler {
016        
017        private static final Pattern REQUEST_PATTERN = Pattern.compile("\\d+ - Sending client request on thread .+");
018        private static final Pattern RESPONSE_PATTERN = Pattern.compile("\\d+ - Received server response on thread .+"); 
019        
020        private static final Pattern RATE_LIMIT_PREFIX_PATTERN = Pattern.compile("^\\d+ < RateLimit-Limit: ");
021//      private static final Pattern RATE_LIMIT_OBSERVED_PREFIX_PATTERN = Pattern.compile("^\\d+ < RateLimit-Observed: ");
022        private static final Pattern RATE_LIMIT_REMAINING_PREFIX_PATTERN = Pattern.compile("^\\d+ < RateLimit-Remaining: ");
023        private static final Pattern RATE_LIMIT_RESET_PREFIX_PATTERN = Pattern.compile("^\\d+ < RateLimit-Reset: ");
024        private static final Pattern DATE_PREFIX_PATTERN = Pattern.compile("^\\d+ < Date: ");
025        
026        private int conservatism = 3;
027        private int factor = 10;
028        
029        // Client-provided rate window
030        private long clientRateLimitWindow = -1;
031        private int clientRateLimit = -1;
032        
033        public ThrottlingHandler() {
034        }
035
036        /**
037         * Client rate limit is enforced before sending a request.
038         * @param clientRateLimitWindow Client rate window in milliseconds. Client rate limit is enforced if this value and clientRateLimit are positive. 
039         * @param clientRateLimit Client rate limit per rate window. Client rate limit is enforced if this value and clientRateLimitWindow are positive.
040         */
041        public ThrottlingHandler(long clientRateLimitWindow, int clientRateLimit) {
042                
043                this.clientRateLimitWindow = clientRateLimitWindow;
044                this.clientRateLimit = clientRateLimit;
045        }
046        
047        public int getConservatism() {
048                return conservatism;
049        }
050        
051        /**
052         * Number in seconds to add to the server rate window when computing throttling sleep interval. 
053         * Non-zero conservatism ensures that not all rate is used before the window ends. The default value is 3.
054         * @param conservatism
055         */
056        public void setConservatism(int conservatism) {
057                this.conservatism = conservatism;
058        }
059        
060        public int getFactor() {
061                return factor;
062        }
063        
064        /**
065         * Determines when to start throttling. 
066         * Throttling starts when the remaining server rate limit is less than rate limit divided by factor.
067         * E.g. if factor is 20 then throttling will start when 95% of the rate has been used. 
068         * The default is 10 - throttling starts after 90% of the rate has been used. 
069         * @param factor
070         */
071        public void setFactor(int factor) {
072                this.factor = factor;
073        }
074        
075        private AtomicLong clientRateLimitWindowEnd = new AtomicLong(); 
076        private AtomicInteger clientRateLimitWindowRequestsRemaining = new AtomicInteger();
077        
078        private AtomicInteger serverRateLimit = new AtomicInteger();
079        private AtomicInteger serverRateLimitRemaining = new AtomicInteger();
080        private AtomicLong serverRateLimitReset = new AtomicLong();
081        private AtomicLong serverResponseTime = new AtomicLong();       
082        
083        @Override
084        public void publish(LogRecord logRecord) {
085                String message = logRecord.getMessage();
086                String lines[] = message.split("\\r?\\n");
087                if (REQUEST_PATTERN.matcher(lines[0]).matches()) {                      
088                        long now = System.currentTimeMillis();
089                        if (clientRateLimit > 0 && clientRateLimitWindow > 0) {                         
090                                long clientRateWindowDelta = now - clientRateLimitWindowEnd.get();
091                                if (clientRateWindowDelta >= 0) {
092                                        // Updating window end and replenishing requests budget
093                                        clientRateLimitWindowEnd.set(now + clientRateLimitWindow);
094                                        clientRateLimitWindowRequestsRemaining.set(clientRateLimit);
095                                }
096                        }                       
097                        
098                        int rateLimit = serverRateLimit.get();
099                        int rateLimitRemaining = serverRateLimitRemaining.get();
100                        long rateLimitReset = serverRateLimitReset.get();
101                        
102                        long toSleep = -1; 
103                        
104                        if (rateLimit > 0 
105                                        && rateLimitRemaining > 0 
106                                        && rateLimitRemaining * factor < rateLimit) { // Start throttling when less than rateLimit / factor remains
107                                
108                                long delta = rateLimitReset - serverResponseTime.get(); // Seconds remaining until reset
109                                delta += conservatism; 
110                                toSleep = delta * 1000 / rateLimitRemaining;
111                        }
112                        
113                        if (clientRateLimit > 0 && clientRateLimitWindow > 0) {
114                                long clientRateWindowDelta = clientRateLimitWindowEnd.get() - now; // Milliseconds until window end (client rate reset)
115                                long clientToSleep = clientRateWindowDelta / Math.max(clientRateLimitWindowRequestsRemaining.decrementAndGet(), 1); // If ran out of rate limit - wait until the end of the window
116                                if (clientToSleep > toSleep) {
117                                        toSleep = clientToSleep;
118                                }
119                        }
120                        
121                        if (toSleep > 0) {
122                                try {
123                                        Thread.sleep(toSleep);
124                                } catch (InterruptedException e) {
125                                        e.printStackTrace();
126                                };
127                        }                                                                                               
128                }
129                                
130                if (RESPONSE_PATTERN.matcher(lines[0]).matches()) {
131                        int rateLimit = -1;
132//                      int rateLimitObserved = -1;
133                        int rateLimitRemaining = -1;
134                        long rateLimitReset = -1;
135                        ZonedDateTime date = null;
136                        
137                        for (String line: lines) {
138                                Matcher rateLimitMatcher = RATE_LIMIT_PREFIX_PATTERN.matcher(line);
139                                if (rateLimitMatcher.find()) {
140                                        rateLimit = Integer.parseInt(line.substring(rateLimitMatcher.end()));
141                                } else {
142                                        Matcher rateLimitRemainingMatcher = RATE_LIMIT_REMAINING_PREFIX_PATTERN.matcher(line);
143                                        if (rateLimitRemainingMatcher.find()) {
144                                                rateLimitRemaining = Integer.parseInt(line.substring(rateLimitRemainingMatcher.end()));
145                                        } else {
146                                                Matcher rateLimitResetMatcher = RATE_LIMIT_RESET_PREFIX_PATTERN.matcher(line);
147                                                if (rateLimitResetMatcher.find()) {
148                                                        rateLimitReset = Long.parseLong(line.substring(rateLimitResetMatcher.end()));
149                                                } else {
150                                                        Matcher dateMatcher = DATE_PREFIX_PATTERN.matcher(line);
151                                                        if (dateMatcher.find()) {
152                                                                String dateString = line.substring(dateMatcher.end());
153                                                                date = ZonedDateTime.parse(dateString, DateTimeFormatter.RFC_1123_DATE_TIME);                                                           
154                                                        } else {
155                                                                // Rate limit observed if needed
156                                                        }
157                                                
158                                                }
159                                        }
160                                }
161                        }
162                        
163                        serverRateLimit.set(rateLimit);
164                        serverRateLimitRemaining.set(rateLimitRemaining);
165                        serverRateLimitReset.set(rateLimitReset);
166                        serverResponseTime.set(date == null ? System.currentTimeMillis() / 1000 : date.toEpochSecond());
167                }
168        }
169        
170        @Override
171        public void flush() {
172                // NOP                          
173        }
174        
175        @Override
176        public void close() throws SecurityException {
177                // NOP                          
178        }
179
180}