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