io: add benchmark for TFRecord reading

This reads a single event file from start to end, parsing the frame of
each TFRecord, but not inspecting the payload at all. Unfortunately, the
TFRecord reading is already ~3× slower than RustBoard’s entire pipeline,
and ~13× slower than RustBoard’s TFRecord reading. :-(

Discussion here is on a 248 MiB event file, the `egraph_edge_cgan_003`
run from a user-supplied log directory:
<tensorflow/tensorboard#766 (comment)>

The effect of the buffer size is a bit strange. As expected, buffering
definitely helps (~3× improvement with default buffer size), and the
improvements taper off as the buffer size increases: 4 KiB and 1 MiB are
about the same. But then in the 4 MiB to 8 MiB range we start to see
sharp improvements: 1 MiB to 4 MiB is no change, but 4 MiB to 8 MiB is
25% faster. The improvements continue even up to 128 or 256 MiB on a
file that’s 248 MiB long. Compare to RustBoard, which sees similar
effects at low buffer sizes but no extra improvements for very large
buffers. (This all running with hot filesystem cache.)

Buffer size sweep:
<https://gist.github.com/wchargin/b73b5af3ef36b88e4e1aacf9a2453ea6>

CPU profiling shows that about 35% of time is spent in `make([]byte)` in
`ExtendTo`, which seems unfortunate, since due to the actual access
patterns we barely overallocate there (overallocating only 12 bytes per
record), so it’s not obvious how to avoid that cost. Another 50% of
total time is spent in `runtime.mallocgc`. And 20% of total time (not
necessarily disjoint) is spent in the `result := TFRecord{...}` final
allocation in `ReadRecord`, which is surprising to me since it just has
two small fields (a slice header and a `uint32`) and they’ve already
been computed above. (Inlining effects?)

Checksum validation is fast when enabled; runtime increases by ~10%.

Author: wchargin

io/cmd/tfrecord_bench/main.go

@@ -0,0 +1,94 @@
+package main
+
+import (
+	"bufio"
+	"flag"
+	"fmt"
+	"io"
+	"os"
+	"runtime"
+	"runtime/pprof"
+
+	tbio "github.com/wchargin/tensorboard-data-server/io"
+)
+
+var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
+var memprofile = flag.String("memprofile", "", "write memory profile to this file")
+
+var bufsize = flag.Int("bufsize", 8192, "bufio.NewReaderSize capacity")
+
+var checksum = flag.Bool("checksum", false, "validate TFRecord payloads against CRC-32")
+
+func main() {
+	flag.Parse()
+	args := flag.Args()
+	if len(args) != 1 {
+		fmt.Fprintf(os.Stderr, "must specify event file path as only argument\n")
+		os.Exit(1)
+	}
+
+	// Profiling structure copied from "go doc runtime/pprof".
+	if *cpuprofile != "" {
+		f, err := os.Create(*cpuprofile)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "creating CPU profile: %v\n", err)
+			os.Exit(1)
+		}
+		pprof.StartCPUProfile(f)
+		defer pprof.StopCPUProfile()
+	}
+
+	ReadAllRecords(args[0])
+
+	if *memprofile != "" {
+		fmt.Println("creating memory profile...")
+		f, err := os.Create(*memprofile)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "creating memory profile: %v\n", err)
+			os.Exit(1)
+		}
+		defer f.Close()
+		runtime.GC()
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			fmt.Fprintf(os.Stderr, "writing memory profile: %v\n", err)
+		}
+		fmt.Println("done with memory profile")
+	}
+}
+
+func ReadAllRecords(fileName string) {
+	rawFile, err := os.Open(fileName)
+	if err != nil {
+		panic(err)
+	}
+	f := bufio.NewReaderSize(rawFile, *bufsize)
+	recordsRead := 0
+	totalPayloadSize := 0
+	for {
+		payloadSize, more := ProcessOneRecord(f)
+		recordsRead++
+		totalPayloadSize += payloadSize
+		if !more {
+			break
+		}
+	}
+	fmt.Printf("all done; read %v records (%v bytes payload)\n", recordsRead, totalPayloadSize)
+}
+
+func ProcessOneRecord(r io.Reader) (payloadSize int, more bool) {
+	var state *tbio.TFRecordState
+	rec, err := tbio.ReadRecord(&state, r)
+	if err == io.EOF {
+		return 0, false
+	}
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "reading record: %v\n", err)
+		os.Exit(1)
+	}
+	if *checksum {
+		if err := rec.Checksum(); err != nil {
+			fmt.Fprintf(os.Stderr, "checksum failure: %v\n", err)
+		}
+	}
+	return len(rec.Data), true
+}