fetch_ml/tests/benchmarks/response_packet_benchmark_test.go

package benchmarks

import (
	"encoding/binary"
	"fmt"
	"testing"

	"github.com/jfraeys/fetch_ml/internal/api"
)

var benchmarkDataPayload = func() []byte {
	buf := make([]byte, 4096)
	for i := range buf {
		buf[i] = byte(i % 251)
	}
	return buf
}()

var benchmarkPackets = []struct {
	name   string
	packet *api.ResponsePacket
}{
	{
		name: "success",
		packet: &api.ResponsePacket{
			PacketType:     api.PacketTypeSuccess,
			Timestamp:      1_732_000_000,
			SuccessMessage: "Job 'benchmark' queued successfully",
		},
	},
	{
		name: "error",
		packet: &api.ResponsePacket{
			PacketType:   api.PacketTypeError,
			Timestamp:    1_732_000_000,
			ErrorCode:    api.ErrorCodeDatabaseError,
			ErrorMessage: "Failed to enqueue task",
			ErrorDetails: "database connection refused",
		},
	},
	{
		name: "data",
		packet: &api.ResponsePacket{
			PacketType:  api.PacketTypeData,
			Timestamp:   1_732_000_000,
			DataType:    "status",
			DataPayload: benchmarkDataPayload,
		},
	},
	{
		name: "progress",
		packet: &api.ResponsePacket{
			PacketType:      api.PacketTypeProgress,
			Timestamp:       1_732_000_000,
			ProgressType:    api.ProgressTypePercentage,
			ProgressValue:   42,
			ProgressTotal:   100,
			ProgressMessage: "running",
		},
	},
}

func BenchmarkResponsePacketSerialize(b *testing.B) {
	for _, variant := range benchmarkPackets {
		variant := variant
		b.Run(variant.name+"/current", func(b *testing.B) {
			benchmarkSerializePacket(b, variant.packet)
		})
		b.Run(variant.name+"/legacy", func(b *testing.B) {
			benchmarkLegacySerializePacket(b, variant.packet)
		})
	}
}

func benchmarkSerializePacket(b *testing.B, packet *api.ResponsePacket) {
	b.Helper()
	b.ReportAllocs()
	for b.Loop() {
		if _, err := packet.Serialize(); err != nil {
			b.Fatalf("serialize failed: %v", err)
		}
	}
}

func benchmarkLegacySerializePacket(b *testing.B, packet *api.ResponsePacket) {
	b.Helper()
	b.ReportAllocs()
	for b.Loop() {
		if _, err := legacySerializePacket(packet); err != nil {
			b.Fatalf("legacy serialize failed: %v", err)
		}
	}
}

func legacySerializePacket(p *api.ResponsePacket) ([]byte, error) {
	var buf []byte

	buf = append(buf, p.PacketType)

	timestampBytes := make([]byte, 8)
	binary.BigEndian.PutUint64(timestampBytes, p.Timestamp)
	buf = append(buf, timestampBytes...)

	switch p.PacketType {
	case api.PacketTypeSuccess:
		buf = append(buf, legacySerializeString(p.SuccessMessage)...)

	case api.PacketTypeError:
		buf = append(buf, p.ErrorCode)
		buf = append(buf, legacySerializeString(p.ErrorMessage)...)
		buf = append(buf, legacySerializeString(p.ErrorDetails)...)

	case api.PacketTypeProgress:
		buf = append(buf, p.ProgressType)
		valueBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(valueBytes, p.ProgressValue)
		buf = append(buf, valueBytes...)

		totalBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(totalBytes, p.ProgressTotal)
		buf = append(buf, totalBytes...)

		buf = append(buf, legacySerializeString(p.ProgressMessage)...)

	case api.PacketTypeStatus:
		buf = append(buf, legacySerializeString(p.StatusData)...)

	case api.PacketTypeData:
		buf = append(buf, legacySerializeString(p.DataType)...)
		buf = append(buf, legacySerializeBytes(p.DataPayload)...)

	case api.PacketTypeLog:
		buf = append(buf, p.LogLevel)
		buf = append(buf, legacySerializeString(p.LogMessage)...)

	default:
		return nil, fmt.Errorf("unknown packet type: %d", p.PacketType)
	}

	return buf, nil
}

func legacySerializeString(s string) []byte {
	length := uint16(len(s))
	buf := make([]byte, 2+len(s))
	binary.BigEndian.PutUint16(buf[:2], length)
	copy(buf[2:], s)
	return buf
}

func legacySerializeBytes(b []byte) []byte {
	length := uint32(len(b))
	buf := make([]byte, 4+len(b))
	binary.BigEndian.PutUint32(buf[:4], length)
	copy(buf[4:], b)
	return buf
}