kshark.HprofWriter.kt Maven / Gradle / Ivy

Go to download
package kshark

import okio.Buffer
import okio.BufferedSink
import okio.Okio
import kshark.GcRoot.Debugger
import kshark.GcRoot.Finalizing
import kshark.GcRoot.InternedString
import kshark.GcRoot.JavaFrame
import kshark.GcRoot.JniGlobal
import kshark.GcRoot.JniLocal
import kshark.GcRoot.JniMonitor
import kshark.GcRoot.MonitorUsed
import kshark.GcRoot.NativeStack
import kshark.GcRoot.ReferenceCleanup
import kshark.GcRoot.StickyClass
import kshark.GcRoot.ThreadBlock
import kshark.GcRoot.ThreadObject
import kshark.GcRoot.Unknown
import kshark.GcRoot.Unreachable
import kshark.GcRoot.VmInternal
import kshark.HprofRecord.HeapDumpEndRecord
import kshark.HprofRecord.HeapDumpRecord.GcRootRecord
import kshark.HprofRecord.HeapDumpRecord.HeapDumpInfoRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.ClassDumpRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.InstanceDumpRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.ObjectArrayDumpRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.BooleanArrayDump
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.ByteArrayDump
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.CharArrayDump
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.DoubleArrayDump
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.FloatArrayDump
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.IntArrayDump
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.LongArrayDump
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord.ShortArrayDump
import kshark.HprofRecord.LoadClassRecord
import kshark.HprofRecord.StackTraceRecord
import kshark.HprofRecord.StringRecord
import kshark.PrimitiveType.BOOLEAN
import kshark.PrimitiveType.BYTE
import kshark.PrimitiveType.CHAR
import kshark.PrimitiveType.DOUBLE
import kshark.PrimitiveType.FLOAT
import kshark.PrimitiveType.INT
import kshark.PrimitiveType.LONG
import kshark.PrimitiveType.SHORT
import kshark.ValueHolder.BooleanHolder
import kshark.ValueHolder.ByteHolder
import kshark.ValueHolder.CharHolder
import kshark.ValueHolder.DoubleHolder
import kshark.ValueHolder.FloatHolder
import kshark.ValueHolder.IntHolder
import kshark.ValueHolder.LongHolder
import kshark.ValueHolder.ReferenceHolder
import kshark.ValueHolder.ShortHolder
import java.io.Closeable
import java.io.File

/**
 * Generates Hprof files.
 *
 * Call [openWriterFor] to obtain a new instance.
 *
 * Call [write] to add records and [close] when you're done.
 */
