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

sync: # file upload improve

Browse source
This commit is contained in:
iamxvbaba 2026-08-03 17:33:49 +08:00
parent 9041abd3d3
commit e70bb0f167
6 changed files with 489 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

101
internal/mtprotoedge/inbound_layer_rpc.go
View file

@ -42,6 +42,7 @@ const (
// 32-MiB per-connection budget and therefore remains admissible by default.
layerRPCAdmissionStaticObjectBytes = 512
layerRPCAdmissionWireFactor = 60
layerRPCAdmissionFlatBytesFactor = 2
layerRPCAdmissionGraphSlack = layerRPCAdmissionStaticObjectBytes * 32
)
@ -164,33 +165,91 @@ func (s *Server) initialLayerRPCAdmissionCursor(ctx context.Context, c *Conn) (l
// Saturation deliberately turns hostile integer-sized inputs into ordinary
// capacity rejection; it must never wrap into a small accepted reservation.
func layerRPCAdmissionReservationSize(wireBytes int) int {
return saturatingLayerRPCAdmissionCharge(
layerRPCAdmissionGraphSlack,
layerRPCAdmissionWireCharge(wireBytes),
)
}
func layerRPCAdmissionWireCharge(wireBytes int) int {
if wireBytes < 0 {
wireBytes = 0
}
maxInt := int(^uint(0) >> 1)
if wireBytes > (maxInt-layerRPCAdmissionGraphSlack)/layerRPCAdmissionWireFactor {
if wireBytes > maxInt/layerRPCAdmissionWireFactor {
return maxInt
}
return wireBytes*layerRPCAdmissionWireFactor + layerRPCAdmissionGraphSlack
return wireBytes * layerRPCAdmissionWireFactor
}
// layerRPCFlatBytesWireCharge keeps the generic factor on fixed TL wire and
// applies a tight copy factor only to a payload which the production Router has
// already proven to be one bounded flat bytes field. The temporary expanded
// buffer remains independently owned by inboundFrameBudget while generated
// admission materializes the request; 2x covers the retained bytes copy plus
// allocator rounding without treating every payload byte as a possible nested
// object/vector node. The per-request graph slack is owned by the base
// reservation and must not be repeated for each nested gzip expansion.
func layerRPCFlatBytesWireCharge(wireBytes, payloadBytes int) int {
if wireBytes < 0 || payloadBytes < 0 || payloadBytes > wireBytes {
return layerRPCAdmissionWireCharge(wireBytes)
}
fixedBytes := wireBytes - payloadBytes
maxInt := int(^uint(0) >> 1)
if fixedBytes > maxInt/layerRPCAdmissionWireFactor {
return maxInt
}
charge := fixedBytes * layerRPCAdmissionWireFactor
if payloadBytes > (maxInt-charge)/layerRPCAdmissionFlatBytesFactor {
return maxInt
}
return charge + payloadBytes*layerRPCAdmissionFlatBytesFactor
}
func saturatingLayerRPCAdmissionCharge(left, right int) int {
if left < 0 {
left = 0
}
if right < 0 {
right = 0
}
maxInt := int(^uint(0) >> 1)
if right > maxInt-left {
return maxInt
}
return left + right
}
func (s *Server) layerRPCExpandedWireCharge(wire []byte) int {
if s != nil {
if sizer, ok := s.layerRPC.(LayerRPCFlatBytesPayloadSizer); ok {
if payloadBytes, flat := sizer.LayerRPCFlatBytesPayloadSize(wire); flat && payloadBytes >= 0 && payloadBytes <= len(wire) {
return layerRPCFlatBytesWireCharge(len(wire), payloadBytes)
}
}
}
return layerRPCAdmissionWireCharge(len(wire))
}
// layerRPCGZIPExpansionBudget bridges transient process-wide expansion memory
// to the durable scheduler charge of one exact typed request. sourceBytes is
// conservative: it retains the original compressed wire plus every successful
// expansion seen across nested envelopes or an authoritative-profile re-decode.
// to the durable scheduler charge of one exact typed request. The original
// compressed wire retains the generic graph charge. Every successful expansion
// adds its own charge; only a handler-proven flat bytes terminal can use the
// tighter payload factor. An authoritative-profile re-decode reuses the largest
// already-held charge instead of double-counting sequential attempts.
type layerRPCGZIPExpansionBudget struct {
server *Server
plan *inboundPlan
reservation *inboundRPCBatchReservation
entry int
baseSourceBytes int
attemptExpandedBytes int
chargedSourceBytes int
server *Server
plan *inboundPlan
reservation *inboundRPCBatchReservation
entry int
baseCharge int
attemptExpandedCharge int
chargedSize int
}
func (b *layerRPCGZIPExpansionBudget) beginAttempt() {
if b != nil {
b.attemptExpandedBytes = 0
b.attemptExpandedCharge = 0
}
}
@ -221,17 +280,19 @@ func (b *layerRPCGZIPExpansionBudget) expand(wire []byte, admissionLimit int) ([
return nil, release, err
}
b.plan.gzipExpandedBytes += len(data)
b.attemptExpandedBytes += len(data)
targetSourceBytes := b.baseSourceBytes + b.attemptExpandedBytes
targetCharge := layerRPCAdmissionReservationSize(targetSourceBytes)
if targetSourceBytes > b.chargedSourceBytes {
b.attemptExpandedCharge = saturatingLayerRPCAdmissionCharge(
b.attemptExpandedCharge,
b.server.layerRPCExpandedWireCharge(data),
)
targetCharge := saturatingLayerRPCAdmissionCharge(b.baseCharge, b.attemptExpandedCharge)
if targetCharge > b.chargedSize {
if err := b.reservation.growEntry(b.entry, targetCharge); err != nil {
if errors.Is(err, ErrInboundRPCQueueFull) {
return nil, release, errors.Join(errLayerRPCGZIPCapacity, err)
}
return nil, release, err
}
b.chargedSourceBytes = targetSourceBytes
b.chargedSize = targetCharge
}
return data, release, nil
}
@ -293,7 +354,9 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
}
expansionBudget := &layerRPCGZIPExpansionBudget{
server: s, plan: plan, reservation: reservation,
entry: reservationIndex, baseSourceBytes: len(item.body), chargedSourceBytes: len(item.body),
entry: reservationIndex,
baseCharge: provisionalSpecs[reservationIndex].size,
chargedSize: provisionalSpecs[reservationIndex].size,
}
expansionBudgets[reservationIndex] = expansionBudget
options := tlprofile.AdmissionOptions{

298
internal/mtprotoedge/layer_admission_budget_test.go
View file

@ -12,6 +12,7 @@ import (
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/tgerr"
"go.uber.org/zap/zaptest"
"github.com/iamxvbaba/td/tlprofile"
@ -85,6 +86,18 @@ func TestLayerRPCAdmissionMaterializationConstants(t *testing.T) {
if got := layerRPCAdmissionReservationSize(-1); got != layerRPCAdmissionGraphSlack {
t.Fatalf("negative wire charge = %d, want fixed slack %d", got, layerRPCAdmissionGraphSlack)
}
if got := layerRPCFlatBytesWireCharge(100, 80); got != 20*layerRPCAdmissionWireFactor+80*layerRPCAdmissionFlatBytesFactor {
t.Fatalf("flat bytes wire charge = %d", got)
}
if got := layerRPCFlatBytesWireCharge(10, 11); got != layerRPCAdmissionWireCharge(10) {
t.Fatalf("invalid flat bytes hint charge = %d, want generic %d", got, layerRPCAdmissionWireCharge(10))
}
if got := layerRPCFlatBytesWireCharge(maxInt, maxInt); got != maxInt {
t.Fatalf("saturating flat bytes wire charge = %d, want max int %d", got, maxInt)
}
if got := saturatingLayerRPCAdmissionCharge(maxInt, 1); got != maxInt {
t.Fatalf("saturating addition = %d, want max int %d", got, maxInt)
}
}
type countingLayerRPCAdmission struct {
@ -196,6 +209,274 @@ func TestLayerRPCAdmissionExpandsTDLibNestedGZIPUnderTransferredBudget(t *testin
}
}
func TestLayerRPCAdmissionAdmitsTDLibUploadParts(t *testing.T) {
tests := []struct {
name string
method string
body bin.Object
withoutGZIP bool
bare bool
}{
{
name: "small_file_part",
method: "upload.saveFilePart",
body: &tg.UploadSaveFilePartRequest{
FileID: 91,
FilePart: 0,
Bytes: make([]byte, 1071),
},
},
{
name: "pixel_9a_big_file_part_gzip",
method: "upload.saveBigFilePart",
body: &tg.UploadSaveBigFilePartRequest{
FileID: 92,
FilePart: 0,
FileTotalParts: 364,
Bytes: make([]byte, 64<<10),
},
},
{
name: "pixel_9a_big_file_part_plain",
method: "upload.saveBigFilePart",
withoutGZIP: true,
body: &tg.UploadSaveBigFilePartRequest{
FileID: 93,
FilePart: 7,
FileTotalParts: 364,
Bytes: make([]byte, 64<<10),
},
},
{
name: "pixel_9a_big_file_part_bare_upload_session",
method: "upload.saveBigFilePart",
bare: true,
body: &tg.UploadSaveBigFilePartRequest{
FileID: 94,
FilePart: 15,
FileTotalParts: 364,
Bytes: make([]byte, 64<<10),
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 31}, sessionID: 831, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
var (
body []byte
expandedBytes int
)
if tc.bare {
body = exactOutboundLayerRPCBody(t, tlprofile.Profile228, tc.body)
} else if tc.withoutGZIP {
body = tdlibWrappedBody(t, tlprofile.Profile228, tc.body)
} else {
body, expandedBytes = tdlibNestedGZIPBody(t, tlprofile.Profile228, tc.body)
}
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if plan.items[0].kind != inboundItemRPC || plan.rpcReservation == nil || len(plan.rpcTasks) != 1 {
t.Fatalf("admitted nested gzip upload plan = kind:%d reservation:%v tasks:%d payload:%v",
plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks), plan.items[0].payload)
}
if got := plan.gzipExpandedBytes; got != expandedBytes {
t.Fatalf("nested gzip upload cumulative expansion = %d, want %d", got, expandedBytes)
}
// Admission alone is insufficient: the prepared wrapper chain must remain
// executable after the temporary gzip expansion buffer has been released.
// With no Files dependency configured, reaching the upload handler has the
// stable terminal NOT_IMPLEMENTED; INPUT_REQUEST_INVALID means the wrapper or
// prepared-call boundary corrupted the otherwise valid request.
requestBody := &bin.Buffer{Buf: append([]byte(nil), body...)}
admitted, err := router.AdmitLayerWithOptions(tlprofile.Profile228, requestBody, tlprofile.AdmissionOptions{
Limits: inboundLayerDecodeLimits,
ExpandGZIP: func(wire []byte, limit int) ([]byte, func(), error) {
return s.decodeGZIPWithGlobalBudgetLimit(&bin.Buffer{Buf: wire}, limit)
},
})
if err != nil {
t.Fatal(err)
}
_, method, err := router.DispatchAdmitted(
rpc.WithUserID(context.Background(), 42),
[8]byte{8, 31},
831,
100,
1,
admitted,
)
if method != tc.method || !tgerr.Is(err, "NOT_IMPLEMENTED") {
t.Fatalf("nested gzip upload dispatch = method:%q err:%v, want %s/NOT_IMPLEMENTED", method, err, tc.method)
}
})
}
}
func TestLayerRPCAdmissionAdmitsTDLibFirstUploadContainer(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 32}, sessionID: 832, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
// A newly opened TDLib upload Session applies its invokeWithLayer /
// initConnection header to every query in the first MTProto container. The
// Pixel 9a trace contains eight gzip-packed 64 KiB saveBigFilePart requests
// in that first container, all with the same known total and distinct
// parts/message IDs. The fixture is an ELF and its first chunks compress to
// roughly 11-14 KiB each, yielding a 129 KiB encrypted write.
plan, legacyGenericCharge := tdlibFirstUploadPlan(t)
defer plan.close()
if legacyGenericCharge <= maxInflightRPCBytes {
t.Fatalf("test fixture generic charge = %d, must exceed old connection ceiling %d", legacyGenericCharge, maxInflightRPCBytes)
}
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if plan.rpcReservation == nil || len(plan.rpcTasks) != len(plan.items) {
t.Fatalf("first upload container reservation/tasks = %v/%d, want retained/%d",
plan.rpcReservation != nil, len(plan.rpcTasks), len(plan.items))
}
if got := plan.rpcReservation.totalSize; got <= 0 || got >= legacyGenericCharge || got > maxInflightRPCBytes {
t.Fatalf("first upload container retained charge = %d, want 0 < charge < legacy %d and <= %d",
got, legacyGenericCharge, maxInflightRPCBytes)
}
for index := range plan.items {
item := &plan.items[index]
if item.kind != inboundItemRPC {
t.Fatalf("first upload container item %d kind = %d, payload = %v", index, item.kind, item.payload)
}
if method := plan.rpcTasks[index].method; method != "upload.saveBigFilePart" {
t.Fatalf("first upload container task %d method = %q, want upload.saveBigFilePart", index, method)
}
}
}
func TestLayerRPCAdmissionTDLibFirstUploadContainerRequiresFlatBytesCapability(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
// The wrapper deliberately exposes exact admission but not the optional
// flat-bytes sizing capability. The edge must therefore retain the generic
// worst-case charge and reject the whole oversized batch atomically.
handler := &countingLayerRPCAdmission{LayerRPCHandler: router}
s := New(Options{DC: 2, LayerRPC: handler, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 33}, sessionID: 833, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
plan, _ := tdlibFirstUploadPlan(t)
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
for index := range plan.items {
if plan.items[index].kind != inboundItemCapacityError {
t.Fatalf("item %d kind = %d, want capacity error", index, plan.items[index].kind)
}
}
if plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
t.Fatalf("generic fallback retained reservation/tasks = %v/%d", plan.rpcReservation != nil, len(plan.rpcTasks))
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("generic fallback leaked connection charge %d", got)
}
if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("generic fallback leaked global budget %d/%d", tasks, bytes)
}
}
func tdlibFirstUploadPlan(t *testing.T) (*inboundPlan, int64) {
t.Helper()
plan := &inboundPlan{items: make([]inboundItem, 8)}
var legacyGenericCharge int64
for index := range plan.items {
payload := make([]byte, 64<<10)
state := uint32(index + 1)
for byteIndex := 0; byteIndex < 13<<10; byteIndex++ {
state ^= state << 13
state ^= state >> 17
state ^= state << 5
payload[byteIndex] = byte(state)
}
body, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.UploadSaveBigFilePartRequest{
FileID: 95,
FilePart: index,
FileTotalParts: 364,
Bytes: payload,
})
plan.items[index] = inboundItem{
kind: inboundItemRPC,
msgID: int64(100 + index*4),
body: body,
}
legacyGenericCharge += int64(layerRPCAdmissionReservationSize(len(body) + expandedBytes))
}
return plan, legacyGenericCharge
}
func TestLayerRPCAdmissionAdmitsTDLibWrappedBindTempAuthKey(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 41}, sessionID: -7940676790771565328, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
request := &tg.AuthBindTempAuthKeyRequest{
PermAuthKeyID: 9179421154451858694,
Nonce: 5318578202586482454,
ExpiresAt: 1785817822,
EncryptedMessage: make([]byte, 104),
}
body := tdlibWrappedBody(t, tlprofile.Profile228, request)
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if plan.items[0].kind != inboundItemRPC || plan.rpcReservation == nil || len(plan.rpcTasks) != 1 {
t.Fatalf("admitted TDLib bind plan = kind:%d reservation:%v tasks:%d payload:%v",
plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks), plan.items[0].payload)
}
requestBody := &bin.Buffer{Buf: append([]byte(nil), body...)}
admitted, err := router.AdmitUnprofiled(requestBody, inboundLayerDecodeLimits)
if err != nil {
t.Fatal(err)
}
result, method, err := router.DispatchAdmitted(
context.Background(),
c.authKeyID,
c.sessionID,
100,
1,
admitted,
)
if err != nil || method != "auth.bindTempAuthKey" || result == nil || !result.WireInvariant() {
t.Fatalf("TDLib wrapped bind dispatch = method:%q result:%T invariant:%v err:%v",
method, result, result != nil && result.WireInvariant(), err)
}
}
func TestLayerRPCAdmissionNestedGZIPGrowFailureRejectsWholeBatch(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
@ -313,7 +594,7 @@ func TestLayerRPCAdmissionNestedGZIPReDecodeReusesMaterializationCharge(t *testi
plan225 := &inboundPlan{}
budget225 := &layerRPCGZIPExpansionBudget{
server: s, plan: plan225, reservation: reservation225,
baseSourceBytes: len(body), chargedSourceBytes: len(body),
baseCharge: initialCharge, chargedSize: initialCharge,
}
options225 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget225.expand}
item225 := inboundItem{msgID: 100, body: body}
@ -327,7 +608,7 @@ func TestLayerRPCAdmissionNestedGZIPReDecodeReusesMaterializationCharge(t *testi
plan227 := &inboundPlan{}
budget227 := &layerRPCGZIPExpansionBudget{
server: s, plan: plan227, reservation: reservation227,
baseSourceBytes: len(body), chargedSourceBytes: len(body),
baseCharge: initialCharge, chargedSize: initialCharge,
}
options227 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget227.expand}
item227 := inboundItem{msgID: 100, body: body}
@ -413,6 +694,11 @@ func TestLayerRPCAdmissionTransfersOriginalReservationToFreshOwner(t *testing.T)
func tdlibNestedGZIPBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) ([]byte, int) {
t.Helper()
terminalWire := exactOutboundLayerRPCBody(t, profile, terminal)
return tdlibWrappedBody(t, profile, &proto.GZIP{Data: terminalWire}), len(terminalWire)
}
func tdlibWrappedBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) []byte {
t.Helper()
request := &tg.InvokeWithLayerRequest{
Layer: int(profile),
Query: &tg.InitConnectionRequest{
@ -423,10 +709,14 @@ func tdlibNestedGZIPBody(t *testing.T, profile tlprofile.Profile, terminal bin.O
SystemLangCode: "en",
LangPack: "",
LangCode: "en",
Query: &proto.GZIP{Data: terminalWire},
Params: &tg.JSONObject{Value: []tg.JSONObjectValue{{
Key: "tz_offset",
Value: &tg.JSONNumber{Value: 8 * 60 * 60},
}}},
Query: terminal,
},
}
return exactLayerRPCBody(t, request), len(terminalWire)
return exactLayerRPCBody(t, request)
}
func TestLayerRPCAdmissionPendingReplayReleasesProvisionalEntry(t *testing.T) {

11
internal/mtprotoedge/server.go
View file

@ -93,6 +93,17 @@ type LayerRPCDefaultProfileOptionsAdmitter interface {
AdmitDefaultLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error)
}
// LayerRPCFlatBytesPayloadSizer is an optional, allocation-free admission
// capability for exact terminal requests whose generated object graph contains
// one already-bounded flat bytes payload. The handler may return ok only after
// proving the complete terminal wire shape and every semantic field cap it
// relies on. The edge still owns the multiplier, graph slack and all process /
// connection budgets; an absent or invalid hint falls back to the conservative
// generic graph charge.
type LayerRPCFlatBytesPayloadSizer interface {
LayerRPCFlatBytesPayloadSize(wire []byte) (payloadBytes int, ok bool)
}
// LayerRPCSessionProfileResolver may restore an exact profile only when it was
// previously proven for this same (auth_key_id, session_id). Auth-key-wide
// device metadata is intentionally ineligible: a client upgrade can reuse its