github.com/gofiber/fiber/v2@v2.47.0/internal/storage/memory/memory.go (about)

     1  // Package memory Is a copy of the storage memory from the external storage packet as a purpose to test the behavior
     2  // in the unittests when using a storages from these packets
     3  package memory
     4  
     5  import (
     6  	"sync"
     7  	"sync/atomic"
     8  	"time"
     9  
    10  	"github.com/gofiber/fiber/v2/utils"
    11  )
    12  
    13  // Storage interface that is implemented by storage providers
    14  type Storage struct {
    15  	mux        sync.RWMutex
    16  	db         map[string]entry
    17  	gcInterval time.Duration
    18  	done       chan struct{}
    19  }
    20  
    21  type entry struct {
    22  	data []byte
    23  	// max value is 4294967295 -> Sun Feb 07 2106 06:28:15 GMT+0000
    24  	expiry uint32
    25  }
    26  
    27  // New creates a new memory storage
    28  func New(config ...Config) *Storage {
    29  	// Set default config
    30  	cfg := configDefault(config...)
    31  
    32  	// Create storage
    33  	store := &Storage{
    34  		db:         make(map[string]entry),
    35  		gcInterval: cfg.GCInterval,
    36  		done:       make(chan struct{}),
    37  	}
    38  
    39  	// Start garbage collector
    40  	utils.StartTimeStampUpdater()
    41  	go store.gc()
    42  
    43  	return store
    44  }
    45  
    46  // Get value by key
    47  func (s *Storage) Get(key string) ([]byte, error) {
    48  	if len(key) <= 0 {
    49  		return nil, nil
    50  	}
    51  	s.mux.RLock()
    52  	v, ok := s.db[key]
    53  	s.mux.RUnlock()
    54  	if !ok || v.expiry != 0 && v.expiry <= atomic.LoadUint32(&utils.Timestamp) {
    55  		return nil, nil
    56  	}
    57  
    58  	return v.data, nil
    59  }
    60  
    61  // Set key with value
    62  func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
    63  	// Ain't Nobody Got Time For That
    64  	if len(key) <= 0 || len(val) <= 0 {
    65  		return nil
    66  	}
    67  
    68  	var expire uint32
    69  	if exp != 0 {
    70  		expire = uint32(exp.Seconds()) + atomic.LoadUint32(&utils.Timestamp)
    71  	}
    72  
    73  	e := entry{val, expire}
    74  	s.mux.Lock()
    75  	s.db[key] = e
    76  	s.mux.Unlock()
    77  	return nil
    78  }
    79  
    80  // Delete key by key
    81  func (s *Storage) Delete(key string) error {
    82  	// Ain't Nobody Got Time For That
    83  	if len(key) <= 0 {
    84  		return nil
    85  	}
    86  	s.mux.Lock()
    87  	delete(s.db, key)
    88  	s.mux.Unlock()
    89  	return nil
    90  }
    91  
    92  // Reset all keys
    93  func (s *Storage) Reset() error {
    94  	ndb := make(map[string]entry)
    95  	s.mux.Lock()
    96  	s.db = ndb
    97  	s.mux.Unlock()
    98  	return nil
    99  }
   100  
   101  // Close the memory storage
   102  func (s *Storage) Close() error {
   103  	s.done <- struct{}{}
   104  	return nil
   105  }
   106  
   107  func (s *Storage) gc() {
   108  	ticker := time.NewTicker(s.gcInterval)
   109  	defer ticker.Stop()
   110  	var expired []string
   111  
   112  	for {
   113  		select {
   114  		case <-s.done:
   115  			return
   116  		case <-ticker.C:
   117  			ts := atomic.LoadUint32(&utils.Timestamp)
   118  			expired = expired[:0]
   119  			s.mux.RLock()
   120  			for id, v := range s.db {
   121  				if v.expiry != 0 && v.expiry <= ts {
   122  					expired = append(expired, id)
   123  				}
   124  			}
   125  			s.mux.RUnlock()
   126  			s.mux.Lock()
   127  			// Double-checked locking.
   128  			// We might have replaced the item in the meantime.
   129  			for i := range expired {
   130  				v := s.db[expired[i]]
   131  				if v.expiry != 0 && v.expiry <= ts {
   132  					delete(s.db, expired[i])
   133  				}
   134  			}
   135  			s.mux.Unlock()
   136  		}
   137  	}
   138  }
   139  
   140  // Return database client
   141  func (s *Storage) Conn() map[string]entry {
   142  	s.mux.RLock()
   143  	defer s.mux.RUnlock()
   144  	return s.db
   145  }