add tracing for sent and received packets

quictrace/tracer.go

@@ -0,0 +1,189 @@
+package quictrace
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/golang/protobuf/proto"
+
+	"github.com/lucas-clemente/quic-go/internal/ackhandler"
+	"github.com/lucas-clemente/quic-go/internal/protocol"
+	"github.com/lucas-clemente/quic-go/internal/wire"
+	"github.com/lucas-clemente/quic-go/quictrace/pb"
+)
+
+type traceEvent struct {
+	connID protocol.ConnectionID
+	ev     Event
+}
+
+// A tracer is used to trace a QUIC connection
+type tracer struct {
+	eventQueue chan traceEvent
+
+	events map[string] /* conn ID */ []traceEvent
+}
+
+var _ Tracer = &tracer{}
+
+// NewTracer creates a new Tracer
+func NewTracer() Tracer {
+	qt := &tracer{
+		eventQueue: make(chan traceEvent, 1<<10),
+		events:     make(map[string][]traceEvent),
+	}
+	go qt.run()
+	return qt
+}
+
+// Trace traces an event
+func (t *tracer) Trace(connID protocol.ConnectionID, ev Event) {
+	t.eventQueue <- traceEvent{connID: connID, ev: ev}
+}
+
+func (t *tracer) run() {
+	for tev := range t.eventQueue {
+		key := string(tev.connID)
+		if _, ok := t.events[key]; !ok {
+			t.events[key] = make([]traceEvent, 0, 10*1<<10)
+		}
+		t.events[key] = append(t.events[key], tev)
+	}
+}
+
+func (t *tracer) GetAllTraces() map[string][]byte {
+	traces := make(map[string][]byte)
+	for connID := range t.events {
+		data, err := t.emitByConnIDAsString(connID)
+		if err != nil {
+			panic(err)
+		}
+		traces[connID] = data
+	}
+	return traces
+}
+
+// Emit emits the serialized protobuf that will be consumed by quic-trace
+func (t *tracer) Emit(connID protocol.ConnectionID) ([]byte, error) {
+	return t.emitByConnIDAsString(string(connID))
+}
+
+func (t *tracer) emitByConnIDAsString(connID string) ([]byte, error) {
+	events, ok := t.events[connID]
+	if !ok {
+		return nil, fmt.Errorf("No trace found for connection ID %s", connID)
+	}
+	trace := &pb.Trace{
+		DestinationConnectionId: []byte{1, 2, 3, 4, 5, 6, 7, 8},
+		SourceConnectionId:      []byte{1, 2, 3, 4, 5, 6, 7, 8},
+		ProtocolVersion:         []byte{0xff, 0, 0, 19},
+		Events:                  make([]*pb.Event, len(events)),
+	}
+	var startTime time.Time
+	for i, ev := range events {
+		event := ev.ev
+		if i == 0 {
+			startTime = event.Time
+		}
+
+		packetNumber := uint64(event.PacketNumber)
+		packetSize := uint64(event.PacketSize)
+
+		trace.Events[i] = &pb.Event{
+			TimeUs:          durationToUs(event.Time.Sub(startTime)),
+			EventType:       getEventType(event.EventType),
+			PacketSize:      &packetSize,
+			PacketNumber:    &packetNumber,
+			TransportState:  getTransportState(event.TransportState),
+			EncryptionLevel: getEncryptionLevel(event.EncryptionLevel),
+			Frames:          getFrames(event.Frames),
+		}
+	}
+	delete(t.events, connID)
+	return proto.Marshal(trace)
+}
+
+func getEventType(evType EventType) *pb.EventType {
+	var t pb.EventType
+	switch evType {
+	case PacketSent:
+		t = pb.EventType_PACKET_SENT
+	case PacketReceived:
+		t = pb.EventType_PACKET_RECEIVED
+	case PacketLost:
+		t = pb.EventType_PACKET_LOST
+	default:
+		panic("unknown event type")
+	}
+	return &t
+}
+
+func getEncryptionLevel(encLevel protocol.EncryptionLevel) *pb.EncryptionLevel {
+	enc := pb.EncryptionLevel_ENCRYPTION_UNKNOWN
+	switch encLevel {
+	case protocol.EncryptionInitial:
+		enc = pb.EncryptionLevel_ENCRYPTION_INITIAL
+	case protocol.EncryptionHandshake:
+		enc = pb.EncryptionLevel_ENCRYPTION_HANDSHAKE
+	case protocol.Encryption1RTT:
+		enc = pb.EncryptionLevel_ENCRYPTION_1RTT
+	}
+	return &enc
+}
+
+func getFrames(wframes []wire.Frame) []*pb.Frame {
+	streamFrameType := pb.FrameType_STREAM
+	ackFrameType := pb.FrameType_ACK
+	var frames []*pb.Frame
+	for _, frame := range wframes {
+		switch f := frame.(type) {
+		case *wire.StreamFrame:
+			streamID := uint64(f.StreamID)
+			offset := uint64(f.Offset)
+			length := uint64(f.DataLen())
+			frames = append(frames, &pb.Frame{
+				FrameType: &streamFrameType,
+				StreamFrameInfo: &pb.StreamFrameInfo{
+					StreamId: &streamID,
+					Offset:   &offset,
+					Length:   &length,
+				},
+			})
+		case *wire.AckFrame:
+			var ackedPackets []*pb.AckBlock
+			for _, ackBlock := range f.AckRanges {
+				firstPacket := uint64(ackBlock.Smallest)
+				lastPacket := uint64(ackBlock.Largest)
+				ackedPackets = append(ackedPackets, &pb.AckBlock{
+					FirstPacket: &firstPacket,
+					LastPacket:  &lastPacket,
+				})
+			}
+			frames = append(frames, &pb.Frame{
+				FrameType: &ackFrameType,
+				AckInfo: &pb.AckInfo{
+					AckDelayUs:   durationToUs(f.DelayTime),
+					AckedPackets: ackedPackets,
+				},
+			})
+		}
+	}
+	return frames
+}
+
+func getTransportState(state *ackhandler.State) *pb.TransportState {
+	bytesInFlight := uint64(state.BytesInFlight)
+	congestionWindow := uint64(state.CongestionWindow)
+	return &pb.TransportState{
+		MinRttUs:      durationToUs(state.MinRTT),
+		SmoothedRttUs: durationToUs(state.SmoothedRTT),
+		LastRttUs:     durationToUs(state.LatestRTT),
+		InFlightBytes: &bytesInFlight,
+		CwndBytes:     &congestionWindow,
+	}
+}
+
+func durationToUs(d time.Duration) *uint64 {
+	dur := uint64(d / 1000)
+	return &dur
+}