Configurable caching of ephemeral prepared statements

[CAFxX]

This is more of as a feature suggestion than anything else (I would have added it to the wiki, but it's readonly): since the driver already creates ephemeral prepared statements when executing queries with interpolateParams=false it could be interesting to allow users to opt-in to have the driver dynamically cache up to a configurable number of automatically-created prepared statements corresponding to the most frequently executed queries.

By default this cache would be off, and users could opt-in e.g. with a DSN parameter like cacheStatements=N, where N is the maximum number of prepared statements that can be in the cache. The driver internally would keep a larger (but bounded) approximate set of the most frequently executed queries, but only the top N statements would be eligible for having the corresponding prepared statement be kept in the cache. The documentation should clearly state that N is per-connection, and that therefore the total number of prepared statements on the server used by this cache can reach up to N*num_open_connections, and that therefore careful configuration is required both on the client (SetMaxOpenConns) and on the server (max_prepared_stmt_count).

The rationale for this is that in our experience many services have a (small) subset of hot queries. If the driver implemented this feature, it would transparently give significant performance benefits - as it would avoid the prepared statement roundtrip at least on the most frequently executed queries without requiring users to manually curate and maintain the list of which statements should be prepared beforehand.

[ljluestc]

package mssql

import (
	"container/list"
	"sync"
)

// StatementCache manages a per-connection cache of prepared statements
type StatementCache struct {
	mu          sync.Mutex
	maxSize     int                    // Maximum number of cached statements
	cache       map[string]*cacheEntry // Query string to cache entry
	lruList     *list.List             // LRU list for eviction
	freqCounter map[string]int         // Approximate frequency counter
	freqMaxSize int                    // Maximum size of frequency counter
}

// cacheEntry holds a prepared statement and its LRU list element
type cacheEntry struct {
	stmt    *Stmt
	element *list.Element
}

// NewStatementCache creates a new cache with the given maximum size
func NewStatementCache(maxSize int) *StatementCache {
	return &StatementCache{
		maxSize:     maxSize,
		cache:       make(map[string]*cacheEntry),
		lruList:     list.New(),
		freqCounter: make(map[string]int),
		freqMaxSize: maxSize * 2, // Track more queries than cached statements
	}
}

// Get retrieves a cached prepared statement for the query, if available
func (sc *StatementCache) Get(query string) *Stmt {
	if sc.maxSize == 0 {
		return nil
	}
	sc.mu.Lock()
	defer sc.mu.Unlock()

	entry, ok := sc.cache[query]
	if !ok {
		return nil
	}
	sc.lruList.MoveToFront(entry.element)
	return entry.stmt
}

// Put attempts to cache a prepared statement for the query
func (sc *StatementCache) Put(query string, stmt *Stmt) {
	if sc.maxSize == 0 {
		stmt.Close()
		return
	}
	sc.mu.Lock()
	defer sc.mu.Unlock()

	// Increment frequency counter
	sc.freqCounter[query]++
	if len(sc.freqCounter) > sc.freqMaxSize {
		// Evict least frequent query
		var minQuery string
		minCount := int(^uint(0) >> 1) // Max int
		for q, count := range sc.freqCounter {
			if count < minCount {
				minCount = count
				minQuery = q
			}
		}
		delete(sc.freqCounter, minQuery)
	}

	// Only cache if query is frequent
	if !sc.isFrequent(query) {
		stmt.Close()
		return
	}

	// Update existing entry
	if entry, ok := sc.cache[query]; ok {
		sc.lruList.MoveToFront(entry.element)
		entry.stmt.Close()
		entry.stmt = stmt
		return
	}

	// Evict oldest if cache is full
	if len(sc.cache) >= sc.maxSize {
		oldest := sc.lruList.Back()
		if oldest != nil {
			oldQuery := oldest.Value.(string)
			if oldEntry, ok := sc.cache[oldQuery]; ok {
				oldEntry.stmt.Close()
				delete(sc.cache, oldQuery)
				sc.lruList.Remove(oldest)
			}
		}
	}

	// Add new entry
	element := sc.lruList.PushFront(query)
	sc.cache[query] = &cacheEntry{stmt: stmt, element: element}
}

// isFrequent checks if the query is among the top N frequent queries
func (sc *StatementCache) isFrequent(query string) bool {
	count := sc.freqCounter[query]
	if len(sc.freqCounter) <= sc.maxSize {
		return true
	}
	// Find Nth highest frequency
	counts := make([]int, 0, len(sc.freqCounter))
	for _, c := range sc.freqCounter {
		counts = append(counts, c)
	}
	// Simple selection sort for top N
	for i := 0; i < sc.maxSize && i < len(counts); i++ {
		maxIdx := i
		for j := i + 1; j < len(counts); j++ {
			if counts[j] > counts[maxIdx] {
				maxIdx = j
			}
		}
		counts[i], counts[maxIdx] = counts[maxIdx], counts[i]
	}
	return count >= counts[sc.maxSize-1]
}

// Close closes all cached statements
func (sc *StatementCache) Close() {
	sc.mu.Lock()
	defer sc.mu.Unlock()
	for query, entry := range sc.cache {
		entry.stmt.Close()
		delete(sc.cache, query)
	}
	sc.lruList.Init()
	sc.freqCounter = make(map[string]int)
}