1 /*
2  * Copyright (C) 2022 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.launcher3.util.rule;
17 
18 import android.os.SystemClock;
19 
20 import androidx.test.InstrumentationRegistry;
21 
22 import org.junit.rules.TestRule;
23 import org.junit.runner.Description;
24 import org.junit.runners.model.Statement;
25 
26 import java.io.BufferedOutputStream;
27 import java.io.File;
28 import java.io.FileOutputStream;
29 import java.io.IOException;
30 import java.io.OutputStreamWriter;
31 import java.text.SimpleDateFormat;
32 import java.util.Date;
33 
34 /**
35  * A rule that generates a file that helps diagnosing cases when the test process was terminated
36  * because the test execution took too long, and tests that ran for too long even without being
37  * terminated. If the process was terminated or the test was long, the test leaves an artifact with
38  * stack traces of all threads, every SAMPLE_INTERVAL_MS. This will help understanding where we
39  * stuck.
40  */
41 public class SamplerRule implements TestRule {
42     private static final int TOO_LONG_TEST_MS = 180000;
43     private static final int SAMPLE_INTERVAL_MS = 3000;
44 
startThread(Description description)45     public static Thread startThread(Description description) {
46         Thread thread =
47                 new Thread() {
48                     @Override
49                     public void run() {
50                         // Write all-threads stack stace every SAMPLE_INTERVAL_MS while the test
51                         // is running.
52                         // After the test finishes, delete that file. If the test process is
53                         // terminated due to timeout, the trace file won't be deleted.
54                         final File file = getFile();
55 
56                         final long startTime = SystemClock.elapsedRealtime();
57                         try (OutputStreamWriter outputStreamWriter =
58                                      new OutputStreamWriter(
59                                              new BufferedOutputStream(
60                                                      new FileOutputStream(file)))) {
61                             writeSamples(outputStreamWriter);
62                         } catch (IOException | InterruptedException e) {
63                             // Simply suppressing the exceptions, nothing to do here.
64                         } finally {
65                             // If the process is not killed, then there was no test timeout, and
66                             // we are not interested in the trace file, unless the test ran too
67                             // long.
68                             if (SystemClock.elapsedRealtime() - startTime < TOO_LONG_TEST_MS) {
69                                 file.delete();
70                             }
71                         }
72                     }
73 
74                     private File getFile() {
75                         final String strDate = new SimpleDateFormat("HH:mm:ss").format(new Date());
76 
77                         final String descStr = description.getTestClass().getSimpleName() + "."
78                                 + description.getMethodName();
79                         return artifactFile(
80                                 "ThreadStackSamples-" + strDate + "-" + descStr + ".txt");
81                     }
82 
83                     private void writeSamples(OutputStreamWriter writer)
84                             throws IOException, InterruptedException {
85                         int count = 0;
86                         while (true) {
87                             writer.write(
88                                     "#"
89                                             + (count++)
90                                             + " =============================================\r\n");
91                             for (StackTraceElement[] stack : getAllStackTraces().values()) {
92                                 writer.write("---------------------\r\n");
93                                 for (StackTraceElement frame : stack) {
94                                     writer.write(frame.toString() + "\r\n");
95                                 }
96                             }
97                             writer.flush();
98 
99                             sleep(SAMPLE_INTERVAL_MS);
100                         }
101                     }
102                 };
103 
104         thread.start();
105         return thread;
106     }
107 
108     @Override
apply(Statement base, Description description)109     public Statement apply(Statement base, Description description) {
110         return new Statement() {
111             @Override
112             public void evaluate() throws Throwable {
113                 final Thread traceThread = startThread(description);
114                 try {
115                     base.evaluate();
116                 } finally {
117                     traceThread.interrupt();
118                     traceThread.join();
119                 }
120             }
121         };
122     }
123 
124     private static File artifactFile(String fileName) {
125         return new File(
126                 InstrumentationRegistry.getInstrumentation().getTargetContext().getFilesDir(),
127                 fileName);
128     }
129 }
130