src/python/bcc/usdt.py

   1 # Copyright 2016 Sasha Goldshtein
   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
  15 from __future__ import print_function
  16 import ctypes as ct
  17 import sys
  18 from .libbcc import lib, _USDT_CB, _USDT_PROBE_CB, \
  19                     bcc_usdt_location, bcc_usdt_argument, \
  20                     BCC_USDT_ARGUMENT_FLAGS
  21
  22 class USDTException(Exception):
  23     pass
  24
  25 class USDTProbeArgument(object):
  26     def __init__(self, argument):
  27         self.signed = argument.size < 0
  28         self.size = abs(argument.size)
  29         self.valid = argument.valid
  30         if self.valid & BCC_USDT_ARGUMENT_FLAGS.CONSTANT != 0:
  31             self.constant = argument.constant
  32         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET != 0:
  33             self.deref_offset = argument.deref_offset
  34         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT != 0:
  35             self.deref_ident = argument.deref_ident
  36         if self.valid & BCC_USDT_ARGUMENT_FLAGS.BASE_REGISTER_NAME != 0:
  37             self.base_register_name = argument.base_register_name
  38         if self.valid & BCC_USDT_ARGUMENT_FLAGS.INDEX_REGISTER_NAME != 0:
  39             self.index_register_name = argument.index_register_name
  40         if self.valid & BCC_USDT_ARGUMENT_FLAGS.SCALE != 0:
  41             self.scale = argument.scale
  42
  43     def _size_prefix(self):
  44         return "%d %s bytes" % \
  45                 (self.size, "signed  " if self.signed else "unsigned")
  46
  47     def _format(self):
  48         # This mimics the logic in cc/usdt_args.cc that gives meaning to the
  49         # various argument settings. A change there will require a change here.
  50         if self.valid & BCC_USDT_ARGUMENT_FLAGS.CONSTANT != 0:
  51             return "%d" % self.constant
  52         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET == 0:
  53             return "%s" % self.base_register_name.decode()
  54         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET != 0 and \
  55            self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT == 0:
  56             if self.valid & BCC_USDT_ARGUMENT_FLAGS.INDEX_REGISTER_NAME != 0:
  57                 index_offset = " + %s" % self.index_register_name.decode()
  58                 if self.valid & BCC_USDT_ARGUMENT_FLAGS.SCALE != 0:
  59                     index_offset += " * %d" % self.scale
  60             else:
  61                 index_offset = ""
  62             sign = '+' if self.deref_offset >= 0 else '-'
  63             return "*(%s %s %d%s)" % (self.base_register_name.decode(),
  64                                       sign, abs(self.deref_offset), index_offset)
  65         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET != 0 and \
  66            self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT != 0 and \
  67            self.valid & BCC_USDT_ARGUMENT_FLAGS.BASE_REGISTER_NAME != 0 and \
  68            self.base_register_name == "ip":
  69             sign = '+' if self.deref_offset >= 0 else '-'
  70             return "*(&%s %s %d)" % (self.deref_ident.decode(),
  71                                      sign, abs(self.deref_offset))
  72         # If we got here, this is an unrecognized case. Doesn't mean it's
  73         # necessarily bad, so just provide the raw data. It just means that
  74         # other tools won't be able to work with this argument.
  75         return "unrecognized argument format, flags %d" % self.valid
  76
  77     def __str__(self):
  78         return "%s @ %s" % (self._size_prefix(), self._format())
  79
  80 class USDTProbeLocation(object):
  81     def __init__(self, probe, index, location):
  82         self.probe = probe
  83         self.index = index
  84         self.num_arguments = probe.num_arguments
  85         self.address = location.address
  86         self.bin_path = location.bin_path
  87
  88     def __str__(self):
  89         return "%s 0x%x" % (self.bin_path.decode(), self.address)
  90
  91     def get_argument(self, index):
  92         arg = bcc_usdt_argument()
  93         res = lib.bcc_usdt_get_argument(self.probe.context, self.probe.provider,
  94                                         self.probe.name,
  95                                         self.index, index, ct.byref(arg))
  96         if res != 0:
  97             raise USDTException(
  98                     "error retrieving probe argument %d location %d" %
  99                     (index, self.index))
 100         return USDTProbeArgument(arg)
 101
 102 class USDTProbe(object):
 103     def __init__(self, context, probe):
 104         self.context = context
 105         self.provider = probe.provider
 106         self.name = probe.name
 107         self.bin_path = probe.bin_path
 108         self.semaphore = probe.semaphore
 109         self.num_locations = probe.num_locations
 110         self.num_arguments = probe.num_arguments
 111
 112     def __str__(self):
 113         return "%s:%s [sema 0x%x]" % \
 114                (self.provider.decode(), self.name.decode(), self.semaphore)
 115
 116     def short_name(self):
 117         return "%s:%s" % (self.provider.decode(), self.name.decode())
 118
 119     def get_location(self, index):
 120         loc = bcc_usdt_location()
 121         res = lib.bcc_usdt_get_location(self.context, self.provider, self.name,
 122                                         index, ct.byref(loc))
 123         if res != 0:
 124             raise USDTException("error retrieving probe location %d" % index)
 125         return USDTProbeLocation(self, index, loc)
 126
 127 class USDT(object):
 128     def __init__(self, pid=None, path=None):
 129         if pid and pid != -1:
 130             self.pid = pid
 131             if path:
 132                 self.context = lib.bcc_usdt_new_frompid(pid, path.encode('ascii'))
 133             else:
 134                 self.context = lib.bcc_usdt_new_frompid(pid, ct.c_char_p(0))
 135             if self.context == None:
 136                 raise USDTException("USDT failed to instrument PID %d" % pid)
 137         elif path:
 138             self.path = path
 139             self.context = lib.bcc_usdt_new_frompath(path.encode('ascii'))
 140             if self.context == None:
 141                 raise USDTException("USDT failed to instrument path %s" % path)
 142         else:
 143             raise USDTException(
 144                     "either a pid or a binary path must be specified")
 145
 146     def __del__(self):
 147         lib.bcc_usdt_close(self.context)
 148
 149     def enable_probe(self, probe, fn_name):
 150         probe_parts = probe.split(":", 1)
 151         if len(probe_parts) == 1:
 152             ret = lib.bcc_usdt_enable_probe(
 153                 self.context, probe.encode('ascii'), fn_name.encode('ascii'))
 154         else:
 155             (provider_name, probe_name) = probe_parts
 156             ret = lib.bcc_usdt_enable_fully_specified_probe(
 157                 self.context, provider_name.encode('ascii'), probe_name.encode('ascii'),
 158                 fn_name.encode('ascii'))
 159
 160         if ret != 0:
 161             raise USDTException(
 162 """Failed to enable USDT probe '%s':
 163 the specified pid might not contain the given language's runtime,
 164 or the runtime was not built with the required USDT probes. Look
 165 for a configure flag similar to --with-dtrace or --enable-dtrace.
 166 To check which probes are present in the process, use the tplist tool.
 167 """ % probe)
 168
 169     def enable_probe_or_bail(self, probe, fn_name):
 170         try:
 171             self.enable_probe(probe, fn_name)
 172         except USDTException as e:
 173             print(e, file=sys.stderr)
 174             sys.exit(1)
 175
 176     def get_context(self):
 177         return self.context
 178
 179     def get_text(self):
 180         ctx_array = (ct.c_void_p * 1)()
 181         ctx_array[0] = ct.c_void_p(self.context)
 182         return lib.bcc_usdt_genargs(ctx_array, 1).decode()
 183
 184     def get_probe_arg_ctype(self, probe_name, arg_index):
 185         probe_parts = probe_name.split(":", 1)
 186         if len(probe_parts) == 1:
 187             return lib.bcc_usdt_get_probe_argctype(
 188                 self.context, probe_name.encode('ascii'), arg_index).decode()
 189         else:
 190             (provider_name, probe) = probe_parts
 191             return lib.bcc_usdt_get_fully_specified_probe_argctype(
 192                 self.context, provider_name.encode('ascii'), probe.encode('ascii'), arg_index).decode()
 193
 194     def enumerate_probes(self):
 195         probes = []
 196         def _add_probe(probe):
 197             probes.append(USDTProbe(self.context, probe.contents))
 198
 199         lib.bcc_usdt_foreach(self.context, _USDT_CB(_add_probe))
 200         return probes
 201
 202     # This is called by the BPF module's __init__ when it realizes that there
 203     # is a USDT context and probes need to be attached.
 204     def attach_uprobes(self, bpf, attach_usdt_ignore_pid):
 205         probes = self.enumerate_active_probes()
 206         for (binpath, fn_name, addr, pid) in probes:
 207             if attach_usdt_ignore_pid:
 208                 pid = -1
 209             bpf.attach_uprobe(name=binpath.decode(), fn_name=fn_name.decode(),
 210                               addr=addr, pid=pid)
 211
 212     def enumerate_active_probes(self):
 213         probes = []
 214         def _add_probe(binpath, fn_name, addr, pid):
 215             probes.append((binpath, fn_name, addr, pid))
 216
 217         lib.bcc_usdt_foreach_uprobe(self.context, _USDT_PROBE_CB(_add_probe))
 218         return probes