remark42/backend/app/rest/cache/mongo.go

package cache

import (
	"log"
	"time"

	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"
	"github.com/go-pkgz/mongo"
	"github.com/go-pkgz/repeater"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/pkg/errors"
)

type mongoCache struct {
	connection   *mongo.Connection
	postFlushFn  func()
	maxKeys      int
	maxValueSize int
	maxCacheSize int64
}

const cacheCollection = "cache"

type mongoDoc struct {
	SiteID string   `bson:"site"`
	Key    string   `bson:"key"`
	Scopes []string `bson:"scopes,omitempty"`
	Data   []byte   `bson:"data"`
}

// NewMongoCache makes mongoCache implementation
func NewMongoCache(connection *mongo.Connection, options ...Option) (LoadingCache, error) {
	log.Printf("[INFO] make mongo cache with %s", connection)
	res := &mongoCache{
		connection:   connection,
		postFlushFn:  func() {},
		maxKeys:      1000,
		maxValueSize: 0,
	}
	for _, opt := range options {
		if err := opt(res); err != nil {
			return nil, errors.Wrap(err, "failed to set cache option")
		}
	}
	if err := res.prepare(); err != nil {
		return nil, err
	}
	return res, nil
}

// Get is loading cache method to get value by key or load via fn if not found
func (m *mongoCache) Get(key Key, fn func() ([]byte, error)) (data []byte, err error) {

	d := mongoDoc{}

	// repeat find from cache with small delay to avoid mgo random error
	rep := repeater.NewDefault(5, 10*time.Millisecond)
	mgErr := rep.Do(func() error {
		return m.connection.WithCustomCollection(cacheCollection, func(coll *mgo.Collection) error {
			return coll.Find(bson.M{"site": key.siteID, "key": key.id}).One(&d)
		})
	}, mgo.ErrNotFound)
	if mgErr == nil { // cached result found
		return d.Data, nil
	}

	if data, err = fn(); err != nil {
		return data, err
	}

	if mgErr != mgo.ErrNotFound { // some other error in mgo query, don't try to update cache
		log.Printf("[WARN] unexpected mgo error %+v", mgErr)
		return data, err
	}

	if !m.allowed(data) {
		return data, nil
	}

	d = mongoDoc{
		SiteID: key.siteID,
		Key:    key.id,
		Data:   data,
		Scopes: key.scopes,
	}
	err = m.connection.WithCustomCollection(cacheCollection, func(coll *mgo.Collection) error {
		_, e := coll.Upsert(bson.M{"site": key.siteID, "key": key.id}, bson.M{"$set": d})
		return e
	})
	if err != nil {
		return nil, errors.Wrapf(err, "can't set cached value for %+v", key)
	}

	if m.maxKeys > 0 {
		err = m.cleanup(key.siteID)
	}

	return data, errors.Wrap(err, "failed to cleanup cached records")
}

func (m *mongoCache) cleanup(siteID string) (err error) {
	ids := []struct {
		ID bson.ObjectId `bson:"_id"`
	}{}

	err = m.connection.WithCustomCollection(cacheCollection, func(coll *mgo.Collection) error {
		n, countErr := coll.Find(bson.M{"site": siteID}).Count()
		if countErr != nil {
			return countErr
		}
		if countErr == nil && n > m.maxKeys {
			if findErr := coll.Find(bson.M{"site": siteID}).Sort("+id").Limit(n - m.maxKeys).All(&ids); findErr == nil {
				bsonIDs := []bson.ObjectId{}
				for _, id := range ids {
					bsonIDs = append(bsonIDs, id.ID)
				}
				_, removalErr := coll.RemoveAll(bson.M{"_id": bson.M{"$in": bsonIDs}})
				return removalErr
			}
		}
		return nil
	})
	return err
}

// Flush clears cache and calls postFlushFn async
func (m *mongoCache) Flush(req FlusherRequest) {
	err := m.connection.WithCustomCollection(cacheCollection, func(coll *mgo.Collection) error {
		q := bson.M{"site": req.siteID}
		if len(req.scopes) > 0 {
			q["scopes"] = bson.M{"$in": req.scopes}
		}
		_, e := coll.RemoveAll(q)
		return e
	})

	if err == nil && m.postFlushFn != nil {
		m.postFlushFn()
	}
}

// prepare collections with all indexes
func (m *mongoCache) prepare() error {
	errs := new(multierror.Error)
	return m.connection.WithCustomCollection(cacheCollection, func(coll *mgo.Collection) error {
		errs = multierror.Append(errs, coll.EnsureIndexKey("site", "key"))
		errs = multierror.Append(errs, coll.EnsureIndexKey("site", "scopes"))
		return errors.Wrapf(errs.ErrorOrNil(), "can't create index for %s", cacheCollection)
	})
}

func (m *mongoCache) allowed(data []byte) bool {
	if m.maxValueSize > 0 && len(data) >= m.maxValueSize {
		return false
	}
	return true
}

func (m *mongoCache) setMaxValSize(max int) error {
	m.maxValueSize = max
	if max <= 0 {
		return errors.Errorf("negative size for MaxValSize, %d", max)
	}
	return nil
}

func (m *mongoCache) setMaxKeys(max int) error {
	m.maxKeys = max
	if max <= 0 {
		return errors.Errorf("negative size for MaxKeys, %d", max)
	}
	return nil
}

func (m *mongoCache) setMaxCacheSize(max int64) error {
	m.maxCacheSize = max
	if max <= 0 {
		return errors.Errorf("negative size or MaxCacheSize, %d", max)
	}
	return nil
}

func (m *mongoCache) setPostFlushFn(postFlushFn func()) error {
	m.postFlushFn = postFlushFn
	return nil
}