chore: Add mutex for session storage in kv package

This commit is contained in:
divyam234 2024-06-04 11:06:04 +05:30
parent e786597e94
commit c500e0f8e0

View file

@ -3,6 +3,7 @@ package kv
import ( import (
"context" "context"
"errors" "errors"
"sync"
"github.com/gotd/td/telegram" "github.com/gotd/td/telegram"
) )
@ -10,6 +11,7 @@ import (
type Session struct { type Session struct {
kv KV kv KV
key string key string
mu sync.Mutex
} }
func NewSession(kv KV, key string) telegram.SessionStorage { func NewSession(kv KV, key string) telegram.SessionStorage {
@ -17,6 +19,8 @@ func NewSession(kv KV, key string) telegram.SessionStorage {
} }
func (s *Session) LoadSession(_ context.Context) ([]byte, error) { func (s *Session) LoadSession(_ context.Context) ([]byte, error) {
s.mu.Lock()
defer s.mu.Unlock()
b, err := s.kv.Get(s.key) b, err := s.kv.Get(s.key)
if err != nil { if err != nil {
@ -29,5 +33,7 @@ func (s *Session) LoadSession(_ context.Context) ([]byte, error) {
} }
func (s *Session) StoreSession(_ context.Context, data []byte) error { func (s *Session) StoreSession(_ context.Context, data []byte) error {
s.mu.Lock()
defer s.mu.Unlock()
return s.kv.Set(s.key, data) return s.kv.Set(s.key, data)
} }