main

square/leakcanary

Last updated at: 29/12/2023 09:39

StreamingHprofReader.kt

TLDR

The provided file, StreamingHprofReader.kt, is a Kotlin class that reads Hprof files. It provides methods to read Hprof records from the source and listen to specific record tags.

Methods

readRecords

Obtains a new source to read all Hprof records from and calls the provided listener for each record that matches one of the provided record tags. It returns the number of bytes read from the source.

Classes

Class StreamingHprofReader

This class reads the entire content of a Hprof source. It provides the readRecords method to read specific records based on tags. It also has two companion methods to create instances of StreamingHprofReader for a given Hprof file or source provider.

package shark

import java.io.File
import okio.Source
import shark.HprofRecordTag.CLASS_DUMP
import shark.HprofRecordTag.HEAP_DUMP
import shark.HprofRecordTag.HEAP_DUMP_END
import shark.HprofRecordTag.HEAP_DUMP_INFO
import shark.HprofRecordTag.HEAP_DUMP_SEGMENT
import shark.HprofRecordTag.INSTANCE_DUMP
import shark.HprofRecordTag.LOAD_CLASS
import shark.HprofRecordTag.OBJECT_ARRAY_DUMP
import shark.HprofRecordTag.PRIMITIVE_ARRAY_DUMP
import shark.HprofRecordTag.PRIMITIVE_ARRAY_NODATA
import shark.HprofRecordTag.ROOT_DEBUGGER
import shark.HprofRecordTag.ROOT_FINALIZING
import shark.HprofRecordTag.ROOT_INTERNED_STRING
import shark.HprofRecordTag.ROOT_JAVA_FRAME
import shark.HprofRecordTag.ROOT_JNI_GLOBAL
import shark.HprofRecordTag.ROOT_JNI_LOCAL
import shark.HprofRecordTag.ROOT_JNI_MONITOR
import shark.HprofRecordTag.ROOT_MONITOR_USED
import shark.HprofRecordTag.ROOT_NATIVE_STACK
import shark.HprofRecordTag.ROOT_REFERENCE_CLEANUP
import shark.HprofRecordTag.ROOT_STICKY_CLASS
import shark.HprofRecordTag.ROOT_THREAD_BLOCK
import shark.HprofRecordTag.ROOT_THREAD_OBJECT
import shark.HprofRecordTag.ROOT_UNKNOWN
import shark.HprofRecordTag.ROOT_UNREACHABLE
import shark.HprofRecordTag.ROOT_VM_INTERNAL
import shark.HprofRecordTag.STACK_FRAME
import shark.HprofRecordTag.STACK_TRACE
import shark.HprofRecordTag.STRING_IN_UTF8
import shark.PrimitiveType.Companion.REFERENCE_HPROF_TYPE
import shark.PrimitiveType.INT
import shark.StreamingHprofReader.Companion.readerFor

/**
 * Reads the entire content of a Hprof source in one fell swoop.
 * Call [readerFor] to obtain a new instance.
 */
