Parser performance improvements, round 2

demo_packet.go

@@ -60,21 +60,46 @@ func (ms pendingMessages) Less(i, j int) bool {
 // multiple inner packets from a single CDemoPacket. This is the main structure
 // that contains all other data types in the demo file.
 func (p *Parser) onCDemoPacket(m *dota.CDemoPacket) error {
+	return p.processDemoPacket(m.GetData())
+}
+
+// processDemoPacket is the core of the CDemoPacket handler. It takes the raw
+// packet payload so the fast envelope path (envelope_fast.go) can call it
+// without materializing a proto message.
+func (p *Parser) processDemoPacket(data []byte) error {
 	// Reuse a parser-level buffer to store pending messages. Messages are read
 	// first as pending messages then sorted before dispatch. onCDemoPacket is
 	// never re-entrant (it is dispatched only via callByDemoType, never nested
 	// within a callByPacketType call), so a single reused backing array is safe
 	// and avoids a heap allocation per embedded message.
 	ms := p.pendingMsgBuf[:0]
 
+	// The inner stream is bit-shifted after the leading 6-bit readUBitVar, so
+	// message bodies almost never sit on a byte boundary and must be copied out.
+	// Carve those copies from a single reused arena sized to the packet instead
+	// of allocating per message: the buffers only live until dispatch below (the
+	// protobuf unmarshal copies what it keeps), so reusing the arena across
+	// packets is safe for the same reason reusing pendingMsgBuf is. Message
+	// headers take space too, so the payload total always fits in len(data).
+	if cap(p.packetArena) < len(data) {
+		p.packetArena = make([]byte, 0, len(data))
+	}
+	arena := p.packetArena[:0]
+
 	// Read all messages from the buffer. Messages are packed serially as
 	// {type, size, data}. We keep reading until until less than a byte remains.
-	r := newReader(m.GetData())
+	r := newReader(data)
 	for r.remBytes() > 0 {
 		t := int32(r.readUBitVar())
 		size := r.readVarUint32()
-		buf := r.readBytes(size)
-		ms = append(ms, pendingMessage{p.Tick, t, buf})
+		start := len(arena)
+		end := start + int(size)
+		if end > cap(arena) {
+			_panicf("onCDemoPacket: message size %d exceeds packet buffer", size)
+		}
+		arena = arena[:end]
+		r.readBytesInto(arena[start:end])
+		ms = append(ms, pendingMessage{p.Tick, t, arena[start:end:end]})
 	}
 
 	// Sort messages to ensure dependencies are met. For example, we need to
@@ -83,10 +108,12 @@ func (p *Parser) onCDemoPacket(m *dota.CDemoPacket) error {
 	// and avoids the reflection allocations of sort.Sort's interface path.
 	sort.Stable(ms)
 
-	// Dispatch messages in order, stopping on handler error.
+	// Dispatch messages in order, stopping on handler error. dispatchPacket
+	// takes the fast envelope path for hot internal-only message types and
+	// falls back to the full protobuf callback path otherwise.
 	var err error
 	for i := range ms {
-		if err = p.Callbacks.callByPacketType(ms[i].t, ms[i].buf); err != nil {
+		if err = p.dispatchPacket(ms[i].t, ms[i].buf); err != nil {
 			break
 		}
 	}
@@ -110,7 +137,7 @@ func (p *Parser) onCDemoFullPacket(m *dota.CDemoFullPacket) error {
 
 	// Then the CDemoPacket.
 	if m.Packet != nil {
-		if err := p.onCDemoPacket(m.GetPacket()); err != nil {
+		if err := p.processDemoPacket(m.GetPacket().GetData()); err != nil {
 			return err
 		}
 	}

entity.go

@@ -176,6 +176,9 @@ func (e *Entity) GetIndex() int32 {
 
 // FindEntity finds a given Entity by index
 func (p *Parser) FindEntity(index int32) *Entity {
+	if index < 0 || int(index) >= len(p.entities) {
+		return nil
+	}
 	return p.entities[index]
 }
 
@@ -207,7 +210,7 @@ func (p *Parser) FindEntityByHandle(handle uint64) *Entity {
 func (p *Parser) FilterEntity(fb func(*Entity) bool) []*Entity {
 	entities := make([]*Entity, 0, 0)
 	for _, et := range p.entities {
-		if fb(et) {
+		if et != nil && fb(et) {
 			entities = append(entities, et)
 		}
 	}
@@ -223,18 +226,25 @@ type entityOpTuple struct {
 
 // Internal Callback for OnCSVCMsg_PacketEntities.
 func (p *Parser) onCSVCMsg_PacketEntities(m *dota.CSVCMsg_PacketEntities) error {
+	return p.processPacketEntities(m.GetEntityData(), m.GetUpdatedEntries(), m.GetLegacyIsDelta())
+}
+
+// processPacketEntities is the core of the PacketEntities handler. It takes
+// the three envelope fields the parser actually uses so the fast envelope
+// path (envelope_fast.go) can call it without materializing a proto message.
+func (p *Parser) processPacketEntities(entityData []byte, updatedEntries int32, isDelta bool) error {
 	r := &p.entityReader
-	r.reset(m.GetEntityData())
+	r.reset(entityData)
 
 	var index = int32(-1)
-	var updates = int(m.GetUpdatedEntries())
+	var updates = int(updatedEntries)
 	var cmd uint32
 	var classId int32
 	var serial int32
 	var e *Entity
 	var op EntityOp
 
-	if !m.GetLegacyIsDelta() {
+	if !isDelta {
 		if p.entityFullPackets > 0 {
 			return nil
 		}

envelope_fast.go

@@ -0,0 +1,251 @@
+package manta
+
+import (
+	"github.com/dotabuff/manta/dota"
+	"google.golang.org/protobuf/encoding/protowire"
+)
+
+// Hand-rolled envelope decoders for the hottest internal messages.
+//
+// The reflective protobuf unmarshal allocates the message, a pointer per
+// scalar field, and — dominating the profile — a fresh copy of every bytes
+// field (CDemoPacket.data, CSVCMsg_PacketEntities.entity_data,
+// CSVCMsg_UpdateStringTable.string_data effectively copy the whole replay a
+// second time). The parser's internal handlers only need a few scalar fields
+// plus the payload, and consume the payload synchronously, so when a message
+// type has no user callbacks registered we can decode the envelope by hand
+// with protowire and alias the payload instead of copying it.
+//
+// Gating: NewParser registers exactly one internal handler per relevant list
+// before returning, and users can only register afterwards, so
+// len(list) == 1 means "internal handler only". Any user registration makes
+// the list longer and permanently reverts that type to the full protobuf
+// path, where the user-visible message owns its own copies as before.
+//
+// Aliasing lifetimes: entity_data and string_data alias the demo-packet arena
+// (p.packetArena), which is stable for the duration of the packet's dispatch
+// loop; both are fully consumed before the handler returns (string-table
+// values are copied out by readBitsAsBytes/readString). CDemoPacket.data
+// aliases the outer-message buffer, which is stable until the next
+// readOuterMessage; processDemoPacket copies message bodies into the arena
+// before dispatching.
+
+// dispatchDemo routes an outer demo message, taking the fast envelope path
+// for CDemoPacket/CDemoSignonPacket when only the internal handler is
+// registered.
+func (p *Parser) dispatchDemo(t int32, buf []byte) error {
+	switch t {
+	case int32(dota.EDemoCommands_DEM_Packet):
+		if len(p.Callbacks.onCDemoPacket) == 1 {
+			return p.fastDemoPacket(buf)
+		}
+	case int32(dota.EDemoCommands_DEM_SignonPacket):
+		if len(p.Callbacks.onCDemoSignonPacket) == 1 {
+			return p.fastDemoPacket(buf)
+		}
+	}
+	return p.Callbacks.callByDemoType(t, buf)
+}
+
+// dispatchPacket routes an embedded packet message, taking the fast envelope
+// path for the hot types when only the internal handler is registered.
+func (p *Parser) dispatchPacket(t int32, buf []byte) error {
+	switch t {
+	case int32(dota.NET_Messages_net_Tick):
+		if len(p.Callbacks.onCNETMsg_Tick) == 1 {
+			return p.fastNetTick(buf)
+		}
+	case int32(dota.SVC_Messages_svc_UpdateStringTable):
+		if len(p.Callbacks.onCSVCMsg_UpdateStringTable) == 1 {
+			return p.fastUpdateStringTable(buf)
+		}
+	case int32(dota.SVC_Messages_svc_PacketEntities):
+		if len(p.Callbacks.onCSVCMsg_PacketEntities) == 1 {
+			return p.fastPacketEntities(buf)
+		}
+	case int32(dota.EBaseGameEvents_GE_Source1LegacyGameEvent):
+		if len(p.Callbacks.onCMsgSource1LegacyGameEvent) == 1 && p.skipGameEvent(buf) {
+			return nil
+		}
+	}
+	return p.Callbacks.callByPacketType(t, buf)
+}
+
+// skipGameEvent peeks the eventid (field 2, varint) of a
+// CMsgSource1LegacyGameEvent and reports whether the event is a known type
+// with no registered handlers, in which case the full unmarshal (a message
+// plus a reflect.New per key) can be skipped. The internal handler returns
+// nil for exactly that case; unknown event ids fall back to the full path so
+// its error behaviour is preserved.
+func (p *Parser) skipGameEvent(buf []byte) bool {
+	for len(buf) > 0 {
+		num, typ, n := protowire.ConsumeTag(buf)
+		if n < 0 {
+			return false
+		}
+		buf = buf[n:]
+		if num == 2 && typ == protowire.VarintType {
+			v, n := protowire.ConsumeVarint(buf)
+			if n < 0 {
+				return false
+			}
+			name, ok := p.gameEventNames[int32(v)]
+			if !ok {
+				return false
+			}
+			return p.gameEventHandlers[name] == nil
+		}
+		n = protowire.ConsumeFieldValue(num, typ, buf)
+		if n < 0 {
+			return false
+		}
+		buf = buf[n:]
+	}
+	return false
+}
+
+// fastDemoPacket decodes the CDemoPacket envelope: data = field 3 (bytes).
+func (p *Parser) fastDemoPacket(buf []byte) error {
+	var data []byte
+	for len(buf) > 0 {
+		num, typ, n := protowire.ConsumeTag(buf)
+		if n < 0 {
+			return _errorf("fastDemoPacket: invalid tag")
+		}
+		buf = buf[n:]
+		if num == 3 && typ == protowire.BytesType {
+			v, n := protowire.ConsumeBytes(buf)
+			if n < 0 {
+				return _errorf("fastDemoPacket: invalid data field")
+			}
+			data = v
+			buf = buf[n:]
+			continue
+		}
+		n = protowire.ConsumeFieldValue(num, typ, buf)
+		if n < 0 {
+			return _errorf("fastDemoPacket: invalid field %d", num)
+		}
+		buf = buf[n:]
+	}
+	return p.processDemoPacket(data)
+}
+
+// fastNetTick decodes the CNETMsg_Tick envelope: tick = field 1 (varint).
+func (p *Parser) fastNetTick(buf []byte) error {
+	var tick uint64
+	for len(buf) > 0 {
+		num, typ, n := protowire.ConsumeTag(buf)
+		if n < 0 {
+			return _errorf("fastNetTick: invalid tag")
+		}
+		buf = buf[n:]
+		if num == 1 && typ == protowire.VarintType {
+			v, n := protowire.ConsumeVarint(buf)
+			if n < 0 {
+				return _errorf("fastNetTick: invalid tick field")
+			}
+			tick = v
+			buf = buf[n:]
+			continue
+		}
+		n = protowire.ConsumeFieldValue(num, typ, buf)
+		if n < 0 {
+			return _errorf("fastNetTick: invalid field %d", num)
+		}
+		buf = buf[n:]
+	}
+	p.NetTick = uint32(tick)
+	return nil
+}
+
+// fastUpdateStringTable decodes the CSVCMsg_UpdateStringTable envelope:
+// table_id = 1 (varint), num_changed_entries = 2 (varint),
+// string_data = 3 (bytes).
+func (p *Parser) fastUpdateStringTable(buf []byte) error {
+	var tableId, numChanged int32
+	var data []byte
+	for len(buf) > 0 {
+		num, typ, n := protowire.ConsumeTag(buf)
+		if n < 0 {
+			return _errorf("fastUpdateStringTable: invalid tag")
+		}
+		buf = buf[n:]
+		switch {
+		case num == 1 && typ == protowire.VarintType:
+			v, n := protowire.ConsumeVarint(buf)
+			if n < 0 {
+				return _errorf("fastUpdateStringTable: invalid table_id")
+			}
+			tableId = int32(v)
+			buf = buf[n:]
+		case num == 2 && typ == protowire.VarintType:
+			v, n := protowire.ConsumeVarint(buf)
+			if n < 0 {
+				return _errorf("fastUpdateStringTable: invalid num_changed_entries")
+			}
+			numChanged = int32(v)
+			buf = buf[n:]
+		case num == 3 && typ == protowire.BytesType:
+			v, n := protowire.ConsumeBytes(buf)
+			if n < 0 {
+				return _errorf("fastUpdateStringTable: invalid string_data")
+			}
+			data = v
+			buf = buf[n:]
+		default:
+			n = protowire.ConsumeFieldValue(num, typ, buf)
+			if n < 0 {
+				return _errorf("fastUpdateStringTable: invalid field %d", num)
+			}
+			buf = buf[n:]
+		}
+	}
+	return p.processUpdateStringTable(tableId, numChanged, data)
+}
+
+// fastPacketEntities decodes the CSVCMsg_PacketEntities envelope:
+// updated_entries = 2 (varint), legacy_is_delta = 3 (varint bool),
+// entity_data = 7 (bytes).
+func (p *Parser) fastPacketEntities(buf []byte) error {
+	var updatedEntries int32
+	var isDelta bool
+	var entityData []byte
+	for len(buf) > 0 {
+		num, typ, n := protowire.ConsumeTag(buf)
+		if n < 0 {
+			return _errorf("fastPacketEntities: invalid tag")
+		}
+		buf = buf[n:]
+		switch {
+		case num == 2 && typ == protowire.VarintType:
+			v, n := protowire.ConsumeVarint(buf)
+			if n < 0 {
+				return _errorf("fastPacketEntities: invalid updated_entries")
+			}
+			updatedEntries = int32(v)
+			buf = buf[n:]
+		case num == 3 && typ == protowire.VarintType:
+			v, n := protowire.ConsumeVarint(buf)
+			if n < 0 {
+				return _errorf("fastPacketEntities: invalid legacy_is_delta")
+			}
+			isDelta = v != 0
+			buf = buf[n:]
+		case num == 7 && typ == protowire.BytesType:
+			v, n := protowire.ConsumeBytes(buf)
+			if n < 0 {
+				return _errorf("fastPacketEntities: invalid entity_data")
+			}
+			entityData = v
+			buf = buf[n:]
+		default:
+			n = protowire.ConsumeFieldValue(num, typ, buf)
+			if n < 0 {
+				return _errorf("fastPacketEntities: invalid field %d", num)
+			}
+			buf = buf[n:]
+		}
+	}
+	return p.processPacketEntities(entityData, updatedEntries, isDelta)
+}