telegabber/telegram/cache/cache.go

package cache

import (
	"sync"

	"github.com/zelenin/go-tdlib/client"
)

// 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
	chatsLock    sync.Mutex
	usersLock    sync.Mutex
	statusesLock 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{},
	}
}

// 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
}

// 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,
	}
}

// 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
}