From 656b01dba79d8c06d288c39c1e54245fd444e760 Mon Sep 17 00:00:00 2001
From: iamxvbaba <28732408+iamxvbaba@users.noreply.github.com>
Date: Mon, 3 Aug 2026 17:33:49 +0800
Subject: [PATCH] perf(rpc): sync eliminate upload preflight allocations

---
 internal/rpc/request_preflight.go      | 14 +++------
 internal/rpc/request_preflight_test.go | 41 ++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/internal/rpc/request_preflight.go b/internal/rpc/request_preflight.go
index 10a039d0..13f86437 100644
--- a/internal/rpc/request_preflight.go
+++ b/internal/rpc/request_preflight.go
@@ -21,12 +21,6 @@ type requestVectorPolicy struct {
 	tooLong      func() error
 }
 
-type rpcPreflightWire interface {
-	WireSize() int
-	ByteAt(offset int) (byte, error)
-	Uint32At(offset int) (uint32, error)
-}
-
 type rawRPCPreflightWire []byte
 
 func (w rawRPCPreflightWire) WireSize() int { return len(w) }
@@ -159,7 +153,7 @@ func preflightRPCRequest(id uint32, b *bin.Buffer) error {
 	return preflightRPCWire(id, rawRPCPreflightWire(b.Buf))
 }
 
-func preflightRPCWire(id uint32, wire rpcPreflightWire) error {
+func preflightRPCWire(id uint32, wire rawRPCPreflightWire) error {
 	if wire == nil {
 		return inputRequestInvalidErr()
 	}
@@ -209,7 +203,7 @@ func (r *Router) LayerRPCFlatBytesPayloadSize(wire []byte) (int, bool) {
 	return payloadBytes, true
 }
 
-func preflightFixedVector(wire rpcPreflightWire, policy requestVectorPolicy) error {
+func preflightFixedVector(wire rawRPCPreflightWire, policy requestVectorPolicy) error {
 	if policy.vectorOffset < 4 || policy.minElemBytes <= 0 || wire.WireSize() < policy.vectorOffset+8 {
 		return inputRequestInvalidErr()
 	}
@@ -240,7 +234,7 @@ func preflightFixedVector(wire rpcPreflightWire, policy requestVectorPolicy) err
 	return nil
 }
 
-func preflightUploadPart(wire rpcPreflightWire, bytesOffset int, big bool) error {
+func preflightUploadPart(wire rawRPCPreflightWire, bytesOffset int, big bool) error {
 	if wire.WireSize() < bytesOffset {
 		return inputRequestInvalidErr()
 	}
@@ -280,7 +274,7 @@ func preflightUploadPart(wire rpcPreflightWire, bytesOffset int, big bool) error
 
 // tlBytesSizeAt parses a TL bytes prefix without copying the payload. encoded
 // includes prefix, payload and 4-byte padding.
-func tlBytesSizeAt(wire rpcPreflightWire, offset int) (n, encoded int, err error) {
+func tlBytesSizeAt(wire rawRPCPreflightWire, offset int) (n, encoded int, err error) {
 	if offset < 0 || offset >= wire.WireSize() {
 		return 0, 0, fmt.Errorf("bytes prefix out of range")
 	}
diff --git a/internal/rpc/request_preflight_test.go b/internal/rpc/request_preflight_test.go
index 71c1ed6d..6dc89b84 100644
--- a/internal/rpc/request_preflight_test.go
+++ b/internal/rpc/request_preflight_test.go
@@ -149,6 +149,47 @@ func TestLayerRPCFlatBytesPayloadSizeRequiresCompleteLegalUpload(t *testing.T) {
 	}
 }
 
+func TestLayerRPCFlatBytesPayloadSizeDoesNotAllocate(t *testing.T) {
+	router := New(Config{}, Deps{}, zaptest.NewLogger(t), clock.System)
+	wire := uploadPartRequest(
+		tg.UploadSaveBigFilePartRequestTypeID,
+		20,
+		7,
+		364,
+		appfiles.MaxUploadPartBytes,
+	)
+	var payloadBytes int
+	var ok bool
+	allocs := testing.AllocsPerRun(1000, func() {
+		payloadBytes, ok = router.LayerRPCFlatBytesPayloadSize(wire)
+	})
+	if !ok || payloadBytes != appfiles.MaxUploadPartBytes {
+		t.Fatalf("flat payload = %d/%v", payloadBytes, ok)
+	}
+	if allocs != 0 {
+		t.Fatalf("flat payload preflight allocations = %v, want 0", allocs)
+	}
+}
+
+func BenchmarkLayerRPCFlatBytesPayloadSize(b *testing.B) {
+	router := New(Config{}, Deps{}, zaptest.NewLogger(b), clock.System)
+	wire := uploadPartRequest(
+		tg.UploadSaveBigFilePartRequestTypeID,
+		20,
+		7,
+		364,
+		appfiles.MaxUploadPartBytes,
+	)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		payloadBytes, ok := router.LayerRPCFlatBytesPayloadSize(wire)
+		if !ok || payloadBytes != appfiles.MaxUploadPartBytes {
+			b.Fatalf("flat payload = %d/%v", payloadBytes, ok)
+		}
+	}
+}
+
 func fixedVectorRequest(id uint32, policy requestVectorPolicy, count int) []byte {
 	raw := make([]byte, policy.vectorOffset+8+count*policy.minElemBytes)
 	binary.LittleEndian.PutUint32(raw[0:4], id)