astra/owpengram-server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

perf(rpc): sync eliminate upload preflight allocations

Browse source
This commit is contained in:
iamxvbaba 2026-08-03 17:33:49 +08:00
parent e70bb0f167
commit 656b01dba7
2 changed files with 45 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

14
internal/rpc/request_preflight.go
View file

@ -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")
}

41
internal/rpc/request_preflight_test.go
View file

@ -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)