add profiling hooks

main.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"runtime/pprof"
 	"strings"
 )
 
@@ -40,9 +41,13 @@ func wrap(err error, t string, args ...interface{}) error {
 }
 
 type options struct {
-	b bool   // bzip compression flag
-	v bool   // verbose flag
-	f string // input file
+	b          bool   // bzip compression flag
+	v          bool   // verbose flag
+	f          string // input file
+	memprofile string
+	cpuprofile string
+	messages   bool // dump messages or no
+	packets    bool // dump packets or no
 }
 
 func (o options) input() (io.Reader, error) {
@@ -63,19 +68,62 @@ func (o options) input() (io.Reader, error) {
 	return r, nil
 }
 
+func memprofile(dest string) {
+	fmt.Println("writing mem profile to", dest)
+	w, err := os.Create(dest)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "unable to open memprofile output %s: %v\n", dest, err)
+	} else {
+		defer w.Close()
+		err := pprof.WriteHeapProfile(w)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "unable to write heap profile: %v\n", err)
+		}
+	}
+}
+
+func cpuprofile(dest string) func() {
+	fmt.Println("writing cpu profile to", dest)
+	w, err := os.Create(dest)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "unable to open cpuprofile output %s: %v\n", dest, err)
+		return func() {}
+	} else {
+		pprof.StartCPUProfile(w)
+		return func() {
+			pprof.StopCPUProfile()
+			w.Close()
+		}
+	}
+}
+
 func main() {
 	var opts options
 	flag.BoolVar(&opts.b, "b", false, "input is expected to be bzip-compressed")
 	flag.BoolVar(&opts.v, "v", false, "verbose mode")
 	flag.StringVar(&opts.f, "f", "--", "input file to be used. -- means stdin")
+	flag.BoolVar(&opts.messages, "messages", false, "dump top-level messages to stdout")
+	flag.BoolVar(&opts.packets, "packets", false, "dump packets to stdout")
+	flag.StringVar(&opts.memprofile, "memprofile", "", "memory profile destination")
+	flag.StringVar(&opts.cpuprofile, "cpuprofile", "", "cpu profile destination")
 	flag.Parse()
 
+	if opts.memprofile != "" {
+		defer memprofile(opts.memprofile)
+	}
+
+	if opts.cpuprofile != "" {
+		defer cpuprofile(opts.cpuprofile)()
+	}
+
 	r, err := opts.input()
 	if err != nil {
 		bail(1, "input error: %v", err)
 	}
 
 	p := newParser(r)
+	p.dumpMessages = opts.messages
+	p.dumpPackets = opts.packets
 	if err := p.start(); err != nil {
 		bail(1, "parse error: %v", err)
 	}

message.go

@@ -17,12 +17,12 @@ type message struct {
 
 func (m message) String() string {
 	if len(m.body) > 30 {
-		return fmt.Sprintf("{cmd: %v tick: %v compressed: %t size: %d body): %q...}", m.cmd, m.tick, m.compressed, len(m.body), m.body[:27])
+		return fmt.Sprintf("{cmd: %v tick: %v compressed: %t size: %d body: %q...}", m.cmd, m.tick, m.compressed, len(m.body), m.body[:27])
 	}
 	return fmt.Sprintf("{cmd: %v tick: %v compressed: %t size: %d body: %q}", m.cmd, m.tick, m.compressed, len(m.body), m.body)
 }
 
-func (m *message) check() error {
+func (m *message) check(dump bool) error {
 	if m.cmd != dota.EDemoCommands_DEM_Packet {
 		return fmt.Errorf("wrong command type in openPacket: %v", m.cmd)
 	}
@@ -40,6 +40,8 @@ func (m *message) check() error {
 	if err := proto.Unmarshal(m.body, packet); err != nil {
 		return wrap(err, "onPacket unable to unmarshal message body")
 	}
-	fmt.Printf("\t{in: %d out: %d data: %x}\n", packet.GetSequenceIn(), packet.GetSequenceOutAck(), packet.GetData()[:8])
+	if dump {
+		fmt.Printf("{in: %d out: %d data: %x}\n", packet.GetSequenceIn(), packet.GetSequenceOutAck(), packet.GetData()[:8])
+	}
 	return nil
 }

parser.go

@@ -14,10 +14,13 @@ type parser struct {
 
 	// re-useable scratch buffer. Contents never guaranteed to be clean.
 	scratch []byte
+
+	dumpMessages bool
+	dumpPackets  bool
 }
 
 func newParser(r io.Reader) *parser {
-	br := bufio.NewReader(r)
+	br := bufio.NewReaderSize(r, 1<<16)
 	return &parser{source: br, scratch: make([]byte, 1<<10)}
 }
 
@@ -41,10 +44,12 @@ func (p *parser) run() error {
 		if err != nil {
 			return wrap(err, "read message error in run loop")
 		}
-		fmt.Println(msg)
+		if p.dumpMessages {
+			fmt.Println(msg)
+		}
 		switch msg.cmd {
 		case dota.EDemoCommands_DEM_Packet:
-			if err := msg.check(); err != nil {
+			if err := msg.check(p.dumpPackets); err != nil {
 				fmt.Printf("error: %v\n", err)
 			}
 		}