third_party/pigweed/repo/pw_allocator/py/pw_allocator/heap_viewer.py

   1 # Copyright 2020 The Pigweed Authors
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
   4 # use this file except in compliance with the License. You may obtain a copy of
   5 # the License at
   6 #
   7 #     https://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, WITHOUT
  11 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12 # License for the specific language governing permissions and limitations under
  13 # the License.
  14 """Heap visualizer of ASCII characters."""
  15
  16 import argparse
  17 import sys
  18 import math
  19 import logging
  20 from typing import Optional
  21 from dataclasses import dataclass
  22 import coloredlogs  # type: ignore
  23
  24
  25 @dataclass
  26 class HeapBlock:
  27     """Building blocks for memory chunk allocated at heap."""
  28     size: int
  29     mem_offset: int
  30     next: Optional['HeapBlock'] = None
  31
  32
  33 @dataclass
  34 class HeapUsage:
  35     """Contains a linked list of allocated HeapBlocks."""
  36     begin: HeapBlock = HeapBlock(0, 0)
  37
  38     def add_block(self, block):
  39         cur_block = self.begin.next
  40         prev_block = self.begin
  41         while cur_block is not None:
  42             if cur_block.mem_offset == block.mem_offset:
  43                 return
  44             if cur_block.mem_offset < block.mem_offset:
  45                 prev_block = cur_block
  46                 cur_block = cur_block.next
  47             else:
  48                 block.next = cur_block
  49                 prev_block.next = block
  50                 return
  51         prev_block.next = block
  52
  53     def remove_block(self, address):
  54         cur_block = self.begin.next
  55         prev_block = self.begin
  56         while cur_block is not None:
  57             if cur_block.mem_offset == address:
  58                 prev_block.next = cur_block.next
  59                 return
  60             if cur_block.mem_offset < address:
  61                 prev_block = cur_block
  62                 cur_block = cur_block.next
  63             else:
  64                 return
  65
  66
  67 def add_parser_arguments(parser):
  68     parser.add_argument('--dump-file',
  69                         help=('dump file that contains a list of malloc and '
  70                               'free instructions. The format should be as '
  71                               'follows: "m <size> <address>" on a line for '
  72                               'each malloc called and "f <address>" on a line '
  73                               'for each free called.'),
  74                         required=True)
  75
  76     parser.add_argument('--heap-low-address',
  77                         help=('lower address of the heap.'),
  78                         type=lambda x: int(x, 0),
  79                         required=True)
  80
  81     parser.add_argument('--heap-high-address',
  82                         help=('higher address of the heap.'),
  83                         type=lambda x: int(x, 0),
  84                         required=True)
  85
  86     parser.add_argument('--poison-enabled',
  87                         help=('if heap poison is enabled or not.'),
  88                         default=False,
  89                         action='store_true')
  90
  91     parser.add_argument('--pointer-size',
  92                         help=('size of pointer on the machine.'),
  93                         default=4,
  94                         type=lambda x: int(x, 0))
  95
  96
  97 _LEFT_HEADER_CHAR = '['
  98 _RIGHT_HEADER_CHAR = ']'
  99 _USED_CHAR = '*'
 100 _FREE_CHAR = ' '
 101 _CHARACTERS_PER_LINE = 64
 102 _BYTES_PER_CHARACTER = 4
 103 _LOG = logging.getLogger(__name__)
 104
 105
 106 def _exit_due_to_file_not_found():
 107     _LOG.critical('Dump file location is not provided or dump file is not '
 108                   'found. Please specify a valid file in the argument.')
 109     sys.exit(1)
 110
 111
 112 def _exit_due_to_bad_heap_info():
 113     _LOG.critical(
 114         'Heap low/high address is missing or invalid. Please put valid '
 115         'addresses in the argument.')
 116     sys.exit(1)
 117
 118
 119 def visualize(dump_file=None,
 120               heap_low_address=None,
 121               heap_high_address=None,
 122               poison_enabled=False,
 123               pointer_size=4):
 124     """Visualization of heap usage."""
 125     # TODO(pwbug/236): Add standarized mechanisms to produce dump file and read
 126     # heap information from dump file.
 127     aligned_bytes = pointer_size
 128     header_size = pointer_size * 2
 129
 130     try:
 131         if heap_high_address < heap_low_address:
 132             _exit_due_to_bad_heap_info()
 133         heap_size = heap_high_address - heap_low_address
 134     except TypeError:
 135         _exit_due_to_bad_heap_info()
 136
 137     if poison_enabled:
 138         poison_offset = pointer_size
 139     else:
 140         poison_offset = 0
 141
 142     try:
 143         allocation_dump = open(dump_file, 'r')
 144     except (FileNotFoundError, TypeError):
 145         _exit_due_to_file_not_found()
 146
 147     heap_visualizer = HeapUsage()
 148     # Parse the dump file.
 149     for line in allocation_dump:
 150         info = line[:-1].split(' ')
 151         if info[0] == 'm':
 152             # Add a HeapBlock when malloc is called
 153             block = HeapBlock(
 154                 int(math.ceil(float(info[1]) / aligned_bytes)) * aligned_bytes,
 155                 int(info[2], 0) - heap_low_address)
 156             heap_visualizer.add_block(block)
 157         elif info[0] == 'f':
 158             # Remove the HeapBlock when free is called
 159             heap_visualizer.remove_block(int(info[1], 0) - heap_low_address)
 160
 161     # next_block indicates the nearest HeapBlock that hasn't finished
 162     # printing.
 163     next_block = heap_visualizer.begin
 164     if next_block.next is None:
 165         next_mem_offset = heap_size + header_size + poison_offset + 1
 166         next_size = 0
 167     else:
 168         next_mem_offset = next_block.next.mem_offset
 169         next_size = next_block.next.size
 170
 171     # Flags to indicate status of the 4 bytes going to be printed.
 172     is_left_header = False
 173     is_right_header = False
 174     is_used = False
 175
 176     # Print overall heap information
 177     _LOG.info('%-40s%-40s', f'The heap starts at {hex(heap_low_address)}.',
 178               f'The heap ends at {hex(heap_high_address)}.')
 179     _LOG.info('%-40s%-40s', f'Heap size is {heap_size // 1024}k bytes.',
 180               f'Heap is aligned by {aligned_bytes} bytes.')
 181     if poison_offset != 0:
 182         _LOG.info(
 183             'Poison is enabled %d bytes before and after the usable '
 184             'space of each block.', poison_offset)
 185     else:
 186         _LOG.info('%-40s', 'Poison is disabled.')
 187     _LOG.info(
 188         '%-40s', 'Below is the visualization of the heap. '
 189         'Each character represents 4 bytes.')
 190     _LOG.info('%-40s', f"    '{_FREE_CHAR}' indicates free space.")
 191     _LOG.info('%-40s', f"    '{_USED_CHAR}' indicates used space.")
 192     _LOG.info(
 193         '%-40s', f"    '{_LEFT_HEADER_CHAR}' indicates header or "
 194         'poisoned space before the block.')
 195     _LOG.info('%-40s', f"    '{_RIGHT_HEADER_CHAR}' poisoned space after "
 196               'the block.')
 197     print()
 198
 199     # Go over the heap space where there will be 64 characters each line.
 200     for line_base_address in range(0, heap_size, _CHARACTERS_PER_LINE *
 201                                    _BYTES_PER_CHARACTER):
 202         # Print the heap address of the current line.
 203         sys.stdout.write(f"{' ': <13}"
 204                          f'{hex(heap_low_address + line_base_address)}'
 205                          f"{f' (+{line_base_address}):': <12}")
 206         for line_offset in range(0,
 207                                  _CHARACTERS_PER_LINE * _BYTES_PER_CHARACTER,
 208                                  _BYTES_PER_CHARACTER):
 209             # Determine if the current 4 bytes is used, unused, or is a
 210             # header.
 211             # The case that we have went over the previous block and will
 212             # turn to the next block.
 213             current_address = line_base_address + line_offset
 214             if current_address == next_mem_offset + next_size + poison_offset:
 215                 next_block = next_block.next
 216                 # If this is the last block, set nextMemOffset to be over
 217                 # the last byte of heap so that the rest of the heap will
 218                 # be printed out as unused.
 219                 # Otherwise set the next HeapBlock allocated.
 220                 if next_block.next is None:
 221                     next_mem_offset = (heap_size + header_size +
 222                                        poison_offset + 1)
 223                     next_size = 0
 224                 else:
 225                     next_mem_offset = next_block.next.mem_offset
 226                     next_size = next_block.next.size
 227
 228             # Determine the status of the current 4 bytes.
 229             if (next_mem_offset - header_size - poison_offset <=
 230                     current_address < next_mem_offset):
 231                 is_left_header = True
 232                 is_right_header = False
 233                 is_used = False
 234             elif (next_mem_offset <= current_address <
 235                   next_mem_offset + next_size):
 236                 is_left_header = False
 237                 is_right_header = False
 238                 is_used = True
 239             elif (next_mem_offset + next_size <= current_address <
 240                   next_mem_offset + next_size + poison_offset):
 241                 is_left_header = False
 242                 is_right_header = True
 243                 is_used = False
 244             else:
 245                 is_left_header = False
 246                 is_right_header = False
 247                 is_used = False
 248
 249             if is_left_header:
 250                 sys.stdout.write(_LEFT_HEADER_CHAR)
 251             elif is_right_header:
 252                 sys.stdout.write(_RIGHT_HEADER_CHAR)
 253             elif is_used:
 254                 sys.stdout.write(_USED_CHAR)
 255             else:
 256                 sys.stdout.write(_FREE_CHAR)
 257         sys.stdout.write('\n')
 258
 259     allocation_dump.close()
 260
 261
 262 def main():
 263     """A python script to visualize heap usage given a dump file."""
 264     parser = argparse.ArgumentParser(description=main.__doc__)
 265     add_parser_arguments(parser)
 266     # Try to use pw_cli logs, else default to something reasonable.
 267     try:
 268         import pw_cli.log  # pylint: disable=import-outside-toplevel
 269         pw_cli.log.install()
 270     except ImportError:
 271         coloredlogs.install(level='INFO',
 272                             level_styles={
 273                                 'debug': {
 274                                     'color': 244
 275                                 },
 276                                 'error': {
 277                                     'color': 'red'
 278                                 }
 279                             },
 280                             fmt='%(asctime)s %(levelname)s | %(message)s')
 281     visualize(**vars(parser.parse_args()))
 282
 283
 284 if __name__ == "__main__":
 285     main()