class HprofWriter private constructor(
  private val sink: BufferedSink,
  val hprofHeader: HprofHeader
) : Closeable {

  @Deprecated(
    "Replaced by HprofWriter.hprofHeader.identifierByteSize",
    ReplaceWith("hprofHeader.identifierByteSize")
  )
  val identifierByteSize: Int
    get() = hprofHeader.identifierByteSize

  @Deprecated(
    "Replaced by HprofWriter.hprofHeader.version",
    ReplaceWith("hprofHeader.version")
  )
  val hprofVersion: Hprof.HprofVersion
    get() = Hprof.HprofVersion.valueOf(hprofHeader.version.name)

  private val workBuffer = Buffer()

  /**
   * Appends a [HprofRecord] to the heap dump. If [record] is a [HprofRecord.HeapDumpRecord] then
   * it will not be written to an in memory buffer and written to file only when the next a record
   * that is not a [HprofRecord.HeapDumpRecord] is written or when [close] is called.
   */
  fun write(record: HprofRecord) {
    sink.write(record)
  }

  /**
   * Helper method for creating a [ByteArray] for [InstanceDumpRecord.fieldValues] from a
   * list of [ValueHolder].
   */
  fun valuesToBytes(values: List): ByteArray {
    val valuesBuffer = Buffer()
    values.forEach { value ->
      valuesBuffer.writeValue(value)
    }
    return valuesBuffer.readByteArray()
  }

  /**
   * Flushes to disk all [HprofRecord.HeapDumpRecord] that are currently written to the in memory
   * buffer, then closes the file.
   */
  override fun close() {
    sink.flushHeapBuffer()
    sink.close()
  }

  private fun BufferedSink.writeValue(wrapper: ValueHolder) {
    when (wrapper) {
      is ReferenceHolder -> writeId(wrapper.value)
      is BooleanHolder -> writeBoolean(wrapper.value)
      is CharHolder -> write(charArrayOf(wrapper.value))
      is FloatHolder -> writeFloat(wrapper.value)
      is DoubleHolder -> writeDouble(wrapper.value)
      is ByteHolder -> writeByte(wrapper.value.toInt())
      is ShortHolder -> writeShort(wrapper.value.toInt())
      is IntHolder -> writeInt(wrapper.value)
      is LongHolder -> writeLong(wrapper.value)
    }
  }

  @Suppress("LongMethod")
  private fun BufferedSink.write(record: HprofRecord) {
    when (record) {
      is StringRecord -> {
        writeNonHeapRecord(HprofRecordTag.STRING_IN_UTF8.tag) {
          writeId(record.id)
          writeUtf8(record.string)
        }
      }
      is LoadClassRecord -> {
        writeNonHeapRecord(HprofRecordTag.LOAD_CLASS.tag) {
          writeInt(record.classSerialNumber)
          writeId(record.id)
          writeInt(record.stackTraceSerialNumber)
          writeId(record.classNameStringId)
        }
      }
      is StackTraceRecord -> {
        writeNonHeapRecord(HprofRecordTag.STACK_TRACE.tag) {
          writeInt(record.stackTraceSerialNumber)
          writeInt(record.threadSerialNumber)
          writeInt(record.stackFrameIds.size)
          writeIdArray(record.stackFrameIds)
        }
      }
      is GcRootRecord -> {
        with(workBuffer) {
          when (val gcRoot = record.gcRoot) {
            is Unknown -> {
              writeByte(HprofRecordTag.ROOT_UNKNOWN.tag)
              writeId(gcRoot.id)
            }
            is JniGlobal -> {
              writeByte(
                HprofRecordTag.ROOT_JNI_GLOBAL.tag
              )
              writeId(gcRoot.id)
              writeId(gcRoot.jniGlobalRefId)
            }
            is JniLocal -> {
              writeByte(HprofRecordTag.ROOT_JNI_LOCAL.tag)
              writeId(gcRoot.id)
              writeInt(gcRoot.threadSerialNumber)
              writeInt(gcRoot.frameNumber)
            }
            is JavaFrame -> {
              writeByte(HprofRecordTag.ROOT_JAVA_FRAME.tag)
              writeId(gcRoot.id)
              writeInt(gcRoot.threadSerialNumber)
              writeInt(gcRoot.frameNumber)
            }
            is NativeStack -> {
              writeByte(HprofRecordTag.ROOT_NATIVE_STACK.tag)
              writeId(gcRoot.id)
              writeInt(gcRoot.threadSerialNumber)
            }
            is StickyClass -> {
              writeByte(HprofRecordTag.ROOT_STICKY_CLASS.tag)
              writeId(gcRoot.id)
            }
            is ThreadBlock -> {
              writeByte(HprofRecordTag.ROOT_THREAD_BLOCK.tag)
              writeId(gcRoot.id)
              writeInt(gcRoot.threadSerialNumber)
            }
            is MonitorUsed -> {
              writeByte(HprofRecordTag.ROOT_MONITOR_USED.tag)
              writeId(gcRoot.id)
            }
            is ThreadObject -> {
              writeByte(HprofRecordTag.ROOT_THREAD_OBJECT.tag)
              writeId(gcRoot.id)
              writeInt(gcRoot.threadSerialNumber)
              writeInt(gcRoot.stackTraceSerialNumber)
            }
            is ReferenceCleanup -> {
              writeByte(HprofRecordTag.ROOT_REFERENCE_CLEANUP.tag)
              writeId(gcRoot.id)
            }
            is VmInternal -> {
              writeByte(HprofRecordTag.ROOT_VM_INTERNAL.tag)
              writeId(gcRoot.id)
            }
            is JniMonitor -> {
              writeByte(HprofRecordTag.ROOT_JNI_MONITOR.tag)
              writeId(gcRoot.id)
              writeInt(gcRoot.stackTraceSerialNumber)
              writeInt(gcRoot.stackDepth)
            }
            is InternedString -> {
              writeByte(HprofRecordTag.ROOT_INTERNED_STRING.tag)
              writeId(gcRoot.id)
            }
            is Finalizing -> {
              writeByte(HprofRecordTag.ROOT_FINALIZING.tag)
              writeId(gcRoot.id)
            }
            is Debugger -> {
              writeByte(HprofRecordTag.ROOT_DEBUGGER.tag)
              writeId(gcRoot.id)
            }
            is Unreachable -> {
              writeByte(HprofRecordTag.ROOT_UNREACHABLE.tag)
              writeId(gcRoot.id)
            }
          }
        }
      }
      is ClassDumpRecord -> {
        with(workBuffer) {
          writeByte(HprofRecordTag.CLASS_DUMP.tag)
          writeId(record.id)
          writeInt(record.stackTraceSerialNumber)
          writeId(record.superclassId)
          writeId(record.classLoaderId)
          writeId(record.signersId)
          writeId(record.protectionDomainId)
          // reserved
          writeId(0)
          // reserved
          writeId(0)
          writeInt(record.instanceSize)
          // Not writing anything in the constant pool
          val constantPoolCount = 0
          writeShort(constantPoolCount)
          writeShort(record.staticFields.size)
          record.staticFields.forEach { field ->
            writeId(field.nameStringId)
            writeByte(field.type)
            writeValue(field.value)
          }
          writeShort(record.fields.size)
          record.fields.forEach { field ->
            writeId(field.nameStringId)
            writeByte(field.type)
          }
        }
      }
      is InstanceDumpRecord -> {
        with(workBuffer) {
          writeByte(HprofRecordTag.INSTANCE_DUMP.tag)
          writeId(record.id)
          writeInt(record.stackTraceSerialNumber)
          writeId(record.classId)
          writeInt(record.fieldValues.size)
          write(record.fieldValues)
        }
      }
      is ObjectArrayDumpRecord -> {
        with(workBuffer) {
          writeByte(HprofRecordTag.OBJECT_ARRAY_DUMP.tag)
          writeId(record.id)
          writeInt(record.stackTraceSerialNumber)
          writeInt(record.elementIds.size)
          writeId(record.arrayClassId)
          writeIdArray(record.elementIds)
        }
      }
      is PrimitiveArrayDumpRecord -> {
        with(workBuffer) {
          writeByte(HprofRecordTag.PRIMITIVE_ARRAY_DUMP.tag)
          writeId(record.id)
          writeInt(record.stackTraceSerialNumber)

          when (record) {
            is BooleanArrayDump -> {
              writeInt(record.array.size)
              writeByte(BOOLEAN.hprofType)
              write(record.array)
            }
            is CharArrayDump -> {
              writeInt(record.array.size)
              writeByte(CHAR.hprofType)
              write(record.array)
            }
            is FloatArrayDump -> {
              writeInt(record.array.size)
              writeByte(FLOAT.hprofType)
              write(record.array)
            }
            is DoubleArrayDump -> {
              writeInt(record.array.size)
              writeByte(DOUBLE.hprofType)
              write(record.array)
            }
            is ByteArrayDump -> {
              writeInt(record.array.size)
              writeByte(BYTE.hprofType)
              write(record.array)
            }
            is ShortArrayDump -> {
              writeInt(record.array.size)
              writeByte(SHORT.hprofType)
              write(record.array)
            }
            is IntArrayDump -> {
              writeInt(record.array.size)
              writeByte(INT.hprofType)
              write(record.array)
            }
            is LongArrayDump -> {
              writeInt(record.array.size)
              writeByte(LONG.hprofType)
              write(record.array)
            }
          }
        }
      }
      is HeapDumpInfoRecord -> {
        with(workBuffer) {
          writeByte(HprofRecordTag.HEAP_DUMP_INFO.tag)
          writeInt(record.heapId)
          writeId(record.heapNameStringId)
        }
      }
      is HeapDumpEndRecord -> {
        throw IllegalArgumentException("HprofWriter automatically emits HeapDumpEndRecord")
      }
    }
  }

  private fun BufferedSink.writeDouble(value: Double) {
    writeLong(value.toBits())
  }

  private fun BufferedSink.writeFloat(value: Float) {
    writeInt(value.toBits())
  }

  private fun BufferedSink.writeBoolean(value: Boolean) {
    writeByte(if (value) 1 else 0)
  }

  private fun BufferedSink.writeIdArray(array: LongArray) {
    array.forEach { writeId(it) }
  }

  private fun BufferedSink.write(array: BooleanArray) {
    array.forEach { writeByte(if (it) 1 else 0) }
  }

  private fun BufferedSink.write(array: CharArray) {
    writeString(String(array), Charsets.UTF_16BE)
  }

  private fun BufferedSink.write(array: FloatArray) {
    array.forEach { writeFloat(it) }
  }

  private fun BufferedSink.write(array: DoubleArray) {
    array.forEach { writeDouble(it) }
  }

  private fun BufferedSink.write(array: ShortArray) {
    array.forEach { writeShort(it.toInt()) }
  }

  private fun BufferedSink.write(array: IntArray) {
    array.forEach { writeInt(it) }
  }

  private fun BufferedSink.write(array: LongArray) {
    array.forEach { writeLong(it) }
  }

  private fun BufferedSink.writeNonHeapRecord(
    tag: Int,
    block: BufferedSink.() -> Unit
  ) {
    flushHeapBuffer()
    workBuffer.block()
    writeTagHeader(tag, workBuffer.size())
    writeAll(workBuffer)
  }

  private fun BufferedSink.flushHeapBuffer() {
    if (workBuffer.size() > 0) {
      writeTagHeader(HprofRecordTag.HEAP_DUMP.tag, workBuffer.size())
      writeAll(workBuffer)
      writeTagHeader(HprofRecordTag.HEAP_DUMP_END.tag, 0)
    }
  }

  private fun BufferedSink.writeTagHeader(
    tag: Int,
    length: Long
  ) {
    writeByte(tag)
    // number of microseconds since the time stamp in the header
    writeInt(0)
    writeInt(length.toInt())
  }

  private fun BufferedSink.writeId(id: Long) {
    when (hprofHeader.identifierByteSize) {
      1 -> writeByte(id.toInt())
      2 -> writeShort(id.toInt())
      4 -> writeInt(id.toInt())
      8 -> writeLong(id)
      else -> throw IllegalArgumentException("ID Length must be 1, 2, 4, or 8")
    }
  }

  companion object {

    fun openWriterFor(
      hprofFile: File,
      hprofHeader: HprofHeader = HprofHeader()
    ): HprofWriter {
      return openWriterFor(Okio.buffer(Okio.sink(hprofFile.outputStream())), hprofHeader)
    }

    fun openWriterFor(
      hprofSink: BufferedSink,
      hprofHeader: HprofHeader = HprofHeader()
    ): HprofWriter {
      hprofSink.writeUtf8(hprofHeader.version.versionString)
      hprofSink.writeByte(0)
      hprofSink.writeInt(hprofHeader.identifierByteSize)
      hprofSink.writeLong(hprofHeader.heapDumpTimestamp)
      return HprofWriter(hprofSink, hprofHeader)
    }

    @Deprecated(
      "Replaced by HprofWriter.openWriterFor()",
      ReplaceWith(
        "kshark.HprofWriter.openWriterFor(hprofFile)"
      )
    )
    fun open(
      hprofFile: File,
      identifierByteSize: Int = 4,
      hprofVersion: Hprof.HprofVersion = Hprof.HprofVersion.ANDROID
    ): HprofWriter = openWriterFor(
      hprofFile,
      HprofHeader(
        version = HprofVersion.valueOf(hprofVersion.name),
        identifierByteSize = identifierByteSize
      )
    )
  }
}