Rename pscache to stmtcache

2019-08-24 19:46:14 -05:00
parent beba629bb5
commit bcd6b9244a
3 changed files with 9 additions and 9 deletions
@@ -0,0 +1,111 @@
+package stmtcache
+
+import (
+	"container/list"
+	"context"
+	"fmt"
+	"sync/atomic"
+
+	"github.com/jackc/pgconn"
+)
+
+var lruCacheCount uint64
+
+// LRUCache implements cache with a Least Recently Used (LRU) cache.
+type LRUCache struct {
+	conn         *pgconn.PgConn
+	mode         int
+	cap          int
+	prepareCount int
+	m            map[string]*list.Element
+	l            *list.List
+	psNamePrefix string
+}
+
+// NewLRUCache creates a new LRUCache. mode is either ModePrepare or ModeDescribe. cap is the maximum size of the cache.
+func NewLRUCache(conn *pgconn.PgConn, mode int, cap int) *LRUCache {
+	mustBeValidMode(mode)
+	mustBeValidCap(cap)
+
+	n := atomic.AddUint64(&lruCacheCount, 1)
+
+	return &LRUCache{
+		conn:         conn,
+		mode:         mode,
+		cap:          cap,
+		m:            make(map[string]*list.Element),
+		l:            list.New(),
+		psNamePrefix: fmt.Sprintf("lrupsc_%d", n),
+	}
+}
+
+// Get returns the prepared statement description for sql preparing or describing the sql on the server as needed.
+func (c *LRUCache) Get(ctx context.Context, sql string) (*pgconn.PreparedStatementDescription, error) {
+	if el, ok := c.m[sql]; ok {
+		c.l.MoveToFront(el)
+		return el.Value.(*pgconn.PreparedStatementDescription), nil
+	}
+
+	if c.l.Len() == c.cap {
+		err := c.removeOldest(ctx)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	psd, err := c.prepare(ctx, sql)
+	if err != nil {
+		return nil, err
+	}
+
+	el := c.l.PushFront(psd)
+	c.m[sql] = el
+
+	return psd, nil
+}
+
+// Clear removes all entries in the cache. Any prepared statements will be deallocated from the PostgreSQL session.
+func (c *LRUCache) Clear(ctx context.Context) error {
+	for c.l.Len() > 0 {
+		err := c.removeOldest(ctx)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// Len returns the number of cached prepared statement descriptions.
+func (c *LRUCache) Len() int {
+	return c.l.Len()
+}
+
+// Cap returns the maximum number of cached prepared statement descriptions.
+func (c *LRUCache) Cap() int {
+	return c.cap
+}
+
+// Mode returns the mode of the cache (ModePrepare or ModeDescribe)
+func (c *LRUCache) Mode() int {
+	return c.mode
+}
+
+func (c *LRUCache) prepare(ctx context.Context, sql string) (*pgconn.PreparedStatementDescription, error) {
+	var name string
+	if c.mode == ModePrepare {
+		name = fmt.Sprintf("%s_%d", c.psNamePrefix, c.prepareCount)
+		c.prepareCount += 1
+	}
+
+	return c.conn.Prepare(ctx, name, sql, nil)
+}
+
+func (c *LRUCache) removeOldest(ctx context.Context) error {
+	oldest := c.l.Back()
+	c.l.Remove(oldest)
+	if c.mode == ModePrepare {
+		return c.conn.Exec(ctx, fmt.Sprintf("deallocate %s", oldest.Value.(*pgconn.PreparedStatementDescription).Name)).Close()
+	}
+	return nil
+}
@@ -0,0 +1,113 @@
+package stmtcache_test
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgconn"
+	"github.com/jackc/pgconn/stmtcache"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestLRUCacheModePrepare(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	conn, err := pgconn.Connect(ctx, os.Getenv("PGX_TEST_CONN_STRING"))
+	require.NoError(t, err)
+	defer conn.Close(ctx)
+
+	cache := stmtcache.NewLRUCache(conn, stmtcache.ModePrepare, 2)
+	require.EqualValues(t, 0, cache.Len())
+	require.EqualValues(t, 2, cache.Cap())
+	require.EqualValues(t, stmtcache.ModePrepare, cache.Mode())
+
+	psd, err := cache.Get(ctx, "select 1")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+	require.EqualValues(t, 1, cache.Len())
+	require.ElementsMatch(t, []string{"select 1"}, fetchServerStatements(t, ctx, conn))
+
+	psd, err = cache.Get(ctx, "select 1")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+	require.EqualValues(t, 1, cache.Len())
+	require.ElementsMatch(t, []string{"select 1"}, fetchServerStatements(t, ctx, conn))
+
+	psd, err = cache.Get(ctx, "select 2")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+	require.EqualValues(t, 2, cache.Len())
+	require.ElementsMatch(t, []string{"select 1", "select 2"}, fetchServerStatements(t, ctx, conn))
+
+	psd, err = cache.Get(ctx, "select 3")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+	require.EqualValues(t, 2, cache.Len())
+	require.ElementsMatch(t, []string{"select 2", "select 3"}, fetchServerStatements(t, ctx, conn))
+
+	err = cache.Clear(ctx)
+	require.NoError(t, err)
+	require.EqualValues(t, 0, cache.Len())
+	require.Empty(t, fetchServerStatements(t, ctx, conn))
+}
+
+func TestLRUCacheModeDescribe(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	conn, err := pgconn.Connect(ctx, os.Getenv("PGX_TEST_CONN_STRING"))
+	require.NoError(t, err)
+	defer conn.Close(ctx)
+
+	cache := stmtcache.NewLRUCache(conn, stmtcache.ModeDescribe, 2)
+	require.EqualValues(t, 0, cache.Len())
+	require.EqualValues(t, 2, cache.Cap())
+	require.EqualValues(t, stmtcache.ModeDescribe, cache.Mode())
+
+	psd, err := cache.Get(ctx, "select 1")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+	require.EqualValues(t, 1, cache.Len())
+	require.Empty(t, fetchServerStatements(t, ctx, conn))
+
+	psd, err = cache.Get(ctx, "select 1")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+	require.EqualValues(t, 1, cache.Len())
+	require.Empty(t, fetchServerStatements(t, ctx, conn))
+
+	psd, err = cache.Get(ctx, "select 2")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+	require.EqualValues(t, 2, cache.Len())
+	require.Empty(t, fetchServerStatements(t, ctx, conn))
+
+	psd, err = cache.Get(ctx, "select 3")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+	require.EqualValues(t, 2, cache.Len())
+	require.Empty(t, fetchServerStatements(t, ctx, conn))
+
+	err = cache.Clear(ctx)
+	require.NoError(t, err)
+	require.EqualValues(t, 0, cache.Len())
+	require.Empty(t, fetchServerStatements(t, ctx, conn))
+}
+
+func fetchServerStatements(t testing.TB, ctx context.Context, conn *pgconn.PgConn) []string {
+	result := conn.ExecParams(ctx, `select statement from pg_prepared_statements`, nil, nil, nil, nil).Read()
+	require.NoError(t, result.Err)
+	var statements []string
+	for _, r := range result.Rows {
+		statements = append(statements, string(r[0]))
+	}
+	return statements
+}
@@ -0,0 +1,52 @@
+// Package stmtcache is a cache that can be used to implement lazy prepared statements.
+package stmtcache
+
+import (
+	"context"
+
+	"github.com/jackc/pgconn"
+)
+
+const (
+	ModePrepare  = iota // Cache should prepare named statements.
+	ModeDescribe        // Cache should prepare the anonymous prepared statement to only fetch the description of the statement.
+)
+
+// Cache prepares and caches prepared statement descriptions.
+type Cache interface {
+	// Get returns the prepared statement description for sql preparing or describing the sql on the server as needed.
+	Get(ctx context.Context, sql string) (*pgconn.PreparedStatementDescription, error)
+
+	// Clear removes all entries in the cache. Any prepared statements will be deallocated from the PostgreSQL session.
+	Clear(ctx context.Context) error
+
+	// Len returns the number of cached prepared statement descriptions.
+	Len() int
+
+	// Cap returns the maximum number of cached prepared statement descriptions.
+	Cap() int
+
+	// Mode returns the mode of the cache (ModePrepare or ModeDescribe)
+	Mode() int
+}
+
+// New returns the preferred cache implementation for mode and cap. mode is either ModePrepare or ModeDescribe. cap is
+// the maximum size of the cache.
+func New(conn *pgconn.PgConn, mode int, cap int) Cache {
+	mustBeValidMode(mode)
+	mustBeValidCap(cap)
+
+	return NewLRUCache(conn, mode, cap)
+}
+
+func mustBeValidMode(mode int) {
+	if mode != ModePrepare && mode != ModeDescribe {
+		panic("mode must be ModePrepare or ModeDescribe")
+	}
+}
+
+func mustBeValidCap(cap int) {
+	if cap < 1 {
+		panic("cache must have cap of >= 1")
+	}
+}