Merge pull request #1388 from mattn/stmt-cache-lru

evict least-recently-used stmt when cache is full

sqlite3.go

@@ -451,9 +451,14 @@ type SQLiteConn struct {
 	txlock         string
 	funcs          []*functionInfo
 	aggregators    []*aggInfo
-	stmtCache      map[string][]*SQLiteStmt
+	// Prepared-statement cache. The slice is allocated at Open with a
-	stmtCacheSize  int
+	// fixed capacity equal to the configured cache size; cap bounds the
-	stmtCacheCount int
+	// cache, len is the live count, and entries are ordered LRU-first
+	// (index 0 is the oldest, the tail is most recently put). Access
+	// requires mu; stmtCacheEnabled is immutable after Open and is the
+	// only field safe to read without the lock.
+	stmtCache        []*SQLiteStmt
+	stmtCacheEnabled bool
 }
 
 // SQLiteTx implements driver.Tx.
@@ -1623,9 +1628,10 @@ func (d *SQLiteDriver) Open(dsn string) (driver.Conn, error) {
 	//
 
 	// Create connection to SQLite
-	conn := &SQLiteConn{db: db, loc: loc, txlock: txlock, stmtCacheSize: stmtCacheSize}
+	conn := &SQLiteConn{db: db, loc: loc, txlock: txlock}
 	if stmtCacheSize > 0 {
-		conn.stmtCache = make(map[string][]*SQLiteStmt)
+		conn.stmtCache = make([]*SQLiteStmt, 0, stmtCacheSize)
+		conn.stmtCacheEnabled = true
 	}
 
 	// Password Cipher has to be registered before authentication
@@ -1919,7 +1925,7 @@ func (c *SQLiteConn) dbConnOpen() bool {
 }
 
 func (c *SQLiteConn) takeCachedStmt(query string) *SQLiteStmt {
-	if c == nil || query == "" || c.stmtCacheSize <= 0 {
+	if c == nil || query == "" || !c.stmtCacheEnabled {
 		return nil
 	}
 
@@ -1929,21 +1935,25 @@ func (c *SQLiteConn) takeCachedStmt(query string) *SQLiteStmt {
 	if c.db == nil {
 		return nil
 	}
-	stmts := c.stmtCache[query]
+	// Scan from the MRU end (tail) so that a stmt put just before is
-	if len(stmts) == 0 {
+	// found immediately.
-		return nil
+	for i := len(c.stmtCache) - 1; i >= 0; i-- {
+		s := c.stmtCache[i]
+		if s.cacheKey != query {
+			continue
+		}
+		n := len(c.stmtCache)
+		copy(c.stmtCache[i:n-1], c.stmtCache[i+1:n])
+		c.stmtCache[n-1] = nil
+		c.stmtCache = c.stmtCache[:n-1]
+		// The stmt was marked closed by Close before being cached, and
+		// cls may have been set if Query opened it; reset both so the
+		// caller gets a stmt equivalent to a fresh Prepare.
+		s.closed = false
+		s.cls = false
+		return s
 	}
-	s := stmts[len(stmts)-1]
+	return nil
-	if len(stmts) == 1 {
-		delete(c.stmtCache, query)
-	} else {
-		c.stmtCache[query] = stmts[:len(stmts)-1]
-	}
-	c.stmtCacheCount--
-	s.closed = false
-	s.cls = false
-	s.t = ""
-	return s
 }
 
 func (c *SQLiteConn) putCachedStmt(s *SQLiteStmt) bool {
@@ -1954,32 +1964,46 @@ func (c *SQLiteConn) putCachedStmt(s *SQLiteStmt) bool {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	if c.db == nil || c.stmtCacheCount >= c.stmtCacheSize {
+	if c.db == nil {
 		return false
 	}
 	rv := C._sqlite3_reset_clear(s.s)
 	if rv != C.SQLITE_ROW && rv != C.SQLITE_OK && rv != C.SQLITE_DONE {
 		return false
 	}
-	c.stmtCache[s.cacheKey] = append(c.stmtCache[s.cacheKey], s)
+	// If full, finalize the LRU entry at index 0 and shift left; the
-	c.stmtCacheCount++
+	// freed tail slot is immediately reused by the append below.
+	if len(c.stmtCache) == cap(c.stmtCache) {
+		finalizeCachedStmt(c.stmtCache[0])
+		copy(c.stmtCache, c.stmtCache[1:])
+		c.stmtCache = c.stmtCache[:len(c.stmtCache)-1]
+	}
+	c.stmtCache = append(c.stmtCache, s)
 	return true
 }
 
 func (c *SQLiteConn) closeCachedStmtsLocked() {
-	for key, stmts := range c.stmtCache {
+	for i, s := range c.stmtCache {
-		for _, s := range stmts {
+		c.stmtCache[i] = nil
-			if s == nil || s.s == nil {
+		finalizeCachedStmt(s)
-				continue
-			}
-			runtime.SetFinalizer(s, nil)
-			C.sqlite3_finalize(s.s)
-			s.s = nil
-			s.c = nil
-		}
-		delete(c.stmtCache, key)
 	}
-	c.stmtCacheCount = 0
+	c.stmtCache = c.stmtCache[:0]
+}
+
+// finalizeCachedStmt tears down a stmt that was sitting in the connection's
+// stmt cache. The caller must hold c.mu. It is safe to pass a nil stmt or a
+// stmt whose handle has already been released.
+func finalizeCachedStmt(s *SQLiteStmt) {
+	if s == nil {
+		return
+	}
+	runtime.SetFinalizer(s, nil)
+	if s.s != nil {
+		C.sqlite3_finalize(s.s)
+		s.s = nil
+	}
+	s.c = nil
+	s.closed = true
 }
 
 // Prepare the query string. Return a new statement.
@@ -2014,7 +2038,7 @@ func (c *SQLiteConn) prepareWithCache(ctx context.Context, query string) (driver
 		return nil, err
 	}
 	ss := stmt.(*SQLiteStmt)
-	if ss.t == "" {
+	if ss.t == "" && c.stmtCacheEnabled {
 		ss.cacheKey = query
 	}
 	return ss, nil

sqlite3_stmt_cache_bench_test.go

@@ -0,0 +1,88 @@
+// Copyright (C) 2019 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+//go:build cgo
+// +build cgo
+
+package sqlite3
+
+import (
+	"context"
+	"database/sql/driver"
+	"fmt"
+	"testing"
+)
+
+// BenchmarkStmtCache measures the stmt cache hit / miss / eviction paths by
+// cycling through a fixed set of queries under various cache sizes. It is
+// intended for comparing cache behavior changes, not for absolute numbers.
+func BenchmarkStmtCache(b *testing.B) {
+	cases := []struct {
+		name      string
+		cacheSize int
+		keyCount  int
+	}{
+		{"off", 0, 1},                 // baseline: no cache
+		{"size4_keys1_hit", 4, 1},     // trivial hit path
+		{"size4_keys4_hit", 4, 4},     // all queries fit, always hit
+		{"size4_keys8_evict", 4, 8},   // working set > cache: miss + eviction
+		{"size16_keys8_hit", 16, 8},   // all queries fit in larger cache
+		{"size16_keys32_evict", 16, 32}, // working set >> cache
+	}
+	for _, tc := range cases {
+		b.Run(tc.name, func(b *testing.B) {
+			dsn := ":memory:"
+			if tc.cacheSize > 0 {
+				dsn = fmt.Sprintf(":memory:?_stmt_cache_size=%d", tc.cacheSize)
+			}
+			d := SQLiteDriver{}
+			conn, err := d.Open(dsn)
+			if err != nil {
+				b.Fatal(err)
+			}
+			defer conn.Close()
+			c := conn.(*SQLiteConn)
+
+			queries := make([]string, tc.keyCount)
+			for i := range queries {
+				// Distinct literal forces a distinct prepared statement.
+				queries[i] = fmt.Sprintf("SELECT %d", i+1)
+			}
+
+			ctx := context.Background()
+			// Warm up: exercise each query at least once so the cache (if any)
+			// reaches steady state before timing begins.
+			for _, q := range queries {
+				rows, err := c.query(ctx, q, nil)
+				if err != nil {
+					b.Fatal(err)
+				}
+				drainRows(b, rows)
+			}
+
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				q := queries[i%len(queries)]
+				rows, err := c.query(ctx, q, nil)
+				if err != nil {
+					b.Fatal(err)
+				}
+				drainRows(b, rows)
+			}
+		})
+	}
+}
+
+func drainRows(b *testing.B, rows driver.Rows) {
+	b.Helper()
+	dest := make([]driver.Value, len(rows.Columns()))
+	for {
+		if err := rows.Next(dest); err != nil {
+			break
+		}
+	}
+	rows.Close()
+}

sqlite3_stmt_cache_test.go

@@ -0,0 +1,153 @@
+// Copyright (C) 2019 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+//go:build cgo
+// +build cgo
+
+package sqlite3
+
+import (
+	"context"
+	"testing"
+)
+
+// TestStmtCacheLRUEviction verifies that when the prepared-statement cache is
+// full, the least-recently-used entry is evicted to make room for a new one.
+// Without eviction, the first N queries to enter the cache would squat on
+// every slot forever and any subsequently-prepared query (even a hot one)
+// would never benefit from caching.
+func TestStmtCacheLRUEviction(t *testing.T) {
+	d := SQLiteDriver{}
+	conn, err := d.Open(":memory:?_stmt_cache_size=2")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	c := conn.(*SQLiteConn)
+	ctx := context.Background()
+
+	prepareAndClose := func(q string) {
+		t.Helper()
+		stmt, err := c.prepareWithCache(ctx, q)
+		if err != nil {
+			t.Fatalf("prepareWithCache(%q): %v", q, err)
+		}
+		if err := stmt.Close(); err != nil {
+			t.Fatalf("Close(%q): %v", q, err)
+		}
+	}
+
+	q1 := "SELECT 1"
+	q2 := "SELECT 2"
+	q3 := "SELECT 3"
+
+	// Fill the cache with q1 and q2.
+	prepareAndClose(q1)
+	prepareAndClose(q2)
+	if got, want := len(c.stmtCache), 2; got != want {
+		t.Fatalf("after filling: len(stmtCache) = %d, want %d", got, want)
+	}
+	if cacheCount(c, q1) != 1 || cacheCount(c, q2) != 1 {
+		t.Fatalf("after filling: expected q1 and q2 cached, got %#v", cacheKeys(c))
+	}
+
+	// Insert q3. q1 is the oldest entry and should be evicted.
+	prepareAndClose(q3)
+	if got, want := len(c.stmtCache), 2; got != want {
+		t.Fatalf("after q3: len(stmtCache) = %d, want %d", got, want)
+	}
+	if cacheCount(c, q1) != 0 {
+		t.Fatalf("after q3: q1 should have been evicted, cache=%#v", cacheKeys(c))
+	}
+	if cacheCount(c, q2) != 1 || cacheCount(c, q3) != 1 {
+		t.Fatalf("after q3: expected q2 and q3 cached, got %#v", cacheKeys(c))
+	}
+
+	// Touching q2 should make q3 the oldest (the entry at index 0).
+	prepareAndClose(q2)
+	if len(c.stmtCache) == 0 || c.stmtCache[0].cacheKey != q3 {
+		var head string
+		if len(c.stmtCache) > 0 {
+			head = c.stmtCache[0].cacheKey
+		}
+		t.Fatalf("after touching q2: expected q3 at stmtCache[0] (LRU), got %q", head)
+	}
+
+	// Insert q1 again. Now q3 should be evicted (q2 is newer).
+	prepareAndClose(q1)
+	if cacheCount(c, q3) != 0 {
+		t.Fatalf("after reinserting q1: q3 should have been evicted, cache=%#v", cacheKeys(c))
+	}
+	if cacheCount(c, q1) != 1 || cacheCount(c, q2) != 1 {
+		t.Fatalf("after reinserting q1: expected q1 and q2 cached, got %#v", cacheKeys(c))
+	}
+	if got, want := len(c.stmtCache), 2; got != want {
+		t.Fatalf("after reinserting q1: len(stmtCache) = %d, want %d", got, want)
+	}
+
+	// Sanity-check: no dangling entries past len(stmtCache).
+	tail := c.stmtCache[:cap(c.stmtCache)]
+	for i := len(c.stmtCache); i < len(tail); i++ {
+		if tail[i] != nil {
+			t.Fatalf("stmtCache tail slot %d = %p, expected nil", i, tail[i])
+		}
+	}
+}
+
+// TestStmtCacheReuseReturnsSameHandle verifies that a cached prepare reuses
+// the underlying sqlite3_stmt rather than preparing a fresh one.
+func TestStmtCacheReuseReturnsSameHandle(t *testing.T) {
+	d := SQLiteDriver{}
+	conn, err := d.Open(":memory:?_stmt_cache_size=4")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	c := conn.(*SQLiteConn)
+	ctx := context.Background()
+
+	const q = "SELECT 42"
+	stmt1, err := c.prepareWithCache(ctx, q)
+	if err != nil {
+		t.Fatal(err)
+	}
+	h1 := stmt1.(*SQLiteStmt).s
+	if err := stmt1.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	stmt2, err := c.prepareWithCache(ctx, q)
+	if err != nil {
+		t.Fatal(err)
+	}
+	h2 := stmt2.(*SQLiteStmt).s
+	if err := stmt2.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	if h1 != h2 {
+		t.Fatalf("expected cached prepare to reuse sqlite3_stmt handle, got %p vs %p", h1, h2)
+	}
+}
+
+func cacheKeys(c *SQLiteConn) map[string]int {
+	out := make(map[string]int)
+	for _, s := range c.stmtCache {
+		out[s.cacheKey]++
+	}
+	return out
+}
+
+func cacheCount(c *SQLiteConn, q string) int {
+	n := 0
+	for _, s := range c.stmtCache {
+		if s.cacheKey == q {
+			n++
+		}
+	}
+	return n
+}