tgbot_misaka_5882f7/msgcache.go

package main

import (
	"context"
	"fmt"

	"git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/utils"
	"github.com/dgraph-io/ristretto"
	"github.com/eko/gocache/lib/v4/cache"
	gocache_lib "github.com/eko/gocache/lib/v4/store"
	ristretto_store "github.com/eko/gocache/store/ristretto/v4"
	"github.com/go-errors/errors"
	tele "gopkg.in/telebot.v3"
)

var (
	msgCache *cache.Cache[tele.Message]
	ctxVoid  = context.Background()
)

func initMsgCache() error {
	ristrettoCache, err := ristretto.NewCache(&ristretto.Config{
		NumCounters: 100_000,
		MaxCost:     20 << 20, // 20 MiB
		BufferItems: 64,
	})
	if err != nil {
		return err
	}
	ristrettoStore := ristretto_store.NewRistretto(ristrettoCache)

	msgCache = cache.New[tele.Message](ristrettoStore)
	return nil
}

func cacheMessage(msg *tele.Message) error {
	if msg == nil || msg.Chat == nil {
		return errors.New("failed to get message ID")
	}

	msgId := fmt.Sprintf("%d:%d", msg.Chat.ID, msg.ID)
	return msgCache.Set(ctxVoid, msgId, *msg)
}

func getCachedMessageById(id string) (*tele.Message, error) {
	msg, err := msgCache.Get(ctxVoid, id)
	if err != nil {
		if _, ok := err.(*gocache_lib.NotFound); ok {
			return nil, nil
		}
		return nil, errors.Wrap(err, 0)
	}
	return &msg, err
}

func getCachedMessage(msg *tele.Message) (*tele.Message, error) {
	if msg == nil || msg.Chat == nil {
		return nil, errors.New("failed to get message ID")
	}

	msgId := fmt.Sprintf("%d:%d", msg.Chat.ID, msg.ID)
	return getCachedMessageById(msgId)
}

func getCachedThread(msg *tele.Message) ([]*tele.Message, error) {
	if msg == nil {
		return nil, errors.New("empty message given")
	}

	threadR := []*tele.Message{msg}
	currentMsg := msg
	for {
		if currentMsg.ReplyTo == nil {
			break
		}

		parentMsg, err := getCachedMessage(currentMsg.ReplyTo)
		if err != nil {
			return nil, err
		}
		if parentMsg == nil {
			parentMsg = currentMsg.ReplyTo
		}
		threadR = append(threadR, parentMsg)
		currentMsg = parentMsg
	}

	return utils.Reverse(threadR), nil
}