HprofWriterHelper.kt

TLDR

The HprofWriterHelper.kt file is a Kotlin file that defines a class called HprofWriterHelper which is used for writing Hprof files. It provides methods for creating class definitions, instances, and various types of arrays in the Hprof file. The file also includes utility methods for dumping the generated Hprof data to a file or to bytes.

Classes

HprofWriterHelper

The HprofWriterHelper class is used for writing Hprof files. It provides methods for creating class definitions, instances, arrays, and references in the Hprof file. It also implements the Closeable interface to allow closing the writer.

Inner classes

• InstanceAndClassDefinition: Represents a class definition and instance definition. It provides methods for defining fields and static fields for class definitions, and instance fields for instance definitions.
• ClassDefinition: Represents a class definition. It provides a method for defining static fields.

Methods

clazz

This method is used for creating a class definition in the Hprof file. It takes parameters like className, superclassId, staticFields, and fields to define the class. It returns the ID of the created class.

stringRecord

This method is used for creating a string record in the Hprof file. It takes a name parameter and returns the created string record.

gcRoot

This method is used for creating a garbage root reference in the Hprof file. It takes a gcRoot parameter and writes the corresponding record to the file.

arrayClass

This method is used for creating an array class definition in the Hprof file. It takes a className parameter and returns the ID of the created array class.

string

This method is used for creating a string instance in the Hprof file. It takes a string parameter and returns the corresponding reference.

keyedWeakReference

This method is used for creating a keyed weak reference instance in the Hprof file. It takes a referentInstanceId parameter and returns the corresponding reference.

instance

This method is used for creating a general instance in the Hprof file. It takes a classId parameter and fields parameter and returns the corresponding reference.

instanceAndClassDefinition

This inner class represents a combined instance and class definition. It provides methods for defining fields for both class and instance definitions.

dump

The dump function is used for dumping the generated Hprof data to a file. It takes an optional hprofHeader parameter and a block parameter of the type HprofWriterHelper.() -> Unit that defines the operations to perform on the Hprof file.

dumpToBytes

The dumpToBytes function is used for dumping the generated Hprof data to bytes. It takes an optional hprofHeader parameter and a block parameter of the type HprofWriterHelper.() -> Unit that defines the operations to perform on the Hprof file.

asHprofBytes

This extension function is used for converting a list of Hprof records to bytes. It returns a DualSourceProvider object that encapsulates the byte array.

package shark

import okio.Buffer
import shark.GcRoot.StickyClass
import shark.HprofRecord.HeapDumpRecord.GcRootRecord
import shark.HprofRecord.HeapDumpRecord.ObjectRecord.ClassDumpRecord
import shark.HprofRecord.HeapDumpRecord.ObjectRecord.ClassDumpRecord.FieldRecord
import shark.HprofRecord.HeapDumpRecord.ObjectRecord.ClassDumpRecord.StaticFieldRecord
import shark.HprofRecord.HeapDumpRecord.ObjectRecord.InstanceDumpRecord
import shark.HprofRecord.HeapDumpRecord.ObjectRecord.ObjectArrayDumpRecord
import shark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.CharArrayDump
import shark.HprofRecord.LoadClassRecord
import shark.HprofRecord.StringRecord
import shark.PrimitiveType.BOOLEAN
import shark.PrimitiveType.BYTE
import shark.PrimitiveType.CHAR
import shark.PrimitiveType.DOUBLE
import shark.PrimitiveType.FLOAT
import shark.PrimitiveType.INT
import shark.PrimitiveType.LONG
import shark.PrimitiveType.SHORT
import shark.ValueHolder.BooleanHolder
import shark.ValueHolder.ByteHolder
import shark.ValueHolder.CharHolder
import shark.ValueHolder.DoubleHolder
import shark.ValueHolder.FloatHolder
import shark.ValueHolder.IntHolder
import shark.ValueHolder.LongHolder
import shark.ValueHolder.ReferenceHolder
import shark.ValueHolder.ShortHolder
import java.io.Closeable
import java.io.File
import java.util.UUID
import kotlin.random.Random
import kotlin.reflect.KClass
import shark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.LongArrayDump

