github.com/jackc/pgx/v5@v5.5.5/internal/stmtcache/lru_cache.go (about)

     1  package stmtcache
     2  
     3  import (
     4  	"container/list"
     5  
     6  	"github.com/jackc/pgx/v5/pgconn"
     7  )
     8  
     9  // LRUCache implements Cache with a Least Recently Used (LRU) cache.
    10  type LRUCache struct {
    11  	cap          int
    12  	m            map[string]*list.Element
    13  	l            *list.List
    14  	invalidStmts []*pgconn.StatementDescription
    15  }
    16  
    17  // NewLRUCache creates a new LRUCache. cap is the maximum size of the cache.
    18  func NewLRUCache(cap int) *LRUCache {
    19  	return &LRUCache{
    20  		cap: cap,
    21  		m:   make(map[string]*list.Element),
    22  		l:   list.New(),
    23  	}
    24  }
    25  
    26  // Get returns the statement description for sql. Returns nil if not found.
    27  func (c *LRUCache) Get(key string) *pgconn.StatementDescription {
    28  	if el, ok := c.m[key]; ok {
    29  		c.l.MoveToFront(el)
    30  		return el.Value.(*pgconn.StatementDescription)
    31  	}
    32  
    33  	return nil
    34  
    35  }
    36  
    37  // Put stores sd in the cache. Put panics if sd.SQL is "". Put does nothing if sd.SQL already exists in the cache or
    38  // sd.SQL has been invalidated and HandleInvalidated has not been called yet.
    39  func (c *LRUCache) Put(sd *pgconn.StatementDescription) {
    40  	if sd.SQL == "" {
    41  		panic("cannot store statement description with empty SQL")
    42  	}
    43  
    44  	if _, present := c.m[sd.SQL]; present {
    45  		return
    46  	}
    47  
    48  	// The statement may have been invalidated but not yet handled. Do not readd it to the cache.
    49  	for _, invalidSD := range c.invalidStmts {
    50  		if invalidSD.SQL == sd.SQL {
    51  			return
    52  		}
    53  	}
    54  
    55  	if c.l.Len() == c.cap {
    56  		c.invalidateOldest()
    57  	}
    58  
    59  	el := c.l.PushFront(sd)
    60  	c.m[sd.SQL] = el
    61  }
    62  
    63  // Invalidate invalidates statement description identified by sql. Does nothing if not found.
    64  func (c *LRUCache) Invalidate(sql string) {
    65  	if el, ok := c.m[sql]; ok {
    66  		delete(c.m, sql)
    67  		c.invalidStmts = append(c.invalidStmts, el.Value.(*pgconn.StatementDescription))
    68  		c.l.Remove(el)
    69  	}
    70  }
    71  
    72  // InvalidateAll invalidates all statement descriptions.
    73  func (c *LRUCache) InvalidateAll() {
    74  	el := c.l.Front()
    75  	for el != nil {
    76  		c.invalidStmts = append(c.invalidStmts, el.Value.(*pgconn.StatementDescription))
    77  		el = el.Next()
    78  	}
    79  
    80  	c.m = make(map[string]*list.Element)
    81  	c.l = list.New()
    82  }
    83  
    84  // GetInvalidated returns a slice of all statement descriptions invalidated since the last call to RemoveInvalidated.
    85  func (c *LRUCache) GetInvalidated() []*pgconn.StatementDescription {
    86  	return c.invalidStmts
    87  }
    88  
    89  // RemoveInvalidated removes all invalidated statement descriptions. No other calls to Cache must be made between a
    90  // call to GetInvalidated and RemoveInvalidated or RemoveInvalidated may remove statement descriptions that were
    91  // never seen by the call to GetInvalidated.
    92  func (c *LRUCache) RemoveInvalidated() {
    93  	c.invalidStmts = nil
    94  }
    95  
    96  // Len returns the number of cached prepared statement descriptions.
    97  func (c *LRUCache) Len() int {
    98  	return c.l.Len()
    99  }
   100  
   101  // Cap returns the maximum number of cached prepared statement descriptions.
   102  func (c *LRUCache) Cap() int {
   103  	return c.cap
   104  }
   105  
   106  func (c *LRUCache) invalidateOldest() {
   107  	oldest := c.l.Back()
   108  	sd := oldest.Value.(*pgconn.StatementDescription)
   109  	c.invalidStmts = append(c.invalidStmts, sd)
   110  	delete(c.m, sd.SQL)
   111  	c.l.Remove(oldest)
   112  }