Make the chats/users cache thread-safe

6 files changed, 99 insertions, 22 deletions

diff --git a/telegram/cache/cache.go b/telegram/cache/cache.go
new file mode 100644
index 0000000..d1d16f9
--- /dev/null
+++ b/telegram/cache/cache.go
@@ -0,0 +1,84 @@
+package cache
+
+import (
+	"sync"
+
+	"github.com/zelenin/go-tdlib/client"
+)
+
+// Cache allows operating the chats and users cache in
+// a thread-safe manner
+type Cache struct {
+	chats     map[int64]*client.Chat
+	users     map[int32]*client.User
+	chatsLock sync.Mutex
+	usersLock sync.Mutex
+}
+
+// NewCache initializes a cache
+func NewCache() *Cache {
+	return &Cache{
+		chats: map[int64]*client.Chat{},
+		users: map[int32]*client.User{},
+	}
+}
+
+// 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.chats {
+		keys = append(keys, id)
+	}
+	return keys
+}
+
+// UsersKeys grabs user ids synchronously to avoid lockups
+// while they are used
+func (cache *Cache) UsersKeys() []int32 {
+	cache.usersLock.Lock()
+	defer cache.usersLock.Unlock()
+
+	var keys []int32
+	for id := range cache.users {
+		keys = append(keys, id)
+	}
+	return keys
+}
+
+// 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.chats[id]
+	return chat, ok
+}
+
+// GetUser retrieves user by id if it's present in the cache
+func (cache *Cache) GetUser(id int32) (*client.User, bool) {
+	cache.usersLock.Lock()
+	defer cache.usersLock.Unlock()
+
+	user, ok := cache.users[id]
+	return user, ok
+}
+
+// SetChat stores a chat in the cache
+func (cache *Cache) SetChat(id int64, chat *client.Chat) {
+	cache.chatsLock.Lock()
+	defer cache.chatsLock.Unlock()
+
+	cache.chats[id] = chat
+}
+
+// SetUser stores a user in the cache
+func (cache *Cache) SetUser(id int32, user *client.User) {
+	cache.usersLock.Lock()
+	defer cache.usersLock.Unlock()
+
+	cache.users[id] = user
+}
diff --git a/telegram/client.go b/telegram/client.go
index 2acd75e..dde15f2 100644
--- a/telegram/client.go
+++ b/telegram/client.go
@@ -9,6 +9,7 @@ import (
 
 	"dev.narayana.im/narayana/telegabber/config"
 	"dev.narayana.im/narayana/telegabber/persistence"
+	"dev.narayana.im/narayana/telegabber/telegram/cache"
 
 	"github.com/zelenin/go-tdlib/client"
 	"gosrc.io/xmpp"
@@ -24,11 +25,6 @@ var logConstants = map[string]int32{
 	":all":     1023,
 }
 
-type cache struct {
-	chats map[int64]*client.Chat
-	users map[int32]*client.User
-}
-
 func stringToLogConstant(c string) int32 {
 	level, ok := logConstants[c]
 	if !ok {
@@ -51,7 +47,7 @@ type Client struct {
 	jid     string
 	Session *persistence.Session
 	content *config.TelegramContentConfig
-	cache   *cache
+	cache   *cache.Cache
 	online  bool
 
 	locks clientLocks
@@ -109,11 +105,8 @@ func NewClient(conf config.TelegramConfig, jid string, component *xmpp.Component
 		jid:        jid,
 		Session:    session,
 		content:    &conf.Content,
-		cache: &cache{
-			chats: map[int64]*client.Chat{},
-			users: map[int32]*client.User{},
-		},
-		options: options,
-		locks:   clientLocks{},
+		cache:      cache.NewCache(),
+		options:    options,
+		locks:      clientLocks{},
 	}, nil
 }
diff --git a/telegram/commands.go b/telegram/commands.go
index 533c627..558652d 100644
--- a/telegram/commands.go
+++ b/telegram/commands.go
@@ -206,7 +206,7 @@ func (c *Client) ProcessTransportCommand(cmdline string) string {
 			return notOnline
 		}
 
-		for id := range c.cache.chats {
+		for _, id := range c.cache.ChatsKeys() {
 			c.unsubscribe(id)
 		}
 
diff --git a/telegram/connect.go b/telegram/connect.go
index cb9936a..f367370 100644
--- a/telegram/connect.go
+++ b/telegram/connect.go
@@ -135,7 +135,7 @@ func (c *Client) Disconnect() {
 	log.Warn("Disconnecting from Telegram network...")
 
 	// we're offline (unsubscribe if logout)
-	for id := range c.cache.chats {
+	for _, id := range c.cache.ChatsKeys() {
 		gateway.SendPresence(
 			c.xmpp,
 			c.jid,
@@ -200,7 +200,7 @@ func (c *Client) interactor() {
 				Limit:       chatsLimit,
 			})
 			if err != nil {
-				log.Error("Could not retrieve chats")
+				log.Errorf("Could not retrieve chats: %v", err)
 			}
 
 			gateway.SendPresence(c.xmpp, c.jid, gateway.SPStatus("Logged in "+c.Session.Login))
diff --git a/telegram/handlers.go b/telegram/handlers.go
index aea1844..c79f765 100644
--- a/telegram/handlers.go
+++ b/telegram/handlers.go
@@ -108,7 +108,7 @@ func (c *Client) updateHandler() {
 
 // new user discovered
 func (c *Client) updateUser(update *client.UpdateUser) {
-	c.cache.users[update.User.Id] = update.User
+	c.cache.SetUser(update.User.Id, update.User)
 	show, status := c.userStatusToText(update.User.Status)
 	go c.processStatusUpdate(int64(update.User.Id), status, show)
 }
@@ -133,7 +133,7 @@ func (c *Client) updateNewChat(update *client.UpdateNewChat) {
 		}
 	}
 
-	c.cache.chats[update.Chat.Id] = update.Chat
+	c.cache.SetChat(update.Chat.Id, update.Chat)
 
 	var isChannel = false
 	if update.Chat.Type.ChatTypeType() == client.TypeChatTypeSupergroup {
diff --git a/telegram/utils.go b/telegram/utils.go
index cef5167..0a0a3fc 100644
--- a/telegram/utils.go
+++ b/telegram/utils.go
@@ -58,7 +58,7 @@ func (c *Client) GetContactByID(id int64, chat *client.Chat) (*client.Chat, *cli
 
 	if id <= math.MaxInt32 && id >= math.MinInt32 {
 		userID := int32(id)
-		user, ok = c.cache.users[userID]
+		user, ok = c.cache.GetUser(userID)
 		if !ok && userID > 0 {
 			user, err = c.client.GetUser(&client.GetUserRequest{
 				UserId: userID,
@@ -67,11 +67,11 @@ func (c *Client) GetContactByID(id int64, chat *client.Chat) (*client.Chat, *cli
 				return nil, nil, err
 			}
 
-			c.cache.users[userID] = user
+			c.cache.SetUser(userID, user)
 		}
 	}
 
-	cacheChat, ok = c.cache.chats[id]
+	cacheChat, ok = c.cache.GetChat(id)
 	if !ok {
 		if chat == nil {
 			cacheChat, err = c.client.GetChat(&client.GetChatRequest{
@@ -86,9 +86,9 @@ func (c *Client) GetContactByID(id int64, chat *client.Chat) (*client.Chat, *cli
 				return nil, nil, err
 			}
 
-			c.cache.chats[id] = cacheChat
+			c.cache.SetChat(id, cacheChat)
 		} else {
-			c.cache.chats[id] = chat
+			c.cache.SetChat(id, chat)
 		}
 	}
 	if chat == nil {