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 16import enum 17from typing import BinaryIO, Dict, Generator, List, Type, Union 18from typing import Generic, Tuple, TypeVar, get_type_hints 19 20from perfetto.trace_uri_resolver import util 21 22TraceUri = str 23TraceGenerator = Generator[bytes, None, None] 24TraceContent = Union[BinaryIO, TraceGenerator] 25_T = TypeVar('_T') 26 27 28@dc.dataclass 29class ConstraintClass(Generic[_T]): 30 31 class Op(enum.Enum): 32 EQ = '=' 33 NE = '!=' 34 LE = '<=' 35 GE = '>=' 36 GT = '>' 37 LT = '<' 38 39 def __str__(self): 40 return self.value 41 42 value: _T 43 op: Op = Op.EQ 44 45 46Constraint = Union[_T, ConstraintClass[_T]] 47ConstraintWithList = Union[Constraint[_T], Constraint[List[_T]]] 48 49 50class TraceUriResolver: 51 """"Resolves a trace URI (e.g. 'ants:trace_id=1234') into a list of traces. 52 53 This class can be subclassed to provide a pluggable mechanism for looking 54 up traces using URI strings. 55 56 For example: 57 class CustomTraceResolver(TraceUriResolver): 58 PREFIX = 'custom' 59 60 def __init__(self, build_branch: List[str] = None, id: str = None): 61 self.build_branch = build_branch 62 self.id = id 63 self.db = init_db() 64 65 def resolve(self): 66 traces = self.db.lookup( 67 id=self.id, build_branch=self.build_branch)['path'] 68 return [ 69 TraceUriResolver.Result( 70 trace=t['path'], 71 args={'iteration': t['iteration'], 'device': t['device']} 72 ) 73 for t in traces 74 ] 75 76 Trace resolvers can be passed to trace processor directly: 77 with TraceProcessor(CustomTraceResolver(id='abcdefg')) as tp: 78 tp.query('select * from slice') 79 80 Alternatively, a trace addesses can be passed: 81 config = TraceProcessorConfig( 82 resolver_registry=ResolverRegistry(resolvers=[CustomTraceResolver]) 83 ) 84 with TraceProcessor('custom:id=abcdefg', config=config) as tp: 85 tp.query('select * from slice') 86 """ 87 88 # Subclasses should set PREFIX to match the trace address prefix they 89 # want to handle. 90 PREFIX: str = None 91 92 @dc.dataclass 93 class Result: 94 # TraceUri is present here because it allows recursive lookups (i.e. 95 # a resolver which returns a path to a trace). 96 trace: Union[TraceUri, TraceContent] 97 98 # metadata allows additional key-value pairs to be provided which are 99 # associated for trace. For example, test names and iteration numbers 100 # could be provivded for traces originating from lab tests. 101 metadata: Dict[str, str] 102 103 def __init__(self, 104 trace: Union[TraceUri, TraceContent], 105 metadata: Dict[str, str] = dict()): 106 self.trace = trace 107 self.metadata = metadata 108 109 def resolve(self) -> List['TraceUriResolver.Result']: 110 """Resolves a list of traces. 111 112 Subclasses should implement this method and resolve the parameters 113 specified in the constructor to a list of traces. 114 """ 115 raise Exception('resolve is unimplemented for this resolver') 116 117 @classmethod 118 def from_trace_uri(cls: Type['TraceUriResolver'], 119 uri: TraceUri) -> 'TraceUriResolver': 120 """Creates a resolver from a URI. 121 122 URIs have the form: 123 android_ci:day=2021-01-01;devices=blueline,crosshatch;key>=value 124 125 This is converted to a dictionary of the form: 126 {'day': '2021-01-01', 'id': ['blueline', 'crosshatch'], 127 'key': ConstraintClass('value', Op.GE)} 128 129 and passed as kwargs to the constructor of the trace resolver (see class 130 documentation for info). 131 132 Generally, sublcasses should not override this method as the standard 133 trace address format should work for most usecases. Instead, simply 134 define your constructor with the parameters you expect to see in the 135 trace address. 136 """ 137 return cls(**_args_dict_from_uri(uri, get_type_hints(cls.__init__))) 138 139 140def _read_op(arg_str: str, op_start_ind: int) -> ConstraintClass.Op: 141 """Parse operator from string. 142 143 Given string and an expected start index for operator it returns Op object or 144 raises error if operator was not found. 145 146 For example: 147 _read_op('a>4', 1) returns Op.GE 148 _read_op('a>4', 0) raises ValueError 149 _read_op('a>4', 3) raises ValueError 150 """ 151 first = arg_str[op_start_ind] if op_start_ind < len(arg_str) else None 152 second = arg_str[op_start_ind + 153 1] if op_start_ind + 1 < len(arg_str) else None 154 Op = ConstraintClass.Op 155 if first == '>': 156 return Op.GE if second == '=' else Op.GT 157 elif first == '<': 158 return Op.LE if second == '=' else Op.LT 159 elif first == '!' and second == '=': 160 return Op.NE 161 elif first == '=': 162 return Op.EQ 163 raise ValueError('Could not find valid operator in uri arg_str: ' + arg_str) 164 165 166def _parse_arg(arg_str: str) -> Tuple[str, ConstraintClass.Op, str]: 167 """Parse argument string and return a tuple (key, operator, value). 168 169 Given a string like 'branch_num>=4000', it returns a tuple ('branch_num', 170 Op.GE,'4000'). Raises ValueError exceptions in case ill formed arg_str is 171 passed like '>30', 'key>', 'key', 'key--31' 172 """ 173 op_start_ind = 0 174 for ind, c in enumerate(arg_str): 175 if not c.isalnum() and c != '_': 176 op_start_ind = ind 177 break 178 if op_start_ind == 0: 179 raise ValueError('Could not find valid key in arg_str: ' + arg_str) 180 key = arg_str[:op_start_ind] 181 op = _read_op(arg_str, op_start_ind) 182 value = arg_str[op_start_ind + len(str(op)):] 183 if not value: 184 raise ValueError('Empty value in trace uri arg_str: ' + arg_str) 185 return (key, op, value) 186 187 188def _args_dict_from_uri(uri: str, 189 type_hints) -> Dict[str, ConstraintWithList[str]]: 190 """Creates an the args dictionary from a trace URI. 191 192 URIs have the form: 193 android_ci:day=2021-01-01;devices=blueline,crosshatch;key>=value 194 195 This is converted to a dictionary of the form: 196 {'day': '2021-01-01', 'id': ['blueline', 'crosshatch'], 197 'key': ConstraintClass('value', Op.GE)} 198 """ 199 _, args_str = util.parse_trace_uri(uri) 200 if not args_str: 201 return {} 202 203 args_lst = args_str.split(';') 204 args_dict = dict() 205 for arg in args_lst: 206 (key, op, value) = _parse_arg(arg) 207 lst = value.split(',') 208 if len(lst) > 1: 209 args_dict[key] = lst 210 else: 211 args_dict[key] = value 212 213 if key not in type_hints: 214 if op != ConstraintClass.Op.EQ: 215 raise ValueError(f'{key} only supports "=" operator') 216 continue 217 have_constraint = False 218 type_hint = type_hints[key] 219 type_args = type_hint.__args__ if hasattr(type_hint, '__args__') else () 220 for type_arg in type_args: 221 type_origin = type_arg.__origin__ if hasattr(type_arg, 222 '__origin__') else None 223 if type_origin is ConstraintClass: 224 have_constraint = True 225 break 226 if not have_constraint and op != ConstraintClass.Op.EQ: 227 raise ValueError('Operator other than "=" passed to argument which ' 228 'does not have constraint type: ' + arg) 229 if have_constraint: 230 args_dict[key] = ConstraintClass(args_dict[key], op) 231 return args_dict 232