github.com/gofiber/fiber/v2@v2.47.0/middleware/limiter/manager.go (about)

     1  package limiter
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/gofiber/fiber/v2"
     8  	"github.com/gofiber/fiber/v2/internal/memory"
     9  )
    10  
    11  // go:generate msgp
    12  // msgp -file="manager.go" -o="manager_msgp.go" -tests=false -unexported
    13  type item struct {
    14  	currHits int
    15  	prevHits int
    16  	exp      uint64
    17  }
    18  
    19  //msgp:ignore manager
    20  type manager struct {
    21  	pool    sync.Pool
    22  	memory  *memory.Storage
    23  	storage fiber.Storage
    24  }
    25  
    26  func newManager(storage fiber.Storage) *manager {
    27  	// Create new storage handler
    28  	manager := &manager{
    29  		pool: sync.Pool{
    30  			New: func() interface{} {
    31  				return new(item)
    32  			},
    33  		},
    34  	}
    35  	if storage != nil {
    36  		// Use provided storage if provided
    37  		manager.storage = storage
    38  	} else {
    39  		// Fallback too memory storage
    40  		manager.memory = memory.New()
    41  	}
    42  	return manager
    43  }
    44  
    45  // acquire returns an *entry from the sync.Pool
    46  func (m *manager) acquire() *item {
    47  	return m.pool.Get().(*item) //nolint:forcetypeassert // We store nothing else in the pool
    48  }
    49  
    50  // release and reset *entry to sync.Pool
    51  func (m *manager) release(e *item) {
    52  	e.prevHits = 0
    53  	e.currHits = 0
    54  	e.exp = 0
    55  	m.pool.Put(e)
    56  }
    57  
    58  // get data from storage or memory
    59  func (m *manager) get(key string) *item {
    60  	var it *item
    61  	if m.storage != nil {
    62  		it = m.acquire()
    63  		raw, err := m.storage.Get(key)
    64  		if err != nil {
    65  			return it
    66  		}
    67  		if raw != nil {
    68  			if _, err := it.UnmarshalMsg(raw); err != nil {
    69  				return it
    70  			}
    71  		}
    72  		return it
    73  	}
    74  	if it, _ = m.memory.Get(key).(*item); it == nil { //nolint:errcheck // We store nothing else in the pool
    75  		it = m.acquire()
    76  		return it
    77  	}
    78  	return it
    79  }
    80  
    81  // set data to storage or memory
    82  func (m *manager) set(key string, it *item, exp time.Duration) {
    83  	if m.storage != nil {
    84  		if raw, err := it.MarshalMsg(nil); err == nil {
    85  			_ = m.storage.Set(key, raw, exp) //nolint:errcheck // TODO: Handle error here
    86  		}
    87  		// we can release data because it's serialized to database
    88  		m.release(it)
    89  	} else {
    90  		m.memory.Set(key, it, exp)
    91  	}
    92  }