feat: sync multilayer td integration
This commit is contained in:
A
parent
20a310f6ca
commit
766c5db992
491 changed files with 26235 additions and 35340 deletions
292
internal/mtprotoedge/rpc_result_budget.go
Normal file
292
internal/mtprotoedge/rpc_result_budget.go
Normal file
|
|
@ -0,0 +1,292 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"hash/maphash"
|
||||
"sync"
|
||||
)
|
||||
|
||||
const rpcResultBudgetShards = 64
|
||||
|
||||
type rpcResultBudgetLimit struct {
|
||||
entries int64
|
||||
bytes int64
|
||||
}
|
||||
|
||||
type rpcResultBudgetUsage struct {
|
||||
entries int64
|
||||
bytes int64
|
||||
pending int64
|
||||
}
|
||||
|
||||
type rpcResultSessionBudgetKey struct {
|
||||
authKeyID [8]byte
|
||||
sessionID int64
|
||||
}
|
||||
|
||||
type rpcResultAuthBudgetShard struct {
|
||||
mu sync.Mutex
|
||||
usage map[[8]byte]rpcResultBudgetUsage
|
||||
}
|
||||
|
||||
type rpcResultSessionBudgetShard struct {
|
||||
mu sync.Mutex
|
||||
usage map[rpcResultSessionBudgetKey]rpcResultBudgetUsage
|
||||
}
|
||||
|
||||
// rpcResultFairBudget accounts one ownership reservation at all three scopes.
|
||||
// A pending owner and its completed result are the same ownership: admission
|
||||
// reserves one entry plus one byte, Put resizes that byte reservation and moves
|
||||
// it to the completed row, while Abort/TTL return the whole reservation.
|
||||
//
|
||||
// Auth and session maps are striped independently. Every operation takes the
|
||||
// auth stripe before the session stripe; global counters remain atomic. This
|
||||
// keeps unrelated auth keys off one process-wide mutex while preserving hard
|
||||
// limits at every hierarchy level.
|
||||
type rpcResultFairBudget struct {
|
||||
seed maphash.Seed
|
||||
globalEntries *rpcResultFlightLimit
|
||||
globalBytes *rpcResultCacheByteBudget
|
||||
authLimit rpcResultBudgetLimit
|
||||
sessionLimit rpcResultBudgetLimit
|
||||
pendingPerAuth int64
|
||||
authShards [rpcResultBudgetShards]rpcResultAuthBudgetShard
|
||||
sessionShards [rpcResultBudgetShards]rpcResultSessionBudgetShard
|
||||
}
|
||||
|
||||
type rpcResultBudgetReservation struct {
|
||||
budget *rpcResultFairBudget
|
||||
key rpcResultCacheKey
|
||||
bytes int
|
||||
pending bool
|
||||
released bool
|
||||
}
|
||||
|
||||
func newRPCResultFairBudget(
|
||||
seed maphash.Seed,
|
||||
globalEntries *rpcResultFlightLimit,
|
||||
globalBytes *rpcResultCacheByteBudget,
|
||||
authLimit rpcResultBudgetLimit,
|
||||
sessionLimit rpcResultBudgetLimit,
|
||||
pendingPerAuth int,
|
||||
) *rpcResultFairBudget {
|
||||
b := &rpcResultFairBudget{
|
||||
seed: seed,
|
||||
globalEntries: globalEntries,
|
||||
globalBytes: globalBytes,
|
||||
authLimit: authLimit,
|
||||
sessionLimit: sessionLimit,
|
||||
pendingPerAuth: int64(pendingPerAuth),
|
||||
}
|
||||
for i := range b.authShards {
|
||||
b.authShards[i].usage = make(map[[8]byte]rpcResultBudgetUsage)
|
||||
b.sessionShards[i].usage = make(map[rpcResultSessionBudgetKey]rpcResultBudgetUsage)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *rpcResultFairBudget) reserveOwner(key rpcResultCacheKey) *rpcResultBudgetReservation {
|
||||
return b.reserve(key, 1, true)
|
||||
}
|
||||
|
||||
func (b *rpcResultFairBudget) reserveCompleted(key rpcResultCacheKey, bytes int) *rpcResultBudgetReservation {
|
||||
return b.reserve(key, bytes, false)
|
||||
}
|
||||
|
||||
func (b *rpcResultFairBudget) reserve(key rpcResultCacheKey, bytes int, pending bool) *rpcResultBudgetReservation {
|
||||
if b == nil || b.globalEntries == nil || b.globalBytes == nil || bytes < 1 {
|
||||
return nil
|
||||
}
|
||||
authShard := b.authShard(key.authKeyID)
|
||||
sessionKey := rpcResultSessionBudgetKey{authKeyID: key.authKeyID, sessionID: key.sessionID}
|
||||
sessionShard := b.sessionShard(sessionKey)
|
||||
authShard.mu.Lock()
|
||||
sessionShard.mu.Lock()
|
||||
|
||||
authUsage := authShard.usage[key.authKeyID]
|
||||
sessionUsage := sessionShard.usage[sessionKey]
|
||||
bytes64 := int64(bytes)
|
||||
canReserve := withinRPCResultBudget(authUsage.entries, 1, b.authLimit.entries) &&
|
||||
withinRPCResultBudget(authUsage.bytes, bytes64, b.authLimit.bytes) &&
|
||||
withinRPCResultBudget(sessionUsage.entries, 1, b.sessionLimit.entries) &&
|
||||
withinRPCResultBudget(sessionUsage.bytes, bytes64, b.sessionLimit.bytes)
|
||||
if pending {
|
||||
canReserve = canReserve && withinRPCResultBudget(authUsage.pending, 1, b.pendingPerAuth)
|
||||
}
|
||||
if !canReserve || !b.globalEntries.reserve() {
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
if !b.globalBytes.reserve(bytes) {
|
||||
b.globalEntries.release()
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
authUsage.entries++
|
||||
authUsage.bytes += bytes64
|
||||
sessionUsage.entries++
|
||||
sessionUsage.bytes += bytes64
|
||||
if pending {
|
||||
authUsage.pending++
|
||||
}
|
||||
authShard.usage[key.authKeyID] = authUsage
|
||||
sessionShard.usage[sessionKey] = sessionUsage
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
return &rpcResultBudgetReservation{budget: b, key: key, bytes: bytes, pending: pending}
|
||||
}
|
||||
|
||||
func withinRPCResultBudget(used, delta, limit int64) bool {
|
||||
return delta >= 0 && limit > 0 && used >= 0 && used <= limit-delta
|
||||
}
|
||||
|
||||
func (r *rpcResultBudgetReservation) resizeBytes(bytes int) bool {
|
||||
if r == nil || r.budget == nil || r.released || bytes < 1 {
|
||||
return false
|
||||
}
|
||||
if bytes == r.bytes {
|
||||
return true
|
||||
}
|
||||
b := r.budget
|
||||
authShard := b.authShard(r.key.authKeyID)
|
||||
sessionKey := rpcResultSessionBudgetKey{authKeyID: r.key.authKeyID, sessionID: r.key.sessionID}
|
||||
sessionShard := b.sessionShard(sessionKey)
|
||||
authShard.mu.Lock()
|
||||
sessionShard.mu.Lock()
|
||||
authUsage, authOK := authShard.usage[r.key.authKeyID]
|
||||
sessionUsage, sessionOK := sessionShard.usage[sessionKey]
|
||||
if !authOK || !sessionOK || authUsage.entries < 1 || sessionUsage.entries < 1 {
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
panic("mtprotoedge: rpc result budget reservation disappeared during resize")
|
||||
}
|
||||
delta := int64(bytes) - int64(r.bytes)
|
||||
if delta > 0 {
|
||||
if !withinRPCResultBudget(authUsage.bytes, delta, b.authLimit.bytes) ||
|
||||
!withinRPCResultBudget(sessionUsage.bytes, delta, b.sessionLimit.bytes) ||
|
||||
!b.globalBytes.reserve(int(delta)) {
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
} else {
|
||||
if authUsage.bytes < -delta || sessionUsage.bytes < -delta {
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
panic("mtprotoedge: rpc result byte reservation underflow during resize")
|
||||
}
|
||||
}
|
||||
authUsage.bytes += delta
|
||||
sessionUsage.bytes += delta
|
||||
authShard.usage[r.key.authKeyID] = authUsage
|
||||
sessionShard.usage[sessionKey] = sessionUsage
|
||||
r.bytes = bytes
|
||||
if delta < 0 {
|
||||
b.globalBytes.release(int(-delta))
|
||||
}
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
return true
|
||||
}
|
||||
|
||||
func (r *rpcResultBudgetReservation) releasePending() {
|
||||
if r == nil || r.budget == nil || r.released || !r.pending {
|
||||
return
|
||||
}
|
||||
b := r.budget
|
||||
authShard := b.authShard(r.key.authKeyID)
|
||||
authShard.mu.Lock()
|
||||
authUsage, ok := authShard.usage[r.key.authKeyID]
|
||||
if !ok || authUsage.pending < 1 {
|
||||
authShard.mu.Unlock()
|
||||
panic("mtprotoedge: rpc result per-auth pending budget underflow")
|
||||
}
|
||||
authUsage.pending--
|
||||
authShard.usage[r.key.authKeyID] = authUsage
|
||||
r.pending = false
|
||||
authShard.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *rpcResultBudgetReservation) release() {
|
||||
if r == nil || r.budget == nil || r.released {
|
||||
return
|
||||
}
|
||||
b := r.budget
|
||||
authShard := b.authShard(r.key.authKeyID)
|
||||
sessionKey := rpcResultSessionBudgetKey{authKeyID: r.key.authKeyID, sessionID: r.key.sessionID}
|
||||
sessionShard := b.sessionShard(sessionKey)
|
||||
authShard.mu.Lock()
|
||||
sessionShard.mu.Lock()
|
||||
authUsage, authOK := authShard.usage[r.key.authKeyID]
|
||||
sessionUsage, sessionOK := sessionShard.usage[sessionKey]
|
||||
bytes64 := int64(r.bytes)
|
||||
if !authOK || !sessionOK || authUsage.entries < 1 || sessionUsage.entries < 1 ||
|
||||
authUsage.bytes < bytes64 || sessionUsage.bytes < bytes64 ||
|
||||
(r.pending && authUsage.pending < 1) {
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
panic("mtprotoedge: rpc result fair budget underflow")
|
||||
}
|
||||
authUsage.entries--
|
||||
authUsage.bytes -= bytes64
|
||||
sessionUsage.entries--
|
||||
sessionUsage.bytes -= bytes64
|
||||
if r.pending {
|
||||
authUsage.pending--
|
||||
}
|
||||
if authUsage == (rpcResultBudgetUsage{}) {
|
||||
delete(authShard.usage, r.key.authKeyID)
|
||||
} else {
|
||||
authShard.usage[r.key.authKeyID] = authUsage
|
||||
}
|
||||
if sessionUsage == (rpcResultBudgetUsage{}) {
|
||||
delete(sessionShard.usage, sessionKey)
|
||||
} else {
|
||||
sessionShard.usage[sessionKey] = sessionUsage
|
||||
}
|
||||
r.released = true
|
||||
r.pending = false
|
||||
r.bytes = 0
|
||||
b.globalBytes.release(int(bytes64))
|
||||
b.globalEntries.release()
|
||||
sessionShard.mu.Unlock()
|
||||
authShard.mu.Unlock()
|
||||
}
|
||||
|
||||
func (b *rpcResultFairBudget) authSnapshot(authKeyID [8]byte) rpcResultBudgetUsage {
|
||||
if b == nil {
|
||||
return rpcResultBudgetUsage{}
|
||||
}
|
||||
shard := b.authShard(authKeyID)
|
||||
shard.mu.Lock()
|
||||
usage := shard.usage[authKeyID]
|
||||
shard.mu.Unlock()
|
||||
return usage
|
||||
}
|
||||
|
||||
func (b *rpcResultFairBudget) sessionSnapshot(authKeyID [8]byte, sessionID int64) rpcResultBudgetUsage {
|
||||
if b == nil {
|
||||
return rpcResultBudgetUsage{}
|
||||
}
|
||||
key := rpcResultSessionBudgetKey{authKeyID: authKeyID, sessionID: sessionID}
|
||||
shard := b.sessionShard(key)
|
||||
shard.mu.Lock()
|
||||
usage := shard.usage[key]
|
||||
shard.mu.Unlock()
|
||||
return usage
|
||||
}
|
||||
|
||||
func (b *rpcResultFairBudget) authShard(authKeyID [8]byte) *rpcResultAuthBudgetShard {
|
||||
index := maphash.Bytes(b.seed, authKeyID[:]) & (rpcResultBudgetShards - 1)
|
||||
return &b.authShards[index]
|
||||
}
|
||||
|
||||
func (b *rpcResultFairBudget) sessionShard(key rpcResultSessionBudgetKey) *rpcResultSessionBudgetShard {
|
||||
var raw [16]byte
|
||||
copy(raw[:8], key.authKeyID[:])
|
||||
binary.LittleEndian.PutUint64(raw[8:], uint64(key.sessionID))
|
||||
index := maphash.Bytes(b.seed, raw[:]) & (rpcResultBudgetShards - 1)
|
||||
return &b.sessionShards[index]
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue