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.tradefed.invoker.tracing;
17 
18 
19 import java.io.File;
20 import java.util.Collections;
21 import java.util.HashMap;
22 import java.util.Map;
23 
24 import javax.annotation.Nullable;
25 
26 /** Class that helps to manage tracing for each test invocation. */
27 public class TracingLogger {
28 
29     private static final Map<ThreadGroup, ActiveTrace> mPerGroupActiveTrace =
30             Collections.synchronizedMap(new HashMap<ThreadGroup, ActiveTrace>());
31 
32     /**
33      * Creates and register an active trace for an invocation.
34      *
35      * @param pid Current process id
36      * @param tid Current thread id
37      */
createActiveTrace(long pid, long tid)38     public static ActiveTrace createActiveTrace(long pid, long tid) {
39         return createActiveTrace(pid, tid, false);
40     }
41 
createActiveTrace(long pid, long tid, boolean mainProcess)42     public static ActiveTrace createActiveTrace(long pid, long tid, boolean mainProcess) {
43         ThreadGroup group = Thread.currentThread().getThreadGroup();
44         synchronized (mPerGroupActiveTrace) {
45             ActiveTrace trace = new ActiveTrace(pid, tid, mainProcess);
46             mPerGroupActiveTrace.put(group, trace);
47             return trace;
48         }
49     }
50 
51     /** If it exists, returns the current trace of the Tradefed process itself. */
getMainTrace()52     public static ActiveTrace getMainTrace() {
53         synchronized (mPerGroupActiveTrace) {
54             for (ActiveTrace t : mPerGroupActiveTrace.values()) {
55                 if (t != null && t.isMainTradefedProcess()) {
56                     return t;
57                 }
58             }
59             return null;
60         }
61     }
62 
63     private static ThreadLocal<ThreadGroup> sLocal = new ThreadLocal<>();
64 
65     /** Tracks a localized context when using the properties inside the gRPC server */
setLocalGroup(ThreadGroup tg)66     public static void setLocalGroup(ThreadGroup tg) {
67         sLocal.set(tg);
68     }
69 
70     /** Resets the localized context. */
resetLocalGroup()71     public static void resetLocalGroup() {
72         sLocal.remove();
73     }
74 
75     /**
76      * Sets the currently active trace for an invocation.
77      *
78      * @return the previous active trace or {@code null} if there was none.
79      */
80     @Nullable
setActiveTrace(ActiveTrace trace)81     static ActiveTrace setActiveTrace(ActiveTrace trace) {
82         ThreadGroup group = Thread.currentThread().getThreadGroup();
83         synchronized (mPerGroupActiveTrace) {
84             return mPerGroupActiveTrace.put(group, trace);
85         }
86     }
87 
88     /** Returns the current active trace for the invocation, or null if none. */
getActiveTrace()89     public static ActiveTrace getActiveTrace() {
90         ThreadGroup group = Thread.currentThread().getThreadGroup();
91         synchronized (mPerGroupActiveTrace) {
92             if (sLocal.get() != null) {
93                 group = sLocal.get();
94             }
95             return mPerGroupActiveTrace.get(group);
96         }
97     }
98 
getActiveTraceForGroup(ThreadGroup group)99     public static ActiveTrace getActiveTraceForGroup(ThreadGroup group) {
100         if (group == null) {
101             return null;
102         }
103         synchronized (mPerGroupActiveTrace) {
104             return mPerGroupActiveTrace.get(group);
105         }
106     }
107 
108     /** Finalize the tracing and clear the tracking. */
finalizeTrace()109     public static File finalizeTrace() {
110         ThreadGroup group = Thread.currentThread().getThreadGroup();
111         synchronized (mPerGroupActiveTrace) {
112             ActiveTrace trace = mPerGroupActiveTrace.remove(group);
113             if (trace != null) {
114                 return trace.finalizeTracing();
115             }
116         }
117         return null;
118     }
119 }
120