Update mongodb.go

This commit is contained in:
Fenny
2020-11-03 15:21:35 +01:00
parent ec4771e6b6
commit 3e4f83c844

View File

@@ -18,11 +18,11 @@ type Storage struct {
items *sync.Pool items *sync.Pool
} }
type MongoStorage struct { type item struct {
ObjectID primitive.ObjectID `json:"_id,omitempty" bson:"_id,omitempty"` ObjectID primitive.ObjectID `json:"_id,omitempty" bson:"_id,omitempty"`
Key string `json:"key" bson:"key"` Key string `json:"key" bson:"key"`
Value []byte `json:"value" bson:"value"` Value []byte `json:"value" bson:"value"`
Exp time.Time `json:"exp,omitempty" bson:"exp,omitempty"` Expiration time.Time `json:"exp,omitempty" bson:"exp,omitempty"`
} }
// New creates a new MongoDB storage // New creates a new MongoDB storage
@@ -85,7 +85,10 @@ func New(config ...Config) *Storage {
// expired data may exist for some time beyond the 60 second period between runs of the background task. // expired data may exist for some time beyond the 60 second period between runs of the background task.
// more on https://docs.mongodb.com/manual/core/index-ttl/ // more on https://docs.mongodb.com/manual/core/index-ttl/
indexModel := mongo.IndexModel{ indexModel := mongo.IndexModel{
Keys: bson.D{{"exp", 1}}, Keys: bson.D{{
Key: "exp",
Value: 1,
}},
// setting to 0 // setting to 0
// means that documents will remain in the collection // means that documents will remain in the collection
// until they're explicitly deleted or the collection is dropped. // until they're explicitly deleted or the collection is dropped.
@@ -101,7 +104,7 @@ func New(config ...Config) *Storage {
col: col, col: col,
items: &sync.Pool{ items: &sync.Pool{
New: func() interface{} { New: func() interface{} {
return new(MongoStorage) return new(item)
}, },
}, },
} }
@@ -111,20 +114,23 @@ func New(config ...Config) *Storage {
// Get value by key // Get value by key
func (s *Storage) Get(key string) ([]byte, error) { func (s *Storage) Get(key string) ([]byte, error) {
res := s.col.FindOne(context.Background(), bson.M{"key": key}) res := s.col.FindOne(context.Background(), bson.M{"key": key})
result := MongoStorage{} item := s.acquireItem()
if err := res.Err(); err != nil { if err := res.Err(); err != nil {
return nil, err return nil, err
} }
if err := res.Decode(&result); err != nil { if err := res.Decode(&item); err != nil {
return nil, err return nil, err
} }
if !result.Exp.IsZero() && result.Exp.Unix() <= time.Now().Unix() { if !item.Expiration.IsZero() && item.Expiration.Unix() <= time.Now().Unix() {
return nil, nil return nil, nil
} }
// // not safe?
return result.Value, nil // res := item.Val
// s.releaseItem(item)
// return res, nil
return item.Value, nil
} }
// Set key with value, replace if document exits // Set key with value, replace if document exits
@@ -137,7 +143,7 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
item.Value = val item.Value = val
if exp != 0 { if exp != 0 {
item.Exp = time.Now().Add(exp).UTC() item.Expiration = time.Now().Add(exp).UTC()
} }
_, err := s.col.ReplaceOne(context.Background(), filter, item, options.Replace().SetUpsert(true)) _, err := s.col.ReplaceOne(context.Background(), filter, item, options.Replace().SetUpsert(true))
@@ -162,16 +168,16 @@ func (s *Storage) Close() error {
} }
// Acquire item from pool // Acquire item from pool
func (s *Storage) acquireItem() *MongoStorage { func (s *Storage) acquireItem() *item {
return s.items.Get().(*MongoStorage) return s.items.Get().(*item)
} }
// Release item from pool // Release item from pool
func (s *Storage) releaseItem(item *MongoStorage) { func (s *Storage) releaseItem(item *item) {
if item != nil { if item != nil {
item.Key = "" item.Key = ""
item.Value = nil item.Value = nil
item.Exp = time.Time{} item.Expiration = time.Time{}
s.items.Put(item) s.items.Put(item)
} }