1# Copyright (C) 2022 The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15import dataclasses as dc 16from typing import BinaryIO, Dict, Generator, List, Type, Union 17 18from perfetto.trace_uri_resolver import util 19 20TraceUri = str 21TraceGenerator = Generator[bytes, None, None] 22TraceContent = Union[BinaryIO, TraceGenerator] 23 24 25class TraceUriResolver: 26 """"Resolves a trace URI (e.g. 'ants:trace_id=1234') into a list of traces. 27 28 This class can be subclassed to provide a pluggable mechanism for looking 29 up traces using URI strings. 30 31 For example: 32 class CustomTraceResolver(TraceResolver): 33 PREFIX = 'custom' 34 35 def __init__(self, build_branch: List[str] = None, id: str = None): 36 self.build_branch = build_branch 37 self.id = id 38 self.db = init_db() 39 40 def resolve(self): 41 traces = self.db.lookup( 42 id=self.id, build_branch=self.build_branch)['path'] 43 return [ 44 TraceResolver.Result( 45 trace=t['path'], 46 args={'iteration': t['iteration'], 'device': t['device']} 47 ) 48 for t in traces 49 ] 50 51 Trace resolvers can be passed to trace processor directly: 52 with TraceProcessor(CustomTraceResolver(id='abcdefg')) as tp: 53 tp.query('select * from slice') 54 55 Alternatively, a trace addesses can be passed: 56 config = TraceProcessorConfig( 57 resolver_registry=ResolverRegistry(resolvers=[CustomTraceResolver]) 58 ) 59 with TraceProcessor('custom:id=abcdefg', config=config) as tp: 60 tp.query('select * from slice') 61 """ 62 63 # Subclasses should set PREFIX to match the trace address prefix they 64 # want to handle. 65 PREFIX: str = None 66 67 @dc.dataclass 68 class Result: 69 # TraceUri is present here because it allows recursive lookups (i.e. 70 # a resolver which returns a path to a trace). 71 trace: Union[TraceUri, TraceContent] 72 73 # metadata allows additional key-value pairs to be provided which are 74 # associated for trace. For example, test names and iteration numbers 75 # could be provivded for traces originating from lab tests. 76 metadata: Dict[str, str] 77 78 def __init__(self, 79 trace: Union[TraceUri, TraceContent], 80 metadata: Dict[str, str] = dict()): 81 self.trace = trace 82 self.metadata = metadata 83 84 def resolve(self) -> List['TraceUriResolver.Result']: 85 """Resolves a list of traces. 86 87 Subclasses should implement this method and resolve the parameters 88 specified in the constructor to a list of traces.""" 89 raise Exception("resolve is unimplemented for this resolver") 90 91 @classmethod 92 def from_trace_uri(cls: Type['TraceUriResolver'], 93 uri: TraceUri) -> 'TraceUriResolver': 94 """Creates a resolver from a URI. 95 96 URIs have the form: 97 android_ci:day=2021-01-01;devices=blueline,crosshatch 98 99 This is converted to a dictionary of the form: 100 {'day': '2021-01-01', 'id': ['blueline', 'crosshatch']} 101 102 and passed as kwargs to the constructor of the trace resolver (see class 103 documentation for info). 104 105 Generally, sublcasses should not override this method as the standard 106 trace address format should work for most usecases. Instead, simply 107 define your constructor with the parameters you expect to see in the 108 trace address.""" 109 return cls(**_args_dict_from_uri(uri)) 110 111 112def _args_dict_from_uri(uri: str) -> Dict[str, str]: 113 """Creates an the args dictionary from a trace URI. 114 115 URIs have the form: 116 android_ci:day=2021-01-01;devices=blueline,crosshatch 117 118 This is converted to a dictionary of the form: 119 {'day': '2021-01-01', 'id': ['blueline', 'crosshatch']} 120 """ 121 _, args_str = util.parse_trace_uri(uri) 122 if not args_str: 123 return {} 124 125 args_lst = args_str.split(';') 126 args_dict = dict() 127 for arg in args_lst: 128 (key, value) = arg.split('=') 129 lst = value.split(',') 130 if len(lst) > 1: 131 args_dict[key] = lst 132 else: 133 args_dict[key] = value 134 return args_dict 135