telegabber/telegram/cache/cache.go

package cache

import (
	"sync"

	"github.com/zelenin/go-tdlib/client"
	"gosrc.io/xmpp/stanza"
)

// Status stores formatted data for XMPP presence
type Status struct {
	ID          int64
	XMPP        string
	Description string
}

// Cache allows operating the chats and users cache in
// a thread-safe manner
type Cache struct {
	ownChats     map[int64]*client.Chat
	auxChats     map[int64]*client.Chat
	users        map[int64]*client.User
	statuses     map[int64]*Status
	capsVers     map[int64]string
	verDiscos    map[string]*stanza.DiscoInfo
	chatsLock    sync.Mutex
	usersLock    sync.Mutex
	statusesLock sync.Mutex
	capsVersLock sync.Mutex
}

// NewCache initializes a cache
func NewCache() *Cache {
	return &Cache{
		ownChats:  map[int64]*client.Chat{},
		auxChats:  map[int64]*client.Chat{},
		users:     map[int64]*client.User{},
		statuses:  map[int64]*Status{},
		capsVers:  map[int64]string{},
		verDiscos: map[string]*stanza.DiscoInfo{},
	}
}

// ChatsKeys grabs chat ids synchronously to avoid lockups
// while they are used
func (cache *Cache) ChatsKeys() []int64 {
	cache.chatsLock.Lock()
	defer cache.chatsLock.Unlock()

	var keys []int64
	for id := range cache.ownChats {
		keys = append(keys, id)
	}
	for id := range cache.auxChats {
		keys = append(keys, id)
	}
	return keys
}

// OwnChatsKeys grabs only own chat ids synchronously to avoid lockups
// while they are used
func (cache *Cache) OwnChatsKeys() []int64 {
	cache.chatsLock.Lock()
	defer cache.chatsLock.Unlock()

	var keys []int64
	for id := range cache.ownChats {
		keys = append(keys, id)
	}
	return keys
}

// UsersKeys grabs user ids synchronously to avoid lockups
// while they are used
func (cache *Cache) UsersKeys() []int64 {
	cache.usersLock.Lock()
	defer cache.usersLock.Unlock()

	var keys []int64
	for id := range cache.users {
		keys = append(keys, id)
	}
	return keys
}

// StatusesRange loops through the map in a thread-safe manner
func (cache *Cache) StatusesRange() chan *Status {
	cache.statusesLock.Lock()

	statusChan := make(chan *Status, 1)

	go func() {
		defer func() {
			cache.statusesLock.Unlock()
			close(statusChan)
		}()

		for _, status := range cache.statuses {
			statusChan <- status
		}
	}()

	return statusChan
}

// GetChat retrieves chat by id if it's present in the cache
func (cache *Cache) GetChat(id int64) (*client.Chat, bool) {
	cache.chatsLock.Lock()
	defer cache.chatsLock.Unlock()

	chat, ok := cache.ownChats[id]
	if !ok {
		chat, ok = cache.auxChats[id]
	}
	return chat, ok
}

// GetUser retrieves user by id if it's present in the cache
func (cache *Cache) GetUser(id int64) (*client.User, bool) {
	cache.usersLock.Lock()
	defer cache.usersLock.Unlock()

	user, ok := cache.users[id]
	return user, ok
}

// GetStatus retrieves status by id if it's present in the cache
func (cache *Cache) GetStatus(id int64) (*Status, bool) {
	cache.statusesLock.Lock()
	defer cache.statusesLock.Unlock()

	status, ok := cache.statuses[id]
	return status, ok
}

// GetCapsVer retrieves capabilities verification string by id if it's present in the cache
func (cache *Cache) GetCapsVer(id int64) (string, *stanza.DiscoInfo, bool) {
	cache.capsVersLock.Lock()
	defer cache.capsVersLock.Unlock()

	ver, ok := cache.capsVers[id]
	var di *stanza.DiscoInfo
	if ok {
		di, ok = cache.verDiscos[ver]
	}
	return ver, di, ok
}

// GetVerDisco retrieves disco info by capability verification string if it's present in the cache
func (cache *Cache) GetVerDisco(ver string) (*stanza.DiscoInfo, bool) {
	cache.capsVersLock.Lock()
	defer cache.capsVersLock.Unlock()

	di, ok := cache.verDiscos[ver]
	return di, ok
}

// SetChat stores a chat in the cache
func (cache *Cache) SetChat(id int64, chat *client.Chat, own bool) {
	cache.chatsLock.Lock()
	defer cache.chatsLock.Unlock()

	if own {
		cache.ownChats[id] = chat
		// move from aux to own, but not vice versa
		// (own: true means that presences for the chat are needed
		// for sure, false means just "not necessarily")
		if _, ok := cache.auxChats[id]; ok {
			delete(cache.auxChats, id)
		}
	} else {
		cache.auxChats[id] = chat
	}
}

// SetUser stores a user in the cache
func (cache *Cache) SetUser(id int64, user *client.User) {
	cache.usersLock.Lock()
	defer cache.usersLock.Unlock()

	cache.users[id] = user
}

// SetStatus stores a status in the cache
func (cache *Cache) SetStatus(id int64, show string, status string) {
	cache.statusesLock.Lock()
	defer cache.statusesLock.Unlock()

	cache.statuses[id] = &Status{
		ID:          id,
		XMPP:        show,
		Description: status,
	}
}

// SetCapsVer stores a capabilities verification string in the cache
func (cache *Cache) SetCapsVer(id int64, ver string, di *stanza.DiscoInfo) {
	cache.capsVersLock.Lock()
	defer cache.capsVersLock.Unlock()

	cache.capsVers[id] = ver
	cache.verDiscos[ver] = di
}

// Destruct splits a cached status into show, description and type
func (status *Status) Destruct() (show, description, typ string) {
	show, description = status.XMPP, status.Description
	if show == "unavailable" {
		typ = show
		show = ""
	}
	return
}