class HprofWriterHelper constructor(
  private val writer: HprofWriter
) : Closeable {

  private var lastId = 0L
  private val id: Long
    get() = ++lastId

  private val weakRefKeyRandom = Random(42)

  // Sequence identical for every test run
  private val weakRefKey: String
    get() =
      UUID(weakRefKeyRandom.nextLong(), weakRefKeyRandom.nextLong()).toString()

  private val typeSizes =
    PrimitiveType.byteSizeByHprofType + (PrimitiveType.REFERENCE_HPROF_TYPE to writer.hprofHeader.identifierByteSize)

  private val classDumps = mutableMapOf<Long, ClassDumpRecord>()

  private val objectClassId = clazz(superclassId = 0, className = "java.lang.Object")
  private val objectArrayClassId = arrayClass("java.lang.Object")
  private val stringClassId = clazz(
    className = "java.lang.String", fields = listOf(
    "value" to ReferenceHolder::class,
    "count" to IntHolder::class
  )
  )

  private val referenceClassId = clazz(
    className = "java.lang.ref.Reference",
    fields = listOf(
      "referent" to ReferenceHolder::class
    )
  )

  private val weakReferenceClassId = clazz(
    className = "java.lang.ref.WeakReference",
    superclassId = referenceClassId
  )
  private val keyedWeakReferenceClassId = clazz(
    superclassId = weakReferenceClassId,
    className = "leakcanary.KeyedWeakReference",
    staticFields = listOf("heapDumpUptimeMillis" to LongHolder(30000)),
    fields = listOf(
      "key" to ReferenceHolder::class,
      "name" to ReferenceHolder::class,
      "watchUptimeMillis" to LongHolder::class,
      "retainedUptimeMillis" to LongHolder::class
    )
  )

  fun clazz(
    className: String,
    superclassId: Long = -1L, // -1 defaults to java.lang.Object
    staticFields: List<Pair<String, ValueHolder>> = emptyList(),
    fields: List<Pair<String, KClass<out ValueHolder>>> = emptyList()
  ): Long {
    val classNameRecord = StringRecord(id, className)
    writer.write(classNameRecord)
    val loadClass = LoadClassRecord(1, id, 1, classNameRecord.id)
    writer.write(loadClass)

    val staticFieldRecords = staticFields.map {
      val fieldName = StringRecord(id, it.first)
      writer.write(fieldName)
      StaticFieldRecord(fieldName.id, typeOf(it.second), it.second)
    }

    val fieldRecords = fields.map {
      val fieldName = StringRecord(id, it.first)
      writer.write(fieldName)
      FieldRecord(fieldName.id, typeOf(it.second))
    }

    var instanceSize = fieldRecords.sumBy {
      typeSizes.getValue(it.type)
    }

    var nextUpId = if (superclassId == -1L) objectClassId else superclassId
    while (nextUpId != 0L) {
      val nextUp = classDumps[nextUpId]!!
      instanceSize += nextUp.fields.sumBy {
        typeSizes.getValue(it.type)
      }
      nextUpId = nextUp.superclassId
    }
    val classDump = ClassDumpRecord(
      id = loadClass.id,
      stackTraceSerialNumber = 1,
      superclassId = if (superclassId == -1L) objectClassId else superclassId,
      classLoaderId = 0,
      signersId = 0,
      protectionDomainId = 0,
      instanceSize = instanceSize,
      staticFields = staticFieldRecords,
      fields = fieldRecords
    )
    classDumps[loadClass.id] = classDump
    writer.write(classDump)
    val gcRoot = StickyClass(classDump.id)
    gcRoot(gcRoot)
    return classDump.id
  }

  fun stringRecord(
    name: String
  ): StringRecord {
    val stringRecord = StringRecord(id, name)
    writer.write(stringRecord)
    return stringRecord
  }

  fun clazz(
    classNameRecord: StringRecord,
    superclassId: Long = -1L, // -1 defaults to java.lang.Object
    staticFields: List<Pair<Long, ValueHolder>> = emptyList(),
    fields: List<Pair<Long, KClass<out ValueHolder>>> = emptyList()
  ): Long {
    val loadClass = LoadClassRecord(1, id, 1, classNameRecord.id)
    writer.write(loadClass)

    val staticFieldRecords = staticFields.map {
      StaticFieldRecord(it.first, typeOf(it.second), it.second)
    }

    val fieldRecords = fields.map {
      FieldRecord(it.first, typeOf(it.second))
    }

    var instanceSize = fieldRecords.sumBy {
      typeSizes.getValue(it.type)
    }

    var nextUpId = if (superclassId == -1L) objectClassId else superclassId
    while (nextUpId != 0L) {
      val nextUp = classDumps[nextUpId]!!
      instanceSize += nextUp.fields.sumBy {
        typeSizes.getValue(it.type)
      }
      nextUpId = nextUp.superclassId
    }
    val classDump = ClassDumpRecord(
      id = loadClass.id,
      stackTraceSerialNumber = 1,
      superclassId = if (superclassId == -1L) objectClassId else superclassId,
      classLoaderId = 0,
      signersId = 0,
      protectionDomainId = 0,
      instanceSize = instanceSize,
      staticFields = staticFieldRecords,
      fields = fieldRecords
    )
    classDumps[loadClass.id] = classDump
    writer.write(classDump)
    val gcRoot = StickyClass(classDump.id)
    gcRoot(gcRoot)
    return classDump.id
  }

  fun gcRoot(gcRoot: GcRoot) {
    val gcRootRecord = GcRootRecord(gcRoot = gcRoot)
    writer.write(gcRootRecord)
  }

  fun arrayClass(className: String): Long {
    return clazz(className = "$className[]")
  }

  fun string(
    string: String
  ): ReferenceHolder {
    return instance(
      stringClassId,
      fields = listOf(string.charArrayDump, IntHolder(string.length))
    )
  }

  fun keyedWeakReference(
    referentInstanceId: ReferenceHolder
  ): ReferenceHolder {
    val referenceKey = string(weakRefKey)
    return instance(
      classId = keyedWeakReferenceClassId,
      fields = listOf(
        referenceKey,
        string("its lifecycle has ended"),
        LongHolder(5000),
        LongHolder(20000),
        ReferenceHolder(referentInstanceId.value)
      )
    )
  }

  fun instance(
    classId: Long,
    fields: List<ValueHolder> = emptyList()
  ): ReferenceHolder {
    val instanceDump = InstanceDumpRecord(
      id = id,
      stackTraceSerialNumber = 1,
      classId = classId,
      fieldValues = writer.valuesToBytes(fields)
    )
    writer.write(instanceDump)
    return ReferenceHolder(instanceDump.id)
  }

  inner class InstanceAndClassDefinition {
    val field = LinkedHashMap<String, ValueHolder>()
    val staticField = LinkedHashMap<String, ValueHolder>()
  }

  inner class ClassDefinition {
    val staticField = LinkedHashMap<String, ValueHolder>()
  }

  infix fun String.watchedInstance(block: InstanceAndClassDefinition.() -> Unit): ReferenceHolder {
    val instance = this.instance(block)
    keyedWeakReference(instance)
    return instance
  }

  infix fun String.instance(block: InstanceAndClassDefinition.() -> Unit): ReferenceHolder {
    val definition = InstanceAndClassDefinition()
    block(definition)

    val classFields = definition.field.map {
      it.key to it.value::class
    }

    val staticFields = definition.staticField.map { it.key to it.value }

    val instanceFields = definition.field.map { it.value }

    return instance(clazz(this, fields = classFields, staticFields = staticFields), instanceFields)
  }

  infix fun String.clazz(block: ClassDefinition.() -> Unit): Long {
    val definition = ClassDefinition()
    block(definition)

    val staticFields = definition.staticField.map { it.key to it.value }
    return clazz(this, staticFields = staticFields)
  }

  val String.charArrayDump: ReferenceHolder
    get() {
      val arrayDump = CharArrayDump(id, 1, toCharArray())
      writer.write(arrayDump)
      return ReferenceHolder(arrayDump.id)
    }

  fun objectArray(
    vararg elements: ReferenceHolder
  ): ReferenceHolder {
    return objectArrayOf(objectArrayClassId, *elements)
  }

  fun objectArrayOf(
    classId: Long,
    vararg elements: ReferenceHolder
  ): ReferenceHolder {
    return ReferenceHolder(objectArray(classId, elements.map { it.value }.toLongArray()))
  }

  fun objectArray(
    classId: Long,
    array: LongArray
  ): Long {
    val arrayDump = ObjectArrayDumpRecord(id, 1, classId, array)
    writer.write(arrayDump)
    return arrayDump.id
  }

  fun primitiveLongArray(array: LongArray): Long {
    val arrayDump = LongArrayDump(id, 1, array)
    writer.write(arrayDump)
    return arrayDump.id
  }

  private fun typeOf(wrapper: ValueHolder): Int {
    return when (wrapper) {
      is ReferenceHolder -> PrimitiveType.REFERENCE_HPROF_TYPE
      is BooleanHolder -> BOOLEAN.hprofType
      is CharHolder -> CHAR.hprofType
      is FloatHolder -> FLOAT.hprofType
      is DoubleHolder -> DOUBLE.hprofType
      is ByteHolder -> BYTE.hprofType
      is ShortHolder -> SHORT.hprofType
      is IntHolder -> INT.hprofType
      is LongHolder -> LONG.hprofType
    }
  }

  private fun typeOf(wrapperClass: KClass<out ValueHolder>): Int {
    return when (wrapperClass) {
      ReferenceHolder::class -> PrimitiveType.REFERENCE_HPROF_TYPE
      BooleanHolder::class -> BOOLEAN.hprofType
      CharHolder::class -> CHAR.hprofType
      FloatHolder::class -> FLOAT.hprofType
      DoubleHolder::class -> DOUBLE.hprofType
      ByteHolder::class -> BYTE.hprofType
      ShortHolder::class -> SHORT.hprofType
      IntHolder::class -> INT.hprofType
      LongHolder::class -> LONG.hprofType
      else -> throw IllegalArgumentException("Unexpected class $wrapperClass")
    }
  }

  override fun close() {
    writer.close()
  }
}

fun File.dump(block: HprofWriterHelper.() -> Unit) {
  HprofWriterHelper(HprofWriter.openWriterFor(this))
    .use(block)
}

fun dump(
  hprofHeader: HprofHeader = HprofHeader(),
  block: HprofWriterHelper.() -> Unit
): DualSourceProvider {
  val buffer = Buffer()
  HprofWriterHelper(HprofWriter.openWriterFor(buffer, hprofHeader))
    .use(block)
  return ByteArraySourceProvider(buffer.readByteArray())
}

fun dumpToBytes(
  hprofHeader: HprofHeader = HprofHeader(),
  block: HprofWriterHelper.() -> Unit
): ByteArray {
  val buffer = Buffer()
  HprofWriterHelper(HprofWriter.openWriterFor(buffer, hprofHeader))
    .use(block)
  return buffer.readByteArray()
}

fun List<HprofRecord>.asHprofBytes(): DualSourceProvider {
  val buffer = Buffer()
  HprofWriter.openWriterFor(buffer)
    .use { writer ->
      forEach { record ->
        writer.write(record)
      }
    }
  return ByteArraySourceProvider(buffer.readByteArray())
}