1 /*
2  * Copyright (C) 2017 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 package com.android.server.testutils;
17 
18 
19 import static android.util.ExceptionUtils.appendCause;
20 import static android.util.ExceptionUtils.propagate;
21 
22 import android.os.Handler;
23 import android.os.Looper;
24 import android.os.Message;
25 import android.os.SystemClock;
26 import android.util.ArrayMap;
27 
28 import java.util.Map;
29 import java.util.PriorityQueue;
30 import java.util.function.LongSupplier;
31 
32 /**
33  * A test {@link Handler} that stores incoming {@link Message}s and {@link Runnable callbacks}
34  * in a {@link PriorityQueue} based on time, to be manually processed later in a correct order
35  * either all together with {@link #flush}, or only those due at the current time with
36  * {@link #timeAdvance}.
37  *
38  * For the latter use case this also supports providing a custom clock (in a format of a
39  * milliseconds-returning {@link LongSupplier}), that will be used for storing the messages'
40  * timestamps to be posted at, and checked against during {@link #timeAdvance}.
41  *
42  * This allows to test code that uses {@link Handler}'s delayed invocation capabilities, such as
43  * {@link Handler#sendMessageDelayed} or {@link Handler#postDelayed} without resorting to
44  * synchronously {@link Thread#sleep}ing in your test.
45  *
46  * @see OffsettableClock for a useful custom clock implementation to use with this handler
47  */
48 public class TestHandler extends Handler {
49     private static final LongSupplier DEFAULT_CLOCK = SystemClock::uptimeMillis;
50 
51     private final PriorityQueue<MsgInfo> mMessages = new PriorityQueue<>();
52     /**
53      * Map of: {@code message id -> count of such messages currently pending }
54      */
55     // Boxing is ok here - both msg ids and their pending counts tend to be well below 128
56     private final Map<Integer, Integer> mPendingMsgTypeCounts = new ArrayMap<>();
57     private final LongSupplier mClock;
58     private int  mMessageCount = 0;
59 
TestHandler(Callback callback)60     public TestHandler(Callback callback) {
61         this(callback, DEFAULT_CLOCK);
62     }
63 
TestHandler(Callback callback, LongSupplier clock)64     public TestHandler(Callback callback, LongSupplier clock) {
65         this(Looper.getMainLooper(), callback, clock);
66     }
67 
TestHandler(Looper looper, Callback callback, LongSupplier clock)68     public TestHandler(Looper looper, Callback callback, LongSupplier clock) {
69         super(looper, callback);
70         mClock = clock;
71     }
72 
73     @Override
sendMessageAtTime(Message msg, long uptimeMillis)74     public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
75         ++mMessageCount;
76         mPendingMsgTypeCounts.put(msg.what,
77                 mPendingMsgTypeCounts.getOrDefault(msg.what, 0) + 1);
78 
79         // uptimeMillis is an absolute time obtained as SystemClock.uptimeMillis() + offsetMillis
80         // if custom clock is given, recalculate the time with regards to it
81         if (mClock != DEFAULT_CLOCK) {
82             uptimeMillis = uptimeMillis - SystemClock.uptimeMillis() + mClock.getAsLong();
83         }
84 
85         // post a dummy queue entry to keep track of message removal
86         return super.sendMessageAtTime(msg, Long.MAX_VALUE)
87                 && mMessages.add(new MsgInfo(Message.obtain(msg), uptimeMillis, mMessageCount));
88     }
89 
90     /** @see TestHandler */
timeAdvance()91     public void timeAdvance() {
92         long now = mClock.getAsLong();
93         while (!mMessages.isEmpty() && mMessages.peek().sendTime <= now) {
94             dispatch(mMessages.poll());
95         }
96     }
97 
98     /**
99      * Dispatch all messages in order
100      *
101      * @see TestHandler
102      */
flush()103     public void flush() {
104         MsgInfo msg;
105         while ((msg = mMessages.poll()) != null) {
106             dispatch(msg);
107         }
108     }
109 
getPendingMessages()110     public PriorityQueue<MsgInfo> getPendingMessages() {
111         return new PriorityQueue<>(mMessages);
112     }
113 
114     /**
115      * Optionally-overridable to allow deciphering message types
116      *
117      * @see android.util.DebugUtils#valueToString - a handy utility to use when overriding this
118      */
messageToString(Message message)119     protected String messageToString(Message message) {
120         return message.toString();
121     }
122 
dispatch(MsgInfo msg)123     private void dispatch(MsgInfo msg) {
124         int msgId = msg.message.what;
125 
126         if (!hasMessages(msgId)) {
127             // Handler.removeMessages(msgId) must have been called
128             return;
129         }
130 
131         try {
132             Integer pendingMsgCount = mPendingMsgTypeCounts.getOrDefault(msgId, 0);
133             if (pendingMsgCount <= 1) {
134                 removeMessages(msgId);
135             }
136             mPendingMsgTypeCounts.put(msgId, pendingMsgCount - 1);
137 
138             dispatchMessage(msg.message);
139         } catch (Throwable t) {
140             // Append stack trace of this message being posted as a cause for a helpful
141             // test error message
142             throw propagate(appendCause(t, msg.postPoint));
143         } finally {
144             msg.message.recycle();
145         }
146     }
147 
148     public class MsgInfo implements Comparable<MsgInfo> {
149         public final Message message;
150         public final long sendTime;
151         public final int mMessageOrder;
152         public final RuntimeException postPoint;
153 
MsgInfo(Message message, long sendTime, int messageOrder)154         private MsgInfo(Message message, long sendTime, int messageOrder) {
155             this.message = message;
156             this.sendTime = sendTime;
157             this.postPoint = new RuntimeException("Message originated from here:");
158             mMessageOrder = messageOrder;
159         }
160 
161         @Override
compareTo(MsgInfo o)162         public int compareTo(MsgInfo o) {
163             final int result = Long.compare(sendTime, o.sendTime);
164             return result != 0 ? result : Integer.compare(mMessageOrder, o.mMessageOrder);
165         }
166 
167         @Override
toString()168         public String toString() {
169             return "MsgInfo{" +
170                     "message =" + messageToString(message)
171                     + ", sendTime =" + sendTime
172                     + ", mMessageOrder =" + mMessageOrder
173                     + '}';
174         }
175     }
176 }
177