github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/utils/cache/cache.go (about)

     1  package cache
     2  
     3  import "go.etcd.io/bbolt"
     4  
     5  type Cache struct {
     6  	db *bbolt.DB
     7  
     8  	bucketName []byte
     9  }
    10  
    11  func NewCache(db *bbolt.DB, bucketName string) *Cache {
    12  	return &Cache{
    13  		db:         db,
    14  		bucketName: []byte(bucketName),
    15  	}
    16  }
    17  
    18  func (c *Cache) Get(k []byte) (v []byte) {
    19  	if c.db == nil {
    20  		return nil
    21  	}
    22  
    23  	_ = c.db.View(func(tx *bbolt.Tx) error {
    24  		bk := tx.Bucket(c.bucketName)
    25  		if bk == nil {
    26  			return nil
    27  		}
    28  
    29  		vv := bk.Get(k)
    30  		if vv != nil {
    31  			v = make([]byte, len(vv))
    32  			copy(v, vv)
    33  		}
    34  		return nil
    35  	})
    36  
    37  	return v
    38  }
    39  
    40  func (c *Cache) Put(k, v []byte) {
    41  	if c.db == nil {
    42  		return
    43  	}
    44  
    45  	_ = c.db.Batch(func(tx *bbolt.Tx) error {
    46  		bk, err := tx.CreateBucketIfNotExists(c.bucketName)
    47  		if err != nil {
    48  			return err
    49  		}
    50  
    51  		return bk.Put(k, v)
    52  	})
    53  }
    54  
    55  func (c *Cache) Delete(k ...[]byte) {
    56  	if c.db == nil {
    57  		return
    58  	}
    59  	_ = c.db.Batch(func(tx *bbolt.Tx) error {
    60  		b := tx.Bucket(c.bucketName)
    61  		if b == nil {
    62  			return nil
    63  		}
    64  
    65  		for _, kk := range k {
    66  			if kk == nil {
    67  				continue
    68  			}
    69  
    70  			if err := b.Delete(kk); err != nil {
    71  				return err
    72  			}
    73  		}
    74  
    75  		return nil
    76  	})
    77  }