Merge pull request #39 from 0xPolygonID/optimizations

Cache schema documents from http

.github/workflows/tests-c.yaml

@@ -14,8 +14,6 @@ jobs:
     strategy:
       matrix:
         containers:
-          - 1.19.13-bullseye
-          - 1.20.8-bullseye
           - 1.21.1-bullseye
     runs-on: ubuntu-22.04
     container: golang:${{matrix.containers}}

.github/workflows/tests.yaml

@@ -11,8 +11,6 @@ jobs:
     strategy:
       matrix:
         containers:
-          - 1.19.13-bullseye
-          - 1.20.8-bullseye
           - 1.21.1-bullseye
     runs-on: ubuntu-20.04
     container: golang:${{ matrix.containers }}

badger_cache_engine.go

@@ -0,0 +1,138 @@
+package c_polygonid
+
+import (
+	"encoding/json"
+	"errors"
+	"log/slog"
+	"time"
+
+	"github.com/dgraph-io/badger/v4"
+	"github.com/iden3/go-schema-processor/v2/loaders"
+	"github.com/piprate/json-gold/ld"
+)
+
+type cachedRemoteDocument struct {
+	RemoteDocument *ld.RemoteDocument
+	ExpireTime     time.Time
+}
+
+type badgerCacheEngine struct {
+	embedDocs map[string]*ld.RemoteDocument
+}
+
+func (m *badgerCacheEngine) Get(
+	key string) (*ld.RemoteDocument, time.Time, error) {
+
+	if m.embedDocs != nil {
+		doc, ok := m.embedDocs[key]
+		if ok {
+			return doc, time.Now().Add(time.Hour), nil
+		}
+	}
+
+	db, cleanup, err := getCacheDB()
+	if err != nil {
+		slog.Error("can't get cache database", "err", err)
+		return nil, time.Time{}, loaders.ErrCacheMiss
+	}
+	defer cleanup()
+
+	var value []byte
+
+	err = db.View(func(txn *badger.Txn) error {
+		entry, err := txn.Get([]byte(key))
+		if err != nil {
+			return err
+		}
+		value, err = entry.ValueCopy(nil)
+		return err
+	})
+	if errors.Is(err, badger.ErrKeyNotFound) {
+		return nil, time.Time{}, loaders.ErrCacheMiss
+	} else if err != nil {
+		slog.Error("error getting remote document from cache",
+			"err", err)
+		return nil, time.Time{}, loaders.ErrCacheMiss
+	}
+
+	var doc cachedRemoteDocument
+	err = json.Unmarshal(value, &doc)
+	if err != nil {
+		slog.Error("error unmarshalling cached document",
+			"err", err)
+		return nil, time.Time{}, loaders.ErrCacheMiss
+	}
+
+	return doc.RemoteDocument, doc.ExpireTime, nil
+}
+
+func (m *badgerCacheEngine) Set(key string, doc *ld.RemoteDocument,
+	expireTime time.Time) error {
+
+	if m.embedDocs != nil {
+		// if we have the document in the embedded cache, do not overwrite it
+		// with the new value.
+		_, ok := m.embedDocs[key]
+		if ok {
+			return nil
+		}
+	}
+
+	db, cleanup, err := getCacheDB()
+	if err != nil {
+		slog.Error("can't get cache database", "err", err)
+		return nil
+	}
+	defer cleanup()
+
+	value, err := json.Marshal(cachedRemoteDocument{
+		RemoteDocument: doc,
+		ExpireTime:     expireTime,
+	})
+	if err != nil {
+		slog.Error("error marshalling cached document", "err", err)
+		return nil
+	}
+
+	err = db.Update(func(txn *badger.Txn) error {
+		return txn.Set([]byte(key), value)
+	})
+	if err != nil {
+		slog.Error("error storing document to BadgerDB", "err", err)
+	}
+	return nil
+}
+
+type badgerCacheEngineOption func(*badgerCacheEngine) error
+
+func withEmbeddedDocumentBytes(u string, doc []byte) badgerCacheEngineOption {
+	return func(engine *badgerCacheEngine) error {
+		if engine.embedDocs == nil {
+			engine.embedDocs = make(map[string]*ld.RemoteDocument)
+		}
+
+		var rd = &ld.RemoteDocument{DocumentURL: u}
+		err := json.Unmarshal(doc, &rd.Document)
+		if err != nil {
+			return err
+		}
+
+		engine.embedDocs[u] = rd
+		return nil
+	}
+}
+
+func newBadgerCacheEngine(
+	opts ...badgerCacheEngineOption) (loaders.CacheEngine, error) {
+
+	e := &badgerCacheEngine{}
+
+	for _, opt := range opts {
+		err := opt(e)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return e, nil
+}

badger_cache_engine_test.go

@@ -0,0 +1,38 @@
+package c_polygonid
+
+import (
+	"testing"
+	"time"
+
+	"github.com/iden3/go-schema-processor/v2/loaders"
+	"github.com/piprate/json-gold/ld"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetPubRemoteDocument(t *testing.T) {
+	flushCacheDB()
+
+	cacheEng, err := newBadgerCacheEngine()
+	require.NoError(t, err)
+
+	key := "123"
+	doc1, expireTime1, err := cacheEng.Get(key)
+	require.EqualError(t, err, loaders.ErrCacheMiss.Error())
+	require.Nil(t, doc1)
+	require.True(t, expireTime1.IsZero())
+
+	doc := &ld.RemoteDocument{
+		DocumentURL: "123",
+		Document:    map[string]any{"one": float64(1)},
+		ContextURL:  "456",
+	}
+	expireTime := time.Now().Add(time.Hour)
+
+	err = cacheEng.Set(key, doc, expireTime)
+	require.NoError(t, err)
+
+	doc2, expireTime2, err := cacheEng.Get(key)
+	require.NoError(t, err)
+	require.Equal(t, doc, doc2)
+	require.True(t, expireTime2.Equal(expireTime))
+}

cache.go

@@ -0,0 +1,113 @@
+package c_polygonid
+
+import (
+	"log/slog"
+	"os"
+	"path"
+	"sync"
+
+	"github.com/dgraph-io/badger/v4"
+)
+
+var globalDB *badger.DB
+var dbCnt int
+var dbCond = sync.NewCond(&sync.Mutex{})
+
+func CleanCache() (err error) {
+	dbCond.L.Lock()
+	for dbCnt != 0 {
+		dbCond.Wait()
+	}
+	defer dbCond.L.Unlock()
+
+	db, err := openDB()
+	if err != nil {
+		return err
+	}
+	defer func() {
+		err2 := db.Close()
+		if err2 != nil {
+			if err == nil {
+				err = err2
+			} else {
+				slog.Error("failed to close db", "err", err2)
+			}
+		}
+	}()
+
+	return db.DropAll()
+}
+
+func getCacheDB() (*badger.DB, func(), error) {
+	dbCond.L.Lock()
+	defer dbCond.L.Unlock()
+
+	err := maybeOpenDB()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// Close the DB only once per calling getCacheDB().
+	var once sync.Once
+
+	releaseDB := func() {
+		once.Do(func() {
+			dbCond.L.Lock()
+			defer dbCond.L.Unlock()
+
+			dbCnt--
+			if dbCnt == 0 {
+				err2 := globalDB.Close()
+				if err2 != nil {
+					slog.Error("failed to close db", "err", err2)
+				}
+				globalDB = nil
+			}
+
+			dbCond.Broadcast()
+		})
+	}
+
+	return globalDB, releaseDB, nil
+}
+
+// If globalDB is nil, open a new DB, assign it to globalDB.
+// Also increment dbCnt.
+// DANGER: This function is not thread-safe and should be called only when
+// dbCond.L is locked. Use getCacheDB() instead.
+func maybeOpenDB() error {
+	if globalDB != nil {
+		dbCnt++
+		return nil
+	}
+
+	db, err := openDB()
+	if err != nil {
+		return err
+	}
+
+	globalDB = db
+	dbCnt = 1
+
+	return nil
+}
+
+func openDB() (*badger.DB, error) {
+	badgerPath, err := getBadgerPath()
+	if err != nil {
+		return nil, err
+	}
+
+	opts := badger.DefaultOptions(badgerPath)
+
+	return badger.Open(opts)
+}
+
+func getBadgerPath() (string, error) {
+	cachePath, err := os.UserCacheDir()
+	if err != nil {
+		return "", err
+	}
+	cachePath = path.Join(cachePath, "c-polygonid-cache")
+	return cachePath, nil
+}