pkg/memstore: Rename from pkg/storage/memstore

1 files changed, 127 insertions, 0 deletions

diff --git a/pkg/memstore/memstore.go b/pkg/memstore/memstore.go
new file mode 100644
index 0000000..f1421c2
--- /dev/null
+++ b/pkg/memstore/memstore.go
@@ -0,0 +1,127 @@
+// Package memstore implements in-memory storage for atlas.
+package memstore
+
+import (
+	"bytes"
+	"compress/gzip"
+	"crypto/sha256"
+	"io"
+	"strings"
+	"sync"
+
+	"github.com/pg9182/atlas/pkg/api/api0"
+)
+
+// AccountStore stores accounts in-memory.
+type AccountStore struct {
+	accounts sync.Map
+}
+
+// NewPdataStore creates a new MemoryPdataStore.
+func NewAccountStore() *AccountStore {
+	return &AccountStore{}
+}
+
+func (m *AccountStore) GetUIDsByUsername(username string) ([]uint64, error) {
+	var uids []uint64
+	if username != "" {
+		m.accounts.Range(func(_, v any) bool {
+			if u := v.(api0.Account); strings.EqualFold(u.Username, username) {
+				uids = append(uids, u.UID)
+			}
+			return true
+		})
+	}
+	return uids, nil
+}
+
+func (m *AccountStore) GetAccount(uid uint64) (*api0.Account, error) {
+	v, ok := m.accounts.Load(uid)
+	if !ok {
+		return nil, nil
+	}
+	a := v.(api0.Account)
+	return &a, nil
+}
+
+func (m *AccountStore) SaveAccount(a *api0.Account) error {
+	if a != nil {
+		m.accounts.Store(a.UID, *a)
+	}
+	return nil
+}
+
+// PdataStore stores pdata in-memory, with optional compression.
+type PdataStore struct {
+	gzip  bool
+	pdata sync.Map
+}
+
+type pdataStoreEntry struct {
+	Hash [sha256.Size]byte
+	Data []byte
+}
+
+// NewPdataStore creates a new MemoryPdataStore.
+func NewPdataStore(compress bool) *PdataStore {
+	return &PdataStore{
+		gzip: compress,
+	}
+}
+
+func (m *PdataStore) GetPdataHash(uid uint64) ([sha256.Size]byte, bool, error) {
+	v, ok := m.pdata.Load(uid)
+	if !ok {
+		return [sha256.Size]byte{}, ok, nil
+	}
+	return v.(pdataStoreEntry).Hash, ok, nil
+}
+
+func (m *PdataStore) GetPdataCached(uid uint64, sha [sha256.Size]byte) ([]byte, bool, error) {
+	v, ok := m.pdata.Load(uid)
+	if !ok {
+		return nil, ok, nil
+	}
+	e := v.(pdataStoreEntry)
+	if sha != [sha256.Size]byte{} && sha == e.Hash {
+		return nil, ok, nil
+	}
+	var b []byte
+	if m.gzip {
+		r, err := gzip.NewReader(bytes.NewReader(e.Data))
+		if err != nil {
+			return nil, ok, err
+		}
+		b, err = io.ReadAll(r)
+		if err != nil {
+			return nil, ok, err
+		}
+	} else {
+		b = make([]byte, len(e.Data))
+		copy(b, e.Data)
+	}
+	return b, ok, nil
+}
+
+func (m *PdataStore) SetPdata(uid uint64, buf []byte) error {
+	var b []byte
+	if m.gzip {
+		var f bytes.Buffer
+		w := gzip.NewWriter(&f)
+		if _, err := w.Write(buf); err != nil {
+			return err
+		}
+		if err := w.Close(); err != nil {
+			return err
+		}
+		b = f.Bytes()
+	} else {
+		b = make([]byte, len(buf))
+		copy(b, buf)
+	}
+	m.pdata.Store(uid, pdataStoreEntry{
+		Hash: sha256.Sum256(buf),
+		Data: b,
+	})
+	return nil
+}