• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * Copyright (C) 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package android.os.profiling;
18 
19 import android.annotation.IntDef;
20 import android.annotation.Nullable;
21 import android.os.Bundle;
22 import android.os.Environment;
23 import android.os.Handler;
24 import android.os.ProfilingManager;
25 import android.os.ProfilingResult;
26 import android.os.RateLimiterRecordsWrapper;
27 import android.provider.DeviceConfig;
28 import android.util.AtomicFile;
29 import android.util.Log;
30 import android.util.SparseIntArray;
31 
32 import com.android.internal.annotations.GuardedBy;
33 import com.android.internal.annotations.VisibleForTesting;
34 
35 import java.io.File;
36 import java.io.FileOutputStream;
37 import java.io.IOException;
38 import java.lang.annotation.Retention;
39 import java.lang.annotation.RetentionPolicy;
40 import java.util.ArrayDeque;
41 import java.util.Queue;
42 import java.util.concurrent.atomic.AtomicBoolean;
43 
44 public class RateLimiter {
45     private static final String TAG = RateLimiter.class.getSimpleName();
46     private static final boolean DEBUG = false;
47 
48     private static final String RATE_LIMITER_STORE_DIR = "profiling_rate_limiter_store";
49     private static final String RATE_LIMITER_INFO_FILE = "profiling_rate_limiter_info";
50 
51     private static final long TIME_HOUR_MS = 60 * 60 * 1000;
52     private static final long TIME_DAY_MS = 24 * 60 * 60 * 1000;
53     private static final long TIME_WEEK_MS = 7 * 24 * 60 * 60 * 1000;
54 
55     private static final int DEFAULT_MAX_COST_SYSTEM_HOUR = 2;
56     private static final int DEFAULT_MAX_COST_PROCESS_HOUR = 1;
57     private static final int DEFAULT_MAX_COST_SYSTEM_DAY = 5;
58     private static final int DEFAULT_MAX_COST_PROCESS_DAY = 2;
59     private static final int DEFAULT_MAX_COST_SYSTEM_WEEK = 15;
60     private static final int DEFAULT_MAX_COST_PROCESS_WEEK = 3;
61     private static final int DEFAULT_COST_PER_SESSION = 1;
62 
63     public static final int RATE_LIMIT_RESULT_ALLOWED = 0;
64     public static final int RATE_LIMIT_RESULT_BLOCKED_PROCESS = 1;
65     public static final int RATE_LIMIT_RESULT_BLOCKED_SYSTEM = 2;
66 
67     private final Object mLock = new Object();
68 
69     private long mPersistToDiskFrequency;
70 
71     /** To be disabled for testing only. */
72     @GuardedBy("mLock")
73     private boolean mRateLimiterDisabled = false;
74 
75     /** Collection of run costs and entries from the last hour. */
76     @VisibleForTesting
77     public final EntryGroupWrapper mPastRunsHour;
78 
79     /** Collection of run costs and entries from the last day. */
80     @VisibleForTesting
81     public final EntryGroupWrapper mPastRunsDay;
82 
83     /** Collection of run costs and entries from the last week. */
84     @VisibleForTesting
85     public final EntryGroupWrapper mPastRunsWeek;
86 
87     private int mCostJavaHeapDump;
88     private int mCostHeapProfile;
89     private int mCostStackSampling;
90     private int mCostSystemTrace;
91 
92     private final HandlerCallback mHandlerCallback;
93 
94     private Runnable mPersistRunnable = null;
95     private boolean mPersistScheduled = false;
96 
97     private long mLastPersistedTimestampMs;
98 
99     /**
100      * The path to the directory which includes the historical rate limiter data file as specified
101      * in {@link #mPersistFile}.
102      */
103     @VisibleForTesting
104     public File mPersistStoreDir;
105 
106     /** The historical rate limiter data file, persisted in the storage. */
107     @VisibleForTesting
108     public File mPersistFile;
109 
110     @VisibleForTesting
111     public AtomicBoolean mDataLoaded = new AtomicBoolean();
112 
113     @IntDef(value = {
114         RATE_LIMIT_RESULT_ALLOWED,
115         RATE_LIMIT_RESULT_BLOCKED_PROCESS,
116         RATE_LIMIT_RESULT_BLOCKED_SYSTEM,
117     })
118     @Retention(RetentionPolicy.SOURCE)
119     @interface RateLimitResult {}
120 
121     /**
122      * @param handlerCallback Callback for rate limiter to obtain a {@link Handler} to schedule
123      *                        work such as persisting to storage.
124      */
RateLimiter(HandlerCallback handlerCallback)125     public RateLimiter(HandlerCallback handlerCallback) {
126         mHandlerCallback = handlerCallback;
127 
128         DeviceConfig.Properties properties = DeviceConfigHelper.getAllRateLimiterProperties();
129 
130         mPastRunsHour = new EntryGroupWrapper(
131                 properties.getInt(DeviceConfigHelper.MAX_COST_SYSTEM_1_HOUR,
132                         DEFAULT_MAX_COST_SYSTEM_HOUR),
133                 properties.getInt(DeviceConfigHelper.MAX_COST_PROCESS_1_HOUR,
134                         DEFAULT_MAX_COST_PROCESS_HOUR),
135                 TIME_HOUR_MS);
136         mPastRunsDay = new EntryGroupWrapper(
137                 properties.getInt(DeviceConfigHelper.MAX_COST_SYSTEM_24_HOUR,
138                         DEFAULT_MAX_COST_SYSTEM_DAY),
139                 properties.getInt(DeviceConfigHelper.MAX_COST_PROCESS_24_HOUR,
140                         DEFAULT_MAX_COST_PROCESS_DAY),
141                 TIME_DAY_MS);
142         mPastRunsWeek = new EntryGroupWrapper(
143                 properties.getInt(DeviceConfigHelper.MAX_COST_SYSTEM_7_DAY,
144                         DEFAULT_MAX_COST_SYSTEM_WEEK),
145                 properties.getInt(DeviceConfigHelper.MAX_COST_PROCESS_7_DAY,
146                         DEFAULT_MAX_COST_PROCESS_WEEK),
147                 TIME_WEEK_MS);
148 
149         mCostJavaHeapDump = properties.getInt(DeviceConfigHelper.COST_JAVA_HEAP_DUMP,
150                 DEFAULT_COST_PER_SESSION);
151         mCostHeapProfile = properties.getInt(DeviceConfigHelper.COST_HEAP_PROFILE,
152                 DEFAULT_COST_PER_SESSION);
153         mCostStackSampling = properties.getInt(DeviceConfigHelper.COST_STACK_SAMPLING,
154                 DEFAULT_COST_PER_SESSION);
155         mCostSystemTrace = properties.getInt(DeviceConfigHelper.COST_SYSTEM_TRACE,
156                 DEFAULT_COST_PER_SESSION);
157 
158         mPersistToDiskFrequency = properties.getLong(
159                 DeviceConfigHelper.PERSIST_TO_DISK_FREQUENCY_MS, 0);
160         mLastPersistedTimestampMs = System.currentTimeMillis();
161 
162         // Get initial value for whether rate limiter should be enforcing or if it should always
163         // allow profiling requests. This is used for (automated and manual) testing only.
164         synchronized (mLock) {
165             mRateLimiterDisabled = DeviceConfigHelper.getTestBoolean(
166                     DeviceConfigHelper.RATE_LIMITER_DISABLE_PROPERTY, false);
167         }
168 
169         setupFromPersistedData();
170     }
171 
isProfilingRequestAllowed(int uid, int profilingType, @Nullable Bundle params)172     public @RateLimitResult int isProfilingRequestAllowed(int uid,
173             int profilingType, @Nullable Bundle params) {
174         synchronized (mLock) {
175             if (mRateLimiterDisabled) {
176                 // Rate limiter is disabled for testing, approve request and don't store cost.
177                 Log.w(TAG, "Rate limiter disabled, request allowed.");
178                 return RATE_LIMIT_RESULT_ALLOWED;
179             }
180             if (!mDataLoaded.get()) {
181                 // Requests are rejected before rate limiter data is loaded or if data load fails.
182                 Log.e(TAG, "Data loading in progress or failed, request denied.");
183                 return RATE_LIMIT_RESULT_BLOCKED_SYSTEM;
184             }
185             final int cost = getCostForProfiling(profilingType);
186             final long currentTimeMillis = System.currentTimeMillis();
187             int status = mPastRunsHour.isProfilingAllowed(uid, cost, currentTimeMillis);
188             if (status == RATE_LIMIT_RESULT_ALLOWED) {
189                 status = mPastRunsDay.isProfilingAllowed(uid, cost, currentTimeMillis);
190             }
191             if (status == RATE_LIMIT_RESULT_ALLOWED) {
192                 status = mPastRunsWeek.isProfilingAllowed(uid, cost, currentTimeMillis);
193             }
194             if (status == RATE_LIMIT_RESULT_ALLOWED) {
195                 mPastRunsHour.add(uid, cost, currentTimeMillis);
196                 mPastRunsDay.add(uid, cost, currentTimeMillis);
197                 mPastRunsWeek.add(uid, cost, currentTimeMillis);
198                 maybePersistToDisk();
199                 return RATE_LIMIT_RESULT_ALLOWED;
200             }
201             return status;
202         }
203     }
204 
getCostForProfiling(int profilingType)205     private int getCostForProfiling(int profilingType) {
206         switch (profilingType) {
207             case ProfilingManager.PROFILING_TYPE_JAVA_HEAP_DUMP:
208                 return mCostJavaHeapDump;
209             case ProfilingManager.PROFILING_TYPE_HEAP_PROFILE:
210                 return mCostHeapProfile;
211             case ProfilingManager.PROFILING_TYPE_STACK_SAMPLING:
212                 return mCostStackSampling;
213             case ProfilingManager.PROFILING_TYPE_SYSTEM_TRACE:
214                 return mCostSystemTrace;
215             default:
216                 return Integer.MAX_VALUE;
217         }
218     }
219 
220     /**
221      * This method is meant to be called every time a profiling record is added to the history.
222      * - If persist frequency is set to 0, it will immediately persist the records to disk.
223      * - If a persist is already scheduled, it will do nothing.
224      * - If the last records persist occurred longer ago than the persist frequency, it will
225      *      persist immediately.
226      * - In all other cases, it will schedule a persist event at persist frequency after the last
227      *      persist event.
228      */
maybePersistToDisk()229     void maybePersistToDisk() {
230         if (mPersistScheduled) {
231             // We're already waiting on a scheduled persist job, do nothing.
232             return;
233         }
234 
235         if (mPersistToDiskFrequency == 0
236                 || (System.currentTimeMillis() - mLastPersistedTimestampMs
237                         >= mPersistToDiskFrequency)) {
238             // If persist frequency is 0 or if it's already been longer than persist frequency since
239             // the last persist then persist immediately.
240             persistToDisk();
241         } else {
242             // Schedule the persist job.
243             if (mPersistRunnable == null) {
244                 mPersistRunnable = new Runnable() {
245                     @Override
246                     public void run() {
247                         persistToDisk();
248                         mPersistScheduled = false;
249                     }
250                 };
251             }
252             mPersistScheduled = true;
253             long persistDelay = mLastPersistedTimestampMs + mPersistToDiskFrequency
254                     - System.currentTimeMillis();
255             mHandlerCallback.obtainHandler().postDelayed(mPersistRunnable, persistDelay);
256         }
257     }
258 
259     /**
260      * Clean up records and persist to disk.
261      *
262      * Skips if {@link mPersistFile} is not accessible to write to.
263      */
persistToDisk()264     public void persistToDisk() {
265         // Check if file exists
266         try {
267             if (mPersistFile == null) {
268                 // Try again to create the necessary files.
269                 if (!setupPersistFiles()) {
270                     // No file, nowhere to save.
271                     if (DEBUG) Log.d(TAG, "Failed setting up persist files so nowhere to save to.");
272                     return;
273                 }
274             }
275 
276             if (!mPersistFile.exists()) {
277                 // File doesn't exist, try to create it.
278                 mPersistFile.createNewFile();
279             }
280         } catch (Exception e) {
281             if (DEBUG) Log.d(TAG, "Exception accessing persisted records", e);
282             return;
283         }
284 
285         // Clean up old records to reduce extraneous writes
286         mPastRunsWeek.cleanUpOldRecords();
287 
288         // Generate proto for records. We only persist week records as this contains all smaller
289         // time ranges.
290         RateLimiterRecordsWrapper outerWrapper = RateLimiterRecordsWrapper.newBuilder()
291                 .setRecords(mPastRunsWeek.toProto())
292                 .build();
293 
294         // Write to disk
295         byte[] protoBytes = outerWrapper.toByteArray();
296         AtomicFile persistFile = new AtomicFile(mPersistFile);
297         FileOutputStream out = null;
298         try {
299             out = persistFile.startWrite();
300             out.write(protoBytes);
301             persistFile.finishWrite(out);
302         } catch (IOException e) {
303             if (DEBUG) Log.d(TAG, "Exception writing records", e);
304             persistFile.failWrite(out);
305         }
306     }
307 
308     /**
309      * Load initial records data from disk and marks rate limiter ready to use if it's in an
310      * acceptable state.
311      */
312     @VisibleForTesting
setupFromPersistedData()313     public void setupFromPersistedData() {
314         // Setup persist files
315         try {
316             if (!setupPersistFiles()) {
317                 // If setup directory and file was unsuccessful then we won't be able to persist
318                 // records, return and leave feature disabled entirely.
319                 if (DEBUG) Log.d(TAG, "Failed to setup persist directory/files. Feature disabled.");
320                 mDataLoaded.set(false);
321                 return;
322             }
323         } catch (SecurityException e) {
324             // Can't access files.
325             if (DEBUG) Log.d(TAG, "Failed to setup persist directory/files. Feature disabled.", e);
326             mDataLoaded.set(false);
327             return;
328         }
329 
330         // Check if file exists
331         try {
332             if (!mPersistFile.exists()) {
333                 // No file, nothing to load. This is an expected state for before the feature has
334                 // ever been used so mark ready to use and return.
335                 if (DEBUG) Log.d(TAG, "Persist file does not exist, skipping load from disk.");
336                 mDataLoaded.set(true);
337                 return;
338             }
339         } catch (SecurityException e) {
340             // Can't access file.
341             if (DEBUG) Log.d(TAG, "Exception accessing persist file", e);
342             mDataLoaded.set(false);
343             return;
344         }
345 
346         // Read the file
347         AtomicFile persistFile = new AtomicFile(mPersistFile);
348         byte[] bytes;
349         try {
350             bytes = persistFile.readFully();
351         } catch (IOException e) {
352             if (DEBUG) Log.d(TAG, "Exception reading persist file", e);
353             // We already handled no file case above and empty file would not result in exception
354             // so this is a problem reading the file. Attempt remediation.
355             if (handleBadFile()) {
356                 // Successfully remediated bad state! Mark ready to use.
357                 mDataLoaded.set(true);
358             } else {
359                 // Failed to remediate bad state. Feature disabled.
360                 mDataLoaded.set(false);
361             }
362             // Return either way as {@link handleBadFile} handles the entirety of remediating the
363             // bad state and the remainder of this method is no longer applicable.
364             return;
365         }
366         if (bytes.length == 0) {
367             // Empty file, nothing to load. This is an expected state for before the feature
368             // persists so mark ready to use and return.
369             if (DEBUG) Log.d(TAG, "Persist file is empty, skipping load from disk.");
370             mDataLoaded.set(true);
371             return;
372         }
373 
374         // Parse file bytes to proto
375         RateLimiterRecordsWrapper outerWrapper;
376         try {
377             outerWrapper = RateLimiterRecordsWrapper.parseFrom(bytes);
378         } catch (Exception e) {
379             // Failed to parse. Attempt remediation.
380             if (DEBUG) Log.d(TAG, "Error parsing proto from persisted bytes", e);
381             if (handleBadFile()) {
382                 // Successfully remediated bad state! Mark ready to use.
383                 mDataLoaded.set(true);
384             } else {
385                 // Failed to remediate bad state. Feature disabled.
386                 mDataLoaded.set(false);
387             }
388             // Return either way as {@link handleBadFile} handles the entirety of remediating the
389             // bad state and the remainder of this method is no longer applicable.
390             return;
391         }
392 
393         // Populate in memory records stores
394         RateLimiterRecordsWrapper.EntryGroupWrapper weekGroupWrapper = outerWrapper.getRecords();
395         final long currentTimeMillis = System.currentTimeMillis();
396         for (int i = 0; i < weekGroupWrapper.getEntriesCount(); i++) {
397             RateLimiterRecordsWrapper.EntryGroupWrapper.Entry entry =
398                     weekGroupWrapper.getEntries(i);
399             // Check if this timestamp fits the time range for each records collection.
400             if (entry.getTimestamp() > currentTimeMillis - mPastRunsHour.mTimeRangeMs) {
401                 mPastRunsHour.add(entry.getUid(), entry.getCost(), entry.getTimestamp());
402             }
403             if (entry.getTimestamp() > currentTimeMillis - mPastRunsDay.mTimeRangeMs) {
404                 mPastRunsDay.add(entry.getUid(), entry.getCost(), entry.getTimestamp());
405             }
406             if (entry.getTimestamp() > currentTimeMillis - mPastRunsWeek.mTimeRangeMs) {
407                 mPastRunsWeek.add(entry.getUid(), entry.getCost(), entry.getTimestamp());
408             }
409         }
410 
411         // Success!
412         mDataLoaded.set(true);
413     }
414 
415     /**
416      * Handle a bad persist file - this can be a file that can't be read or can't be parsed.
417      *
418      * This case is handled by attempting to delete and recreate the persist file. If this is
419      * successful, it adds some fake records to make up for potentially lost records.
420      *
421      * If the bad file is successfully remediated then RateLimiter is ready to use and no further
422      * initialization is needed.
423      *
424      * @return whether the bad file state has been successfully remediated.
425      */
426     @VisibleForTesting
handleBadFile()427     public boolean handleBadFile() {
428         if (mPersistFile == null) {
429             // This should not happen, if there is no file how can it have been determined to be
430             // bad?
431             if (DEBUG) Log.d(TAG, "Attempted to remediate a bad file but the file doesn't exist.");
432             return false;
433         }
434 
435         try {
436             // Delete the bad file, we won't likely have better luck reading it a second time.
437             mPersistFile.delete();
438             if (DEBUG) Log.d(TAG, "Deleted persist file which could not be parsed.");
439         } catch (SecurityException e) {
440             // Can't delete file so we can't recover from this state.
441             if (DEBUG) Log.d(TAG, "Failed to delete persist file", e);
442             return false;
443         }
444 
445         try {
446             if (!setupPersistFiles()) {
447                 // If setup files was unsuccessful then we won't be able to persist files.
448                 if (DEBUG) Log.d(TAG, "Failed to setup persist directory/files. Feature disabled.");
449                 return false;
450             }
451             mPersistFile.createNewFile();
452             if (!mPersistFile.exists()) {
453                 // If creating the file failed then we won't be able to persist.
454                 if (DEBUG) Log.d(TAG, "Failed to create persist file. Feature disabled.");
455                 return false;
456             }
457         } catch (SecurityException | IOException e) {
458             // Can't access/setup files.
459             if (DEBUG) Log.d(TAG, "Failed to setup persist directory/files. Feature disabled.", e);
460             return false;
461         }
462 
463         // If we made it this far then we have successfully deleted the bad file and created a new
464         // useable one - the feature is now ready to be used!
465         // However, we may have lost some records from the bad file, so add some fake records for
466         // the current time with a very high cost, this effectively disables the feature for the
467         // duration of rate limiting (1 week) to err on the cautious side regarding the potentially
468         // lost records.
469         final long timestamp = System.currentTimeMillis();
470         mPastRunsHour.add(-1 /*fake uid*/, Integer.MAX_VALUE, timestamp);
471         mPastRunsDay.add(-1 /*fake uid*/, Integer.MAX_VALUE, timestamp);
472         mPastRunsWeek.add(-1 /*fake uid*/, Integer.MAX_VALUE, timestamp);
473 
474         // Now persist the fake records.
475         maybePersistToDisk();
476 
477         // Finally, return true as we successfully remediated the bad file state.
478         return true;
479     }
480 
481     /** Update the disable rate limiter flag if present in the provided properties. */
maybeUpdateRateLimiterDisabled(DeviceConfig.Properties properties)482     public void maybeUpdateRateLimiterDisabled(DeviceConfig.Properties properties) {
483         synchronized (mLock) {
484             mRateLimiterDisabled = properties.getBoolean(
485                     DeviceConfigHelper.RATE_LIMITER_DISABLE_PROPERTY, mRateLimiterDisabled);
486         }
487     }
488 
489     /**
490      * Update DeviceConfig set configuration values if present in the provided properties, leaving
491      * not present values unchanged.
492      */
maybeUpdateConfigs(DeviceConfig.Properties properties)493     public void maybeUpdateConfigs(DeviceConfig.Properties properties) {
494         // If the field is not present in the changed properties then we want the value to stay the
495         // same, so use the current value as the default in the properties.get.
496         mPersistToDiskFrequency = properties.getLong(
497                 DeviceConfigHelper.PERSIST_TO_DISK_FREQUENCY_MS, mPersistToDiskFrequency);
498         mCostJavaHeapDump = properties.getInt(DeviceConfigHelper.COST_JAVA_HEAP_DUMP,
499                 mCostJavaHeapDump);
500         mCostHeapProfile = properties.getInt(DeviceConfigHelper.COST_HEAP_PROFILE,
501                 mCostHeapProfile);
502         mCostStackSampling = properties.getInt(DeviceConfigHelper.COST_STACK_SAMPLING,
503                 mCostStackSampling);
504         mCostSystemTrace = properties.getInt(DeviceConfigHelper.COST_SYSTEM_TRACE,
505                 mCostSystemTrace);
506 
507         // For max cost values, set a invalid default value and pass through to each group wrapper
508         // to determine whether to update values.
509         mPastRunsHour.maybeUpdateMaxCosts(
510                 properties.getInt(DeviceConfigHelper.MAX_COST_SYSTEM_1_HOUR, -1),
511                 properties.getInt(DeviceConfigHelper.MAX_COST_PROCESS_1_HOUR, -1));
512         mPastRunsDay.maybeUpdateMaxCosts(
513                 properties.getInt(DeviceConfigHelper.MAX_COST_SYSTEM_24_HOUR, -1),
514                 properties.getInt(DeviceConfigHelper.MAX_COST_PROCESS_24_HOUR, -1));
515         mPastRunsWeek.maybeUpdateMaxCosts(
516                 properties.getInt(DeviceConfigHelper.MAX_COST_SYSTEM_7_DAY, -1),
517                 properties.getInt(DeviceConfigHelper.MAX_COST_PROCESS_7_DAY, -1));
518     }
519 
statusToResult(@ateLimitResult int resultStatus)520     static int statusToResult(@RateLimitResult int resultStatus) {
521         switch (resultStatus) {
522             case RATE_LIMIT_RESULT_BLOCKED_PROCESS:
523                 return ProfilingResult.ERROR_FAILED_RATE_LIMIT_PROCESS;
524             case RATE_LIMIT_RESULT_BLOCKED_SYSTEM:
525                 return ProfilingResult.ERROR_FAILED_RATE_LIMIT_SYSTEM;
526             default:
527                 return ProfilingResult.ERROR_UNKNOWN;
528         }
529     }
530 
531     /**
532      * Create the directory and initialize the file variable for persisting records.
533      *
534      * @return Whether the files were successfully created.
535      */
536     @VisibleForTesting
setupPersistFiles()537     public boolean setupPersistFiles() throws SecurityException {
538         File dataDir = Environment.getDataDirectory();
539         File systemDir = new File(dataDir, "system");
540         mPersistStoreDir = new File(systemDir, RATE_LIMITER_STORE_DIR);
541         if (createDir(mPersistStoreDir)) {
542             mPersistFile = new File(mPersistStoreDir, RATE_LIMITER_INFO_FILE);
543             return true;
544         }
545         return false;
546     }
547 
createDir(File dir)548     private static boolean createDir(File dir) throws SecurityException {
549         if (dir.mkdir()) {
550             return true;
551         }
552 
553         if (dir.exists()) {
554             return dir.isDirectory();
555         }
556 
557         return false;
558     }
559 
560     public static final class EntryGroupWrapper {
561         private final Object mLock = new Object();
562 
563         @GuardedBy("mLock")
564         final Queue<CollectionEntry> mEntries;
565         // uid indexed
566         final SparseIntArray mPerUidCost;
567         final long mTimeRangeMs;
568 
569         int mMaxCost;
570         int mMaxCostPerUid;
571         int mTotalCost;
572 
EntryGroupWrapper(int maxCost, int maxPerUidCost, final long timeRangeMs)573         EntryGroupWrapper(int maxCost, int maxPerUidCost, final long timeRangeMs) {
574             synchronized (mLock) {
575                 mMaxCost = maxCost;
576                 mMaxCostPerUid = maxPerUidCost;
577                 mTimeRangeMs = timeRangeMs;
578                 mEntries = new ArrayDeque<>();
579                 mPerUidCost = new SparseIntArray();
580             }
581         }
582 
583         /** Update max per system and process costs if values are valid (>=0). */
maybeUpdateMaxCosts(int maxCost, int maxPerUidCost)584         public void maybeUpdateMaxCosts(int maxCost, int maxPerUidCost) {
585             synchronized (mLock) {
586                 if (maxCost >= 0) {
587                     mMaxCost = maxCost;
588                 }
589                 if (maxPerUidCost >= 0) {
590                     mMaxCostPerUid = maxPerUidCost;
591                 }
592             }
593         }
594 
595         /** Add a record and update cached costs accordingly. */
add(final int uid, final int cost, final long timestamp)596         public void add(final int uid, final int cost, final long timestamp) {
597             synchronized (mLock) {
598                 mTotalCost += cost;
599                 final int index = mPerUidCost.indexOfKey(uid);
600                 if (index < 0) {
601                     mPerUidCost.put(uid, cost);
602                 } else {
603                     mPerUidCost.put(uid, mPerUidCost.valueAt(index) + cost);
604                 }
605                 mEntries.offer(new CollectionEntry(uid, cost, timestamp));
606             }
607         }
608 
609         /**
610          * Clean up the queue by removing entries that are too old.
611          */
cleanUpOldRecords()612         public void cleanUpOldRecords() {
613             removeOlderThan(System.currentTimeMillis() - mTimeRangeMs);
614         }
615 
616         /**
617          * Clean up the queue by removing entries that are too old.
618          *
619          * @param olderThanTimestamp timestamp to remove record which are older than.
620          */
removeOlderThan(final long olderThanTimestamp)621         public void removeOlderThan(final long olderThanTimestamp) {
622             synchronized (mLock) {
623                 while (mEntries.peek() != null
624                         && mEntries.peek().mTimestamp <= olderThanTimestamp) {
625                     final CollectionEntry entry = mEntries.poll();
626                     if (entry == null) {
627                         return;
628                     }
629                     mTotalCost -= entry.mCost;
630                     if (mTotalCost < 0) {
631                         mTotalCost = 0;
632                     }
633                     final int index = mPerUidCost.indexOfKey(entry.mUid);
634                     if (index >= 0) {
635                         mPerUidCost.setValueAt(index, Math.max(0,
636                                 mPerUidCost.valueAt(index) - entry.mCost));
637                     }
638                 }
639             }
640         }
641 
642         /**
643          * Check if the requested profiling is allowed by the limits of this collection after
644          * ensuring the collection is up to date.
645          *
646          * @param uid of requesting process
647          * @param cost calculated perf cost of running this query
648          * @param currentTimeMillis cache time and keep consistent across checks
649          * @return status indicating whether request is allowed, or which rate limiting applied to
650          *         deny it.
651          */
isProfilingAllowed(final int uid, final int cost, final long currentTimeMillis)652         @RateLimitResult int isProfilingAllowed(final int uid, final int cost,
653                 final long currentTimeMillis) {
654             synchronized (mLock) {
655                 removeOlderThan(currentTimeMillis - mTimeRangeMs);
656                 if (mTotalCost + cost > mMaxCost) {
657                     return RATE_LIMIT_RESULT_BLOCKED_SYSTEM;
658                 }
659                 if (mPerUidCost.get(uid, 0) + cost > mMaxCostPerUid) {
660                     return RATE_LIMIT_RESULT_BLOCKED_PROCESS;
661                 }
662                 return RATE_LIMIT_RESULT_ALLOWED;
663             }
664         }
665 
toProto()666         RateLimiterRecordsWrapper.EntryGroupWrapper toProto() {
667             synchronized (mLock) {
668                 RateLimiterRecordsWrapper.EntryGroupWrapper.Builder builder =
669                         RateLimiterRecordsWrapper.EntryGroupWrapper.newBuilder();
670 
671                 CollectionEntry[] entries = mEntries.toArray(new CollectionEntry[mEntries.size()]);
672                 for (int i = 0; i < entries.length; i++) {
673                     builder.addEntries(entries[i].toProto());
674                 }
675 
676                 return builder.build();
677             }
678         }
679 
populateFromProto(RateLimiterRecordsWrapper.EntryGroupWrapper group)680         void populateFromProto(RateLimiterRecordsWrapper.EntryGroupWrapper group) {
681             synchronized (mLock) {
682                 final long currentTimeMillis = System.currentTimeMillis();
683                 for (int i = 0; i < group.getEntriesCount(); i++) {
684                     RateLimiterRecordsWrapper.EntryGroupWrapper.Entry entry = group.getEntries(i);
685                     if (entry.getTimestamp() > currentTimeMillis - mTimeRangeMs) {
686                         add(entry.getUid(), entry.getCost(), entry.getTimestamp());
687                     }
688                 }
689             }
690         }
691 
692         /** Get a copied array of the backing data. */
getEntriesCopy()693         public CollectionEntry[] getEntriesCopy() {
694             synchronized (mLock) {
695                 CollectionEntry[] array = new CollectionEntry[mEntries.size()];
696                 array = mEntries.toArray(array);
697                 return array.clone();
698             }
699         }
700     }
701 
702     public static final class CollectionEntry {
703         public final int mUid;
704         public final int mCost;
705         public final Long mTimestamp;
706 
CollectionEntry(final int uid, final int cost, final Long timestamp)707         CollectionEntry(final int uid, final int cost, final Long timestamp) {
708             mUid = uid;
709             mCost = cost;
710             mTimestamp = timestamp;
711         }
712 
toProto()713         RateLimiterRecordsWrapper.EntryGroupWrapper.Entry toProto() {
714             return RateLimiterRecordsWrapper.EntryGroupWrapper.Entry.newBuilder()
715                     .setUid(mUid)
716                     .setCost(mCost)
717                     .setTimestamp(mTimestamp)
718                     .build();
719         }
720     }
721 
722     public interface HandlerCallback {
723 
724         /** Obtain a handler to schedule persisting records to disk. */
obtainHandler()725         Handler obtainHandler();
726     }
727 }
728 
729