class StreamingHprofReader private constructor(
  private val sourceProvider: StreamingSourceProvider,
  private val header: HprofHeader
) {

  /**
   * Obtains a new source to read all hprof records from and calls [listener] back for each record
   * that matches one of the provided [recordTags].
   *
   * @return the number of bytes read from the source
   */
  @Suppress("ComplexMethod", "NestedBlockDepth")
  fun readRecords(
    recordTags: Set<HprofRecordTag>,
    listener: OnHprofRecordTagListener
  ): Long {
    return sourceProvider.openStreamingSource().use { source ->
      val reader = HprofRecordReader(header, source)
      reader.skip(header.recordsPosition)

      // Local ref optimizations
      val intByteSize = INT.byteSize
      val identifierByteSize = reader.sizeOf(REFERENCE_HPROF_TYPE)

      while (!source.exhausted()) {
        // type of the record
        val tag = reader.readUnsignedByte()

        // number of microseconds since the time stamp in the header
        reader.skip(intByteSize)

        // number of bytes that follow and belong to this record
        val length = reader.readUnsignedInt()

        when (tag) {
          STRING_IN_UTF8.tag -> {
            if (STRING_IN_UTF8 in recordTags) {
              listener.onHprofRecord(STRING_IN_UTF8, length, reader)
            } else {
              reader.skip(length)
            }
          }
          LOAD_CLASS.tag -> {
            if (LOAD_CLASS in recordTags) {
              listener.onHprofRecord(LOAD_CLASS, length, reader)
            } else {
              reader.skip(length)
            }
          }
          STACK_FRAME.tag -> {
            if (STACK_FRAME in recordTags) {
              listener.onHprofRecord(STACK_FRAME, length, reader)
            } else {
              reader.skip(length)
            }
          }
          STACK_TRACE.tag -> {
            if (STACK_TRACE in recordTags) {
              listener.onHprofRecord(STACK_TRACE, length, reader)
            } else {
              reader.skip(length)
            }
          }
          HEAP_DUMP.tag, HEAP_DUMP_SEGMENT.tag -> {
            val heapDumpStart = reader.bytesRead
            var previousTag = 0
            var previousTagPosition = 0L
            while (reader.bytesRead - heapDumpStart < length) {
              val heapDumpTagPosition = reader.bytesRead
              val heapDumpTag = reader.readUnsignedByte()
              when (heapDumpTag) {
                ROOT_UNKNOWN.tag -> {
                  if (ROOT_UNKNOWN in recordTags) {
                    listener.onHprofRecord(ROOT_UNKNOWN, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }
                ROOT_JNI_GLOBAL.tag -> {
                  if (ROOT_JNI_GLOBAL in recordTags) {
                    listener.onHprofRecord(ROOT_JNI_GLOBAL, -1, reader)
                  } else {
                    reader.skip(identifierByteSize + identifierByteSize)
                  }
                }
                ROOT_JNI_LOCAL.tag -> {
                  if (ROOT_JNI_LOCAL in recordTags) {
                    listener.onHprofRecord(ROOT_JNI_LOCAL, -1, reader)
                  } else {
                    reader.skip(identifierByteSize + intByteSize + intByteSize)
                  }
                }

                ROOT_JAVA_FRAME.tag -> {
                  if (ROOT_JAVA_FRAME in recordTags) {
                    listener.onHprofRecord(ROOT_JAVA_FRAME, -1, reader)
                  } else {
                    reader.skip(identifierByteSize + intByteSize + intByteSize)
                  }
                }

                ROOT_NATIVE_STACK.tag -> {
                  if (ROOT_NATIVE_STACK in recordTags) {
                    listener.onHprofRecord(ROOT_NATIVE_STACK, -1, reader)
                  } else {
                    reader.skip(identifierByteSize + intByteSize)
                  }
                }

                ROOT_STICKY_CLASS.tag -> {
                  if (ROOT_STICKY_CLASS in recordTags) {
                    listener.onHprofRecord(ROOT_STICKY_CLASS, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }
                ROOT_THREAD_BLOCK.tag -> {
                  if (ROOT_THREAD_BLOCK in recordTags) {
                    listener.onHprofRecord(ROOT_THREAD_BLOCK, -1, reader)
                  } else {
                    reader.skip(identifierByteSize + intByteSize)
                  }
                }

                ROOT_MONITOR_USED.tag -> {
                  if (ROOT_MONITOR_USED in recordTags) {
                    listener.onHprofRecord(ROOT_MONITOR_USED, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }

                ROOT_THREAD_OBJECT.tag -> {
                  if (ROOT_THREAD_OBJECT in recordTags) {
                    listener.onHprofRecord(ROOT_THREAD_OBJECT, -1, reader)
                  } else {
                    reader.skip(identifierByteSize + intByteSize + intByteSize)
                  }
                }

                ROOT_INTERNED_STRING.tag -> {
                  if (ROOT_INTERNED_STRING in recordTags) {
                    listener.onHprofRecord(ROOT_INTERNED_STRING, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }

                ROOT_FINALIZING.tag -> {
                  if (ROOT_FINALIZING in recordTags) {
                    listener.onHprofRecord(ROOT_FINALIZING, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }

                ROOT_DEBUGGER.tag -> {
                  if (ROOT_DEBUGGER in recordTags) {
                    listener.onHprofRecord(ROOT_DEBUGGER, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }

                ROOT_REFERENCE_CLEANUP.tag -> {
                  if (ROOT_REFERENCE_CLEANUP in recordTags) {
                    listener.onHprofRecord(ROOT_REFERENCE_CLEANUP, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }

                ROOT_VM_INTERNAL.tag -> {
                  if (ROOT_VM_INTERNAL in recordTags) {
                    listener.onHprofRecord(ROOT_VM_INTERNAL, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }

                ROOT_JNI_MONITOR.tag -> {
                  if (ROOT_JNI_MONITOR in recordTags) {
                    listener.onHprofRecord(ROOT_JNI_MONITOR, -1, reader)
                  } else {
                    reader.skip(identifierByteSize + intByteSize + intByteSize)
                  }
                }

                ROOT_UNREACHABLE.tag -> {
                  if (ROOT_UNREACHABLE in recordTags) {
                    listener.onHprofRecord(ROOT_UNREACHABLE, -1, reader)
                  } else {
                    reader.skip(identifierByteSize)
                  }
                }
                CLASS_DUMP.tag -> {
                  if (CLASS_DUMP in recordTags) {
                    listener.onHprofRecord(CLASS_DUMP, -1, reader)
                  } else {
                    reader.skipClassDumpRecord()
                  }
                }
                INSTANCE_DUMP.tag -> {
                  if (INSTANCE_DUMP in recordTags) {
                    listener.onHprofRecord(INSTANCE_DUMP, -1, reader)
                  } else {
                    reader.skipInstanceDumpRecord()
                  }
                }

                OBJECT_ARRAY_DUMP.tag -> {
                  if (OBJECT_ARRAY_DUMP in recordTags) {
                    listener.onHprofRecord(OBJECT_ARRAY_DUMP, -1, reader)
                  } else {
                    reader.skipObjectArrayDumpRecord()
                  }
                }

                PRIMITIVE_ARRAY_DUMP.tag -> {
                  if (PRIMITIVE_ARRAY_DUMP in recordTags) {
                    listener.onHprofRecord(PRIMITIVE_ARRAY_DUMP, -1, reader)
                  } else {
                    reader.skipPrimitiveArrayDumpRecord()
                  }
                }

                PRIMITIVE_ARRAY_NODATA.tag -> {
                  throw UnsupportedOperationException("$PRIMITIVE_ARRAY_NODATA cannot be parsed")
                }

                HEAP_DUMP_INFO.tag -> {
                  if (HEAP_DUMP_INFO in recordTags) {
                    listener.onHprofRecord(HEAP_DUMP_INFO, -1, reader)
                  } else {
                    reader.skipHeapDumpInfoRecord()
                  }
                }
                else -> throw IllegalStateException(
                  "Unknown tag ${
                    "0x%02x".format(
                      heapDumpTag
                    )
                  } at $heapDumpTagPosition after ${
                    "0x%02x".format(
                      previousTag
                    )
                  } at $previousTagPosition"
                )
              }
              previousTag = heapDumpTag
              previousTagPosition = heapDumpTagPosition
            }
          }
          HEAP_DUMP_END.tag -> {
            if (HEAP_DUMP_END in recordTags) {
              listener.onHprofRecord(HEAP_DUMP_END, length, reader)
            }
          }
          else -> {
            reader.skip(length)
          }
        }
      }
      reader.bytesRead
    }
  }

  companion object {

    /**
     * Creates a [StreamingHprofReader] for the provided [hprofFile]. [hprofHeader] will be read from
     * [hprofFile] unless you provide it.
     */
    fun readerFor(
      hprofFile: File,
      hprofHeader: HprofHeader = HprofHeader.parseHeaderOf(hprofFile)
    ): StreamingHprofReader {
      val sourceProvider = FileSourceProvider(hprofFile)
      return readerFor(sourceProvider, hprofHeader)
    }

    /**
     * Creates a [StreamingHprofReader] that will call [StreamingSourceProvider.openStreamingSource]
     * on every [readRecords] to obtain a [Source] to read the hprof data from. Before reading the
     * hprof records, [StreamingHprofReader] will skip [HprofHeader.recordsPosition] bytes.
     */
    fun readerFor(
      hprofSourceProvider: StreamingSourceProvider,
      hprofHeader: HprofHeader = hprofSourceProvider.openStreamingSource()
        .use { HprofHeader.parseHeaderOf(it) }
    ): StreamingHprofReader {
      return StreamingHprofReader(hprofSourceProvider, hprofHeader)
    }
  }
}