♻️ to sync pool

This commit is contained in:
Wei Lun
2020-11-03 21:10:15 +08:00
parent 12648efd73
commit d393ee607a

View File

@@ -2,6 +2,7 @@ package mongodb
import ( import (
"context" "context"
"sync"
"time" "time"
"go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson"
@@ -12,8 +13,9 @@ import (
// Storage interface that is implemented by storage providers // Storage interface that is implemented by storage providers
type Storage struct { type Storage struct {
db *mongo.Database db *mongo.Database
col *mongo.Collection col *mongo.Collection
items *sync.Pool
} }
type MongoStorage struct { type MongoStorage struct {
@@ -94,10 +96,16 @@ func New(config ...Config) *Storage {
panic(err) panic(err)
} }
return &Storage{ store := &Storage{
db: db, db: db,
col: col, col: col,
items: &sync.Pool{
New: func() interface{} {
return new(MongoStorage)
},
},
} }
return store
} }
// Get value by key // Get value by key
@@ -112,7 +120,7 @@ func (s *Storage) Get(key string) ([]byte, error) {
return nil, err return nil, err
} }
if result.Exp.Unix() > 0 && result.Exp.Unix() <= time.Now().Unix() { if !result.Exp.IsZero() && result.Exp.Unix() <= time.Now().Unix() {
return nil, nil return nil, nil
} }
@@ -124,15 +132,16 @@ func (s *Storage) Get(key string) ([]byte, error) {
// document will be remove automatically if exp is set, based on MongoDB TTL Indexes // document will be remove automatically if exp is set, based on MongoDB TTL Indexes
func (s *Storage) Set(key string, val []byte, exp time.Duration) error { func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
filter := bson.M{"key": key} filter := bson.M{"key": key}
replace := MongoStorage{ item := s.acquireItem()
Key: key, item.Key = key
Value: val, item.Value = val
}
if exp != 0 { if exp != 0 {
replace.Exp = time.Now().Add(exp).UTC() item.Exp = time.Now().Add(exp).UTC()
} }
_, err := s.col.ReplaceOne(context.Background(), filter, replace, options.Replace().SetUpsert(true)) _, err := s.col.ReplaceOne(context.Background(), filter, item, options.Replace().SetUpsert(true))
s.releaseItem(item)
return err return err
} }
@@ -151,3 +160,19 @@ func (s *Storage) Clear() error {
func (s *Storage) Close() error { func (s *Storage) Close() error {
return s.db.Client().Disconnect(context.Background()) return s.db.Client().Disconnect(context.Background())
} }
// Acquire item from pool
func (s *Storage) acquireItem() *MongoStorage {
return s.items.Get().(*MongoStorage)
}
// Release item from pool
func (s *Storage) releaseItem(item *MongoStorage) {
if item != nil {
item.Key = ""
item.Value = nil
item.Exp = time.Time{}
s.items.Put(item)
}
}