github.com/boomhut/fiber/v2@v2.0.0-20230603160335-b65c856e57d3/internal/memory/memory.go (about)

     1  // Package memory Is a slight copy of the memory storage, but far from the storage interface it can not only work with bytes
     2  // but directly store any kind of data without having to encode it each time, which gives a huge speed advantage
     3  package memory
     4  
     5  import (
     6  	"sync"
     7  	"sync/atomic"
     8  	"time"
     9  
    10  	"github.com/boomhut/fiber/v2/utils"
    11  )
    12  
    13  type Storage struct {
    14  	sync.RWMutex
    15  	data map[string]item // data
    16  }
    17  
    18  type item struct {
    19  	// max value is 4294967295 -> Sun Feb 07 2106 06:28:15 GMT+0000
    20  	e uint32      // exp
    21  	v interface{} // val
    22  }
    23  
    24  func New() *Storage {
    25  	store := &Storage{
    26  		data: make(map[string]item),
    27  	}
    28  	utils.StartTimeStampUpdater()
    29  	go store.gc(1 * time.Second)
    30  	return store
    31  }
    32  
    33  // Get value by key
    34  func (s *Storage) Get(key string) interface{} {
    35  	s.RLock()
    36  	v, ok := s.data[key]
    37  	s.RUnlock()
    38  	if !ok || v.e != 0 && v.e <= atomic.LoadUint32(&utils.Timestamp) {
    39  		return nil
    40  	}
    41  	return v.v
    42  }
    43  
    44  // Set key with value
    45  func (s *Storage) Set(key string, val interface{}, ttl time.Duration) {
    46  	var exp uint32
    47  	if ttl > 0 {
    48  		exp = uint32(ttl.Seconds()) + atomic.LoadUint32(&utils.Timestamp)
    49  	}
    50  	i := item{exp, val}
    51  	s.Lock()
    52  	s.data[key] = i
    53  	s.Unlock()
    54  }
    55  
    56  // Delete key by key
    57  func (s *Storage) Delete(key string) {
    58  	s.Lock()
    59  	delete(s.data, key)
    60  	s.Unlock()
    61  }
    62  
    63  // Reset all keys
    64  func (s *Storage) Reset() {
    65  	nd := make(map[string]item)
    66  	s.Lock()
    67  	s.data = nd
    68  	s.Unlock()
    69  }
    70  
    71  func (s *Storage) gc(sleep time.Duration) {
    72  	ticker := time.NewTicker(sleep)
    73  	defer ticker.Stop()
    74  	var expired []string
    75  
    76  	for range ticker.C {
    77  		ts := atomic.LoadUint32(&utils.Timestamp)
    78  		expired = expired[:0]
    79  		s.RLock()
    80  		for key, v := range s.data {
    81  			if v.e != 0 && v.e <= ts {
    82  				expired = append(expired, key)
    83  			}
    84  		}
    85  		s.RUnlock()
    86  		s.Lock()
    87  		// Double-checked locking.
    88  		// We might have replaced the item in the meantime.
    89  		for i := range expired {
    90  			v := s.data[expired[i]]
    91  			if v.e != 0 && v.e <= ts {
    92  				delete(s.data, expired[i])
    93  			}
    94  		}
    95  		s.Unlock()
    96  	}
    97  }