chore: refresh gramsrv public release

internal/mtprotoedge/rpc_result_cache.go

@@ -0,0 +1,156 @@
+package mtprotoedge
+
+import (
+	"container/list"
+	"sync"
+	"time"
+)
+
+const (
+	rpcResultCacheTTL        = 3 * time.Minute
+	rpcResultCacheMaxEntries = 4096
+	rpcResultCacheMaxBytes   = 64 << 20
+)
+
+type rpcResultCacheKey struct {
+	authKeyID [8]byte
+	sessionID int64
+	reqMsgID  int64
+}
+
+type rpcResultCacheEntry struct {
+	key       rpcResultCacheKey
+	encoded   *encodedOutboundMessage
+	size      int
+	expiresAt time.Time
+}
+
+type rpcResultCache struct {
+	mu         sync.Mutex
+	now        func() time.Time
+	ttl        time.Duration
+	maxEntries int
+	maxBytes   int
+	bytes      int
+	order      *list.List
+	byKey      map[rpcResultCacheKey]*list.Element
+}
+
+func newRPCResultCache(now func() time.Time) *rpcResultCache {
+	if now == nil {
+		now = time.Now
+	}
+	return &rpcResultCache{
+		now:        now,
+		ttl:        rpcResultCacheTTL,
+		maxEntries: rpcResultCacheMaxEntries,
+		maxBytes:   rpcResultCacheMaxBytes,
+		order:      list.New(),
+		byKey:      make(map[rpcResultCacheKey]*list.Element),
+	}
+}
+
+func (c *rpcResultCache) Get(authKeyID [8]byte, sessionID, reqMsgID int64) (*encodedOutboundMessage, bool) {
+	if c == nil || reqMsgID == 0 {
+		return nil, false
+	}
+	key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
+	now := c.now()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	elem, ok := c.byKey[key]
+	if !ok {
+		return nil, false
+	}
+	entry := elem.Value.(*rpcResultCacheEntry)
+	if !entry.expiresAt.After(now) {
+		c.removeElement(elem)
+		return nil, false
+	}
+	return cloneEncodedOutboundMessage(entry.encoded), true
+}
+
+func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encoded *encodedOutboundMessage) {
+	if c == nil || reqMsgID == 0 || encoded == nil {
+		return
+	}
+	copied := cloneEncodedOutboundMessage(encoded)
+	if copied == nil {
+		return
+	}
+	size := len(copied.body)
+	if c.maxBytes > 0 && size > c.maxBytes {
+		return
+	}
+
+	key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
+	now := c.now()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.expireLocked(now)
+	if elem, ok := c.byKey[key]; ok {
+		c.removeElement(elem)
+	}
+	entry := &rpcResultCacheEntry{
+		key:       key,
+		encoded:   copied,
+		size:      size,
+		expiresAt: now.Add(c.ttl),
+	}
+	elem := c.order.PushBack(entry)
+	c.byKey[key] = elem
+	c.bytes += size
+	c.trimLocked()
+}
+
+func (c *rpcResultCache) expireLocked(now time.Time) {
+	for elem := c.order.Front(); elem != nil; {
+		next := elem.Next()
+		entry := elem.Value.(*rpcResultCacheEntry)
+		if entry.expiresAt.After(now) {
+			return
+		}
+		c.removeElement(elem)
+		elem = next
+	}
+}
+
+func (c *rpcResultCache) trimLocked() {
+	for c.order.Len() > 0 {
+		tooManyEntries := c.maxEntries > 0 && c.order.Len() > c.maxEntries
+		tooManyBytes := c.maxBytes > 0 && c.bytes > c.maxBytes
+		if !tooManyEntries && !tooManyBytes {
+			return
+		}
+		c.removeElement(c.order.Front())
+	}
+}
+
+func (c *rpcResultCache) removeElement(elem *list.Element) {
+	if elem == nil {
+		return
+	}
+	entry := elem.Value.(*rpcResultCacheEntry)
+	delete(c.byKey, entry.key)
+	c.bytes -= entry.size
+	if c.bytes < 0 {
+		c.bytes = 0
+	}
+	c.order.Remove(elem)
+}
+
+func cloneEncodedOutboundMessage(src *encodedOutboundMessage) *encodedOutboundMessage {
+	if src == nil {
+		return nil
+	}
+	body := append([]byte(nil), src.body...)
+	return &encodedOutboundMessage{
+		body:     body,
+		typeID:   src.typeID,
+		reqMsgID: src.reqMsgID,
+	}
+}